From 408ac63405cec5517ae784e4c1089840a213d3b6 Mon Sep 17 00:00:00 2001 From: Allison Lim Date: Wed, 26 Aug 2026 17:20:56 -0700 Subject: [PATCH 1/2] WIP: reuse draft KV prefixes for fused EAGLE3 --- tensorrt_llm/_torch/pyexecutor/_util.py | 15 + .../_torch/pyexecutor/kv_cache_manager_v2.py | 304 +++++++++++++++++- tensorrt_llm/_torch/pyexecutor/py_executor.py | 42 +++ tensorrt_llm/_torch/speculative/eagle3.py | 83 +++++ 4 files changed, 432 insertions(+), 12 deletions(-) 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..38e902e1bbcd 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @@ -121,6 +121,29 @@ if field_name not in KV_CACHE_ITERATION_STATS_REUSE_FIELDS ) +# TEMPORARY DIAGNOSTIC INSTRUMENTATION for the EAGLE3/MTP draft-side +# block-reuse fix (see scripts/repro.py). Not a product change; safe to +# delete once validated. Uses plain print() (not `logger`) so it is visible +# regardless of TLLM_LOG_LEVEL, including under TLLM_WORKER_USE_SINGLE_PROCESS=1. +_EAGLE3_REPRO_DEBUG = os.environ.get("TLLM_EAGLE3_REPRO_DEBUG", "0") == "1" + + +def _repro_debug(tag: str, **fields) -> None: + if not _EAGLE3_REPRO_DEBUG: + return + kv = " ".join(f"{k}={v}" for k, v in fields.items()) + print(f"[EAGLE3-DEBUG][{tag}] {kv}", flush=True) + + +def _repro_debug_token_hash(tokens) -> str: + """Short stable hash of a reuse lookup key, for diagnostic log lines only + (not used for any actual radix-tree/hashing decision).""" + try: + data = b",".join(str(int(t)).encode() for t in tokens) + except (TypeError, ValueError): + data = repr(list(tokens)).encode() + return hashlib.sha256(data).hexdigest()[:16] + class Role: KEY = DataRole("key") @@ -819,6 +842,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 +2511,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 +2537,39 @@ 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 _EAGLE3_REPRO_DEBUG and not self.is_draft: + _repro_debug( + "TARGET-reuse", + req_id=req.py_request_id, + safe_reused_prefix_applied=req.context_current_position, + target_forward_position_range=( + f"[{req.context_current_position}, " + f"{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 +2586,64 @@ 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 _EAGLE3_REPRO_DEBUG: + target_hit_probe = self.probe_prefix_match_length( + tokens, req.lora_task_id, req.cache_salt + ) + _repro_debug( + "TARGET-draft-paired-reuse-cap", + req_id=req.py_request_id, + target_trie_hit_tokens=target_hit_probe, + draft_trie_hit_tokens=draft_hit, + safe_reused_prefix=min(target_hit_probe, draft_hit), + ) + 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 +2789,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 +2821,55 @@ 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 _EAGLE3_REPRO_DEBUG: + try: + attached_block_ids = self.get_batch_cache_indices( + [req.py_request_id])[0] + attached_block_ids = ( + attached_block_ids.tolist() if hasattr( + attached_block_ids, "tolist") else + list(attached_block_ids)) + except Exception: # noqa: BLE001 - diagnostic only + attached_block_ids = None + _repro_debug( + "DRAFT-reuse-lookup", + req_id=req.py_request_id, + safe_prefix_requested=safe_prefix, + lookup_key_len=( + len(draft_lookup_tokens) + if draft_lookup_tokens is not None else 0 + ), + lookup_key_hash=( + _repro_debug_token_hash(draft_lookup_tokens) + if draft_lookup_tokens is not None else "n/a" + ), + draft_trie_hit_tokens=matched, + draft_context_start=req.context_current_position, + draft_context_chunk=req.context_chunk_size, + attached_block_ids=attached_block_ids, + ) 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}" @@ -2687,6 +2881,20 @@ def _prepare_draft_resources(self, scheduled_batch: ScheduledRequests): + draft_len + self.num_extra_kv_tokens ) + if _EAGLE3_REPRO_DEBUG: + # This manager's OWN (draft-mode) bookkeeping view, not + # necessarily proof of what gets embedded/computed by the + # actual forward -- see DRAFT-forward-ground-truth + # (eagle3.py), logged from attn_metadata/position_ids + # immediately before the real forward, for that. + _repro_debug( + "DRAFT-capacity-range", + req_id=req.py_request_id, + draft_kv_cache_manager_position_range=( + f"[{req.context_current_position}, " + f"{req.context_current_position + req.context_chunk_size})" + ), + ) if not kv_cache.resize(capacity): raise RuntimeError( f"Draft KV cache context resize failed for request " @@ -2727,6 +2935,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 +3737,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,15 +3762,36 @@ 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) + if _EAGLE3_REPRO_DEBUG and self.is_draft: + try: + block_ids = self.get_batch_cache_indices( + [request.py_request_id])[0] + block_ids = (block_ids.tolist() if hasattr( + block_ids, "tolist") else list(block_ids)) + except Exception: # noqa: BLE001 - diagnostic only + block_ids = None + _repro_debug( + "DRAFT-commit", + req_id=request.py_request_id, + committed_tokens=kv_cache.num_committed_tokens, + committed_block_ids=block_ids, + ) if request.context_remaining_length == 0: kv_cache.stop_committing() @@ -4033,7 +4310,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 diff --git a/tensorrt_llm/_torch/speculative/eagle3.py b/tensorrt_llm/_torch/speculative/eagle3.py index d497862e0c8e..36f5e7bf6747 100644 --- a/tensorrt_llm/_torch/speculative/eagle3.py +++ b/tensorrt_llm/_torch/speculative/eagle3.py @@ -1,6 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import os from dataclasses import dataclass, field from typing import TYPE_CHECKING, Dict, List, Optional, Set @@ -27,6 +28,85 @@ if TYPE_CHECKING: from ...llmapi.llm_args import EagleDecodingConfig +# TEMPORARY DIAGNOSTIC INSTRUMENTATION for the EAGLE3/MTP draft-side +# block-reuse fix (see scripts/repro.py). Not a product change; safe to +# delete once validated. Uses plain print() (not `logger`) so it is visible +# regardless of TLLM_LOG_LEVEL, including under TLLM_WORKER_USE_SINGLE_PROCESS=1. +_EAGLE3_REPRO_DEBUG = os.environ.get("TLLM_EAGLE3_REPRO_DEBUG", "0") == "1" + + +def _repro_debug(tag: str, **fields) -> None: + if not _EAGLE3_REPRO_DEBUG: + return + kv = " ".join(f"{k}={v}" for k, v in fields.items()) + print(f"[EAGLE3-DEBUG][{tag}] {kv}", flush=True) + + +def _repro_debug_draft_forward_ground_truth(attn_metadata, inputs, + draft_kv_cache_manager, + step_idx: int) -> None: + """Log ground-truth state immediately before the real draft-model + forward, for the context (prefill) step of one-model (fused) + EAGLE3/MTP-Eagle speculative decoding. + + Called with ``attn_metadata`` already swapped to point at the draft KV + cache manager (inside ``draft_kv_cache_context``), and with + ``inputs["position_ids"]``/``inputs["input_ids"]`` being the *exact* + tensors about to be fed into ``_run_draft_forward`` -- i.e. this reads + what is actually about to be computed, not a prediction from request + bookkeeping. request.context_current_position/context_chunk_size are + deliberately NOT used here: they are a C++-level dual-mode field + (target vs draft, gated on request.use_draft_model) that is only valid + to read while inside the specific request_context(...) scope that wrote + it (see kv_cache_manager_v2.py's _prepare_draft_resources / + _prepare_context_impl comments) -- by the time this forward runs, that + scope has long exited. Everything logged here instead comes from the + draft KV cache manager's own committed-token counter and the tensors + actually bound to this forward. + """ + if not _EAGLE3_REPRO_DEBUG or step_idx != 0 or draft_kv_cache_manager is None: + return + num_contexts = getattr(attn_metadata, "num_contexts", 0) + if num_contexts <= 0: + return + request_ids = list(attn_metadata.request_ids[:num_contexts]) + context_lens = attn_metadata.context_lens[:num_contexts].tolist() + position_ids = inputs.get("position_ids") + tokens_per_block = getattr(draft_kv_cache_manager, "tokens_per_block", None) + offset = 0 + for req_id, num_tokens in zip(request_ids, context_lens): + pos_slice = (position_ids[offset:offset + num_tokens] + if position_ids is not None else None) + kv_cache = draft_kv_cache_manager.kv_cache_map.get(req_id) + matched = kv_cache.num_committed_tokens if kv_cache is not None else None + try: + block_ids = draft_kv_cache_manager.get_batch_cache_indices( + [req_id])[0] + block_ids = (block_ids.tolist() + if hasattr(block_ids, "tolist") else list(block_ids)) + except Exception: # noqa: BLE001 - diagnostic only + block_ids = None + reused_blocks, new_blocks = None, None + if block_ids is not None and matched is not None and tokens_per_block: + n_reused = min( + (matched + tokens_per_block - 1) // tokens_per_block, + len(block_ids)) + reused_blocks = block_ids[:n_reused] + new_blocks = block_ids[n_reused:] + _repro_debug( + "DRAFT-forward-ground-truth", + req_id=req_id, + draft_trie_hit_tokens=matched, + num_input_tokens_this_forward=num_tokens, + actual_forward_position_range=( + f"[{int(pos_slice[0])}, {int(pos_slice[-1]) + 1})" + if pos_slice is not None and len(pos_slice) else "n/a"), + full_block_table=block_ids, + reused_blocks=reused_blocks, + newly_allocated_blocks=new_blocks, + ) + offset += num_tokens + class Eagle3ResourceManager(BaseResourceManager): """ @@ -954,6 +1034,9 @@ def _forward_linear_draft_loop(self, inputs, attn_metadata, spec_metadata, for i in range(runtime_draft_len): if uses_dsa_mtp_metadata: attn_metadata.set_skip_topk(i > 0) + if _EAGLE3_REPRO_DEBUG: + _repro_debug_draft_forward_ground_truth( + attn_metadata, inputs, draft_kv_cache_manager, i) # Run draft model (mode-specific via helper). The helper # passes ``all_rank_num_tokens`` as a kwarg so the draft model # handles save/restore internally (Eagle3DraftModel.forward From d6cfdee75727504324fbd0ef16c2b61252facb25 Mon Sep 17 00:00:00 2001 From: Allison Lim Date: Wed, 26 Aug 2026 17:34:31 -0700 Subject: [PATCH 2/2] Clean up draft KV prefix reuse --- .../_torch/pyexecutor/kv_cache_manager_v2.py | 99 ------------------- tensorrt_llm/_torch/speculative/eagle3.py | 83 ---------------- 2 files changed, 182 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index 38e902e1bbcd..dc8f15134257 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @@ -121,29 +121,6 @@ if field_name not in KV_CACHE_ITERATION_STATS_REUSE_FIELDS ) -# TEMPORARY DIAGNOSTIC INSTRUMENTATION for the EAGLE3/MTP draft-side -# block-reuse fix (see scripts/repro.py). Not a product change; safe to -# delete once validated. Uses plain print() (not `logger`) so it is visible -# regardless of TLLM_LOG_LEVEL, including under TLLM_WORKER_USE_SINGLE_PROCESS=1. -_EAGLE3_REPRO_DEBUG = os.environ.get("TLLM_EAGLE3_REPRO_DEBUG", "0") == "1" - - -def _repro_debug(tag: str, **fields) -> None: - if not _EAGLE3_REPRO_DEBUG: - return - kv = " ".join(f"{k}={v}" for k, v in fields.items()) - print(f"[EAGLE3-DEBUG][{tag}] {kv}", flush=True) - - -def _repro_debug_token_hash(tokens) -> str: - """Short stable hash of a reuse lookup key, for diagnostic log lines only - (not used for any actual radix-tree/hashing decision).""" - try: - data = b",".join(str(int(t)).encode() for t in tokens) - except (TypeError, ValueError): - data = repr(list(tokens)).encode() - return hashlib.sha256(data).hexdigest()[:16] - class Role: KEY = DataRole("key") @@ -2560,16 +2537,6 @@ def _prepare_context_impl(self, req: LlmRequest) -> bool: req.py_draft_target_chunk_end = ( req.context_current_position + req.context_chunk_size ) - if _EAGLE3_REPRO_DEBUG and not self.is_draft: - _repro_debug( - "TARGET-reuse", - req_id=req.py_request_id, - safe_reused_prefix_applied=req.context_current_position, - target_forward_position_range=( - f"[{req.context_current_position}, " - f"{req.context_current_position + req.context_chunk_size})" - ), - ) if req.is_disagg_generation_init_state: # Disagg generation receives prompt KV from the context worker; @@ -2625,17 +2592,6 @@ def _cap_tokens_for_paired_draft_reuse( draft_hit = draft_mgr.probe_prefix_match_length( draft_tokens, req.lora_task_id, req.cache_salt ) - if _EAGLE3_REPRO_DEBUG: - target_hit_probe = self.probe_prefix_match_length( - tokens, req.lora_task_id, req.cache_salt - ) - _repro_debug( - "TARGET-draft-paired-reuse-cap", - req_id=req.py_request_id, - target_trie_hit_tokens=target_hit_probe, - draft_trie_hit_tokens=draft_hit, - safe_reused_prefix=min(target_hit_probe, draft_hit), - ) if draft_hit >= len(tokens): return tokens if draft_hit <= 0: @@ -2843,33 +2799,6 @@ def _prepare_draft_resources(self, scheduled_batch: ScheduledRequests): if chunk_end is not None: req.context_current_position = matched req.context_chunk_size = chunk_end - matched - if _EAGLE3_REPRO_DEBUG: - try: - attached_block_ids = self.get_batch_cache_indices( - [req.py_request_id])[0] - attached_block_ids = ( - attached_block_ids.tolist() if hasattr( - attached_block_ids, "tolist") else - list(attached_block_ids)) - except Exception: # noqa: BLE001 - diagnostic only - attached_block_ids = None - _repro_debug( - "DRAFT-reuse-lookup", - req_id=req.py_request_id, - safe_prefix_requested=safe_prefix, - lookup_key_len=( - len(draft_lookup_tokens) - if draft_lookup_tokens is not None else 0 - ), - lookup_key_hash=( - _repro_debug_token_hash(draft_lookup_tokens) - if draft_lookup_tokens is not None else "n/a" - ), - draft_trie_hit_tokens=matched, - draft_context_start=req.context_current_position, - draft_context_chunk=req.context_chunk_size, - attached_block_ids=attached_block_ids, - ) 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}" @@ -2881,20 +2810,6 @@ def _prepare_draft_resources(self, scheduled_batch: ScheduledRequests): + draft_len + self.num_extra_kv_tokens ) - if _EAGLE3_REPRO_DEBUG: - # This manager's OWN (draft-mode) bookkeeping view, not - # necessarily proof of what gets embedded/computed by the - # actual forward -- see DRAFT-forward-ground-truth - # (eagle3.py), logged from attn_metadata/position_ids - # immediately before the real forward, for that. - _repro_debug( - "DRAFT-capacity-range", - req_id=req.py_request_id, - draft_kv_cache_manager_position_range=( - f"[{req.context_current_position}, " - f"{req.context_current_position + req.context_chunk_size})" - ), - ) if not kv_cache.resize(capacity): raise RuntimeError( f"Draft KV cache context resize failed for request " @@ -3778,20 +3693,6 @@ def try_commit_blocks(self, request: LlmRequest) -> None: # TODO: On a disaggregated prefill server, pass is_end=True for # the last context chunk to improve performance. kv_cache.commit(tokens) - if _EAGLE3_REPRO_DEBUG and self.is_draft: - try: - block_ids = self.get_batch_cache_indices( - [request.py_request_id])[0] - block_ids = (block_ids.tolist() if hasattr( - block_ids, "tolist") else list(block_ids)) - except Exception: # noqa: BLE001 - diagnostic only - block_ids = None - _repro_debug( - "DRAFT-commit", - req_id=request.py_request_id, - committed_tokens=kv_cache.num_committed_tokens, - committed_block_ids=block_ids, - ) if request.context_remaining_length == 0: kv_cache.stop_committing() diff --git a/tensorrt_llm/_torch/speculative/eagle3.py b/tensorrt_llm/_torch/speculative/eagle3.py index 36f5e7bf6747..d497862e0c8e 100644 --- a/tensorrt_llm/_torch/speculative/eagle3.py +++ b/tensorrt_llm/_torch/speculative/eagle3.py @@ -1,7 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -import os from dataclasses import dataclass, field from typing import TYPE_CHECKING, Dict, List, Optional, Set @@ -28,85 +27,6 @@ if TYPE_CHECKING: from ...llmapi.llm_args import EagleDecodingConfig -# TEMPORARY DIAGNOSTIC INSTRUMENTATION for the EAGLE3/MTP draft-side -# block-reuse fix (see scripts/repro.py). Not a product change; safe to -# delete once validated. Uses plain print() (not `logger`) so it is visible -# regardless of TLLM_LOG_LEVEL, including under TLLM_WORKER_USE_SINGLE_PROCESS=1. -_EAGLE3_REPRO_DEBUG = os.environ.get("TLLM_EAGLE3_REPRO_DEBUG", "0") == "1" - - -def _repro_debug(tag: str, **fields) -> None: - if not _EAGLE3_REPRO_DEBUG: - return - kv = " ".join(f"{k}={v}" for k, v in fields.items()) - print(f"[EAGLE3-DEBUG][{tag}] {kv}", flush=True) - - -def _repro_debug_draft_forward_ground_truth(attn_metadata, inputs, - draft_kv_cache_manager, - step_idx: int) -> None: - """Log ground-truth state immediately before the real draft-model - forward, for the context (prefill) step of one-model (fused) - EAGLE3/MTP-Eagle speculative decoding. - - Called with ``attn_metadata`` already swapped to point at the draft KV - cache manager (inside ``draft_kv_cache_context``), and with - ``inputs["position_ids"]``/``inputs["input_ids"]`` being the *exact* - tensors about to be fed into ``_run_draft_forward`` -- i.e. this reads - what is actually about to be computed, not a prediction from request - bookkeeping. request.context_current_position/context_chunk_size are - deliberately NOT used here: they are a C++-level dual-mode field - (target vs draft, gated on request.use_draft_model) that is only valid - to read while inside the specific request_context(...) scope that wrote - it (see kv_cache_manager_v2.py's _prepare_draft_resources / - _prepare_context_impl comments) -- by the time this forward runs, that - scope has long exited. Everything logged here instead comes from the - draft KV cache manager's own committed-token counter and the tensors - actually bound to this forward. - """ - if not _EAGLE3_REPRO_DEBUG or step_idx != 0 or draft_kv_cache_manager is None: - return - num_contexts = getattr(attn_metadata, "num_contexts", 0) - if num_contexts <= 0: - return - request_ids = list(attn_metadata.request_ids[:num_contexts]) - context_lens = attn_metadata.context_lens[:num_contexts].tolist() - position_ids = inputs.get("position_ids") - tokens_per_block = getattr(draft_kv_cache_manager, "tokens_per_block", None) - offset = 0 - for req_id, num_tokens in zip(request_ids, context_lens): - pos_slice = (position_ids[offset:offset + num_tokens] - if position_ids is not None else None) - kv_cache = draft_kv_cache_manager.kv_cache_map.get(req_id) - matched = kv_cache.num_committed_tokens if kv_cache is not None else None - try: - block_ids = draft_kv_cache_manager.get_batch_cache_indices( - [req_id])[0] - block_ids = (block_ids.tolist() - if hasattr(block_ids, "tolist") else list(block_ids)) - except Exception: # noqa: BLE001 - diagnostic only - block_ids = None - reused_blocks, new_blocks = None, None - if block_ids is not None and matched is not None and tokens_per_block: - n_reused = min( - (matched + tokens_per_block - 1) // tokens_per_block, - len(block_ids)) - reused_blocks = block_ids[:n_reused] - new_blocks = block_ids[n_reused:] - _repro_debug( - "DRAFT-forward-ground-truth", - req_id=req_id, - draft_trie_hit_tokens=matched, - num_input_tokens_this_forward=num_tokens, - actual_forward_position_range=( - f"[{int(pos_slice[0])}, {int(pos_slice[-1]) + 1})" - if pos_slice is not None and len(pos_slice) else "n/a"), - full_block_table=block_ids, - reused_blocks=reused_blocks, - newly_allocated_blocks=new_blocks, - ) - offset += num_tokens - class Eagle3ResourceManager(BaseResourceManager): """ @@ -1034,9 +954,6 @@ def _forward_linear_draft_loop(self, inputs, attn_metadata, spec_metadata, for i in range(runtime_draft_len): if uses_dsa_mtp_metadata: attn_metadata.set_skip_topk(i > 0) - if _EAGLE3_REPRO_DEBUG: - _repro_debug_draft_forward_ground_truth( - attn_metadata, inputs, draft_kv_cache_manager, i) # Run draft model (mode-specific via helper). The helper # passes ``all_rank_num_tokens`` as a kwarg so the draft model # handles save/restore internally (Eagle3DraftModel.forward