diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 5835d677b6f0..3f45d8d39947 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -2092,6 +2092,21 @@ def build_managers(self, original_max_seq_len, estimating_kv_cache, kv_cache_config_override=draft_build_kv_cache_config) + # One-model (fused) draft: the draft forward shares the target's + # own request object and compute range this iteration, so the + # target manager's first-chunk reuse lookup can be safely bounded + # by what the draft manager's own trie can back (see + # KVCacheManagerV2._cap_tokens_for_paired_draft_reuse). Two-model + # (separate draft engine, self._draft_model_engine is not None, + # handled in the branch above) is deliberately never wired here: + # its draft LlmRequest objects and token stream are independently + # tracked, so the target's request object is the wrong key for + # that manager's trie. + if isinstance(kv_cache_manager, + KVCacheManagerV2) and isinstance( + draft_kv_cache_manager, KVCacheManagerV2): + kv_cache_manager._paired_draft_kv_cache_manager = ( + draft_kv_cache_manager) # Encoder-decoder cross-attention pool cross_kv_cache_manager = None diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index 43f7d9773876..dc8f15134257 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @@ -819,6 +819,19 @@ def __init__( ) self.is_draft = is_draft + # Set by the creator (_util.py) for a *target* manager paired with a + # one-model (fused) draft V2 KV cache manager that keys its own + # prefix-reuse trie on the EAGLE3-shifted token sequence (see + # _draft_reuse_tokens). When set, first-chunk context prep bounds how + # much of a request's prefix this manager may attach as reused to + # what the paired draft manager's own trie can also back with valid + # KV, via a non-attaching probe -- see _cap_tokens_for_paired_draft_reuse. + # None for the draft manager itself, and for any target manager not + # paired with a one-model draft (e.g. two-model/external-drafter + # configurations, which must not be driven through this coupling -- + # see the docstring on _cap_tokens_for_paired_draft_reuse). + self._paired_draft_kv_cache_manager: Optional["KVCacheManagerV2"] = None + # Retained so consumers (e.g. CUDAGraphRunner.preallocate_padding_dummies) # can distinguish the throwaway estimation-phase managers from the # final ones: the estimation cache is sized with no headroom for @@ -2475,6 +2488,7 @@ def _prepare_context_impl(self, req: LlmRequest) -> bool: tokens = self._augment_tokens_for_block_reuse( all_tokens, req, end=len(all_tokens) - 1 ) + tokens = self._cap_tokens_for_paired_draft_reuse(req, tokens) else: tokens = None kv_cache = self._create_kv_cache( @@ -2500,6 +2514,29 @@ def _prepare_context_impl(self, req: LlmRequest) -> bool: req.set_prepopulated_prompt_len( kv_cache.num_committed_tokens, self.tokens_per_block ) + # req.context_current_position is a C++-level DUAL-MODE field + # (mContextCurrentPositionTarget / mContextCurrentPositionDraft, + # selected by req.use_draft_model) -- see llmRequest.h + # getContextCurrentPosition/setContextCurrentPosition. This + # write, running here in target mode (not inside + # request_context(True, ...)), only ever touches the TARGET + # side. A paired one-model draft manager's own + # _prepare_draft_resources runs *inside* + # request_context(True, ...), so a plain read of + # req.context_current_position there resolves to the + # separate, independently-tracked DRAFT-side field -- which + # nothing above ever writes. Stash the final (already + # paired-capped, if applicable) target-side value as an + # ordinary Python attribute (bypasses the dual-mode C++ + # property) so the draft manager can read the real number. + req.py_draft_reuse_safe_prefix = req.context_current_position + # Same reasoning: stash the end of THIS iteration's target + # chunk (also target-mode-only otherwise) so a paired draft + # manager can size its own [safe_prefix, chunk_end) split + # without being able to read the target-mode field directly. + req.py_draft_target_chunk_end = ( + req.context_current_position + req.context_chunk_size + ) if req.is_disagg_generation_init_state: # Disagg generation receives prompt KV from the context worker; @@ -2516,6 +2553,53 @@ def _prepare_context_impl(self, req: LlmRequest) -> bool: ) return self._resume_and_restore(req.py_request_id, kv_cache) + def _cap_tokens_for_paired_draft_reuse( + self, req: LlmRequest, tokens: Sequence[TokenIdExt] + ) -> Sequence[TokenIdExt]: + """Bound a first-chunk target reuse lookup key to what a paired + one-model draft V2 KV cache manager's own trie can also back with + valid draft KV, via a non-attaching probe of the draft trie. + + Called *before* this manager's real (attaching) lookup runs, so the + target never commits/skips more of the prefix than the draft can + actually serve this iteration -- avoiding ever needing to "recompute" + target hidden states for an already-committed/shared prefix region + (not supported, and would be a target fallback/recompute path). + Combined with capping the draft's own real attach to + ``req.context_current_position`` in ``_prepare_draft_resources`` + (which, after this runs, already reflects + ``min(target_trie_hit_tokens, draft_trie_hit_tokens)``), this makes + both managers agree on the same safe reused prefix without either + side ever attaching more than the other can validate. + + A no-op (returns *tokens* unchanged) when: this manager has no + paired draft manager (only wired for one-model/fused draft configs -- + see ``_paired_draft_kv_cache_manager``); the paired manager has block + reuse disabled; *req* is a dummy/warmup request (never populates the + draft trie); *req* has speculative decoding disabled (no draft KV + cache will exist for it); or *tokens* is already empty. + """ + draft_mgr = self._paired_draft_kv_cache_manager + if ( + draft_mgr is None + or not draft_mgr.enable_block_reuse + or req.is_dummy + or getattr(req, "py_disable_speculative_decoding", False) + or not len(tokens) + ): + return tokens + draft_tokens = draft_mgr._draft_reuse_tokens(req) + draft_hit = draft_mgr.probe_prefix_match_length( + draft_tokens, req.lora_task_id, req.cache_salt + ) + if draft_hit >= len(tokens): + return tokens + if draft_hit <= 0: + # No draft-side reuse at all: match the enable_block_reuse=False + # convention (None, not an empty sequence) for the real lookup. + return None + return tokens[:draft_hit] + def resize_context(self, req: LlmRequest, num_tokens: int) -> bool: """Resize KV cache to cover context_current_position + num_tokens. @@ -2661,10 +2745,28 @@ def _prepare_draft_resources(self, scheduled_batch: ScheduledRequests): for req in scheduled_batch.context_requests: kv_cache = self.kv_cache_map.get(req.py_request_id) if kv_cache is None: + # req.context_current_position is a C++-level DUAL-MODE + # field (mContextCurrentPositionTarget vs + # mContextCurrentPositionDraft, selected by + # req.use_draft_model -- see llmRequest.h). We are inside + # request_context(True, scheduled_batch) here, so a plain + # read of req.context_current_position resolves to the + # DRAFT side, which nothing else ever writes before this + # point -- it is NOT the target's skip boundary. The + # target manager's own _prepare_context_impl stashes its + # final (already paired-capped, if applicable) value as + # a plain Python attribute for exactly this reason; see + # the comment there. + safe_prefix = getattr(req, "py_draft_reuse_safe_prefix", 0) + draft_lookup_tokens = None + if self.enable_block_reuse and not req.is_dummy and safe_prefix > 0: + draft_lookup_tokens = self._draft_reuse_tokens( + req, end=safe_prefix + ) kv_cache = self._create_kv_cache( req.py_request_id, req.lora_task_id, - None, + draft_lookup_tokens, cache_salt=req.cache_salt, is_dummy=req.is_dummy, ) @@ -2675,7 +2777,28 @@ def _prepare_draft_resources(self, scheduled_batch: ScheduledRequests): # slots free up, before the request runs any spec-dec # forward that needs the mirror. continue - kv_cache.stop_committing() + if not self.enable_block_reuse or req.is_dummy: + kv_cache.stop_committing() + else: + # Else: leave committing open. try_commit_blocks + # (called for this manager from py_executor.py, + # mirroring the target manager's own + # update_context_resources call, and running + # *outside* request_context -- i.e. in target mode, + # so it reads the target-side field, which the + # target manager already advances correctly) + # commits whatever this draft forward actually + # computes into the trie for a future request to hit. + matched = kv_cache.num_committed_tokens + # Reflect the real, validated split into this + # manager's OWN (draft-mode) position/chunk fields -- + # still request_context(True, ...) here, so this + # writes mContextCurrentPositionDraft/ + # mContextChunkSizeDraft, not the target's fields. + chunk_end = getattr(req, "py_draft_target_chunk_end", None) + if chunk_end is not None: + req.context_current_position = matched + req.context_chunk_size = chunk_end - matched if not self._resume_and_restore(req.py_request_id, kv_cache): raise RuntimeError( f"Failed to resume draft KV cache for request {req.py_request_id}" @@ -2727,6 +2850,40 @@ def _reuse_token_source(self, req: LlmRequest) -> Sequence[int]: return req.get_tokens_view(DEFAULT_BEAM_INDEX) return req.get_tokens(DEFAULT_BEAM_INDEX) + def _draft_reuse_tokens( + self, req: LlmRequest, start: int = 0, end: int | None = None + ) -> Sequence[TokenIdExt]: + """EAGLE3-transformed token sequence used as this draft manager's own + radix-tree key -- both to look up (probe or attach) a draft KV + prefix committed by an earlier request with the same prompt, and to + commit newly-computed draft KV under that same key + (try_commit_blocks). + + EAGLE3/MTP-Eagle drafts consume the target's own token stream shifted + left by one position (draft position ``p`` is fed + ``target_tokens[p + 1]``; see ``_prepare_context_input_ids`` in + ``speculative/interface.py`` and ``get_draft_model_prompt`` in + ``speculative/model_drafter.py``, which apply the same shift when + building the tokens actually fed to the draft model). Multimodal + content digests are spliced in first (via + ``_augment_tokens_for_block_reuse``, over the *unshifted* token + positions, matching where ``req.multimodal_positions`` are defined) + so the shift only ever reorders already-content-addressed entries. + + ``start``/``end`` are positions in the *shifted* (draft) sequence. + When ``end`` is None, it defaults to one past the last known + position: the final element of the full shifted sequence is the + position whose token is not yet known (the next token to be + sampled), the same "last token cannot be recovered" convention the + target uses for its own first-chunk lookup key. + """ + all_tokens = self._reuse_token_source(req) + augmented = self._augment_tokens_for_block_reuse(all_tokens, req) + shifted = augmented[1:] + if end is None: + end = len(shifted) - 1 + return shifted[start:end] + def _augment_tokens_for_block_reuse( self, tokens: Sequence[int], req: LlmRequest, start: int = 0, end: int | None = None ) -> Sequence[TokenIdExt]: @@ -3495,9 +3652,23 @@ def release_resources( return requests def try_commit_blocks(self, request: LlmRequest) -> None: - should_block_reuse = ( - self.enable_block_reuse and not self.is_draft and not request.is_dummy_request - ) + """Commit this manager's own newly-computed KV for *request* into its + prefix-reuse trie, from ``num_committed_tokens`` up to + ``request.context_current_position``. + + Works for both the target manager and a draft V2 manager + (``self.is_draft``): the draft manager keys its trie on the + EAGLE3-shifted token sequence (``_draft_reuse_tokens``) instead of + the raw target tokens, since that is what its own KV positions + actually hold. Callers are responsible for only invoking this with a + *request* object whose ``context_current_position``/token stream is + valid for *this* manager: for a one-model (fused) draft manager that + is the same shared request object the target manager uses (their + compute ranges are identical by construction); a two-model (separate + draft engine) draft manager must not be driven through this path + with the target's own request object. + """ + should_block_reuse = self.enable_block_reuse and not request.is_dummy_request if not should_block_reuse: return @@ -3506,12 +3677,19 @@ def try_commit_blocks(self, request: LlmRequest) -> None: return if request.context_current_position > kv_cache.num_committed_tokens: - tokens = self._augment_tokens_for_block_reuse( - self._reuse_token_source(request), - request, - start=kv_cache.num_committed_tokens, - end=request.context_current_position, - ) + if self.is_draft: + tokens = self._draft_reuse_tokens( + request, + start=kv_cache.num_committed_tokens, + end=request.context_current_position, + ) + else: + tokens = self._augment_tokens_for_block_reuse( + self._reuse_token_source(request), + request, + start=kv_cache.num_committed_tokens, + end=request.context_current_position, + ) # TODO: On a disaggregated prefill server, pass is_end=True for # the last context chunk to improve performance. kv_cache.commit(tokens) @@ -4033,7 +4211,10 @@ def probe_prefix_match_length(self, input_tokens, lora_task_id=None, cache_salt= """ if not self.enable_block_reuse: return 0 - if not input_tokens: + # len(), not `not input_tokens`: input_tokens may be a zero-copy + # numpy int32 view (get_tokens_view), whose truth value is ambiguous + # for more than one element. + if len(input_tokens) == 0: return 0 salt_int = self._derive_reuse_salt(cache_salt) return self.impl.probe_reuse( diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index c3cd9e97ea37..b79795568885 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -3234,6 +3234,7 @@ def _handle_executed_batch(self, # handling can terminate the request. self.kv_cache_manager.update_context_resources( scheduled_requests) + self._commit_draft_context_kv_blocks(scheduled_requests) if self.kv_cache_transceiver: finished_ctx_reqs = scheduled_requests.context_requests_last_chunk self._send_kv_async(finished_ctx_reqs) @@ -4370,6 +4371,7 @@ def _executor_loop(self): # handling can terminate the request. self.kv_cache_manager.update_context_resources( scheduled_batch) + self._commit_draft_context_kv_blocks(scheduled_batch) self._send_kv_async(scheduled_batch.all_requests()) self._handle_canceled_requests() @@ -5249,6 +5251,7 @@ def _executor_loop_overlap(self): and scheduled_batch.context_requests): self.kv_cache_manager.update_context_resources( scheduled_batch) + self._commit_draft_context_kv_blocks(scheduled_batch) if self.previous_batch is not None and should_process_previous_batch: self._commit_kv_cache_stats( @@ -7778,6 +7781,45 @@ def _update_request_states_tp(self, scheduled_requests: ScheduledRequests): else: request.state = LlmRequestState.GENERATION_IN_PROGRESS + def _commit_draft_context_kv_blocks( + self, scheduled_requests: ScheduledRequests) -> None: + """Commit this iteration's newly-computed draft context KV into the + draft V2 KV cache manager's own prefix-reuse trie, mirroring + ``self.kv_cache_manager.update_context_resources()``'s commit of the + target manager's blocks. + + Only applies to a *fused* (one-model) draft: ``self.drafter is + None`` there (see ``SpeculativeDecodingMode.has_spec_drafter()`` -- + EAGLE3/MTP one-model modes don't register a ``ModelDrafter``, + because the draft forward is embedded directly in the target + engine's own forward rather than driven by a separate drafter + object). In that case ``scheduled_requests.context_requests`` are the + *same* request objects the target manager just used -- their + ``context_current_position`` already reflects exactly the token + range this iteration's fused forward computed for both target and + draft (the safe common reused prefix, once + ``KVCacheManagerV2._cap_tokens_for_paired_draft_reuse`` capped it), + so it is valid to pass straight to the draft manager's + ``try_commit_blocks``. + + A two-model (separate draft engine) drafter builds its own draft + ``LlmRequest`` objects with independently-tracked positions and a + differently-transformed token stream (see ``model_drafter.py``); the + target's own request objects would be the wrong key for that + manager, so this deliberately does nothing when ``self.drafter is + not None``. That path is not modified by this change. + """ + if self.drafter is not None: + return + draft_kv_cache_manager = self.resource_manager.resource_managers.get( + ResourceManagerType.DRAFT_KV_CACHE_MANAGER) + if draft_kv_cache_manager is None or not getattr( + draft_kv_cache_manager, "enable_block_reuse", False): + return + for request in scheduled_requests.context_requests: + if request.state != LlmRequestState.GENERATION_COMPLETE: + draft_kv_cache_manager.try_commit_blocks(request) + @nvtx_range("_update_request_states") def _update_request_states(self, scheduled_requests: ScheduledRequests): cp_config = self.dist.cp_config