[https://nvbugs/6713231][fix] Guard Whisper cross-KV reuse without encoder tokens

cpp/tensorrt_llm/batch_manager/capacityScheduler.cpp

@@ -37,6 +37,14 @@ using kv_cache_manager::BlockKeyHasher;
 namespace
 {
 
+/// Feature-driven encoders (e.g. Whisper) carry no encoder token ids, so the request has no
+/// encoder unique tokens to key cross-KV blocks on and cross prefix reuse cannot apply to it.
+bool hasEncoderUniqueTokens(LlmRequest const& req)
+{
+    auto const& encoderUniqueTokens = req.getEncoderUniqueTokens();
+    return encoderUniqueTokens.has_value() && encoderUniqueTokens.value() != nullptr;
+}
+
 std::tuple<std::unordered_set<BlockKey, BlockKeyHasher>, std::unordered_set<BlockKey, BlockKeyHasher>>
 prefillWithChunkedContextsAlreadyExecuting(RequestList const& activeRequests,
     kv_cache_manager::BaseKVCacheManager const& kvCacheManager,
@@ -59,7 +67,7 @@ prefillWithChunkedContextsAlreadyExecuting(RequestList const& activeRequests,
                     newlyContributedContextBlocks.insert(summary.firstNewBlock.value());
                 }
             }
-            if (crossKvCacheManager && crossKvCacheManager->isEnableBlockReuse())
+            if (crossKvCacheManager && crossKvCacheManager->isEnableBlockReuse() && hasEncoderUniqueTokens(*req))
             {
                 auto uniqueTokens = *(req->getEncoderUniqueTokens().value());
                 auto summary = crossKvCacheManager->analyzePrefixReuse(uniqueTokens, *req);
@@ -344,12 +352,20 @@ std::tuple<RequestVector, RequestVector> GuaranteedNoEvictScheduler::impl(
                         if (crossKvCacheManager && crossKvCacheManager->isEnableBlockReuse()
                             && !crossKvCacheManager->getBlockManager().isVariableWindow())
                         {
-                            auto uniqueTokens = *(req->getEncoderUniqueTokens().value());
-                            crossSummary = crossKvCacheManager->analyzePrefixReuse(uniqueTokens, *req);
+                            if (hasEncoderUniqueTokens(*req))
+                            {
+                                auto uniqueTokens = *(req->getEncoderUniqueTokens().value());
+                                crossSummary = crossKvCacheManager->analyzePrefixReuse(uniqueTokens, *req);
+                            }
+                            else
+                            {
+                                // Nothing to look up: an empty summary means "no reusable cross blocks".
+                                crossSummary = kv_cache_manager::PrefixReuseSummary{};
+                            }
                         }
                     }
                     else if (isEncoderInit && crossKvCacheManager && crossKvCacheManager->isEnableBlockReuse()
-                        && !crossKvCacheManager->getBlockManager().isVariableWindow())
+                        && !crossKvCacheManager->getBlockManager().isVariableWindow() && hasEncoderUniqueTokens(*req))
                     {
                         // Encoder admission only needs the cross summary for reuse ordering.
                         auto uniqueTokens = *(req->getEncoderUniqueTokens().value());

cpp/tests/unit_tests/batch_manager/capacitySchedulerTest.cpp

@@ -2608,6 +2608,21 @@ std::shared_ptr<LlmRequest> createEncoderInitRequest(
     EXPECT_EQ(req->getState(), LlmRequestState::kENCODER_INIT);
     return req;
 }
+
+std::shared_ptr<LlmRequest> createFeatureEncoderInitRequest(
+    int32_t promptLen, int32_t maxNewTokens, int32_t encoderInputLen, int32_t encoderOutputLen, uint64_t reqId)
+{
+    auto inputTokens = VecTokens(promptLen, 1);
+    auto executorReq = tensorrt_llm::executor::Request(inputTokens, maxNewTokens);
+    auto encoderInputFeatures = tensorrt_llm::executor::Tensor::cpu(
+        tensorrt_llm::executor::DataType::kFP32, {encoderInputLen, /*featureSize=*/1});
+    executorReq.setEncoderInputFeatures(std::move(encoderInputFeatures));
+    executorReq.setEncoderOutputLength(encoderOutputLen);
+    auto req = std::make_shared<LlmRequest>(reqId, executorReq);
+    EXPECT_EQ(req->getState(), LlmRequestState::kENCODER_INIT);
+    EXPECT_FALSE(req->getEncoderUniqueTokens().has_value());
+    return req;
+}
 } // namespace
 
 // GuaranteedNoEvict: a single encoder-init request is admitted without
@@ -2688,3 +2703,45 @@ TEST_F(CapacitySchedulerTest, EncoderInitDoesNotConsumeCrossPool)
         EXPECT_EQ(crossKvCacheManager->getNumFreeBlocks(), crossFreeBefore) << "policy=" << static_cast<int>(policy);
     }
 }
