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

cpp/include/tensorrt_llm/batch_manager/llmRequest.h

@@ -1169,6 +1169,19 @@ class GenericLlmRequest
         return mReturnContextLogits;
     }
 
+    //! \brief Allocate the whole prompt afresh instead of claiming a cached prefix.
+    //! \details Outputs computed per prompt token, such as context logits, exist only for tokens the
+    //!          forward pass computes. The request's blocks are still stored for reuse by others.
+    void setSkipPrefixReuse(bool const skipPrefixReuse)
+    {
+        mSkipPrefixReuse = skipPrefixReuse;
+    }
+
+    [[nodiscard]] bool getSkipPrefixReuse() const
+    {
+        return mSkipPrefixReuse;
+    }
+
     void setReturnGenerationLogits(bool const returnGenerationLogits)
     {
         TLLM_CHECK_WITH_INFO(!(mIsStreaming && mSamplingConfig.getBeamWidth() > 1 && returnGenerationLogits),
@@ -1810,6 +1823,7 @@ class GenericLlmRequest
     bool mReturnContextLogits;
     bool mReturnGenerationLogits;
     bool mReturnLogProbs{false};
+    bool mSkipPrefixReuse{false};
     TensorPtr mContextLogitsHost;    // [mPromptLen, vocabSizePadded]
     TensorPtr mGenerationLogitsHost; // [beamSize, mMaxNewTokens, vocabSizePadded]
 

cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp

@@ -1418,12 +1418,16 @@ PrefixReuseSummary BlockManager::analyzePrefixReuse(
 PrefixReuseSummary WindowBlockManager::analyzePrefixReuse(
     VecUniqueTokens const& uniqueTokens, LlmRequest const& llmRequest) const
 {
+    PrefixReuseSummary summary;
+    if (llmRequest.getSkipPrefixReuse())
+    {
+        return summary;
+    }
+
     auto blockedUniqueTokens
         = chopVectorIntoBlocks<UniqueToken>(uniqueTokens, uniqueTokens.size(), mTokensPerBlock, false);
     auto blockKeys = buildBlockKeys(blockedUniqueTokens, llmRequest);
 
-    PrefixReuseSummary summary;
-
     std::lock_guard<std::recursive_mutex> lock(mLookupTree->getMutex());
     auto reuseMatches = findReusableBlockMatches(
         blockKeys, /*enablePartialReuse=*/false, /*copyOnPartialReuse=*/false, std::numeric_limits<SizeType32>::max());
@@ -1617,8 +1621,10 @@ WindowBlockManager::ClaimResult WindowBlockManager::claimMatchingBlocks(Generati
     result.numSharedContextBlocks
         = (beamWidth > 1 && !isShareLastContextBlock) ? numContextBlocks - 1 : numContextBlocks;
     result.shareLastContextBlockAmongBeams = result.numSharedContextBlocks == numContextBlocks;
-    auto reuseMatches = findReusableBlockMatches(
-        result.blockKeys, mEnablePartialReuse, mCopyOnPartialReuse, sequence.getCurrentPrepopulatedPromptLen());
+    auto reuseMatches = llmRequest.getSkipPrefixReuse()
+        ? ReuseMatchResult{}
+        : findReusableBlockMatches(
+            result.blockKeys, mEnablePartialReuse, mCopyOnPartialReuse, sequence.getCurrentPrepopulatedPromptLen());
     result.totalMatchedTokens = reuseMatches.totalMatchedTokens;
 
     for (int bi = 0; bi < result.numSharedContextBlocks && bi < static_cast<int>(reuseMatches.matches.size()); ++bi)
@@ -1963,7 +1969,7 @@ SizeType32 WindowBlockManager::onboardAndAllocateBlocks(
     }
 
     SizeType32 numConnectorMatchedTokens = 0;
-    if (mKvCacheConnectorManager && !llmRequest.isDummyRequest())
+    if (mKvCacheConnectorManager && !llmRequest.isDummyRequest() && !llmRequest.getSkipPrefixReuse())
     {
         numConnectorMatchedTokens
             = mKvCacheConnectorManager->getNumNewMatchedTokens(llmRequest, sequence.getCurrentPrepopulatedPromptLen());

cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp

@@ -200,6 +200,7 @@ void initBindings(nb::module_& m)
             nb::arg("kv_tokens_per_block"))
         .def_prop_rw(
             "estimated_reusable_tokens", &GenLlmReq::getEstimatedReusableTokens, &GenLlmReq::setEstimatedReusableTokens)
+        .def_prop_rw("skip_prefix_reuse", &GenLlmReq::getSkipPrefixReuse, &GenLlmReq::setSkipPrefixReuse)
         .def_prop_rw(
             "expect_snapshot_points", &GenLlmReq::getExpectedSnapshotPoints, &GenLlmReq::setExpectedSnapshotPoints)
         .def_prop_rw("context_phase_params", &GenLlmReq::getContextPhaseParams, &GenLlmReq::setContextPhaseParams)

cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp

@@ -8536,6 +8536,51 @@ TEST_F(KVCacheManagerTest, AddSequenceBatchLeavesOneFinalContextTokenAfterReuse)
     EXPECT_NO_THROW(static_cast<void>(mgr->removeSequence(requestId, req)));
 }
 
+TEST_F(KVCacheManagerTest, AddSequenceBatchSkipPrefixReuseClaimsNothingButStoresItsBlocks)
+{
+    auto const stream = std::make_shared<tr::CudaStream>();
+    auto mgr = makeBatchTestKVCacheManager(stream);
+    auto constexpr beamWidth = 1;
+    auto constexpr promptLen = 9;
+    auto const cachedTokens = std::make_shared<VecTokens>(VecTokens{0, 1, 2, 3, 4, 5, 6, 7, 8});
+    auto const freshTokens = std::make_shared<VecTokens>(VecTokens{9, 10, 11, 12, 13, 14, 15, 16, 17});
+
+    auto makeRequest = [&](LlmRequest::RequestIdType requestId, std::shared_ptr<VecTokens> const& tokens)
+    {
+        return std::make_shared<LlmRequest>(
+            requestId, /*maxNewTokens=*/2, tokens, tle::SamplingConfig{beamWidth}, /*isStreaming=*/false);
+    };
+    auto addThenRemove = [&](std::shared_ptr<LlmRequest> const& req)
+    {
+        mgr->addSequenceBatch({{{req->mRequestId, promptLen, beamWidth}}}, {std::ref(*req)});
+        auto const prepopulated = req->getPrepopulatedPromptLen();
+        tensorrt_llm::testing::KvCacheManagerTestUtil::simulatePrefillCompletion(*req);
+        (void) mgr->removeSequence(req->mRequestId, req);
+        return prepopulated;
+    };
+
+    // A cached prompt is neither credited to the scheduler nor claimed.
+    seedAndRelease(*mgr, /*reqId=*/0, cachedTokens);
+    auto skipping = makeRequest(1, cachedTokens);
+    skipping->setSkipPrefixReuse(true);
+    auto const summary = mgr->analyzePrefixReuse(skipping->getUniqueTokens(0), *skipping);
+    EXPECT_EQ(summary.reusableBlocksAll, 0);
+    EXPECT_EQ(summary.reusableBlocksAllocated, 0);
+    EXPECT_FALSE(summary.firstNewBlock.has_value());
+    mgr->addSequenceBatch({{{1, promptLen, beamWidth}}}, {std::ref(*skipping)});
+    EXPECT_EQ(skipping->getPrepopulatedPromptLen(), 0);
+    EXPECT_EQ(skipping->getContextCurrentPosition(), 0);
+    EXPECT_EQ(skipping->getReusedBlocksPerRequest(), 0);
+    tensorrt_llm::testing::KvCacheManagerTestUtil::simulatePrefillCompletion(*skipping);
+    (void) mgr->removeSequence(1, skipping);
+
+    // The skipping request's own blocks are stored, so a later request reuses them.
+    auto freshSkipping = makeRequest(2, freshTokens);
+    freshSkipping->setSkipPrefixReuse(true);
+    EXPECT_EQ(addThenRemove(freshSkipping), 0);
+    EXPECT_EQ(addThenRemove(makeRequest(3, freshTokens)), promptLen - 1);
+}
+
 TEST_F(KVCacheManagerTest, AddSequenceBatchPreservesDraftTokensOnFinalContextAfterReuse)
 {
     auto const stream = std::make_shared<tr::CudaStream>();
@@ -10530,6 +10575,33 @@ std::vector<std::vector<SizeType32>> driveDecode(
 }
 } // namespace
 
+// A request that skips prefix reuse is not offered an externally matched prefix
+// either; a plain request still is.
+TEST_F(KVCacheManagerTest, KvCacheConnector_SkipPrefixReuseIsNotQueried)
+{
+    auto const stream = std::make_shared<tr::CudaStream>();
+    auto connector = std::make_shared<MockKvCacheConnectorManager>(/*numNewMatchedTokens=*/4);
+    auto mgr = makeConnectorTestKVCacheManager(stream, connector);
+
+    auto constexpr beamWidth = 1;
+    auto tokens = std::make_shared<VecTokens>(VecTokens{0, 1, 2, 3, 4, 5, 6, 7, 8});
+    auto const promptLen = static_cast<SizeType32>(tokens->size());
+    tle::SamplingConfig const samplingConfig{beamWidth};
+
+    auto skipping = std::make_shared<LlmRequest>(
+        LlmRequest::RequestIdType{0}, /*maxNewTokens=*/2, tokens, samplingConfig, /*isStreaming=*/false);
+    skipping->setSkipPrefixReuse(true);
+    mgr->addSequenceBatch({{{skipping->mRequestId, promptLen, beamWidth}}}, {std::ref(*skipping)});
+    EXPECT_EQ(connector->getCallCount(), 0);
+    EXPECT_EQ(skipping->getPrepopulatedPromptLen(), 0);
+
+    auto plain = std::make_shared<LlmRequest>(
+        LlmRequest::RequestIdType{1}, /*maxNewTokens=*/2, tokens, samplingConfig, /*isStreaming=*/false);
+    mgr->addSequenceBatch({{{plain->mRequestId, promptLen, beamWidth}}}, {std::ref(*plain)});
+    EXPECT_EQ(connector->getCallCount(), 1);
+    EXPECT_GT(plain->getPrepopulatedPromptLen(), 0);
+}
+
 // Test 1: connector attached but reports zero external matches. Decode-time
 // boundary allocation must still fire; this is the simplest possible
 // invariant a KvCacheConnector must not break.

tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py

@@ -1224,6 +1224,28 @@ def _spec_recompute_claim_limit(
     return max(0, (req.prompt_len - tail) // tokens_per_block * tokens_per_block)
 
 
+def _skip_prefix_reuse_claim_limit(req: LlmRequest) -> int | None:
+    """Reuse-claim cap for a request that skips prefix reuse.
+
+    Such a request, for example one returning context logits (which exist only
+    for tokens that pass through the forward pass), claims no prefix. 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.
+    """
+    return 0 if req.skip_prefix_reuse else 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 +3868,17 @@ 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 that skip prefix reuse claim nothing.
+        limit = _skip_prefix_reuse_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 +4390,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. Requests
+        # that skip prefix reuse cap to zero as well.
+        limit = _tighter_claim_limit(
+            _skip_prefix_reuse_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/llm_request.py

@@ -978,6 +978,10 @@ def __init__(
         # Router Replay: attach routed_experts to this request's output.
         self.py_return_routed_experts = return_routed_experts
         self.py_return_context_logits = return_context_logits
+        # Context logits exist only for tokens the forward pass computes, so
+        # the KV cache manager allocates the whole prompt instead of claiming
+        # a cached prefix.
+        self.skip_prefix_reuse = return_context_logits
         self.py_return_generation_logits = return_generation_logits
         self.py_return_logits_device_memory = return_logits_device_memory
         self.py_additional_outputs = additional_outputs

tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py

@@ -1099,6 +1099,8 @@ def _reuse_claim_limit(self, req: LlmRequest) -> int | None:
         draft_manager = self._joint_draft_manager
         if draft_manager is None:
             return None
+        if req.skip_prefix_reuse:
+            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)