diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/spec_worker.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/spec_worker.py index 277d1f07ecf3..a35aa480082e 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/spec_worker.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp16ep1/spec_worker.py @@ -61,6 +61,9 @@ class KimiK3DFlashWorker(DFlashWorker): # inputs as trtllm::k3_spec_accept, keep the target logits vocabulary-sharded and write the drafter's context # K/V with trtllm::k3_ctx_kv. k3_decode = False + # The generation-step preparation below reads _ctx_block_tables, not the + # rebased view, so every drafter slot numbers from 0. + _ctx_reads_rebased_view = False def _mask_token_id(self, draft_model) -> int: """The drafter's mask token id, resolved once.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/spec_worker.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/spec_worker.py index 277d1f07ecf3..a35aa480082e 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/spec_worker.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_mxfp4__sm_100__tp16_moetp4ep4/spec_worker.py @@ -61,6 +61,9 @@ class KimiK3DFlashWorker(DFlashWorker): # inputs as trtllm::k3_spec_accept, keep the target logits vocabulary-sharded and write the drafter's context # K/V with trtllm::k3_ctx_kv. k3_decode = False + # The generation-step preparation below reads _ctx_block_tables, not the + # rebased view, so every drafter slot numbers from 0. + _ctx_reads_rebased_view = False def _mask_token_id(self, draft_model) -> int: """The drafter's mask token id, resolved once.""" diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index 42b621bde4c1..b0189844f34b 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -762,6 +762,24 @@ def _extend_swa_windows_for_reuse( return retention_windows +def _draft_slot_trail_tokens( + spec_config: Optional["DecodingBaseConfig"], *, is_draft: bool, tokens_per_block: int +) -> int: + """SWA retention a DFlash/DSpark draft pool keeps below its windows. + + The drafter seeds a slot that holds no context yet at the first position of + the page holding the request's position (``DFlashWorker._seed_generation_slot``), + so the slot's positions trail the request's, which the mirror's history + follows, by up to a page less one token. The far edge of the drafter's + window reaches that far below the mirror's window; one more page of + retention keeps every entry it reads on a page the request holds. + """ + if not is_draft or spec_config is None: + return 0 + mode = spec_config.spec_dec_mode + return int(tokens_per_block) if mode.is_dflash() or mode.is_dspark() else 0 + + def _derive_request_windows_from_layer_types( model_config: ModelConfigPython, local_layer_ids: Sequence[int], @@ -2021,6 +2039,15 @@ def __init__( from tensorrt_llm._torch.speculative import draft_prompt_lookahead, get_num_extra_kv_tokens self.num_extra_kv_tokens = get_num_extra_kv_tokens(spec_config) + # The draft request view (LlmRequest.use_draft_model) reads the context + # position and chunk a separate draft engine advances for its own + # prefill; a one-engine external drafter (DFlash, DSpark, PARD, + # DraftTarget one-model) has no such engine, so those fields stay at + # position 0 / chunk = prompt length and its mirror follows the + # target's chunking instead. + self._mirror_uses_draft_view = not ( + is_draft and spec_config is not None and spec_config.spec_dec_mode.is_external_drafter() + ) # Mirror V1: expose max_draft_len so the native disagg AuxBuffer # (_make_aux_buffer's getattr fallback) is sized for MTP/spec decoding. self.max_draft_len = spec_config.max_draft_len if spec_config is not None else 0 @@ -2086,6 +2113,14 @@ def __init__( self.reuse_match_backoff, max_seq_len, ) + # The drafter's own windows come from its model config, so this widens + # only what the pool retains (see _draft_slot_trail_tokens). + self.draft_slot_trail_tokens = _draft_slot_trail_tokens( + spec_config, is_draft=is_draft, tokens_per_block=tokens_per_block + ) + self.max_attention_window_vec = _extend_swa_windows_for_reuse( + self.max_attention_window_vec, self.draft_slot_trail_tokens, max_seq_len + ) event_window_size = max( self.max_seq_len if window_size is None else int(window_size) @@ -6011,16 +6046,26 @@ def _draft_pool_diagnostic(self) -> str: byte budget, and the capacity scheduler admits on the TARGET pool alone (`scheduler_v2` touches `draft_kv_cache_manager` only to suspend/free). A draft pool smaller in tokens than the target cannot backpressure -- it can - only raise, and the raise kills every rank. + only raise (see _mirror_resize_error), and the raise kills every rank. The first question is always how big the draft pool was and how full, so the message answers it rather than leaving post-hoc arithmetic over the split log. """ - live = sum(c.capacity for c in self.kv_cache_map.values()) + # num_blocks counts block ordinals up to the capacity, including the + # stale ones a sliding window left without a page, so these are logical + # blocks, not pages held. + capacity = sum(c.capacity for c in self.kv_cache_map.values()) + blocks = sum(int(c.num_blocks) for c in self.kv_cache_map.values()) + per_cache = ", ".join( + f"{req_id}: capacity={c.capacity} history={c.history_length} " + f"logical_blocks={c.num_blocks}" + for req_id, c in self.kv_cache_map.items() + ) return ( - f" [draft pool: {len(self.kv_cache_map)} live caches holding " - f"{live} tokens, gpu_max_tokens={self._gpu_max_tokens}]" + f" [draft pool: {len(self.kv_cache_map)} live caches, capacity {capacity} tokens " + f"in {blocks} logical blocks ({per_cache}), " + f"windows={self.max_attention_window_vec}, gpu_max_tokens={self._gpu_max_tokens}]" ) def _prepare_draft_resources(self, scheduled_batch: ScheduledRequests): @@ -6032,8 +6077,11 @@ def _prepare_draft_resources(self, scheduled_batch: ScheduledRequests): scheduler pre-creates context caches to coordinate the reusable prefix. """ # Joint V2 reuse deliberately shares the target request progress. A - # separate draft engine still uses the legacy draft request view. - with request_context(not self.enable_joint_kv_cache_reuse, scheduled_batch): + # separate draft engine still uses the legacy draft request view; a + # one-engine external drafter never advances that view, so its mirror + # reads the target's position and chunk (see _mirror_uses_draft_view). + use_draft_view = self._mirror_uses_draft_view and not self.enable_joint_kv_cache_reuse + with request_context(use_draft_view, scheduled_batch): for req in scheduled_batch.context_requests: kv_cache = self._mirror_draft_kv_cache(req) if kv_cache is None: @@ -6066,12 +6114,44 @@ def _prepare_draft_resources(self, scheduled_batch: ScheduledRequests): + draft_len + self.num_extra_kv_tokens ) - if not kv_cache.resize(capacity): - raise RuntimeError( - f"Draft KV cache context resize failed for request " - f"{req.py_request_id}: could not resize to {capacity} tokens" - f"{self._draft_pool_diagnostic()}" - ) + # None keeps the cache's current history, as every drafter + # that reads the draft request view (MTP, Eagle, two-engine) + # needs: that view's position is not where its layers write. + history = None + if not self._mirror_uses_draft_view: + if req.is_disagg_generation_init_state: + # The transferred prompt KV is the target's only, so + # the mirror holds no drafter K/V below the prompt end. + # The mirror and the drafter share absolute positions: + # the drafter (DFlash/DSpark, the only separate draft + # pool a disaggregated worker keeps) seeds this + # request's slot at the prompt end and reads nothing + # below it (DFlashSpecMetadata.prepare). The whole + # prompt is therefore history, which a sliding window + # trims to its tail plus the generation headroom, and + # which update_resources keeps advancing absolutely. + history = max( + kv_cache.history_length, + req.context_current_position + req.context_chunk_size, + ) + elif ( + kv_cache.capacity >= req.context_current_position + or self.draft_slot_trail_tokens > 0 + ): + # The chunk start is history for a mirror that already + # holds [0, chunk start): sliding-window layers keep + # pages only from history + 1 - window (AttnLifeCycle:: + # getStaleRange), so the mirror holds its window plus + # the chunk in flight rather than the whole prompt so + # far. So it is for a DFlash/DSpark mirror first seen + # past position 0 (a prefix this worker did not + # compute): that drafter stores the chunk at its target + # position over a base hiding the prefix + # (DFlashWorker._ctx_store_start). Any other mirror + # first seen there keeps its history, i.e. pages from 0. + history = max(kv_cache.history_length, req.context_current_position) + if not kv_cache.resize(capacity, history): + raise self._mirror_resize_error(req, "context", capacity) self._after_shrink(req.py_request_id, kv_cache) for req in scheduled_batch.generation_requests: @@ -6095,11 +6175,25 @@ def _prepare_draft_resources(self, scheduled_batch: ScheduledRequests): continue new_cap = self._required_gen_capacity(req, kv_cache.capacity) if not kv_cache.resize(new_cap): - raise RuntimeError( - f"Draft KV cache generation resize failed for request " - f"{req.py_request_id}: could not resize to {new_cap} tokens" - f"{self._draft_pool_diagnostic()}" - ) + raise self._mirror_resize_error(req, "generation", new_cap) + + def _mirror_resize_error(self, req: LlmRequest, phase: str, capacity: int) -> RuntimeError: + """The error for a draft mirror that could not grow to ``capacity``. + + `KvCache.resize` leaves the cache as it was when the pool has no page to + give, and the cache as it was cannot hold what the drafter writes next: + the DFlash block decode appends the block's own K/V past the stored + context through the mirror's block table, so a mirror that did not grow + (a fresh one has no page at all) hands that write a table entry with no + page behind it (``BAD_PAGE_INDEX`` on the python V2 backend, another + request's block 0 on the cpp one). There is no per-request way to skip + drafting, so the failure stays explicit. The message carries the pool + diagnostic. + """ + return RuntimeError( + f"Draft KV cache {phase} resize failed for request {req.py_request_id}: " + f"could not resize to {capacity} tokens{self._draft_pool_diagnostic()}" + ) def _reuse_token_source(self, req: LlmRequest) -> Sequence[int]: """Return a zero-copy beam-0 token view for block reuse.""" @@ -6967,7 +7061,12 @@ def release_resources( if not success: release_resources(req, free_draft_resources=True) return None - success = draft_kv_cache.resize(dummy_capacity) + # Same history hint as the target cache above: a sliding-window + # draft pool then backs only the dummy's window, not its whole + # length. The hint exceeds the window only for the + # longest-length warmup dummies, which run before any request + # holds a draft page. + success = draft_kv_cache.resize(dummy_capacity, history_length=history_hint) if not success: release_resources(req, free_draft_resources=True) return None @@ -7572,6 +7671,13 @@ def get_cache_size_per_token( backoff, max_seq_len, ) + attention_windows = _extend_swa_windows_for_reuse( + attention_windows, + _draft_slot_trail_tokens( + spec_config, is_draft=is_draft, tokens_per_block=tokens_per_block + ), + max_seq_len, + ) _, generation_capacity_headroom = _get_generation_kv_capacity( spec_config, is_draft=is_draft ) @@ -7712,7 +7818,8 @@ def _get_per_request_layer_tokens( return None # ``covered_windows`` are the windows ``get_cache_size_per_token`` # charges as fixed per-request retention in the affine intercept; keep - # them in lockstep with its accounting (reuse backoff included). + # them in lockstep with its accounting (reuse backoff and the drafter's + # trailing page included). if ( kv_cache_config is not None and kv_cache_config.enable_block_reuse @@ -7722,6 +7829,13 @@ def _get_per_request_layer_tokens( backoff = draft_prompt_lookahead(spec_config) or 0 covered_windows = _extend_swa_windows_for_reuse(covered_windows, backoff, max_seq_len) + covered_windows = _extend_swa_windows_for_reuse( + covered_windows, + _draft_slot_trail_tokens( + spec_config, is_draft=is_draft, tokens_per_block=tokens_per_block + ), + max_seq_len, + ) if is_draft: # Draft mirrors are reserved at the request's full context length # (measured: a 41k-token context charges 41k tokens per mirror diff --git a/tensorrt_llm/_torch/speculative/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py index fb85eafd3cd8..e6a6f957bdbe 100644 --- a/tensorrt_llm/_torch/speculative/dflash.py +++ b/tensorrt_llm/_torch/speculative/dflash.py @@ -17,7 +17,7 @@ import weakref from collections import deque from dataclasses import dataclass -from typing import TYPE_CHECKING, Callable, List, Optional +from typing import TYPE_CHECKING, Callable, List, Optional, Tuple import torch from torch import nn @@ -516,6 +516,10 @@ def dflash_allocated_ctx_limit( ) -> torch.Tensor: """Context length a request may advertise, given the pages it holds. + ``block_counts`` is each row's extent: one past its last backed entry + (``decode_draft_block_table``), so positions are absolute even when a + sliding window released the low entries. + Leaves ``block_size`` positions of room: the MLA path writes the block's own latents at ``ctx_len .. ctx_len + block_size``, so stopping at exactly the allocation puts the first of them on an unallocated page, whose table entry @@ -531,6 +535,51 @@ def dflash_allocated_ctx_limit( return (block_counts * page_size - block_size).clamp(min=0) +def dflash_windows_span_a_page(windows: Optional[List[Optional[int]]], page_size: int) -> bool: + """Whether every attention window of a draft pool spans at least a page. + + ``None`` (for the list or an entry) is a layer without a window, which + keeps every page it was given. + """ + return all(window >= page_size for window in windows or () if window is not None) + + +def dflash_rebased_ctx_rows(rows: torch.Tensor, base_pages: torch.Tensor) -> torch.Tensor: + """Block-table rows shifted left by ``base_pages`` entries each. + + Entry ``j`` of the result names the page entry ``j + base`` of the input + row named, so a reader that numbers positions from ``base * page_size`` + reaches the same physical pages as one numbering from 0 and never + addresses an entry below ``base``. Entries shifted in from past the row's + end repeat the last entry; a reader never dereferences them because its + sequence length is shortened by the same ``base * page_size``. + """ + width = rows.size(1) + index = torch.arange(width, device=rows.device).unsqueeze(0) + base_pages.view(-1, 1) + return rows.gather(1, index.clamp_(max=width - 1)) + + +def dflash_rebased_ctx_view( + rows: torch.Tensor, base_pages: torch.Tensor, num_ctx: torch.Tensor, page_size: int +) -> Tuple[torch.Tensor, torch.Tensor]: + """Block-table rows and context lengths the block decode reads for slots + whose first ``base_pages`` pages hold no drafter K/V (see + ``DFlashWorker._ctx_base_pages``). + + The context length is shortened by the base, so the decode never reads + below it. The decode also writes the block's own K/V through the returned + rows, at rebased positions ``num_ctx .. num_ctx + block_size``, so entry 0 + of a row has to name a page the request holds. While the context is + shorter than the base (the first tokens stored inside the rounded-up page) + it reads nothing and the row is shifted only to the page holding the + context length, which is the page those tokens were written to. From + ``base * page_size`` on the shift is the base itself. + """ + rebase_pages = torch.minimum(base_pages, num_ctx // page_size) + num_ctx = (num_ctx - base_pages * page_size).clamp(min=0) + return dflash_rebased_ctx_rows(rows, rebase_pages), num_ctx + + @dataclass class DFlashSpecMetadata(SpecMetadata): """Metadata for DFlash speculative decoding. @@ -602,32 +651,48 @@ def prepare(self): worker = getattr(self, "_dflash_worker", None) if worker is not None and worker._ctx_buf_inited: current = set(self.request_ids) - evicted = {} + # Slot lengths (and bases) written in one batch: evicted slots are + # emptied, newly assigned ones seeded below. + updates, bases = {}, {} for rid in list(worker._req_to_slot.keys()): if rid not in current: slot = worker._req_to_slot.pop(rid) - evicted[slot] = 0 + updates[slot] = 0 worker._req_ctx_pos.pop(rid, None) worker._free_slots.append(slot) - worker._write_ctx_len(evicted) # A disagg generation worker receives prompt KV instead of # prefilling, so _store_prefill_context -- the only other assigner -- - # never runs and every request would land on the dummy slot. + # never runs and every request would land on the dummy slot. The + # same holds for a request whose slot was evicted above. Each such + # slot starts empty at its row's absolute position, on a page + # boundary (_seed_generation_slot). num_contexts = max(0, len(self.request_ids) - self.num_generations) - for rid in self.request_ids[num_contexts:]: + row_starts = None + for row, rid in enumerate(self.request_ids[num_contexts:], start=num_contexts): if ( rid != ATTENTION_DP_DUMMY_REQUEST_ID and rid < worker._graph_dummy_id_floor and rid not in worker._req_to_slot ): - if worker._assign_slot(rid) is None: + slot = worker._assign_slot(rid, updates=updates, bases=bases) + if slot is None: # Not fatal -- spec dec is lossless, so the dummy-slot # fallback only costs acceptance. Warn so it is not silent. logger.warning( f"DFlash: no free drafter slot for generation request {rid}; " "it shares the dummy slot and will draft poorly." ) + continue + if self.host_position_ids is None or self.seq_lens is None: + continue + if row_starts is None: + row_starts = [0] + for seq_len in self.seq_lens: + row_starts.append(row_starts[-1] + int(seq_len)) + position = int(self.host_position_ids[row_starts[row]]) + worker._seed_generation_slot(slot, position, updates, bases) + worker._write_ctx_len(updates, bases) # Route unknown request IDs (cuda graph padding or warmup dummies) # to dummy slot to avoid corrupting real request's context @@ -735,6 +800,14 @@ class DFlashWorker(SpecWorkerBase): Reference: https://arxiv.org/pdf/2602.06036 """ + # Whether the block decode reads _ctx_block_tables_read, the per-slot view + # rebased past _ctx_base_pages. Only then can a slot on the managed pool + # start at an absolute position other than 0 (_ctx_store_start, + # _seed_generation_slot); a subclass with its own generation-step + # preparation that reads _ctx_block_tables directly sets it False and + # numbers every slot from 0. + _ctx_reads_rebased_view = True + def __init__( self, spec_config: "DFlashDecodingConfig", @@ -789,9 +862,22 @@ def __init__( # manager it is replaced by _ctx_block_tables, refreshed per iteration. self._ctx_page_table = None self._ctx_block_tables = None # [max_batch+1, max_blocks_per_seq] int32 + # The block table the block decode reads: each row of _ctx_block_tables + # shifted left by that slot's _ctx_base_pages, so pages below the first + # position the drafter stored are not addressable. + self._ctx_block_tables_read = None + # [max_batch+1] pages below a slot's first stored position: the slot + # holds no drafter K/V there (a prompt this worker did not prefill, or + # a prefix the target reused from an unpaired pool). Host shadow and + # pre-forward snapshots as for _ctx_len. + self._ctx_base_pages = None + self._ctx_base_host = None + self._saved_ctx_base = None + self._saved_ctx_base_host = None self._ctx_block_indptr = None - # Pages the manager handed this request, per batch row. Saturates at - # the table width on the default backend -- dflash_allocated_ctx_limit. + # Extent of the pages the manager handed this request (one past its + # last backed entry), per batch row. Saturates at the table width on + # the default backend -- dflash_allocated_ctx_limit. self._ctx_block_counts = None self._ctx_pool_idx = 0 self._ctx_block_divisor = 1 # encoded offset -> raw pool block index @@ -805,6 +891,11 @@ def __init__( # makes a reused prefix's drafter K/V still addressable. Set in # _lazy_init_ctx_buffers; read by _store_prefill_context. self._ctx_reuse_addressable = False + # Whether a seeded slot may start at the first position of the page + # holding its row's position (_seed_generation_slot): true when every + # window of the bound draft pool spans at least a page. Set in + # _lazy_init_ctx_buffers. + self._ctx_seed_page_aligned = False self._dflash_attention_backend = spec_config.attention_backend # Set by bind_target_lm_head. self._vocab_parallel: Optional[VocabParallelSelection] = None @@ -1146,6 +1237,7 @@ def _init_ctx_block_tables( self._ctx_block_tables = torch.zeros( (num_slots, max_blocks), dtype=torch.int32, device="cuda" ) + self._ctx_block_tables_read = torch.zeros_like(self._ctx_block_tables) self._ctx_block_indptr = torch.arange( 0, (num_slots + 1) * max_blocks, max_blocks, dtype=torch.int32, device="cuda" ) @@ -1216,6 +1308,7 @@ def _lazy_init_ctx_buffers( # pass falling back to the arena still rebinds once the real pool lands. self._ctx_kv_manager = draft_kv_cache_manager self._ctx_block_tables = None + self._ctx_block_tables_read = None self._ctx_block_counts = None self._paged_latent_prepare = False max_batch = spec_metadata.max_num_requests @@ -1272,6 +1365,8 @@ def _lazy_init_ctx_buffers( self._ctx_len = torch.zeros(num_slots, dtype=torch.long, device="cuda") self._ctx_len_host = [0] * num_slots + self._ctx_base_pages = torch.zeros(num_slots, dtype=torch.long, device="cuda") + self._ctx_base_host = [0] * num_slots self._batch_to_slot_buffer = DeviceCopyBuffer( torch.zeros(max_batch, dtype=torch.long, device="cuda"), pin_memory=prefer_pinned(), @@ -1402,6 +1497,15 @@ def _lazy_init_ctx_buffers( self._ctx_reuse_addressable = bool( getattr(draft_kv_cache_manager, "enable_joint_kv_cache_reuse", False) ) + # A slot seeded on a page boundary trails its request by less + # than a page; a windowed pool backs that trail only when it + # keeps a page below its windows (kv_cache_manager_v2's + # draft_slot_trail_tokens; a V1 pool does not). + windows = getattr(draft_kv_cache_manager, "max_attention_window_vec", None) + self._ctx_seed_page_aligned = dflash_windows_span_a_page(windows, page_size) and ( + all(window is None for window in windows or ()) + or getattr(draft_kv_cache_manager, "draft_slot_trail_tokens", 0) >= page_size + ) self._ctx_kv_last_page_len = torch.full( (num_slots,), page_size, dtype=torch.int32, device="cuda" ) @@ -1454,13 +1558,100 @@ def _ctx_paged_index_args(self): indptr = self._ctx_kv_indptr return table.flatten().contiguous(), indptr - def _assign_slot(self, req_id: int, reset: bool = False, updates: Optional[dict] = None): + def _ctx_store_start(self, cur: int, first_pos: int) -> int: + """Position at which a context chunk's drafter K/V is stored. + + ``cur`` is the slot's stored length, ``first_pos`` the chunk's first + target position. They differ only when the target skipped + [0, first_pos) (block reuse, or a prefix computed elsewhere), so just + the tail reaches the capture buffer. + + On the managed pool the block table is indexed by target position, so + the tail is stored at ``first_pos``: always on a pool paired with the + target's reuse, which holds the prefix's drafter K/V below it, and on + an unpaired pool when the worker reads the rebased view, which hides + the empty prefix. Otherwise (the private arena, or an unpaired pool + read from position 0) every request numbers from 0, where the tail is + all the context there is. + """ + if ( + cur == 0 + and first_pos > 0 + and self._ctx_block_tables is not None + and (self._ctx_reuse_addressable or self._ctx_reads_rebased_view) + ): + return first_pos + return cur + + def _ctx_base_for_store(self, cur: int, start: int) -> Optional[int]: + """Pages the block decode must not read for a slot whose first stored + position is ``start``, or None when the slot numbers from ``cur``. + + A pool paired with the target's reuse holds the prefix's drafter K/V + below ``start``, so it stays readable. An unpaired pool holds no drafter + K/V there: the block decode is handed the slot's block-table row shifted + past those pages (``dflash_rebased_ctx_rows``) and a context length + shortened by the same amount, so a tail shorter than the drafter's + window attends only to itself. ``start`` is rounded up to a page + boundary; the tail tokens inside that page stay stored but unread. + """ + if start == cur or self._ctx_reuse_addressable: + return None + return -(-start // self._ctx_page_size) + + def _seed_generation_slot(self, slot: int, position: int, updates: dict, bases: dict) -> None: + """Record the context length and base of a slot newly assigned to a + generation row whose first token sits at absolute ``position``. + + Bound to the draft KV cache manager, the slot addresses the request's + mirror, whose positions are absolute (the target's). The slot holds no + drafter K/V for the positions before ``position`` -- a disaggregated + generation worker received only the target's prompt KV, and a request + that lost its slot lost its context -- whether or not the pool joins + block reuse: the mirror of such a request never looked a prefix up. + The slot therefore starts with an empty context on a page boundary, + and the block decode reads from its base, never below it: + + * at the first position of the page holding ``position``. The slot's + positions then trail the absolute ones by less than a page, and the + mirror keeps one page below each window for that trail + (``_draft_slot_trail_tokens`` in kv_cache_manager_v2), so every + entry the window reads stays on a page the request holds. Taken + when every window of the pool spans a page + (``_ctx_seed_page_aligned``); the slot reads what it stores from + its first step; + * otherwise at ``position`` itself, with the base rounded up to the + next page: the slot reads nothing until its context reaches it. + + The private arena is slot-addressed and starts at 0, as does a worker + that does not read the rebased view (_ctx_reads_rebased_view). + """ + if self._ctx_block_tables is None or not self._ctx_reads_rebased_view or position <= 0: + updates[slot] = 0 + return + position = min(position, self._max_ctx) + page = self._ctx_page_size + if self._ctx_seed_page_aligned: + bases[slot] = position // page + updates[slot] = bases[slot] * page + else: + bases[slot] = -(-position // page) + updates[slot] = position + + def _assign_slot( + self, + req_id: int, + reset: bool = False, + updates: Optional[dict] = None, + bases: Optional[dict] = None, + ) -> Optional[int]: """Get (or refresh) this request's context slot; None when none is free. ``reset`` recycles an existing slot, which is how a reused request id is - told apart from a continuation. Pass ``updates`` to defer the length - resets into a caller's batch instead of writing _ctx_len per request; - either way the host mirror stays in step with the device tensor. + told apart from a continuation. Pass ``updates`` (and ``bases``, which + then gets the base-page resets) to defer the length resets into a + caller's batch instead of writing _ctx_len per request; either way the + host mirror stays in step with the device tensor. """ def clear(slot: int) -> None: @@ -1468,6 +1659,8 @@ def clear(slot: int) -> None: self._write_ctx_len({slot: 0}) else: updates[slot] = 0 + if bases is not None: + bases[slot] = 0 if reset and req_id in self._req_to_slot: old_slot = self._req_to_slot.pop(req_id) @@ -1590,30 +1783,51 @@ def _rebuild_batch_to_slot(self, request_ids: List[int]) -> None: mapping = [self._req_to_slot.get(rid, self._dummy_slot) for rid in request_ids] self._batch_to_slot_buffer.update(torch.tensor(mapping, dtype=torch.long)) - def _write_ctx_len(self, updates: dict[int, int]) -> None: - """Apply a slot-to-length mapping to _ctx_len in one async scatter.""" - if not updates: + def _write_ctx_len( + self, updates: dict[int, int], base_updates: Optional[dict[int, int]] = None + ) -> None: + """Apply slot-to-length (and slot-to-base-page) mappings in one async + scatter each. A slot reset to length 0 also drops its base pages, and + that reset wins over a base recorded for the same slot in + ``base_updates`` (a chunk's base is recorded before the overflow check + that may reset the slot).""" + if not updates and not base_updates: return + bases = dict(base_updates or {}) + bases.update({slot: 0 for slot, value in updates.items() if value == 0}) for slot, value in updates.items(): # Update host side mirror self._ctx_len_host[slot] = value + for slot, value in bases.items(): + self._ctx_base_host[slot] = value - slots, values = zip(*updates.items()) pinned = prefer_pinned() - slots_pinned = torch.tensor(slots, dtype=torch.long, pin_memory=pinned) - vals_pinned = torch.tensor(values, dtype=torch.long, pin_memory=pinned) - self._ctx_len.index_copy_( - 0, - slots_pinned.to("cuda", non_blocking=True), - vals_pinned.to("cuda", non_blocking=True), - ) + for target, mapping in ((self._ctx_len, updates), (self._ctx_base_pages, bases)): + if not mapping: + continue + slots, values = zip(*mapping.items()) + slots_pinned = torch.tensor(slots, dtype=torch.long, pin_memory=pinned) + vals_pinned = torch.tensor(values, dtype=torch.long, pin_memory=pinned) + target.index_copy_( + 0, + slots_pinned.to(target.device, non_blocking=True), + vals_pinned.to(target.device, non_blocking=True), + ) def _restore_ctx_len_host(self) -> None: """Roll the host shadows back to their pre-forward snapshot.""" if self._saved_ctx_len_host is not None: self._ctx_len_host = list(self._saved_ctx_len_host) + if self._saved_ctx_base_host is not None: + self._ctx_base_host = list(self._saved_ctx_base_host) if self._saved_req_ctx_pos is not None: self._req_ctx_pos = dict(self._saved_req_ctx_pos) + def _restore_ctx_len_device(self) -> None: + """Roll _ctx_len and _ctx_base_pages back to their pre-forward snapshot.""" + self._ctx_len.copy_(self._saved_ctx_len) + if self._saved_ctx_base is not None: + self._ctx_base_pages.copy_(self._saved_ctx_base) + def _store_prefill_context( self, draft_model, @@ -1653,6 +1867,7 @@ def _store_prefill_context( # the same values, but reading it costs a D2H sync per request. host_pos = spec_metadata.host_position_ids ctx_len_updates = {} + ctx_base_updates = {} offset = 0 num_contexts = attn_metadata.num_contexts # A context write addresses pages through row i of the block table, so @@ -1682,7 +1897,12 @@ def _store_prefill_context( # request id reused after completion, a prefill restarted after # preemption -- has to start from a clean slot. reset = self._req_ctx_pos.get(req_id) != first_pos - if self._assign_slot(req_id, reset=reset, updates=ctx_len_updates) is None: + if ( + self._assign_slot( + req_id, reset=reset, updates=ctx_len_updates, bases=ctx_base_updates + ) + is None + ): logger.warning("DFlash: no free slots, skipping context store") self._req_ctx_pos.pop(req_id, None) offset += slen @@ -1691,11 +1911,14 @@ def _store_prefill_context( slot = self._req_to_slot[req_id] cur = ctx_len_updates.get(slot, self._ctx_len_host[slot]) cap = self._max_ctx if ctx_alloc is None else min(self._max_ctx, ctx_alloc[i]) - if cur == 0 and first_pos > 0 and self._ctx_reuse_addressable: - # Reuse let the target skip [0, first_pos), so only the tail - # reaches the capture buffer. The prefix K/V are still in the - # paired pool: precompute_context_kv is per-token in (hidden, p). - cur = first_pos + # A paired pool still holds the prefix K/V below the tail + # (precompute_context_kv is per-token in (hidden, p)); an unpaired + # one does not, so the slot reads from a base above it. + start = self._ctx_store_start(cur, first_pos) + base = self._ctx_base_for_store(cur, start) + if base is not None: + ctx_base_updates[slot] = base + cur = start if cur + slen > cap: # Request-level, like the no-free-slots path above: truncating # would silently draft from a stale prefix, but killing the @@ -1754,7 +1977,7 @@ def _store_prefill_context( self._ctx_v_buf[slot, :, cur:end] = chunk_v.permute(1, 0, 2, 3) offset += slen - self._write_ctx_len(ctx_len_updates) + self._write_ctx_len(ctx_len_updates, ctx_base_updates) def _ensure_spec_dec_state_restored(self, attn_metadata, spec_metadata): # Restore first (in warmup mode kv_lens_cuda was saved and comes back @@ -1774,7 +1997,7 @@ def _ensure_spec_dec_state_restored(self, attn_metadata, spec_metadata): ): # A failed forward must not keep this iteration's in-place # _ctx_len updates: roll back to the pre-forward snapshot. - self._ctx_len.copy_(self._saved_ctx_len) + self._restore_ctx_len_device() self._restore_ctx_len_host() self._ctx_len_restore_pending = False @@ -1840,6 +2063,8 @@ def _forward_impl( # mutate _ctx_len until replay. self._saved_ctx_len = self._ctx_len.clone() self._saved_ctx_len_host = list(self._ctx_len_host) + self._saved_ctx_base = self._ctx_base_pages.clone() + self._saved_ctx_base_host = list(self._ctx_base_host) self._saved_req_ctx_pos = dict(self._req_ctx_pos) self._ctx_len_restore_pending = True @@ -2044,7 +2269,7 @@ def _forward_impl( # Restore context lengths after warmup; real runs keep the updates. if is_warmup: - self._ctx_len.copy_(self._saved_ctx_len) + self._restore_ctx_len_device() self._restore_ctx_len_host() self._ctx_len_restore_pending = False @@ -2393,6 +2618,7 @@ def prepare_1st_drafter_inputs( self._advance_ctx_len(slots, gen_num_accepted_long) num_ctx_per_req_t = self._ctx_len[slots] + ctx_page_table = self._ctx_page_table if self._ctx_block_tables is not None: # The write above clamps columns to the request's allocation, so # a context that outruns it has its tail written over the last @@ -2402,6 +2628,18 @@ def prepare_1st_drafter_inputs( self._ctx_block_counts, self._ctx_page_size, block_size ) num_ctx_per_req_t = torch.minimum(num_ctx_per_req_t, allocated[gen_rows_out]) + # Pages below a slot's base hold no drafter K/V: the read + # numbers from the base instead (dflash_rebased_ctx_view). + # Writes above keep the absolute table. + self._ctx_block_tables_read[gen_rows_out], num_ctx_per_req_t = ( + dflash_rebased_ctx_view( + self._ctx_block_tables[gen_rows_out], + self._ctx_base_pages[slots], + num_ctx_per_req_t, + self._ctx_page_size, + ) + ) + ctx_page_table = self._ctx_block_tables_read noise_embedding = noise_embed_2d query_positions = query_position_ids.long() @@ -2418,6 +2656,11 @@ def prepare_1st_drafter_inputs( slots = torch.empty(0, dtype=torch.long, device="cuda") gen_rows_out = slots bonus = torch.empty(0, dtype=torch.long, device="cuda") + ctx_page_table = ( + self._ctx_block_tables_read + if self._ctx_block_tables is not None + else self._ctx_page_table + ) return { "lilicorr_anchor": lilicorr_anchor, @@ -2433,11 +2676,8 @@ def prepare_1st_drafter_inputs( # token. The dspark Markov chain conditions its first step on it. "first_prev_tokens": bonus, "ctx_kv_cache": self._ctx_kv_buf, - "ctx_page_table": ( - self._ctx_block_tables - if self._ctx_block_tables is not None - else self._ctx_page_table - ), + # The manager-bound table is the read view rebased per slot above. + "ctx_page_table": ctx_page_table, "block_placement": None, } @@ -2480,6 +2720,7 @@ def _prepare_paged_latent_gen_inputs( accepted_tokens, self._batch_to_slot_buffer.tensor[num_contexts : num_contexts + num_gens], self._ctx_len, + self._ctx_base_pages, self._ctx_block_counts, self._ctx_block_tables, attn_metadata._seq_lens_cuda, @@ -2514,6 +2755,15 @@ def _prepare_paged_latent_gen_inputs( prep.store_offsets, prep.store_mask, ) + # The block decode reads each row shifted past the slot's base, as the + # eager path's dflash_rebased_ctx_view; the kernel already placed the + # block and advertised the context in those rebased positions. + gen_rows = torch.arange( + num_contexts, num_contexts + num_gens, dtype=torch.long, device="cuda" + ) + self._ctx_block_tables_read[gen_rows] = dflash_rebased_ctx_rows( + self._ctx_block_tables[gen_rows], prep.rebase_pages + ) return { "lilicorr_anchor": None, "noise_embedding": noise_embedding, @@ -2521,12 +2771,10 @@ def _prepare_paged_latent_gen_inputs( "num_ctx_per_req": prep.num_ctx, "ctx_k_cache": self._ctx_k_buf, "ctx_v_cache": self._ctx_v_buf, - "ctx_cache_batch_idx": torch.arange( - num_contexts, num_contexts + num_gens, dtype=torch.long, device="cuda" - ), + "ctx_cache_batch_idx": gen_rows, "first_prev_tokens": prep.bonus, "ctx_kv_cache": self._ctx_kv_buf, - "ctx_page_table": self._ctx_block_tables, + "ctx_page_table": self._ctx_block_tables_read, "block_placement": DraftBlockPlacement( prep.blk_pages, prep.blk_offsets, prep.blk_seq_lens, prep.blk_kv_bounds ), diff --git a/tensorrt_llm/_torch/speculative/spec_tail_kernels.py b/tensorrt_llm/_torch/speculative/spec_tail_kernels.py index 6705f605a9eb..aa40462f6049 100644 --- a/tensorrt_llm/_torch/speculative/spec_tail_kernels.py +++ b/tensorrt_llm/_torch/speculative/spec_tail_kernels.py @@ -143,20 +143,25 @@ def _decode_block_table_kernel( offs = tl.arange(0, BLOCK) src_row = src_ptr + row.to(tl.int64) * stride_src_row table_row = table_ptr + row.to(tl.int64) * stride_table - count = tl.zeros([BLOCK], dtype=tl.int64) + extent = tl.zeros([BLOCK], dtype=tl.int64) for start in range(0, n_blocks, BLOCK): cols = start + offs mask = cols < n_blocks encoded = tl.load(src_row + cols * stride_src_col, mask=mask, other=-1).to(tl.int64) - count += (encoded >= 0).to(tl.int64) + extent = tl.maximum(extent, tl.where(encoded >= 0, cols.to(tl.int64) + 1, 0)) tl.store(table_row + cols, (tl.maximum(encoded, 0) // divisor).to(tl.int32), mask=mask) - tl.store(counts_ptr + row, tl.sum(count, axis=0)) + tl.store(counts_ptr + row, tl.max(extent, axis=0)) def decode_draft_block_table( src: torch.Tensor, divisor: int, table: torch.Tensor, counts: torch.Tensor ) -> None: - """``counts[i] = (src[i] >= 0).sum()``; ``table[i] = src[i].clamp(min=0) // divisor``. + """``counts[i]`` = one past the last ``j`` with ``src[i, j] >= 0`` (0 when + none); ``table[i] = src[i].clamp(min=0) // divisor``. + + The extent, not the number of backed entries: a sliding-window pool + releases a row's low entries, and positions are indexed from entry 0, so + only the extent bounds the positions a request may address. Args: src: ``[rows, blocks]`` encoded block offsets (any strides). @@ -212,6 +217,7 @@ class PagedLatentPrepare(NamedTuple): blk_offsets: torch.Tensor # [G, block] int64 blk_seq_lens: torch.Tensor # [G] int32, num_ctx + block blk_kv_bounds: torch.Tensor # [G * block] int32, blk_seq_lens per query token + rebase_pages: torch.Tensor # [G] int64, row entries the block decode skips @triton.jit @@ -221,6 +227,7 @@ def _paged_latent_prepare_kernel( stride_accepted, slots_ptr, ctx_len_ptr, + base_pages_ptr, block_counts_ptr, table_ptr, stride_table, @@ -237,6 +244,7 @@ def _paged_latent_prepare_kernel( blk_offsets_ptr, blk_seq_lens_ptr, blk_kv_bounds_ptr, + rebase_ptr, num_gens, first_row, dummy_slot, @@ -289,8 +297,14 @@ def _paged_latent_prepare_kernel( tl.store(ctx_len_ptr + slot, new_len, mask=gm) tl.store(ctx_len_ptr + dummy_slot, tl.zeros([], dtype=tl.int64)) - # Advertised context, capped so the block still fits in the allocation. + # Advertised context, capped so the block still fits in the allocation, + # then rebased past the slot's base pages (dflash_rebased_ctx_view): the + # block decode reads row entries from ``rebase`` on. num_ctx = tl.minimum(new_len, tl.maximum(capacity - BLOCK_SIZE, 0)) + base = tl.load(base_pages_ptr + slot, mask=gm, other=0).to(tl.int64) + rebase = tl.minimum(base, num_ctx // PAGE_SIZE) + num_ctx = tl.maximum(num_ctx - base * PAGE_SIZE, 0) + tl.store(rebase_ptr + g, rebase, mask=gm) tl.store(num_ctx_ptr + g, num_ctx, mask=gm) tl.store( seq_lens_ptr + row, @@ -300,7 +314,8 @@ def _paged_latent_prepare_kernel( # Block-decode operands: where the block's own latents go, and its bounds. bpos = num_ctx[:, None] + j[None, :] - bpage = tl.load(table_row[:, None] + tl.minimum(bpos // PAGE_SIZE, max_blocks - 1), mask=qmask) + bentry = tl.minimum(rebase[:, None] + bpos // PAGE_SIZE, max_blocks - 1) + bpage = tl.load(table_row[:, None] + bentry, mask=qmask) tl.store(blk_pages_ptr + blk, bpage.to(tl.int64), mask=qmask) tl.store(blk_offsets_ptr + blk, bpos % PAGE_SIZE, mask=qmask) seq_len = (num_ctx + BLOCK_SIZE).to(tl.int32) @@ -317,6 +332,7 @@ def dflash_paged_latent_prepare( accepted_tokens: torch.Tensor, slots: torch.Tensor, ctx_len: torch.Tensor, + base_pages: torch.Tensor, block_counts: torch.Tensor, block_tables: torch.Tensor, seq_lens_cuda: torch.Tensor, @@ -340,7 +356,11 @@ def dflash_paged_latent_prepare( accepted_tokens: ``[batch, width]`` accepted tokens, unit inner stride. slots: ``[G]`` context slots. ctx_len: ``[slots]`` int64 context lengths. - block_counts: ``[batch]`` pages held per batch row. + base_pages: ``[slots]`` int64 pages below each slot's first stored + position; ``num_ctx`` and the block placement are rebased past them + and ``rebase_pages`` says by how many row entries. + block_counts: ``[batch]`` extent of the pages held per batch row + (``decode_draft_block_table``). block_tables: ``[batch, max_blocks]`` int32 page tables, unit inner stride. seq_lens_cuda: ``[batch]`` attention sequence lengths. first_row: batch row of generation request 0. @@ -364,6 +384,7 @@ def dflash_paged_latent_prepare( blk_offsets=torch.empty((num_gens, block_size), **i64), blk_seq_lens=torch.empty((num_gens,), dtype=torch.int32, device=device), blk_kv_bounds=torch.empty((num_gens * block_size,), dtype=torch.int32, device=device), + rebase_pages=torch.empty((num_gens,), **i64), ) if num_gens == 0: return out @@ -374,6 +395,7 @@ def dflash_paged_latent_prepare( accepted_tokens.stride(0), slots, ctx_len, + base_pages, block_counts, block_tables, block_tables.stride(0), diff --git a/tests/unittest/_torch/pyexecutor/kv_cache/sharing/test_staging_lender.py b/tests/unittest/_torch/pyexecutor/kv_cache/sharing/test_staging_lender.py index 79717b94f868..b8b5a3deaa12 100644 --- a/tests/unittest/_torch/pyexecutor/kv_cache/sharing/test_staging_lender.py +++ b/tests/unittest/_torch/pyexecutor/kv_cache/sharing/test_staging_lender.py @@ -2704,9 +2704,10 @@ def test_a_windowed_joint_draft_pool_resumes_only_where_its_context_resize_holds ) as draft, ): pools = (target, draft) - for mgr in pools: - assert JOINT_WINDOW in kit.windows(mgr), "a pool has no window: nothing checked" - beg, end = kit.stale_blocks(mgr, kit.windows(mgr).index(JOINT_WINDOW), PAIRED_END) + # A DFlash draft pool retains one page below its window (_draft_slot_trail_tokens). + for mgr, window in ((target, JOINT_WINDOW), (draft, JOINT_WINDOW + TPB)): + assert window in kit.windows(mgr), "a pool has no window: nothing checked" + beg, end = kit.stale_blocks(mgr, kit.windows(mgr).index(window), PAIRED_END) assert end <= beg, "a window released blocks at the fetch's end: nothing checked" scheduler = KVCacheV2Scheduler( max_batch_size=4, diff --git a/tests/unittest/_torch/pyexecutor/kv_cache/test_kvcm2_integration.py b/tests/unittest/_torch/pyexecutor/kv_cache/test_kvcm2_integration.py index 83bd77f0c2f4..f5c47141fa34 100644 --- a/tests/unittest/_torch/pyexecutor/kv_cache/test_kvcm2_integration.py +++ b/tests/unittest/_torch/pyexecutor/kv_cache/test_kvcm2_integration.py @@ -15,6 +15,7 @@ import array import os +from collections.abc import Callable from dataclasses import dataclass, field, replace from types import SimpleNamespace from unittest.mock import Mock, call, patch @@ -33,6 +34,7 @@ KVCacheManagerV2, Role, _extend_swa_windows_for_reuse, + _get_generation_kv_capacity, _KVCacheManagerInitStatus, _swa_branch_snapshot_point, _swa_endpoint_priority, @@ -42,6 +44,7 @@ from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, LlmRequestState from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManager, ResourceManagerType from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests +from tensorrt_llm._torch.speculative.dflash import DFlashSpecMetadata, DFlashWorker from tensorrt_llm.bindings import DataType, SamplingConfig from tensorrt_llm.bindings.BuildInfo import ENABLE_MULTI_DEVICE from tensorrt_llm.bindings.executor import KvCacheRetentionConfig @@ -50,7 +53,10 @@ from tensorrt_llm.conversation_params import ConversationParams from tensorrt_llm.llmapi.llm_args import ( BlockReuseConfig, + DecodingBaseConfig, DFlashDecodingConfig, + DraftTargetDecodingConfig, + DSparkDecodingConfig, Eagle3DecodingConfig, KvCacheConfig, MTPDecodingConfig, @@ -1848,6 +1854,7 @@ def test_draft_manager_keeps_shared_progress_across_context_and_generation() -> manager.is_draft = True manager.enable_block_reuse = True manager.enable_joint_kv_cache_reuse = True + manager._mirror_uses_draft_view = True manager.kv_cache_map = {request.py_request_id: kv_cache} observed_progress_views = [] @@ -1873,6 +1880,865 @@ def resume_and_restore(request_id, current_cache): manager._required_gen_capacity.assert_not_called() +@pytest.mark.parametrize( + ( + "uses_draft_view", + "disagg_gen_init", + "position", + "mirror_capacity", + "mirror_history", + "expected_resize", + "trail_tokens", + ), + [ + # External drafter (target view), aggregated chunked prefill. The + # prompt is 8192 + 2 * 1024 tokens, so the 1024-token chunk at 8192 is + # not the last one and carries no draft tokens. + pytest.param(False, False, 0, 0, 0, (1024, 0), 0, id="first_chunk"), + # A later chunk of a mirror that holds [0, 8192): the chunk start is + # history, so a sliding window drops the pages below it. + pytest.param(False, False, 8192, 8192, 4096, (9216, 8192), 0, id="later_chunk"), + # The history is already ahead of the chunk start; a resize never + # lowers it. + pytest.param(False, False, 8192, 8192, 8704, (9216, 8704), 0, id="history_ahead"), + # A mirror first seen past position 0 (a target prefix this worker did + # not compute). A DFlash/DSpark pool (one with a trailing page) takes + # the chunk start as history: that drafter stores the chunk at its + # target position above the prefix. Any other keeps its history and + # covers [0, chunk end). + pytest.param(False, False, 8192, 0, 0, (9216, 8192), 32, id="fresh_past_zero-dflash"), + pytest.param(False, False, 8192, 0, 0, (9216, None), 0, id="fresh_past_zero-pard"), + # MTP, Eagle and two-engine drafters keep the draft request view: it + # reads position 0 with the whole prompt as its (last) chunk, and the + # resize keeps the cache's history. + pytest.param(True, False, 8192, 8192, 0, (10243, None), 0, id="draft_view"), + # Disaggregated generation init: the mirror holds no drafter K/V for + # the transferred prompt (10240 tokens), so the whole prompt is + # history; the capacity stays absolute, as the drafter's positions. + pytest.param(False, True, 0, 0, 0, (10243, 10240), 0, id="disagg_gen_init-target_view"), + pytest.param(True, True, 0, 0, 0, (10243, None), 0, id="disagg_gen_init-draft_view"), + ], +) +def test_unpaired_draft_mirror_context_resize( + uses_draft_view: bool, + disagg_gen_init: bool, + position: int, + mirror_capacity: int, + mirror_history: int, + expected_resize: tuple, + trail_tokens: int, +) -> None: + """Which request view an unpaired draft mirror reads, and the capacity and + history its context resize asks for.""" + prompt, chunk, draft_len = 8192 + 2 * 1024, 1024, 3 + request = LlmRequest( + request_id=41, + max_new_tokens=4, + input_tokens=[0] * prompt, + sampling_config=SamplingConfig(1), + is_streaming=False, + draft_tokens=[1] * draft_len, + ) + if disagg_gen_init: + request.state = LlmRequestState.DISAGG_GENERATION_INIT + assert request.is_last_context_chunk + else: + request.state = LlmRequestState.CONTEXT_INIT + request.context_current_position = position + request.context_chunk_size = chunk + assert not request.is_last_context_chunk + + kv_cache = Mock(is_active=True, capacity=mirror_capacity, history_length=mirror_history) + views = [] + kv_cache.resize.side_effect = lambda *args: views.append(request.use_draft_model) or True + manager = object.__new__(KVCacheManagerV2) + manager.is_draft = True + manager.enable_joint_kv_cache_reuse = False + manager._mirror_uses_draft_view = uses_draft_view + manager.draft_slot_trail_tokens = trail_tokens + manager.num_extra_kv_tokens = 0 + manager.kv_cache_map = {request.py_request_id: kv_cache} + manager._resume_and_restore = Mock(return_value=True) + + batch = ScheduledRequests() + batch.context_requests_last_chunk.append(request) + manager._prepare_draft_resources(batch) + + kv_cache.resize.assert_called_once_with(*expected_resize) + assert views == [uses_draft_view] + + +@pytest.mark.parametrize("phase", ["context", "generation"]) +def test_unpaired_draft_mirror_that_cannot_grow_raises_with_the_pool_diagnostic(phase: str) -> None: + """When a mirror finds no page to grow into, the iteration fails + explicitly: the mirror left at its size has no storage for the K/V the + drafter writes next, and there is no per-request way to skip drafting. + The error carries the pool diagnostic.""" + position, chunk, draft_len = 8192, 1024, 3 + + def make_request(request_id: int) -> LlmRequest: + request = LlmRequest( + request_id=request_id, + max_new_tokens=4, + input_tokens=[0] * (position + chunk), + sampling_config=SamplingConfig(1), + is_streaming=False, + draft_tokens=[1] * draft_len, + ) + if phase == "context": + request.state = LlmRequestState.CONTEXT_INIT + request.context_current_position = position + request.context_chunk_size = chunk + else: + request.state = LlmRequestState.GENERATION_IN_PROGRESS + return request + + requests = [make_request(51), make_request(52)] + caches = { + 51: Mock(is_active=True, capacity=position, history_length=position, beam_width=1), + 52: Mock(is_active=True, capacity=position, history_length=position, beam_width=1), + } + caches[51].resize.return_value = True + caches[52].resize.return_value = False + manager = object.__new__(KVCacheManagerV2) + manager.is_draft = True + manager.enable_joint_kv_cache_reuse = False + manager._mirror_uses_draft_view = False + manager.num_extra_kv_tokens = 0 + manager.kv_cache_map = dict(caches) + manager._allocated_draft_lens = set() + manager._resume_and_restore = Mock(return_value=True) + manager._required_gen_capacity = Mock(return_value=position + 1 + draft_len) + manager._draft_pool_diagnostic = Mock(return_value=" [draft pool: full]") + manager.kv_cache_type = CacheType.SELF + manager.max_beam_width = 1 + + batch = ScheduledRequests() + if phase == "context": + batch.context_requests_last_chunk.extend(requests) + expected_resize = call(position + chunk + draft_len, position) + expected_capacity = position + chunk + draft_len + else: + batch.generation_requests.extend(requests) + expected_resize = call(position + 1 + draft_len) + expected_capacity = position + 1 + draft_len + with pytest.raises(RuntimeError) as excinfo: + manager._prepare_draft_resources(batch) + + assert caches[51].resize.call_args_list == [expected_resize] + assert caches[52].resize.call_args_list == [expected_resize] + message = str(excinfo.value) + assert f"Draft KV cache {phase} resize failed for request 52" in message + assert f"could not resize to {expected_capacity} tokens" in message + assert "[draft pool: full]" in message + + +def test_exhausted_pool_fails_explicitly_on_an_empty_fresh_mirror() -> None: + """A fresh mirror whose first resize fails holds no page at all. The + manager raises before the forward, with a diagnostic that shows the empty + mirror and how full the pool was.""" + prompt, draft_len = 1024, 3 + request = LlmRequest( + request_id=61, + max_new_tokens=4, + input_tokens=[0] * prompt, + sampling_config=SamplingConfig(1), + is_streaming=False, + draft_tokens=[1] * draft_len, + ) + request.state = LlmRequestState.CONTEXT_INIT + request.context_current_position = 0 + request.context_chunk_size = prompt + + # Two occupants with distinct fields, so a swapped field or a max in + # place of the sum changes the message. + occupant = Mock(is_active=True, capacity=1216, history_length=512, num_blocks=38) + small = Mock(is_active=True, capacity=64, history_length=0, num_blocks=2) + fresh = Mock(is_active=True, capacity=0, history_length=0, num_blocks=0) + fresh.resize.return_value = False + manager = object.__new__(KVCacheManagerV2) + manager.is_draft = True + manager.enable_joint_kv_cache_reuse = False + manager._mirror_uses_draft_view = False + manager.num_extra_kv_tokens = 0 + manager.kv_cache_map = {60: occupant, 62: small, 61: fresh} + manager.max_attention_window_vec = [2048] + manager._gpu_max_tokens = 1216 + manager._resume_and_restore = Mock(return_value=True) + + batch = ScheduledRequests() + batch.context_requests_last_chunk.append(request) + with pytest.raises(RuntimeError) as excinfo: + manager._prepare_draft_resources(batch) + + fresh.resize.assert_called_once_with(prompt + draft_len, 0) + message = str(excinfo.value) + assert ( + f"context resize failed for request 61: could not resize to {prompt + draft_len} tokens" + in message + ) + assert "3 live caches, capacity 1280 tokens in 40 logical blocks" in message + assert "60: capacity=1216 history=512 logical_blocks=38" in message + assert "61: capacity=0 history=0 logical_blocks=0" in message + assert "windows=[2048]" in message + assert "gpu_max_tokens=1216" in message + + +def test_windowed_draft_pool_backs_only_the_window_of_a_generation_dummy() -> None: + """A CUDA-graph warmup dummy is a generation request at the longest length + the target pool can hold. Its draft mirror carries the same history hint + as the target cache, so a windowed draft pool sized for a few mirrors of + window + chunk backs only the dummy's window; without the hint the mirror + would need the dummy's whole length, which such a pool cannot hold.""" + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + torch.cuda.init() + spec = DFlashDecodingConfig(max_draft_len=4, speculative_model="draft") + # 32 KiB per token (8 layers, 8 KV heads of 128, fp16): a mirror of a + # 256-token window plus a 512-token chunk is ~26 MiB; the pool holds four + # of them (~104 MiB), while a 60k-token dummy held in full needs ~1.8 GiB. + num_layers, num_kv_heads, head_dim = 8, 8, 128 + tokens_per_block, window, chunk, max_seq_len, batch_size = 32, 256, 512, 65536, 2 + token_num = 60_000 + draft_config = KvCacheConfig( + enable_block_reuse=False, max_attention_window=[window] * num_layers + ) + bytes_per_token = num_layers * 2 * num_kv_heads * head_dim * 2 + _, headroom = _get_generation_kv_capacity(spec, is_draft=True) + tokens_per_mirror = window + chunk + 2 * headroom + tokens_per_block + pages_per_mirror = -(-tokens_per_mirror // tokens_per_block) + draft_budget = 2 * batch_size * pages_per_mirror * tokens_per_block * bytes_per_token + assert token_num * bytes_per_token > 4 * draft_budget + draft = KVCacheManagerV2( + draft_config.model_copy(update={"max_gpu_total_bytes": draft_budget}), + CacheType.SELF, + num_layers=num_layers, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + tokens_per_block=tokens_per_block, + max_seq_len=max_seq_len, + max_batch_size=batch_size, + max_num_tokens=chunk, + mapping=Mapping(), + dtype=DataType.HALF, + spec_config=spec, + is_draft=True, + ) + # A small windowless target pool (512 B per token) that holds the whole + # dummy, so the target side of add_dummy_requests cannot be what fails. + target = KVCacheManagerV2( + KvCacheConfig(enable_block_reuse=False, max_gpu_total_bytes=128 << 20), + CacheType.SELF, + num_layers=1, + num_kv_heads=1, + head_dim=head_dim, + tokens_per_block=tokens_per_block, + max_seq_len=max_seq_len, + max_batch_size=batch_size, + max_num_tokens=chunk, + mapping=Mapping(), + dtype=DataType.HALF, + spec_config=spec, + ) + stream = torch.cuda.current_stream().cuda_stream + + def live_pages(cache: object) -> int: + return len(cache.get_aggregated_page_indices(0, valid_only=True)) + + requests = None + try: + # The mirror resize without a history hint: the whole dummy on the + # windowed pool, which is more than the pool has. + probe = draft.impl.create_kv_cache() + try: + assert probe.resume(stream) + assert not probe.resize(token_num + draft.num_extra_kv_tokens) + finally: + probe.close() + requests = target.add_dummy_requests( + [7], + token_nums=[token_num], + is_gen=True, + max_num_draft_tokens=spec.max_draft_len, + draft_kv_cache_manager=draft, + ) + assert requests is not None, "the dummy's draft mirror did not fit the windowed pool" + (request,) = requests + target_cache = target.kv_cache_map[request.py_request_id] + mirror = draft.kv_cache_map[request.py_request_id] + assert live_pages(target_cache) >= token_num // tokens_per_block + assert mirror.capacity == target_cache.capacity + assert mirror.history_length == token_num - 1 + assert live_pages(mirror) <= pages_per_mirror + finally: + for request in requests or (): + draft.free_resources(request) + target.free_resources(request) + draft.shutdown() + target.shutdown() + + +def _small_spec_manager( + spec: DecodingBaseConfig, + is_draft: bool, + *, + window: int | None = None, + num_kv_heads: int = 2, + dtype: DataType = DataType.HALF, + max_gpu_total_bytes: int = 256 << 20, + max_seq_len: int = 16384, + max_num_tokens: int = 8192, + tokens_per_block: int = 32, +) -> KVCacheManagerV2: + """A two-layer, head_dim 64 V2 manager for ``spec`` (32-token pages by + default), optionally with a uniform sliding window and no block reuse.""" + num_layers = 2 + config = KvCacheConfig( + enable_block_reuse=False, + max_attention_window=None if window is None else [window] * num_layers, + max_gpu_total_bytes=max_gpu_total_bytes, + ) + return KVCacheManagerV2( + config, + CacheType.SELF, + num_layers=num_layers, + num_kv_heads=num_kv_heads, + head_dim=64, + tokens_per_block=tokens_per_block, + max_seq_len=max_seq_len, + max_batch_size=2, + max_num_tokens=max_num_tokens, + mapping=Mapping(), + dtype=dtype, + spec_config=spec, + is_draft=is_draft, + ) + + +def _backed_ordinals(cache: object) -> set[int]: + """Block ordinals of ``cache`` that have a page behind them.""" + return { + ordinal + for ordinal, page in enumerate(cache.get_base_page_indices(0)) + if page != BAD_PAGE_INDEX + } + + +@pytest.mark.parametrize( + ("make_spec", "external_drafter"), + [ + pytest.param( + lambda: DFlashDecodingConfig(max_draft_len=4, speculative_model="draft"), + True, + id="dflash", + ), + pytest.param( + lambda: DSparkDecodingConfig(max_draft_len=4, speculative_model="draft"), + True, + id="dspark", + ), + pytest.param( + lambda: PARDDecodingConfig(max_draft_len=4, speculative_model="draft"), + True, + id="pard", + ), + pytest.param( + lambda: DraftTargetDecodingConfig(max_draft_len=4, speculative_model="draft"), + True, + id="draft_target_one_model", + ), + pytest.param(lambda: MTPDecodingConfig(num_nextn_predict_layers=1), False, id="mtp"), + pytest.param( + lambda: Eagle3DecodingConfig(max_draft_len=3, speculative_model="draft"), + False, + id="eagle3_one_model", + ), + ], +) +def test_only_an_external_drafters_draft_manager_reads_the_target_view( + make_spec: Callable[[], DecodingBaseConfig], external_drafter: bool +) -> None: + """The draft manager of a one-engine external drafter reads the target's + position and chunk; every other manager keeps the draft request view.""" + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + torch.cuda.init() + spec = make_spec() + assert spec.spec_dec_mode.is_external_drafter() is external_drafter + for is_draft in (False, True): + manager = _small_spec_manager(spec, is_draft, max_gpu_total_bytes=4 << 20) + try: + assert manager._mirror_uses_draft_view is not (is_draft and external_drafter) + finally: + manager.shutdown() + + +@pytest.mark.parametrize( + ("make_spec", "releases_below_window"), + [ + # MTP/Eagle draft layers process the target's chunks through the draft + # pool but the mirror reads the draft request view, which nothing + # advances: it keeps the whole prompt backed. + pytest.param( + lambda: Eagle3DecodingConfig(max_draft_len=3, speculative_model="draft"), + False, + id="eagle3_one_model", + ), + # An external drafter follows the target's chunks and drops what falls + # out of the window. + pytest.param( + lambda: DFlashDecodingConfig(max_draft_len=4, speculative_model="draft"), + True, + id="dflash", + ), + ], +) +def test_windowed_draft_mirror_backs_what_the_next_chunk_attends( + make_spec: Callable[[], DecodingBaseConfig], releases_below_window: bool +) -> None: + """Chunked prefill with block reuse off on a sliding-window draft pool: + after the second chunk's resize, every page the chunk attends (its window + back from the chunk start, through the chunk end) is backed.""" + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + torch.cuda.init() + window, page, prompt, chunk = 256, 32, 8192, 512 + draft = _small_spec_manager(make_spec(), True, window=window) + request = LlmRequest( + request_id=101, + max_new_tokens=8, + input_tokens=[1] * prompt, + sampling_config=SamplingConfig(1), + is_streaming=False, + ) + request.state = LlmRequestState.CONTEXT_INIT + try: + request.context_chunk_size = chunk + batch = ScheduledRequests() + batch.context_requests_chunking.append(request) + draft._prepare_draft_resources(batch) + + request.move_to_next_context_chunk() + request.context_chunk_size = chunk + assert request.context_current_position == chunk + batch = ScheduledRequests() + batch.context_requests_chunking.append(request) + draft._prepare_draft_resources(batch) + + backed = _backed_ordinals(draft.kv_cache_map[request.py_request_id]) + attended = set(range((chunk - window) // page, 2 * chunk // page)) + assert attended <= backed + assert (0 in backed) is not releases_below_window + draft.free_resources(request) + finally: + draft.shutdown() + + +class _PoolBoundDFlashDraftModel: + """Just enough of a DFlash draft model for the worker to bind the draft + KV pool and run its generation-step context store.""" + + block_size = 4 + config = SimpleNamespace(max_position_embeddings=1 << 17) + dflash_attention_backend = "TRTLLM" + hidden_norm = object() + + def __init__(self, num_layers: int, num_kv_heads: int, head_dim: int, hidden: int) -> None: + self.fc = SimpleNamespace(weight=torch.empty(1, dtype=torch.bfloat16, device="cuda")) + self._layers = (num_layers, num_kv_heads, head_dim) + embed_tokens = torch.nn.Embedding(16, hidden).cuda() + self.draft_model_full = SimpleNamespace(model=SimpleNamespace(embed_tokens=embed_tokens)) + + def _build_fused_kv_buffers(self) -> None: + self._num_attn_layers, self._num_kv_heads, self._head_dim = self._layers + self._num_heads = self._num_kv_heads + + def _get_attention_mask_args(self, layer_idx: int) -> tuple: + return True, (-1, -1) + + def validate_block_attention_windows(self) -> frozenset: + return frozenset() + + def project_target_hidden(self, hidden_states: torch.Tensor) -> torch.Tensor: + return hidden_states + + def precompute_context_kv( + self, projected: torch.Tensor, positions: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + num_layers, num_kv_heads, head_dim = self._layers + k = torch.full( + (projected.shape[0], num_layers, num_kv_heads, head_dim), + 3.0, + dtype=torch.bfloat16, + device="cuda", + ) + return k, k.clone() + + +def _disagg_draft_fixture( + spec: DFlashDecodingConfig, + prompt: int, + occupant_prompt: int, + window: int, + tokens_per_block: int = 32, +) -> tuple[KVCacheManagerV2, LlmRequest, LlmRequest]: + """A sliding-window draft pool with room for 8192 tokens (two layers, one + KV head of 64, bf16), an occupant that prefilled ``occupant_prompt`` + tokens on it, and a request in disaggregated generation init whose + transferred prompt is ``prompt`` tokens.""" + num_layers, num_kv_heads, head_dim = 2, 1, 64 + bytes_per_token = num_layers * 2 * num_kv_heads * head_dim * 2 + draft = _small_spec_manager( + spec, + True, + window=window, + num_kv_heads=num_kv_heads, + dtype=DataType.BF16, + max_gpu_total_bytes=8192 * bytes_per_token, + max_seq_len=1 << 17, + max_num_tokens=occupant_prompt, + tokens_per_block=tokens_per_block, + ) + occupant = LlmRequest( + request_id=1, + max_new_tokens=8, + input_tokens=[1] * occupant_prompt, + sampling_config=SamplingConfig(1), + is_streaming=False, + ) + occupant.state = LlmRequestState.CONTEXT_INIT + occupant.context_chunk_size = occupant_prompt + request = LlmRequest( + request_id=2, + max_new_tokens=64, + input_tokens=[1] * prompt, + sampling_config=SamplingConfig(1), + is_streaming=False, + draft_tokens=[1] * spec.max_draft_len, + ) + request.state = LlmRequestState.DISAGG_GENERATION_INIT + return draft, occupant, request + + +def _disagg_generation_step(draft: KVCacheManagerV2, request: LlmRequest) -> None: + """One executor generation step on the draft manager: grow the mirror, + then the post-forward update_resources (nothing accepted, one token).""" + batch = ScheduledRequests() + batch.generation_requests.append(request) + draft._prepare_draft_resources(batch) + request.py_num_accepted_draft_tokens = 0 + request.py_rewind_len = 0 + draft.update_resources(batch) + request.add_new_token(1, 0) + + +def test_gen_init_draft_mirror_survives_generation_updates() -> None: + """A disaggregated generation worker's draft mirror for a prompt far + larger than its sliding-window pool fits at generation init and stays + consistent through generation steps, including update_resources, which + advances the history in absolute positions: the history follows the + request and the pages from the prompt end up to the capacity stay + backed.""" + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + torch.cuda.init() + spec = DFlashDecodingConfig(max_draft_len=3, speculative_model="draft") + prompt, page = 65536, 32 + draft, occupant, request = _disagg_draft_fixture(spec, prompt, 512, window=256) + try: + for req in (occupant, request): + batch = ScheduledRequests() + batch.context_requests_last_chunk.append(req) + draft._prepare_draft_resources(batch) + cache = draft.kv_cache_map[request.py_request_id] + # The context worker's first token arrives with the transferred KV. + request.add_new_token(1, 0) + request.state = LlmRequestState.GENERATION_IN_PROGRESS + for _ in range(3): + _disagg_generation_step(draft, request) + assert cache.history_length == request.max_beam_num_tokens - 2 + assert set(range(prompt // page, -(-cache.capacity // page))) <= _backed_ordinals(cache) + draft.free_resources(request) + draft.free_resources(occupant) + finally: + draft.shutdown() + + +@pytest.mark.parametrize("slot_loss", [False, True], ids=["first_step", "slot_loss"]) +@pytest.mark.parametrize("reuse_addressable", [False, True], ids=["unpaired", "reuse_addressable"]) +def test_disagg_generation_drafter_writes_stay_on_its_own_draft_pages( + monkeypatch: pytest.MonkeyPatch, reuse_addressable: bool, slot_loss: bool +) -> None: + """A disaggregated generation worker receives a prompt far larger than its + sliding-window draft pool, while another request owns the pool's low + pages. After generation steps (update_resources included), the drafter's + step -- DFlashSpecMetadata.prepare() assigning the request a fresh slot, + the context K/V store of the accepted tokens, the block decode's read and + its block append -- must address only pages the request's own mirror + holds. The slot is seeded on the page holding the request's absolute + position and reads only what it stored, whether or not the pool joins + block reuse, and again after the request loses its slot and comes back.""" + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + torch.cuda.init() + monkeypatch.setattr( + "tensorrt_llm._torch.speculative.dflash.validate_dflash_trtllm_gen_runtime", + lambda **kwargs: None, + ) + num_layers, num_kv_heads, head_dim, hidden = 2, 1, 64, 8 + window, prompt, occupant_prompt, max_draft_len = 256, 65536, 512, 3 + spec = DFlashDecodingConfig( + max_draft_len=max_draft_len, + speculative_model="draft", + attention_backend="TRTLLM", + mask_token_id=0, + ) + draft, occupant, request = _disagg_draft_fixture(spec, prompt, occupant_prompt, window) + offsets = torch.zeros( + (getattr(draft, "num_attention_op_pools", draft.num_pools), 2, 2, draft.max_blocks_per_seq), + dtype=torch.int32, + device="cuda", + ) + tokens_per_req = max_draft_len + 1 + try: + for req in (occupant, request): + batch = ScheduledRequests() + batch.context_requests_last_chunk.append(req) + draft._prepare_draft_resources(batch) + request.add_new_token(1, 0) + request.state = LlmRequestState.GENERATION_IN_PROGRESS + for _ in range(2): + _disagg_generation_step(draft, request) + + worker = DFlashWorker(spec, Mapping()) + draft_model = _PoolBoundDFlashDraftModel(num_layers, num_kv_heads, head_dim, hidden) + worker.set_draft_model(draft_model) + attn_metadata = SimpleNamespace( + max_seq_len=1 << 17, + draft_kv_cache_block_offsets=offsets, + num_contexts=0, + num_seqs=1, + num_ctx_tokens=0, + _seq_lens_cuda=torch.zeros(1, dtype=torch.int32, device="cuda"), + _seq_lens=torch.zeros(1, dtype=torch.int32), + ) + worker._lazy_init_ctx_buffers( + draft_model, + SimpleNamespace(max_num_requests=2), + attn_metadata, + draft_kv_cache_manager=draft, + ) + assert worker._ctx_block_tables is not None, "the drafter did not bind the draft pool" + worker._ctx_reuse_addressable = reuse_addressable + # V2 hands out per-layer views over one interleaved pool: a table entry + # names layer 0's slot of a page, and layer l of that page is slot + # entry + l of the same dense view, so view 0 sees every layer's write. + assert worker._ctx_block_divisor == draft.kv_factor + pool = worker._ctx_kv_buf[0] + page = worker._ctx_page_size + + spec_metadata = DFlashSpecMetadata( + max_num_requests=2, max_draft_len=max_draft_len, max_total_draft_tokens=max_draft_len + ) + spec_metadata._dflash_worker = worker + + def prepare(request_ids: list[int]) -> None: + """DFlashSpecMetadata.prepare() for a batch of generation rows at + their current positions.""" + reqs = {request.py_request_id: request, occupant.py_request_id: occupant} + spec_metadata.request_ids = request_ids + spec_metadata.num_generations = len(request_ids) + spec_metadata.seq_lens = [tokens_per_req] * len(request_ids) + spec_metadata.host_position_ids = torch.tensor( + [ + reqs[rid].max_beam_num_tokens - 1 + offset + for rid in request_ids + for offset in range(tokens_per_req) + ], + dtype=torch.int32, + ) + spec_metadata.prepare() + + prepare([request.py_request_id]) + if slot_loss: + # The request misses a batch (its slot is evicted), runs a step + # elsewhere in the meantime, and comes back on a fresh slot. + prepare([]) + _disagg_generation_step(draft, request) + prepare([request.py_request_id]) + batch = ScheduledRequests() + batch.generation_requests.append(request) + draft._prepare_draft_resources(batch) + + def own_slots(req: LlmRequest) -> set[int]: + draft.copy_batch_block_offsets(offsets, [req.py_request_id], 1, 0, 1) + worker._refresh_ctx_block_tables(attn_metadata, 1) + row = worker._ctx_block_tables[0].tolist() + return { + row[ordinal] + layer + for ordinal in _backed_ordinals(draft.kv_cache_map[req.py_request_id]) + for layer in range(num_layers) + } + + occupant_slots = own_slots(occupant) + request_slots = own_slots(request) + assert occupant_slots and request_slots + assert not occupant_slots & request_slots + + pool.fill_(7.0) + before = pool.clone() + captured = torch.randn(tokens_per_req, hidden, device="cuda") + step_metadata = SimpleNamespace( + hidden_size=hidden, + runtime_draft_len=max_draft_len, + captured_hidden_states=None, + get_hidden_states=lambda num_tokens: captured[:num_tokens], + ) + inputs = worker.prepare_1st_drafter_inputs( + input_ids=torch.zeros(tokens_per_req, dtype=torch.long, device="cuda"), + position_ids=torch.zeros(tokens_per_req, dtype=torch.long, device="cuda"), + hidden_states=torch.zeros(tokens_per_req, hidden, device="cuda"), + accepted_tokens=torch.ones((1, tokens_per_req), dtype=torch.long, device="cuda"), + num_accepted_tokens=torch.tensor([tokens_per_req], dtype=torch.int32, device="cuda"), + attn_metadata=attn_metadata, + spec_metadata=step_metadata, + draft_model=draft_model, + total_target_tokens=tokens_per_req, + ) + torch.cuda.synchronize() + written = set((pool != before).flatten(1).any(dim=1).nonzero().flatten().tolist()) + assert written, "the drafter wrote nothing" + assert not written & occupant_slots, "the drafter wrote another request's draft pages" + assert written <= request_slots + + # The block decode reads num_ctx positions and appends the block right + # after them, both through the table row it is handed. + num_ctx = int(inputs["num_ctx_per_req"][0]) + read_row = inputs["ctx_page_table"][0].tolist() + block = worker._compute_block_size + addressed = { + read_row[pos // page] + layer + for pos in range(num_ctx + block) + for layer in range(num_layers) + } + assert addressed <= request_slots, "the block decode addresses another request's pages" + + # The slot started empty on the page holding the request's absolute + # position, every page below it hidden, and reads exactly the context + # it stored this step. + base = (request.max_beam_num_tokens - 1) // page + slot = worker._req_to_slot[request.py_request_id] + assert worker._ctx_seed_page_aligned + assert worker._ctx_base_host[slot] == base + assert worker._ctx_len_host[slot] == base * page + assert num_ctx == tokens_per_req + draft.free_resources(request) + draft.free_resources(occupant) + finally: + draft.shutdown() + + +@pytest.mark.parametrize("lag", [2, 7, 31, 127]) +def test_trailing_drafter_slot_reads_only_pages_its_request_holds( + monkeypatch: pytest.MonkeyPatch, lag: int +) -> None: + """A disaggregated generation worker seeds the drafter slot on the page + holding the request's position, ``lag`` tokens behind it, while the + request's sliding-window draft mirror advances its history with the + request. Over many generation steps (one accepted token each), every + block-table entry the drafter's window reads -- keys from the slot length + + 1 - window, never below the slot's base -- must stay on a page the + request's own mirror holds; a released entry resolves to another + request's page.""" + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + torch.cuda.init() + monkeypatch.setattr( + "tensorrt_llm._torch.speculative.dflash.validate_dflash_trtllm_gen_runtime", + lambda **kwargs: None, + ) + num_layers, num_kv_heads, head_dim, hidden = 2, 1, 64, 8 + window, page, max_draft_len, steps = 256, 128, 3, 600 + spec = DFlashDecodingConfig( + max_draft_len=max_draft_len, + speculative_model="draft", + attention_backend="TRTLLM", + mask_token_id=0, + ) + # The first token and two generation steps put the request at prompt + 2. + prompt = 65536 + lag - 2 + draft, occupant, request = _disagg_draft_fixture(spec, prompt, 512, window, page) + offsets = torch.zeros( + (getattr(draft, "num_attention_op_pools", draft.num_pools), 2, 2, draft.max_blocks_per_seq), + dtype=torch.int32, + device="cuda", + ) + try: + for req in (occupant, request): + batch = ScheduledRequests() + batch.context_requests_last_chunk.append(req) + draft._prepare_draft_resources(batch) + request.add_new_token(1, 0) + request.state = LlmRequestState.GENERATION_IN_PROGRESS + for _ in range(2): + _disagg_generation_step(draft, request) + + worker = DFlashWorker(spec, Mapping()) + draft_model = _PoolBoundDFlashDraftModel(num_layers, num_kv_heads, head_dim, hidden) + worker.set_draft_model(draft_model) + attn_metadata = SimpleNamespace( + max_seq_len=1 << 17, + draft_kv_cache_block_offsets=offsets, + num_contexts=0, + num_seqs=1, + num_ctx_tokens=0, + _seq_lens_cuda=torch.zeros(1, dtype=torch.int32, device="cuda"), + _seq_lens=torch.zeros(1, dtype=torch.int32), + ) + worker._lazy_init_ctx_buffers( + draft_model, + SimpleNamespace(max_num_requests=2), + attn_metadata, + draft_kv_cache_manager=draft, + ) + assert worker._ctx_page_size == page + assert worker._ctx_seed_page_aligned + spec_metadata = DFlashSpecMetadata( + max_num_requests=2, max_draft_len=max_draft_len, max_total_draft_tokens=max_draft_len + ) + spec_metadata._dflash_worker = worker + spec_metadata.request_ids = [request.py_request_id] + spec_metadata.num_generations = 1 + spec_metadata.seq_lens = [max_draft_len + 1] + position = request.max_beam_num_tokens - 1 + spec_metadata.host_position_ids = torch.tensor( + [position + offset for offset in range(max_draft_len + 1)], dtype=torch.int32 + ) + spec_metadata.prepare() + slot = worker._req_to_slot[request.py_request_id] + seed_len, base = int(worker._ctx_len_host[slot]), int(worker._ctx_base_host[slot]) + assert (base, position - seed_len) == (position // page, lag) + + cache = draft.kv_cache_map[request.py_request_id] + unbacked_reads = [] + for step in range(steps): + slot_len = seed_len + step + 1 + lowest = max(base * page, slot_len + 1 - window) + read = {key // page for key in range(lowest, slot_len)} + missing = read - _backed_ordinals(cache) + if missing: + unbacked_reads.append((step, cache.history_length, slot_len, sorted(missing))) + _disagg_generation_step(draft, request) + assert not unbacked_reads, ( + f"{len(unbacked_reads)}/{steps} steps read a released entry; first " + f"(step, history, slot length, entries): {unbacked_reads[:3]}" + ) + draft.free_resources(request) + draft.free_resources(occupant) + finally: + draft.shutdown() + + def _make_publishing_manager(policy: BlockReusePolicy) -> KVCacheManagerV2: """A manager wired just far enough to run the publish/history bookkeeping.""" manager = object.__new__(KVCacheManagerV2) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_copy_skip.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_copy_skip.py index f88ba9687763..d9b18896ce83 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_copy_skip.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_copy_skip.py @@ -23,10 +23,10 @@ def _worker(req_to_slot): _graph_dummy_id_floor=1 << 30, _dummy_slot=DUMMY_SLOT, _batch_to_slot_buffer=DeviceCopyBuffer(torch.zeros(8, dtype=torch.long)), - _write_ctx_len=lambda updates: None, + _write_ctx_len=lambda updates, base_updates=None: None, ) # Generation requests are already assigned in these tests; prepare assigns nothing. - worker._assign_slot = lambda rid: worker._req_to_slot.get(rid) + worker._assign_slot = lambda rid, **kwargs: worker._req_to_slot.get(rid) return worker diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_draft_kv.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_draft_kv.py index 53958cbab57c..d89b92753a90 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_draft_kv.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_draft_kv.py @@ -47,3 +47,297 @@ def test_eagle3_one_model_still_gets_a_separate_draft_kv_cache(): config = EagleDecodingConfig(max_draft_len=3, speculative_model="/tmp/eagle3-draft") assert config.spec_dec_mode.is_eagle3_one_model() assert should_use_separate_draft_kv_cache(config) is True + + +@pytest.mark.parametrize( + ("bound_to_pool", "reuse_addressable", "reads_rebased_view", "cur", "first_pos", "expected"), + [ + # Managed pool: the block table is indexed by target position, so a + # tail that reuse let the target start at first_pos is stored there. + (True, False, True, 0, 8192, 8192), + (True, True, True, 0, 8192, 8192), + # A pool paired with the target's reuse holds the prefix below the + # tail, so the tail goes on top of it even for a worker that reads + # from position 0 (the Kimi K3 workers). + (True, True, False, 0, 8192, 8192), + # An unpaired pool read from position 0 numbers the tail from 0. + (True, False, False, 0, 8192, 0), + # A continuing chunk keeps the slot's stored length. + (True, False, True, 4096, 4096, 4096), + # Private arena: every request numbers from 0. + (False, False, True, 0, 8192, 0), + (False, False, True, 0, 0, 0), + ], +) +def test_dflash_ctx_store_start_follows_the_pool_numbering( + bound_to_pool: bool, + reuse_addressable: bool, + reads_rebased_view: bool, + cur: int, + first_pos: int, + expected: int, +) -> None: + from tensorrt_llm._torch.speculative.dflash import DFlashWorker + + worker = object.__new__(DFlashWorker) + worker._ctx_block_tables = object() if bound_to_pool else None + worker._ctx_reuse_addressable = reuse_addressable + worker._ctx_reads_rebased_view = reads_rebased_view + assert worker._ctx_store_start(cur, first_pos) == expected + + +@pytest.mark.parametrize( + ("cur", "start", "reuse_addressable", "expected"), + [ + # Unpaired pool, tail stored at a reused prefix's end: the pages below + # it hold no drafter K/V and are hidden from the block decode. + (0, 8192, False, 128), + # A start inside a page rounds the base up; the tail tokens in that + # page stay stored but unread. + (0, 8200, False, 129), + # A continuing chunk numbers from the stored length: no base change. + (4096, 4096, False, None), + # A pool paired with the target's reuse holds the prefix's drafter K/V. + (0, 8192, True, None), + ], +) +def test_dflash_ctx_base_for_store( + cur: int, start: int, reuse_addressable: bool, expected: int | None +) -> None: + from tensorrt_llm._torch.speculative.dflash import DFlashWorker + + worker = object.__new__(DFlashWorker) + worker._ctx_page_size = 64 + worker._ctx_reuse_addressable = reuse_addressable + assert worker._ctx_base_for_store(cur, start) == expected + + +def test_dflash_rebased_ctx_rows_hide_the_reused_prefix() -> None: + """A tail stored after a reused prefix on an unpaired pool. The block decode + reads the slot's row rebased to the tail and a context length shortened by + the prefix, so the never-written prefix pages (poisoned here) are not within + reach, and every tail position still resolves to the page and offset the + store wrote it to.""" + import torch + + from tensorrt_llm._torch.speculative.dflash import dflash_rebased_ctx_rows + + page_size, first_pos, tail = 4, 24, 5 + row = torch.tensor([[7, 3, 11, 0, 5, 9, 13, 2, 6, 1, 15, 4, 8, 10, 12, 14]], dtype=torch.int32) + pool = torch.zeros(16, page_size) + pool[row[0, : first_pos // page_size].long()] = float("nan") + for pos in range(first_pos, first_pos + tail): + pool[row[0, pos // page_size], pos % page_size] = pos + + base_pages = torch.tensor([first_pos // page_size]) + view = dflash_rebased_ctx_rows(row, base_pages) + num_ctx = first_pos + tail - base_pages.item() * page_size + assert view.shape == row.shape + readable = view[0, : -(-num_ctx // page_size)].long() + assert torch.isfinite(pool[readable]).all() + for pos in range(num_ctx): + assert pool[view[0, pos // page_size], pos % page_size] == first_pos + pos + # A slot without a base reads its row unchanged. + assert torch.equal(dflash_rebased_ctx_rows(row, torch.zeros(1, dtype=torch.long)), row) + + +@pytest.mark.parametrize( + ("first_pos", "tail"), + [ + # A reused prefix ending inside a 64-token page with a one-token tail: + # the base rounds up to the next page and the context stays shorter + # than it for 55 steps. + (8200, 1), + # One token into a page: the first write offsets overlap the stored + # tail, which sits below the base and is never read. + (8193, 1), + # A page-aligned prefix end is never short of its base. + (8192, 1), + ], +) +def test_dflash_rebased_ctx_view_writes_the_block_into_an_owned_page( + first_pos: int, tail: int +) -> None: + """A tail stored after a reused prefix on an unpaired pool, with the mirror + funded the way the manager funds it: pages up to the one holding the + context plus the draft tokens, the window below it, nothing past it. The + block decode writes the block's own K/V at rebased positions + num_ctx .. num_ctx + block_size through the view it reads, so at every + context length those positions must resolve to a funded entry, including + the steps on which the context is still shorter than the rounded-up base + and the decode reads nothing. Entry ``base`` is unfunded on those steps and + holds 0, the pool page another request owns.""" + import torch + + from tensorrt_llm._torch.speculative.dflash import DFlashWorker, dflash_rebased_ctx_view + + page_size, block_size, draft_len, window_pages = 64, 3, 3, 64 + worker = object.__new__(DFlashWorker) + worker._ctx_page_size = page_size + worker._ctx_reuse_addressable = False + base = worker._ctx_base_for_store(0, first_pos) + assert base == -(-first_pos // page_size) + base_pages = torch.tensor([base]) + + def funded_row(ctx_len: int) -> torch.Tensor: + last = (ctx_len + draft_len - 1) // page_size + row = torch.zeros(1, base + 8, dtype=torch.int32) + row[0, last + 1 - window_pages : last + 1] = ( + torch.arange(last + 1 - window_pages, last + 1, dtype=torch.int32) + 1000 + ) + return row + + blind_steps = 0 + for ctx_len in range(first_pos + tail, (base + 1) * page_size): + row = funded_row(ctx_len) + view, num_ctx = dflash_rebased_ctx_view(row, base_pages, torch.tensor([ctx_len]), page_size) + assert view.shape == row.shape + blind = ctx_len < base * page_size + blind_steps += blind + assert num_ctx.item() == max(0, ctx_len - base * page_size) + if blind: + assert view[0, 0] == row[0, ctx_len // page_size] + if (ctx_len + draft_len - 1) // page_size < base: + # Entry ``base`` is funded only once the draft tokens reach it. + assert row[0, base] == 0 + for pos in range(num_ctx.item() + block_size): + entry = view[0, pos // page_size] + assert entry >= 1000, (ctx_len, pos) + if pos < num_ctx.item(): + # A read position p is absolute position base * page_size + p. + assert entry == row[0, (base * page_size + pos) // page_size] + assert blind_steps == max(0, base * page_size - first_pos - tail) + + +def test_dflash_write_ctx_len_zero_length_reset_wins_over_a_pending_base() -> None: + """_store_prefill_context records a chunk's base pages before the overflow + check that may reset the same slot to length 0 in the same call. The reset + takes precedence: the slot ends with length 0 and base 0 on the host + mirrors and the device tensors alike, while a slot that was not reset keeps + its new length and base.""" + import torch + + from tensorrt_llm._torch.speculative.dflash import DFlashWorker + + worker = object.__new__(DFlashWorker) + worker._ctx_len_host = [0, 512, 64] + worker._ctx_base_host = [0, 8, 1] + worker._ctx_len = torch.tensor(worker._ctx_len_host) + worker._ctx_base_pages = torch.tensor(worker._ctx_base_host) + + # Slot 1: base recorded, then reset by the overflow check. Slot 2: a + # continuing chunk with a new base. + worker._write_ctx_len({1: 0, 2: 96}, {1: 128, 2: 2}) + + assert worker._ctx_len_host == [0, 0, 96] + assert worker._ctx_base_host == [0, 0, 2] + assert worker._ctx_len.tolist() == [0, 0, 96] + assert worker._ctx_base_pages.tolist() == [0, 0, 2] + + +@pytest.mark.parametrize( + ("bound_to_pool", "page_aligned", "reuse_addressable", "position", "expected"), + [ + # A disaggregated generation row: the slot starts empty at the first + # position of the page holding its position and reads from there, + # paired with reuse or not (the mirror of such a request never looked + # a prefix up). + (True, True, False, 8200, (8192, 128)), + (True, True, True, 8200, (8192, 128)), + (True, True, False, 8192, (8192, 128)), + # Clamped to the drafter's context ceiling, base included. + (True, True, False, 1 << 20, (65536, 1024)), + # A pool window shorter than a page may already have released that + # page: the slot starts at the position, reading from the next page. + (True, False, False, 8200, (8200, 129)), + # The private arena is slot-addressed and starts at 0. + (False, True, False, 8200, (0, None)), + ], +) +def test_dflash_generation_slot_seed_is_absolute_with_a_base( + bound_to_pool: bool, + page_aligned: bool, + reuse_addressable: bool, + position: int, + expected: tuple, +) -> None: + """A generation row that gets a fresh slot (a disaggregated generation + worker's first step, or a request whose slot was evicted) starts in the + coordinates of its draft mirror, at a page boundary no further than a page + below its absolute position, with every page below the base hidden from + the block decode.""" + from tensorrt_llm._torch.speculative.dflash import DFlashWorker + + worker = object.__new__(DFlashWorker) + worker._ctx_block_tables = object() if bound_to_pool else None + worker._ctx_reuse_addressable = reuse_addressable + worker._ctx_seed_page_aligned = page_aligned + worker._ctx_page_size = 64 + worker._max_ctx = 65536 + updates, bases = {}, {} + worker._seed_generation_slot(3, position, updates, bases) + assert (updates[3], bases.get(3)) == expected + + +def test_dflash_assign_slot_resets_the_base_of_a_reused_slot() -> None: + """A slot freed and handed out again inside one batched update starts + with base 0, not the previous occupant's base.""" + from collections import deque + + from tensorrt_llm._torch.speculative.dflash import DFlashWorker + + worker = object.__new__(DFlashWorker) + worker._req_to_slot = {7: 0} + worker._free_slots = deque() + updates, bases = {}, {} + assert worker._assign_slot(8, reset=True, updates=updates, bases=bases) is None + assert worker._assign_slot(7, reset=True, updates=updates, bases=bases) == 0 + assert (updates, bases) == ({0: 0}, {0: 0}) + + +def test_dflash_rollback_restores_the_base_pages() -> None: + """A failed or warmup forward rolls the slot lengths back to the + pre-forward snapshot, and the base pages with them, on the device and the + host mirror; a base seeded during the forward does not survive.""" + import torch + + from tensorrt_llm._torch.speculative.dflash import DFlashWorker + + worker = object.__new__(DFlashWorker) + worker._ctx_len_host = [0, 512, 64] + worker._ctx_base_host = [0, 8, 0] + worker._ctx_len = torch.tensor(worker._ctx_len_host) + worker._ctx_base_pages = torch.tensor(worker._ctx_base_host) + worker._req_ctx_pos = {} + worker._saved_ctx_len = worker._ctx_len.clone() + worker._saved_ctx_len_host = list(worker._ctx_len_host) + worker._saved_ctx_base = worker._ctx_base_pages.clone() + worker._saved_ctx_base_host = list(worker._ctx_base_host) + worker._saved_req_ctx_pos = {} + + worker._write_ctx_len({1: 0, 2: 8200}, {2: 129}) + worker._restore_ctx_len_device() + worker._restore_ctx_len_host() + + assert worker._ctx_len.tolist() == [0, 512, 64] + assert worker._ctx_base_pages.tolist() == [0, 8, 0] + assert worker._ctx_len_host == [0, 512, 64] + assert worker._ctx_base_host == [0, 8, 0] + + +@pytest.mark.parametrize( + ("windows", "expected"), + [ + (None, True), + ([None, None], True), + ([None, 4096], True), + ([128, 4096], True), + ([None, 64], False), + ], +) +def test_dflash_windows_span_a_page(windows: list | None, expected: bool) -> None: + """A full-attention layer (window None) keeps every page; only a window + shorter than a page rules out the page-aligned generation seed.""" + from tensorrt_llm._torch.speculative.dflash import dflash_windows_span_a_page + + assert dflash_windows_span_a_page(windows, 128) is expected diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_dummy_slot.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_dummy_slot.py index 269e1de42197..7056c17c3ea7 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_dummy_slot.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_dummy_slot.py @@ -48,6 +48,8 @@ def _worker(lengths): _max_ctx=1000, _ctx_len=torch.tensor(lengths, dtype=torch.long, device="cuda"), _ctx_len_host=list(lengths), + _ctx_base_pages=torch.zeros(len(lengths), dtype=torch.long, device="cuda"), + _ctx_base_host=[0] * len(lengths), _batch_to_slot_buffer=DeviceCopyBuffer(torch.zeros(8, dtype=torch.long, device="cuda")), ) w._write_ctx_len = partial(DFlashWorker._write_ctx_len, w) @@ -109,9 +111,9 @@ def test_prepare_resets_only_evicted_slots(): written = [] write = worker._write_ctx_len - def record(updates): + def record(updates: dict[int, int], base_updates: dict[int, int] | None = None) -> None: written.append(dict(updates)) - write(updates) + write(updates, base_updates) worker._write_ctx_len = record _prepare(worker, [11, 12, 13, FLOOR + 1], num_generations=4) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_guided_rollback.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_guided_rollback.py index 2e29fd57a5cc..7960f49ba0cd 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_guided_rollback.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_guided_rollback.py @@ -93,6 +93,8 @@ def test_dflash_worker_rolls_back_after_verification(monkeypatch): worker._ctx_len = torch.zeros(batch_size, dtype=torch.long, device="cuda") # _lazy_init_ctx_buffers (stubbed here) normally builds the host mirror. worker._ctx_len_host = [0] * batch_size + worker._ctx_base_pages = torch.zeros(batch_size, dtype=torch.long, device="cuda") + worker._ctx_base_host = [0] * batch_size accepted = torch.zeros(batch_size, K + 1, dtype=torch.int32, device="cuda") # Distinct per-row counts: one row accepts the whole draft, the other only # the golden token. Uniform values would let a broken forward that drops or diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_draft_kv_resume_pressure.py b/tests/unittest/_torch/speculative/hw_agnostic/test_draft_kv_resume_pressure.py index dc90385bf0cb..01e1edd11e67 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_draft_kv_resume_pressure.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_draft_kv_resume_pressure.py @@ -45,7 +45,7 @@ class _Req: - def __init__(self, req_id): + def __init__(self, req_id: int) -> None: self.py_request_id = req_id self.lora_task_id = None self.cache_salt = None @@ -54,6 +54,7 @@ def __init__(self, req_id): self.context_current_position = 0 self.context_chunk_size = 8 self.py_draft_tokens = None + self.is_disagg_generation_init_state = False class _Batch: @@ -67,7 +68,7 @@ def all_requests(self): return self.context_requests + self.generation_requests -def _manager(resume_ok): +def _manager(resume_ok: bool) -> tuple[KVCacheManagerV2, MagicMock]: """A draft manager whose resume answers ``resume_ok``, with nothing else real. Built without ``__init__``: the method under test needs only the mirror @@ -78,10 +79,12 @@ def _manager(resume_ok): manager.is_draft = True manager.enable_joint_kv_cache_reuse = False manager.num_extra_kv_tokens = 0 + manager._mirror_uses_draft_view = False # one-engine external drafter: target view manager._allocated_draft_lens = set() cache = MagicMock(name="kv_cache") cache.resize.return_value = True cache.capacity = 0 + cache.history_length = 0 manager._mirror_draft_kv_cache = MagicMock(return_value=cache) manager._resume_and_restore = MagicMock(return_value=resume_ok) manager._required_gen_capacity = MagicMock(return_value=1) diff --git a/tests/unittest/_torch/speculative/test_spec_tail_kernels.py b/tests/unittest/_torch/speculative/test_spec_tail_kernels.py index ffddda1808bc..2c11c160187c 100644 --- a/tests/unittest/_torch/speculative/test_spec_tail_kernels.py +++ b/tests/unittest/_torch/speculative/test_spec_tail_kernels.py @@ -64,7 +64,11 @@ def test_strict_accept_matches_eager(num_ctx, num_gens, draft_len): assert torch.equal(accepted[num_ctx:], ref_accepted[num_ctx:]) -def test_decode_draft_block_table_matches_eager(): +def test_decode_draft_block_table_matches_eager() -> None: + """The decoded table is the clamped, divided offsets; the count is each + row's extent, one past its last backed entry, so a sliding window that + released the low entries leaves the count covering the positions the row + still addresses.""" num_pools, rows, n_blocks, divisor = 2, 9, 3000, 6 gen = torch.Generator(device="cuda").manual_seed(1) src = torch.randint( @@ -77,12 +81,18 @@ def test_decode_draft_block_table_matches_eager(): view = src[1, :rows, 0] view[0] = -torch.randint(1, 1 << 20, (n_blocks,), device="cuda", generator=gen).to(torch.int32) view[1] = 0 + # A windowed row: entries 0..99 released, 100..139 backed, the rest unfunded. + view[2] = -1 + view[2, 100:140] = torch.arange(100, 140, device="cuda", dtype=torch.int32) * divisor assert bool((view < 0).float().mean() > 0.3) and bool((view == 0).float().mean() > 0.2) table = torch.full((rows + 3, n_blocks), -7, dtype=torch.int32, device="cuda") counts = torch.full((rows + 3,), -7, dtype=torch.long, device="cuda") decode_draft_block_table(view, divisor, table, counts) encoded = view.to(torch.int64) - assert torch.equal(counts[:rows], (encoded >= 0).sum(dim=1)) + ordinals = torch.arange(1, n_blocks + 1, device="cuda") + extent = torch.where(encoded >= 0, ordinals, 0).max(dim=1).values + assert torch.equal(counts[:rows], extent) + assert counts[2].item() == 140 expected = encoded.clamp(min=0).div(divisor, rounding_mode="floor").to(torch.int32) assert torch.equal(table[:rows], expected) assert counts[0].item() == 0 and bool((table[0] == 0).all()) @@ -143,12 +153,18 @@ def _paged_worker(draft_len, page_size, max_blocks, num_slots, layers, head_dim, perm = torch.randperm(pages, device="cuda", generator=gen).to(torch.int32) worker._ctx_block_tables = perm[: num_slots * max_blocks].view(num_slots, max_blocks).clone() worker._ctx_block_counts = torch.zeros(num_slots, dtype=torch.long, device="cuda") + worker._ctx_base_pages = torch.zeros(num_slots, dtype=torch.long, device="cuda") + worker._ctx_block_tables_read = torch.zeros_like(worker._ctx_block_tables) worker._resolved_mask_token_id = 3 return worker, storage +@pytest.mark.parametrize("based", [False, True], ids=["no_base", "based"]) @pytest.mark.parametrize("draft_len", [3, 7]) -def test_paged_latent_prepare_matches_eager(draft_len): +def test_paged_latent_prepare_matches_eager(draft_len: int, based: bool) -> None: + """The fused generation-step preparation matches the eager one, including + the read view rebased past each slot's base pages: the advertised context, + the rebased block-table rows and the block placement.""" vocab, hidden, features, layers, head_dim = 100, 64, 128, 3, 96 page_size, max_blocks, num_slots, max_ctx = 16, 6, 10, 70 num_ctx_reqs, num_gens = 1, 6 @@ -166,6 +182,11 @@ def test_paged_latent_prepare_matches_eager(draft_len): worker._batch_to_slot_buffer.tensor[:batch] = slots worker._ctx_len[:5] = torch.tensor([10, 3, 66, 5, 40], device="cuda") worker._ctx_block_counts[:batch] = torch.tensor([0, 3, 6, 2, 0, 1, 1], device="cuda") + if based: + # Slot 2 (row 2) advertises past its base, so its read row shifts by + # it; slot 4 (row 3) advertises less than its base, reads nothing and + # shifts only to the page holding its context. The dummy slot keeps 0. + worker._ctx_base_pages[:5] = torch.tensor([0, 1, 1, 1, 2], device="cuda") num_accepted = torch.tensor([1, 2, width, 3, 1, 2, 1], dtype=torch.int32, device="cuda") accepted = torch.randint( 0, vocab, (batch, width), dtype=torch.int32, device="cuda", generator=gen @@ -193,6 +214,7 @@ def attn_metadata(): def prepare(fused): w = copy.copy(worker) w._ctx_len = worker._ctx_len.clone() + w._ctx_block_tables_read = worker._ctx_block_tables_read.clone() pool_storage = storage.clone() w._ctx_kv_buf = [pool_storage[:, layer] for layer in range(layers)] w._ctx_kv_ptrs = torch.tensor([p.data_ptr() for p in w._ctx_kv_buf], device="cuda") @@ -253,9 +275,13 @@ def prepare(fused): assert bool(shared.any()) and bool(untouched.any()) rows = eager["ctx_cache_batch_idx"] + read_rows = eager["ctx_page_table"].index_select(0, rows) + assert torch.equal(fused["ctx_page_table"].index_select(0, rows), read_rows) + if based: + assert not torch.equal(read_rows, worker._ctx_block_tables.index_select(0, rows)) fixup = _build_mla_block_fixup( eager["num_ctx_per_req"], - worker._ctx_block_tables.index_select(0, rows).long(), + read_rows.long(), draft_len, page_size, )