+
+TEST_F(CapacitySchedulerTest, FeatureEncoderWithoutTokensSkipsCrossPoolReuseAnalysis)
+{
+    SizeType32 const maxNumRequests = 1;
+    SizeType32 const tokensPerBlock = 10;
+    SizeType32 const selfMaxTokens = 100;
+    SizeType32 const crossMaxTokens = 20;
+    int32_t const promptLen = 10;
+    int32_t const encoderInputLen = 20;
+    int32_t const encoderOutputLen = 20;
+
+    auto kvCacheManager = getKvCacheManager(maxNumRequests, tokensPerBlock, selfMaxTokens, selfMaxTokens,
+        /*sinkTokenLength=*/0, /*enableReuse=*/true);
+    auto crossKvCacheManager = getKvCacheManager(maxNumRequests, tokensPerBlock, crossMaxTokens, crossMaxTokens,
+        /*sinkTokenLength=*/0, /*enableReuse=*/true, kv_cache_manager::CacheType::kCROSS);
+    auto peftCacheManager = getPeftCacheManager();
+    auto capacityScheduler
+        = CapacityScheduler(maxNumRequests, CapacitySchedulerPolicy::kGUARANTEED_NO_EVICT, kvCacheManager != nullptr,
+            /*twoStepsLookAhead=*/false, LlmRequestState::kENCODER_INIT, LlmRequestState::kGENERATION_COMPLETE);
+    auto req = createFeatureEncoderInitRequest(
+        promptLen, /*maxNewTokens=*/40, encoderInputLen, encoderOutputLen, /*reqId=*/1);
+    RequestList activeRequests{req};
+
+    auto expectRequestScheduled = [&]()
+    {
+        auto [fittingRequests, fittingDisaggGenInitRequests, pausedRequests]
+            = capacityScheduler(activeRequests, kvCacheManager, peftCacheManager, crossKvCacheManager);
+        ASSERT_EQ(fittingRequests.size(), 1u);
+        EXPECT_EQ(fittingRequests.front()->mRequestId, req->mRequestId);
+        EXPECT_TRUE(fittingDisaggGenInitRequests.empty());
+        EXPECT_TRUE(pausedRequests.empty());
+    };
+
+    // Whisper-like feature encoders have no token IDs to key cross-KV reuse.
+    // Exercise encoder admission, first decoder context, and chunked decoder
+    // context, which use the scheduler's three cross-prefix analysis paths.
+    expectRequestScheduled();
+    req->setState(LlmRequestState::kCONTEXT_INIT);
+    expectRequestScheduled();
+    req->setContextCurrentPosition(1);
+    expectRequestScheduled();
+}

tensorrt_llm/_torch/pyexecutor/_util.py

