[https://nvbugs/5577178][fix] Score every prompt token when prompt logprobs hit the KV cache

tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py

@@ -1224,6 +1224,31 @@ def _spec_recompute_claim_limit(
     return max(0, (req.prompt_len - tail) // tokens_per_block * tokens_per_block)
 
 
+def _context_output_claim_limit(req: LlmRequest) -> int | None:
+    """Reuse-claim cap for a request that returns per-prompt-token outputs.
+
+    Context logits, and the prompt logprobs derived from them, exist only for
+    tokens that pass through the forward pass; cached KV cannot reproduce them.
+    Such a request claims no prefix, so every prompt position is scored. Its
+    computed blocks are still committed, so later requests keep reusing them.
+
+    Returns None for no cap, else 0. Applies to the draft pool too: both pools
+    settle the same request cursor.
+    """
+    if req.py_return_context_logits and not req.is_dummy:
+        return 0
+    return None
+
+
+def _tighter_claim_limit(first: int | None, second: int | None) -> int | None:
+    """The stricter of two reuse-claim caps, where None means no cap."""
+    if first is None:
+        return second
+    if second is None:
+        return first
+    return min(first, second)
+
+
 class KVCacheManagerV2(BaseResourceManager):
     # Sparse managers attach a metadata batch during initialization.
     sparse_metadata_batch: Batch | None = None
@@ -3846,14 +3871,18 @@ def prepare_context(self, req: LlmRequest) -> bool:
         # at the matched endpoint are onboarded; a cursor rewound below that
         # endpoint would read and write spans with no pages behind them). The
         # tail check runs first: partially constructed managers carry only the
-        # class-level zero tail.
-        limit = None
+        # class-level zero tail. Requests returning context logits claim nothing
+        # (see _context_output_claim_limit).
+        limit = _context_output_claim_limit(req)
         if self._spec_recompute_tail:
-            limit = _spec_recompute_claim_limit(
-                req,
-                tail=self._spec_recompute_tail,
-                tokens_per_block=self.tokens_per_block,
-                is_draft=self.is_draft,
+            limit = _tighter_claim_limit(
+                limit,
+                _spec_recompute_claim_limit(
+                    req,
+                    tail=self._spec_recompute_tail,
+                    tokens_per_block=self.tokens_per_block,
+                    is_draft=self.is_draft,
+                ),
             )
         reused = self.prepare_context_cache(req, limit)
         if reused is None:
@@ -4365,12 +4394,16 @@ def _prepare_connector_prefix_reservation(self, req: LlmRequest) -> None:
         # A connector-served prefix starves a hidden-state drafter exactly like
         # local reuse does. Cap the reservation at the same claim limit, so
         # the recomputed tail stays local by construction; a full-re-prefill
-        # tail (-1) caps to zero and releases the reservation below.
-        limit = _spec_recompute_claim_limit(
-            req,
-            tail=self._spec_recompute_tail,
-            tokens_per_block=self.tokens_per_block,
-            is_draft=self.is_draft,
+        # tail (-1) caps to zero and releases the reservation below. A served
+        # prefix has no context logits either, so those requests cap to zero.
+        limit = _tighter_claim_limit(
+            _context_output_claim_limit(req),
+            _spec_recompute_claim_limit(
+                req,
+                tail=self._spec_recompute_tail,
+                tokens_per_block=self.tokens_per_block,
+                is_draft=self.is_draft,
+            ),
         )
         if limit is not None:
             end = min(end, limit)

tensorrt_llm/_torch/pyexecutor/resource_manager.py

@@ -1194,6 +1194,13 @@ def prepare_resources(self, scheduled_batch: ScheduledRequests):
                 scheduled_batch, is_cross=False)
 
             if batch_request_infos:
+                # The chunk each context-logits request was admitted with,
+                # before add_sequence_batch shifts it past any reused prefix.
+                scheduled_chunks = {
+                    req.py_request_id: req.context_chunk_size
+                    for req in batch_llm_requests
+                    if req.py_return_context_logits
+                }
                 self.impl.add_sequence_batch(batch_request_infos,
                                              batch_llm_requests)
                 for req in batch_llm_requests:
@@ -1211,6 +1218,8 @@ def prepare_resources(self, scheduled_batch: ScheduledRequests):
                             req, block_ids)
 
                 self._maybe_apply_spec_recompute_tail(batch_llm_requests)
+                self._maybe_rewind_context_output_requests(
+                    batch_llm_requests, scheduled_chunks)
 
             for req in scheduled_batch.generation_requests:
                 if self.mapping.has_cp_helix():
@@ -1329,6 +1338,54 @@ def _maybe_rewind_reused_context(
                 req.context_current_position = 0
             req.context_chunk_size = prompt_len - target
 
+    def _has_windowed_attention(self) -> bool:
+        # Entries <= 0 are not attention windows (recurrent-state sentinels).
+        return any(window is not None and 0 < window < self.max_seq_len
+                   for window in self.max_attention_window_vec)
+
+    def _maybe_rewind_context_output_requests(
+            self, batch_llm_requests: List[LlmRequest],
+            scheduled_chunks: Dict[int, int]) -> None:
+        """Fully re-prefill cache-hit requests that return context logits.
+
+        Context logits (and the prompt logprobs derived from them) exist only
+        for tokens that pass through the forward pass, so a reused prefix leaves
+        those positions unscored. Rewinding to position 0 keeps the reused
+        blocks and scores every prompt token. The recompute rewrites those
+        shared blocks with the same K/V up to kernel numerics, which other
+        requests reading the prefix may observe. The scheduler admitted these requests without a
+        reuse discount (``clear_context_output_reuse_estimates``), so restoring
+        the admitted chunk from position 0 stays within the token budget.
+
+        A sliding-window layer detaches out-of-window blocks, so the rewound
+        cursor would cross spans with no pages behind them. There the request
+        keeps its reuse, and the worker rejects the incomplete context logits
+        rather than misaligning them.
+        """
+        if self.is_draft or not scheduled_chunks:
+            return
+        rewind = [
+            req for req in batch_llm_requests
+            if req.py_request_id in scheduled_chunks and not req.is_dummy
+            and req.prepopulated_prompt_len > 0
+        ]
+        if not rewind:
+            return
+        if self._has_windowed_attention():
+            logger.warning_once(
+                "Requests returning context logits or prompt logprobs cannot "
+                "bypass KV cache reuse with sliding-window attention on this "
+                "KV cache manager, so cache hits among them will fail. Use "
+                "kv_cache_config.use_kv_cache_manager_v2=True or disable "
+                "block reuse.",
+                key="context_output_rewind_unsupported")
+            return
+        for req in rewind:
+            # The C++ setter only advances the position for nonzero values.
+            req.set_prepopulated_prompt_len(0, self.tokens_per_block)
+            req.context_current_position = 0
+            req.context_chunk_size = scheduled_chunks[req.py_request_id]
+
     def report_batch_to_connector(self,
                                   scheduled_batch: ScheduledRequests) -> None:
         """Report the batch to the KV connector.

tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py

@@ -147,6 +147,19 @@ def drop_decoder_context_requests_waiting_for_encoder_output(
     return filtered_requests
 
 
+def clear_context_output_reuse_estimates(requests: RequestList) -> None:
+    """Charge requests that return context logits for their whole prompt.
+
+    Capacity scheduling credits their cached prefix in ``estimated_reusable_tokens``,
+    but the KV cache manager re-prefills them from position 0 so every prompt token
+    gets a logit (``KVCacheManager._maybe_rewind_context_output_requests``). Leaving
+    the credit would admit more tokens than the forward pass can hold.
+    """
+    for req in requests:
+        if req.py_return_context_logits and req.estimated_reusable_tokens:
+            req.estimated_reusable_tokens = 0
+
+
 def split_encoder_from_decoder_context_requests(
     requests: RequestList,
 ) -> tuple[RequestList, RequestList]:
@@ -542,6 +555,7 @@ def __init__(
     def schedule(
         self, active_requests: RequestList, inflight_request_ids: set[int]
     ) -> tuple[list[LlmRequest], list[LlmRequest], list[LlmRequest]]:
+        clear_context_output_reuse_estimates(active_requests)
         encoder_or_context_requests, generation_requests = self.impl(
             active_requests, inflight_request_ids, self.max_batch_size, self.max_num_tokens
         )
@@ -1053,6 +1067,7 @@ def _can_be_scheduled(self, req: LlmRequest) -> bool:
     def schedule(
         self, active_requests: RequestList, inflight_request_ids: set[int]
     ) -> tuple[RequestList, RequestList, RequestList]:
+        clear_context_output_reuse_estimates(active_requests)
         encoder_requests: RequestList = []
         context_requests: RequestList = []
         generation_requests: RequestList = []

tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py

@@ -1096,9 +1096,13 @@ def _reuse_claim_limit(self, req: LlmRequest) -> int | None:
         Doing this up front rather than claim-then-reconcile keeps the mismatch
         path (free the draft pages and re-claim shallower) off the common case.
         """
+        from ..kv_cache.kv_cache_manager_v2 import _context_output_claim_limit
+
         draft_manager = self._joint_draft_manager
         if draft_manager is None:
             return None
+        if _context_output_claim_limit(req) == 0:
+            return 0
         draft_match = draft_manager.probe_context_reuse(req)
         target_match = self.kv_cache_manager.probe_context_reuse(req)
         if draft_match is None or target_match is None:

tensorrt_llm/executor/base_worker.py

@@ -22,7 +22,7 @@
 import weakref
 from pathlib import Path
 from queue import Empty, Queue
-from typing import TYPE_CHECKING, Dict, List, Literal, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Dict, List, Literal, Optional, Set, Tuple, Union
 
 import torch
 
@@ -133,6 +133,9 @@ def __init__(
         self._results: Dict[int, GenerationResult] = {}
         # mapping: client_id from Proxy -> request_id returned from runtime backend
         self._client_id_to_request_id: Dict[int, int] = {}
+        # Requests already answered with an error by _maybe_wrap_response and
+        # aborted; their remaining engine responses are dropped.
+        self._refused_client_ids: Set[int] = set()
         self._await_response_helper = AwaitResponseHelper(weakref.proxy(self))
         self._backend = None if llm_args is None else llm_args.backend
         self._is_pytorch_backend = self._backend == "pytorch"
@@ -1531,9 +1534,12 @@ def handle_for_worker(self, responses: List[tllm.Response]) -> None:
         async_queues = []
         for response in responses:
             assert response is not None
+            is_error = isinstance(response, ErrorResponse)
+            if not is_error and _drop_refused_response(self.worker, response):
+                continue
             queue = self.worker.return_queue(response.client_id)
 
-            if not response.has_error():
+            if not is_error and not response.has_error():
                 response = _maybe_wrap_response(self.worker, response,
                                                 self.worker._is_pytorch_backend)
 
@@ -1548,7 +1554,8 @@ def handle_for_worker(self, responses: List[tllm.Response]) -> None:
             else:
                 queue.put(response)
 
-            if response.has_error() or response.result.is_final:
+            if (isinstance(response, ErrorResponse) or response.has_error()
+                    or response.result.is_final):
                 self.worker._pop_result(response.client_id)
 
         # Notify the events in bulk for performance.
@@ -1571,6 +1578,8 @@ def handle_for_ipc_batched(self, responses: List[tllm.Response]) -> None:
 
             if isinstance(response, ErrorResponse):
                 pass  # send ErrorResponse directly
+            elif _drop_refused_response(self.worker, response):
+                continue
             elif self.worker._has_background_error():
                 response = self.worker._create_error_response(response)
             elif response.has_error():
@@ -1614,6 +1623,10 @@ def _get_params_for_first_rsp(
     return None, None, None
 
 
+class IncompleteContextLogitsError(RuntimeError):
+    """Context logits do not cover the whole prompt, so prompt logprobs cannot be aligned."""
+
+
 def _compute_pytorch_prompt_logprobs(
         generation_result: GenerationResult,
         response: LlmResponse) -> Optional[LogProbsResult]:
@@ -1628,6 +1641,16 @@ def _compute_pytorch_prompt_logprobs(
             )  # generation logprobs, if requested, is provided directly in response.result.log_probs from the sampler.
     context_logits = response.result.context_logits
     assert context_logits is not None, "context_logits must not be None when prompt_logprobs is requested."
+    # Row t scores prompt token t + 1, so a short tensor (a reused prefix the
+    # forward pass skipped) cannot be aligned; scoring it would attribute the
+    # suffix's logprobs to the leading prompt tokens.
+    num_prompt_tokens = len(
+        generation_result._generation_request.prompt_token_ids)
+    if context_logits.shape[-2] < num_prompt_tokens:
+        raise IncompleteContextLogitsError(
+            f"prompt logprobs need context logits for all {num_prompt_tokens} "
+            f"prompt tokens, got {context_logits.shape[-2]} (a reused KV cache "
+            "prefix is not rescored on this configuration)")
     result = response.result.get_result()
     assert result is not None, "result must not be None when prompt_logprobs is requested."
     # Single element list
@@ -1787,12 +1810,35 @@ def _send_rsp(
         raise ValueError(f"Unknown response type: {response}")
 
 
+def _drop_refused_response(worker, response: tllm.Response) -> bool:
+    """Whether ``response`` belongs to a request already answered with an error.
+
+    Such a request was aborted, so the engine may still emit steps that were in
+    flight and a final cancellation; the client must see none of them.
+    """
+    if response.client_id not in worker._refused_client_ids:
+        return False
+    if response.has_error() or response.result.is_final:
+        worker._refused_client_ids.discard(response.client_id)
+    return True
+
+
 def _maybe_wrap_response(
-        worker,
-        response: tllm.Response,
-        is_pytorch_backend=False) -> Union[tllm.Response, ResponseWrapper]:
+    worker,
+    response: tllm.Response,
+    is_pytorch_backend=False
+) -> Union[tllm.Response, ResponseWrapper, ErrorResponse]:
 
-    logprobs_result = _get_logprobs(worker, response, is_pytorch_backend)
+    try:
+        logprobs_result = _get_logprobs(worker, response, is_pytorch_backend)
+    except IncompleteContextLogitsError as error:
+        # Fail only this request, and stop the engine generating for a client
+        # that has already been answered.
+        if not response.result.is_final:
+            worker._refused_client_ids.add(response.client_id)
+            worker.abort_request(response.client_id)
+        return ErrorResponse(response.client_id, str(error),
+                             response.request_id)
     req_perf_metrics = get_metrics_dict(response)
     if logprobs_result or req_perf_metrics:
         response = ResponseWrapper(response, logprobs_result, req_perf_metrics)

tensorrt_llm/grpc/openengine/formatting.py

@@ -167,12 +167,20 @@ def safe_length(self, text: str) -> int:
         return len(text) - max(self._states, default=0)
 
 
+def _missing_prompt_logprobs(result: Any) -> int:
+    """How many scoreable prompt positions (all but the first) have no logprob."""
+    prompt_logprobs = result.outputs[0].prompt_logprobs if result.outputs else None
+    return max(0, len(result.prompt_token_ids or ()) - 1 - len(prompt_logprobs or ()))
+
+
 def _prompt_output(
     tokenizer: Any, result: Any, candidate_limit: int
 ) -> generation_pb2.PromptOutput:
     token_ids = list(result.prompt_token_ids)
     prompt_logprobs = result.outputs[0].prompt_logprobs if result.outputs else None
     prompt_logprobs = prompt_logprobs or []
+    # The first prompt token has no preceding context, so it never carries a
+    # logprob; entry i scores prompt token i + 1.
     aligned_logprobs = [None, *prompt_logprobs[: max(0, len(token_ids) - 1)]]
     return generation_pb2.PromptOutput(
         tokens=_token_infos(tokenizer, token_ids, aligned_logprobs, candidate_limit)

tensorrt_llm/grpc/openengine/servicer.py

@@ -182,6 +182,13 @@ async def _generate(
             trace_headers, target_dp_rank, affinity_headers = _metadata_from_context(context)
             conversation_params = conversation_params_from_request(request)
             disaggregated_params = disaggregated_params_from_request(request)
+            if (
+                disaggregated_params is not None
+                and disaggregated_params.request_type == "generation_only"
+            ):
+                # Prompt scores are the context_only leg's PromptOutput; this
+                # leg receives the prompt KV and never computes its logits.
+                sampling_params.prompt_logprobs = None
             parent_affinity_id = validate_parent_affinity(
                 self.subagent_affinity_auth_key,
                 request,
@@ -361,6 +368,11 @@ def on_rpc_done(_: asyncio.Task) -> None:
                             retryable=True,
                         )
                         return
+                    if resp.HasField("error") and not result.error:
+                        # An error event ends the stream; unless the request
+                        # already finished, the finally below aborts it.
+                        engine_terminal = result.finished
+                        return
                 if result.error:
                     engine_terminal = True
                     return