@@ -2227,6 +2227,11 @@ def _split_kv_cache_budget_for_draft(
     def _is_encoder_decoder(self) -> bool:
         return self._model_engine.model.model_config.is_encoder_decoder
 
+    def _encoder_input_is_features(self) -> bool:
+        return bool(
+            getattr(self._model_engine.input_processor,
+                    "requires_encoder_features", False))
+
     @staticmethod
     def _get_config_int_attr(config, names: tuple[str, ...]) -> Optional[int]:
         for name in names:
@@ -2354,6 +2359,21 @@ def _split_kv_cache_budget_for_cross(
 
         self_kv_cache_config = base_kv_cache_config.model_copy()
         cross_kv_cache_config = base_kv_cache_config.model_copy()
+        if (base_kv_cache_config.enable_block_reuse
+                and self._encoder_input_is_features()):
+            # Decoder self-KV is conditioned on the encoder output, while
+            # cross-KV is keyed on encoder token ids. Feature-driven encoders
+            # provide neither reusable token ids nor an input discriminator,
+            # so neither pool can safely reuse blocks between requests.
+            logger.info(
+                "Disabling block reuse for the self- and cross-KV caches: "
+                "the encoder takes feature tensors, so requests carry no "
+                "encoder tokens or input discriminator to key cache entries.")
+            self_kv_cache_config.enable_block_reuse = False
+            cross_kv_cache_config.enable_block_reuse = False
+            # The attention backend reads this shared runtime flag rather than
+            # either derived manager config when selecting paged-context FMHA.
+            self._model_engine.attn_runtime_features.cache_reuse = False
         split_any_budget = False
 
         free_fraction = base_kv_cache_config.free_gpu_memory_fraction

tests/integration/defs/llmapi/test_llm_api_pytorch_whisper.py

@@ -116,6 +116,8 @@ def _make_llm(
     cuda_graph_batch_sizes: list[int] | None = None,
     tensor_parallel_size: int = 1,
     encoder_graphs: bool = False,
+    enable_block_reuse: bool = False,
+    use_python_scheduler: bool = True,
 ) -> LLM:
     """Build a Whisper LLM for the test matrix, optionally with encoder CUDA graphs."""
     # CudaGraphConfig captures the decode step; the enc-dec encoder step opts in
@@ -154,7 +156,7 @@ def _make_llm(
         disable_overlap_scheduler=True,  # overlap scheduler unsupported
         enable_chunked_prefill=False,
         kv_cache_config=KvCacheConfig(
-            enable_block_reuse=False,
+            enable_block_reuse=enable_block_reuse,
             free_gpu_memory_fraction=_FREE_GPU_MEMORY_FRACTION,
             cross_kv_cache_fraction=_CROSS_KV_CACHE_FRACTION,
             use_kv_cache_manager_v2=use_kv_cache_manager_v2,
@@ -165,7 +167,7 @@ def _make_llm(
         # 1500 encoder positions every Whisper request produces.
         max_input_len=_ENCODER_OUTPUT_LEN,
         max_num_tokens=2 * _ENCODER_OUTPUT_LEN,
-        scheduler_config=SchedulerConfig(use_python_scheduler=True),
+        scheduler_config=SchedulerConfig(use_python_scheduler=use_python_scheduler),
         tensor_parallel_size=tensor_parallel_size,
         **encoder_kwargs,
         **dtype_kwargs,
@@ -248,6 +250,37 @@ def test_whisper_pytorch_transcribe_end_to_end(monkeypatch):
 ]
 
 
+def test_whisper_pytorch_block_reuse_requested(monkeypatch):
+    """Greedy transcription when KV block reuse is requested.
+
+    Whisper requests carry encoder features, not encoder token ids, so the
+    executor disables reuse for both KV pools and must still admit and run
+    them (https://nvbugs/6713231). Batch 2 co-schedules two
+    encoder-init requests, the shape that reached the unguarded cross-reuse
+    lookup in the C++ capacity scheduler.
+    """
+    monkeypatch.setenv("TLLM_WORKER_USE_SINGLE_PROCESS", "1")
+
+    model_path = _get_whisper_model_path()
+    wave, sample_rate = soundfile.read(_get_audio_path())
+    sampling_params = SamplingParams(temperature=0.0, max_tokens=_MAX_NEW_TOKENS)
+
+    with _make_llm(
+        model_path,
+        enable_block_reuse=True,
+        use_python_scheduler=False,
+    ) as llm:
+        for batch_size in (1, 2):
+            outputs = llm.generate(
+                [_audio_prompt(wave, sample_rate) for _ in range(batch_size)],
+                sampling_params,
+            )
+            for output in outputs:
+                completion = output.outputs[0]
+                assert list(completion.token_ids) == _EXPECTED_GREEDY_OUTPUT_TOKEN_IDS
+                assert _EXPECTED_TRANSCRIPT_FRAGMENT in completion.text.lower()
+
+
 @pytest.mark.parametrize("torch_dtype,cuda_graph_batch_sizes,graphs_captured", _BEAM_SEARCH_CASES)
 def test_whisper_pytorch_beam_search(
     monkeypatch, torch_dtype, cuda_graph_batch_sizes, graphs_captured

tests/integration/test_lists/test-db/l0_h100.yml

@@ -207,6 +207,7 @@ l0_h100:
   - llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_generate_encoder_decoder_end_to_end[bf16-kv-v1-cuda-graph-on-beam2-t5-small]
   - llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_generate_encoder_decoder_end_to_end[bf16-kv-v1-cuda-graph-on-greedy-overlap-t5-small]
   - llmapi/test_llm_api_pytorch_whisper.py::test_whisper_pytorch_transcribe_end_to_end
+  - llmapi/test_llm_api_pytorch_whisper.py::test_whisper_pytorch_block_reuse_requested
   - llmapi/test_llm_api_pytorch_whisper.py::test_whisper_pytorch_feature_combinations[bf16-kv-v2-decoder-graphs-on-greedy]
   - llmapi/test_llm_api_pytorch_whisper.py::test_whisper_pytorch_beam_search[bf16-kv-v1-decoder-graphs-on-beam2]
   - test_e2e.py::test_openai_chat_harmony

tests/integration/test_lists/test-db/l0_l40s.yml

@@ -53,6 +53,7 @@ l0_l40s:
   # The encoder-graphs case also exercises decoder graphs, so it stands in for
   # a decoder-only case rather than adding to it; KV-v2 stays covered on H100.
   - llmapi/test_llm_api_pytorch_whisper.py::test_whisper_pytorch_transcribe_end_to_end
+  - llmapi/test_llm_api_pytorch_whisper.py::test_whisper_pytorch_block_reuse_requested
   - llmapi/test_llm_api_pytorch_whisper.py::test_whisper_pytorch_feature_combinations[bf16-kv-v1-encoder-graphs-on-greedy]
 - condition:
     ranges:

tests/unittest/_torch/executor/kv_cache/test_dual_pool_kv_cache.py

@@ -50,6 +50,14 @@ class _FakeCudaStream:
     cuda_stream = 0
 
 
+class _FakeCompiledModel:
+    """Minimal torch.compile-style wrapper that changes the model's type."""
+
+    def __init__(self, original_model):
+        self._orig_mod = original_model
+        self.model_config = original_model.model_config
+
+
 class _FakeKVCacheManagerCpp:
     def __init__(self, **kwargs):
         self.kwargs = kwargs
@@ -141,6 +149,7 @@ def _make_mock_model_engine(model_config):
     engine.dtype = torch.bfloat16
     engine.is_draft_model = False
     engine.kv_cache_manager_key = ResourceManagerType.KV_CACHE_MANAGER
+    engine.input_processor = SimpleNamespace(requires_encoder_features=False)
     return engine
 
 
@@ -161,6 +170,9 @@ def _make_creator(
     if model_config is None:
         model_config = _make_model_config(is_encoder_decoder=is_enc_dec)
     model_engine = _make_mock_model_engine(model_config)
+    model_engine.attn_runtime_features = SimpleNamespace(
+        cache_reuse=kv_cache_config.enable_block_reuse
+    )
 
     if manager_cls is None:
         manager_cls = (
@@ -290,6 +302,31 @@ def test_split_free_fraction_when_budget_is_zero(self):
         assert self_config.free_gpu_memory_fraction == pytest.approx(0.4)
         assert config.free_gpu_memory_fraction == pytest.approx(0.8)
 
+    def test_wrapped_feature_encoder_disables_reuse_for_both_pools(self):
+        """Feature detection survives a torch.compile-style model wrapper."""
+        config = _make_kv_cache_config(cross_kv_cache_fraction=0.5)
+        creator = _make_creator(config, is_enc_dec=True)
+        creator._model_engine.model = _FakeCompiledModel(creator._model_engine.model)
+        creator._model_engine.input_processor.requires_encoder_features = True
+
+        self_config, cross_config = creator._split_kv_cache_budget_for_cross()
+
+        assert config.enable_block_reuse
+        assert not self_config.enable_block_reuse
+        assert not cross_config.enable_block_reuse
+        assert not creator._model_engine.attn_runtime_features.cache_reuse
+
+    def test_token_encoder_preserves_reuse_for_both_pools(self) -> None:
+        """Token inputs retain reusable identities for both KV pools."""
+        config = _make_kv_cache_config(cross_kv_cache_fraction=0.5)
+        creator = _make_creator(config, is_enc_dec=True)
+
+        self_config, cross_config = creator._split_kv_cache_budget_for_cross()
+
+        assert self_config.enable_block_reuse
+        assert cross_config.enable_block_reuse
+        assert creator._model_engine.attn_runtime_features.cache_reuse
+
     def test_is_encoder_decoder_helper(self):
         dec_config = _make_model_config(is_encoder_decoder=False)
         dec_creator = _make_creator(_make_kv_cache_config(), model_config=dec_config)
@@ -684,6 +721,32 @@ def create_cross_manager(cross_cfg, *_args, **_kwargs):
             (pytest.approx(0.45), expected_split),
         ]
 
+    @pytest.mark.parametrize("use_kv_cache_manager_v2", [False, True])
+    def test_build_managers_disables_feature_encoder_reuse(
+        self, use_kv_cache_manager_v2: bool
+    ) -> None:
+        """Pass reuse-disabled configs to both feature-encoder managers."""
+        config = _make_kv_cache_config(
+            cross_kv_cache_fraction=0.5,
+            max_gpu_total_bytes=8 * (1 << 30),
+            use_kv_cache_manager_v2=use_kv_cache_manager_v2,
+        )
+        creator = _make_creator(config, is_enc_dec=True)
+        creator.configure_kv_cache_capacity = Mock()
+        creator._model_engine.input_processor.requires_encoder_features = True
+        creator._should_create_separate_draft_kv_cache = Mock(return_value=False)
+        creator._create_kv_cache_manager = Mock(return_value=Mock())
+        creator._create_cross_kv_cache_manager = Mock(return_value=Mock())
+
+        creator.build_managers({}, estimating_kv_cache=False)
+
+        self_config = creator._create_kv_cache_manager.call_args.kwargs["kv_cache_config_override"]
+        cross_config = creator._create_cross_kv_cache_manager.call_args.args[0]
+        assert config.enable_block_reuse
+        assert not self_config.enable_block_reuse
+        assert not cross_config.enable_block_reuse
+        assert not creator._model_engine.attn_runtime_features.cache_reuse
+
     def test_build_managers_skips_cross_pool_for_decoder_only(self):
         creator = _make_creator(
             _make_kv_cache_config(

tests/unittest/_torch/modeling/test_modeling_whisper.py

@@ -25,7 +25,11 @@
 import torch
 from transformers import WhisperConfig, WhisperFeatureExtractor
 
-from tensorrt_llm._torch.models.modeling_whisper import WhisperLogMelFrontend
+from tensorrt_llm._torch.models.modeling_whisper import WhisperInputProcessor, WhisperLogMelFrontend
+
+
+def test_whisper_input_processor_requires_encoder_features():
+    assert WhisperInputProcessor.requires_encoder_features
 
 
 def _synthetic_waveform_batch(n_samples: int, seed: int = 1234) -> np.ndarray: