From 3b346ceb00387b9678231827547787d136595e23 Mon Sep 17 00:00:00 2001 From: Simeng Liu Date: Mon, 27 Jul 2026 12:56:11 -0700 Subject: [PATCH 1/4] [https://nvbugs/6441022][perf] Enable CUDA graphs for final single-token contexts Promote eligible final single-token context rows through a temporary decode-shaped execution view after KV preparation. Commit the view only when the existing graph runner finds a matching graph, and retain semantic eager fallback otherwise. Preserve request state, request type, KV ownership, sampling order, and lifecycle updates. Cover context logits, guided decoding, zero-runtime-draft target execution, overlap token sourcing, sparse graph keys, local offload, changed-tail reuse, and TP2 rank-local replay. Tests added: cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp: - KVCacheManagerTest.AddSequenceBatchLeavesOneFinalContextTokenAfterReuse - KVCacheManagerTest.AddSequenceBatchPreservesDraftTokensOnFinalContextAfterReuse - KVCacheManagerTest.AddSequenceBatchPreservesGuidanceAndContextLogitsAfterReuse - KVCacheManagerTest.AddSequenceBatchOnboardsOffloadedPrefixForFinalContextToken - KVCacheManagerTest.AddSequenceBatchLeavesOneFinalMultimodalContextTokenAfterReuse tests/integration/defs/kv_cache/test_final_single_token_context_cuda_graph.py: - test_final_token_reuse_cuda_graph - test_changed_final_token_reuse_cuda_graph - test_final_token_reuse_cuda_graph_tp2 - test_context_logits_after_final_token_reuse - test_guided_decoding_after_final_token_reuse - test_zero_runtime_draft_speculation_after_final_token_reuse tests/unittest/_torch/executor/test_pytorch_model_engine.py: - SingleTokenContextGraphBatchTestCase.test_generation_only_is_identity - SingleTokenContextGraphBatchTestCase.test_eligible_batch_has_independent_lists_and_stable_order - SingleTokenContextGraphBatchTestCase.test_structural_fallbacks_return_semantic_batch - SingleTokenContextGraphBatchTestCase.test_context_shape_and_mode_fallback_matrix - SingleTokenContextGraphBatchTestCase.test_context_logits_use_final_token_graph_candidate - SingleTokenContextGraphBatchTestCase.test_generation_only_request_in_context_list_falls_back - SingleTokenContextGraphBatchTestCase.test_generation_shape_fallback_matrix - SingleTokenContextGraphBatchTestCase.test_mixed_one_and_two_token_contexts_fall_back_together - SingleTokenContextGraphBatchTestCase.test_mrope_delta_is_supported_by_decode_provider - SingleTokenContextGraphBatchTestCase.test_multimodal_context_requires_compatible_decode_token - SingleTokenContextGraphBatchTestCase.test_multimodal_pending_event_is_rechecked - SingleTokenContextGraphBatchTestCase.test_multimodal_decode_compatibility_uses_final_prompt_token - SingleTokenContextGraphBatchTestCase.test_sparse_sequence_mode_uses_promoted_context_cursor - SingleTokenContextGraphBatchTestCase.test_graph_key_forwards_promoted_context_ids - SingleTokenContextGraphBatchTestCase.test_graph_lookup_forwards_promoted_context_ids - SingleTokenContextGraphBatchTestCase.test_forward_commits_candidate_only_on_graph_hit - SingleTokenContextGraphBatchTestCase.test_forward_graph_miss_uses_semantic_eager_batch - SingleTokenContextGraphBatchTestCase.test_zero_runtime_draft_speculation_commits_graph_candidate - SingleTokenContextGraphBatchTestCase.test_zero_runtime_draft_speculation_graph_miss_is_semantic_eager - SingleTokenContextGraphBatchTestCase.test_forward_allows_guided_context_logits_on_graph_hit - SingleTokenContextGraphBatchTestCase.test_multimodal_graph_miss_preserves_semantic_payload - SingleTokenContextGraphBatchTestCase.test_generation_only_forward_does_not_call_new_selector - SingleTokenContextGraphBatchTestCase.test_global_incompatibilities_bypass_candidate_selection - PyTorchModelEngineTestCase.test_promoted_context_uses_prompt_token_during_overlap - PyTorchModelEngineTestCase.test_promoted_context_precedes_speculative_overlap_generation - PyTorchModelEngineTestCase.test_promoted_mrope_context_uses_decode_state_contract - PyTorchModelEngineTestCase.test_kv_cache_manager_with_execution_stream - PyTorchModelEngineTestCase.test_cuda_graph_replay_observes_execution_stream_dependency tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_behavior.py: - test_v2_resume_restores_offsets_only_after_execution_stream_ready Signed-off-by: Simeng Liu --- .../batch_manager/kvCacheManagerTest.cpp | 259 ++++- .../_torch/pyexecutor/cuda_graph_runner.py | 53 +- .../_torch/pyexecutor/model_engine.py | 235 ++++- ...t_final_single_token_context_cuda_graph.py | 520 ++++++++++ .../test_lists/test-db/l0_dgx_h100.yml | 1 + .../test_lists/test-db/l0_h100.yml | 5 + .../executor/test_pytorch_model_engine.py | 963 +++++++++++++++++- .../test_kv_cache_stats_behavior.py | 39 + 8 files changed, 2000 insertions(+), 75 deletions(-) create mode 100644 tests/integration/defs/kv_cache/test_final_single_token_context_cuda_graph.py diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp index 7bacaf9a2575..177424477b8d 100644 --- a/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp @@ -8386,16 +8386,17 @@ TEST_F(KVCacheManagerTest, StaticLinearHybridAllocationTest) // ownership tracking. /////////////////////////////////////////////////////////////////////////////// -// Helper: create a KVCacheManager for batch tests. -// tokensPerBlock=4, 16 primary blocks, block reuse enabled, partial reuse enabled. -static auto makeBatchTestKVCacheManager(std::shared_ptr const& stream) +// Helper: create a KVCacheManager for batch tests. The default pool geometry is +// tokensPerBlock=4, 16 primary blocks, no secondary blocks. Tests that certify +// local offload can opt into secondary blocks without duplicating the manager +// construction used by the ownership tests below. +static auto makeBatchTestKVCacheManager(std::shared_ptr const& stream, + SizeType32 blocksInPrimaryPool = 16, SizeType32 blocksInSecondaryPool = 0) { auto constexpr numLayers = 1; auto constexpr numKvHeads = 1; auto constexpr sizePerHead = 16; auto constexpr tokensPerBlock = 4; - auto constexpr blocksInPrimaryPool = 16; - auto constexpr blocksInSecondaryPool = 0; auto constexpr maxNumSequences = 16; auto constexpr beamWidth = 1; auto constexpr maxAttentionWindow = tokensPerBlock * 8; @@ -8426,6 +8427,254 @@ static void seedAndRelease(KVCacheManager& mgr, LlmRequest::RequestIdType reqId, (void) mgr.removeSequence(reqId, req); } +TEST_F(KVCacheManagerTest, AddSequenceBatchLeavesOneFinalContextTokenAfterReuse) +{ + auto const stream = std::make_shared(); + auto mgr = makeBatchTestKVCacheManager(stream); + auto constexpr beamWidth = 1; + auto constexpr promptLen = 9; + auto constexpr seedRequestId = 0; + auto constexpr requestId = 1; + auto const inputTokens = std::make_shared(VecTokens{0, 1, 2, 3, 4, 5, 6, 7, 8}); + + seedAndRelease(*mgr, seedRequestId, inputTokens); + + auto req = std::make_shared( + LlmRequest::RequestIdType{requestId}, /*maxNewTokens=*/2, inputTokens, tr::SamplingConfig{beamWidth}, false); + auto const initialState = req->getState(); + auto const requestType = req->getLlmRequestType(); + + mgr->addSequenceBatch({{{requestId, promptLen, beamWidth}}}, {std::ref(*req)}); + + EXPECT_EQ(req->getPrepopulatedPromptLen(), promptLen - 1); + EXPECT_EQ(req->getContextCurrentPosition(), promptLen - 1); + EXPECT_EQ(req->getContextRemainingLength(), 1); + EXPECT_EQ(req->getContextChunkSize(), 1); + EXPECT_TRUE(req->isLastContextChunk()); + EXPECT_EQ(req->getState(), initialState); + EXPECT_EQ(req->getLlmRequestType(), requestType); + + tensorrt_llm::testing::KvCacheManagerTestUtil::simulatePrefillCompletion(*req); + EXPECT_NO_THROW(static_cast(mgr->removeSequence(requestId, req))); +} + +TEST_F(KVCacheManagerTest, AddSequenceBatchPreservesDraftTokensOnFinalContextAfterReuse) +{ + auto const stream = std::make_shared(); + auto mgr = makeBatchTestKVCacheManager(stream); + auto constexpr beamWidth = 1; + auto constexpr promptLen = 9; + auto constexpr seedRequestId = 0; + auto constexpr requestId = 1; + auto const inputTokens = std::make_shared(VecTokens{0, 1, 2, 3, 4, 5, 6, 7, 8}); + + seedAndRelease(*mgr, seedRequestId, inputTokens); + + auto req = std::make_shared( + LlmRequest::RequestIdType{requestId}, /*maxNewTokens=*/2, inputTokens, tr::SamplingConfig{beamWidth}, false); + auto const draftTokens = std::make_shared(VecTokens{42}); + req->setDraftTokens(draftTokens); + auto const initialState = req->getState(); + auto const requestType = req->getLlmRequestType(); + + mgr->addSequenceBatch({{{requestId, promptLen, beamWidth}}}, {std::ref(*req)}); + + EXPECT_EQ(req->getPrepopulatedPromptLen(), promptLen - 1); + EXPECT_EQ(req->getContextCurrentPosition(), promptLen - 1); + EXPECT_EQ(req->getContextRemainingLength(), 1); + EXPECT_EQ(req->getContextChunkSize(), 1); + EXPECT_EQ(req->getNumDraftTokens(), 1); + EXPECT_EQ(req->getState(), initialState); + EXPECT_EQ(req->getLlmRequestType(), requestType); + + tensorrt_llm::testing::KvCacheManagerTestUtil::simulatePrefillCompletion(*req); + EXPECT_NO_THROW(static_cast(mgr->removeSequence(requestId, req))); +} + +TEST_F(KVCacheManagerTest, AddSequenceBatchPreservesGuidanceAndContextLogitsAfterReuse) +{ + auto const stream = std::make_shared(); + auto mgr = makeBatchTestKVCacheManager(stream); + auto constexpr promptLen = 9; + auto constexpr reusableLen = promptLen - 1; + auto constexpr beamWidth = 1; + auto const inputTokens = std::make_shared(VecTokens{0, 1, 2, 3, 4, 5, 6, 7, 8}); + tr::SamplingConfig const samplingConfig{beamWidth}; + + auto seedReq = std::make_shared( + LlmRequest::RequestIdType{0}, SizeType32{1}, inputTokens, samplingConfig, /*isStreaming=*/false); + mgr->addSequenceBatch({{{0, promptLen, beamWidth}}}, {std::ref(*seedReq)}); + tensorrt_llm::testing::KvCacheManagerTestUtil::simulatePrefillCompletion(*seedReq); + (void) mgr->removeSequence(0, seedReq); + + tle::OutputConfig outputConfig; + outputConfig.returnContextLogits = true; + auto const guidedParams = tle::GuidedDecodingParams(tle::GuidedDecodingParams::GuideType::kREGEX, R"([0-9]+)"); + tle::Request executorRequest( + *inputTokens, /*maxTokens=*/1, /*streaming=*/false, tle::SamplingConfig{}, outputConfig); + executorRequest.setGuidedDecodingParams(guidedParams); + auto req = std::make_shared(LlmRequest::RequestIdType{1}, executorRequest); + auto const semanticState = req->getState(); + auto const semanticType = req->getLlmRequestType(); + + mgr->addSequenceBatch({{{1, promptLen, beamWidth}}}, {std::ref(*req)}); + + EXPECT_EQ(req->getContextCurrentPosition(), reusableLen); + EXPECT_EQ(req->getContextRemainingLength(), 1); + EXPECT_EQ(req->getContextChunkSize(), 1); + EXPECT_TRUE(req->getReturnContextLogits()); + ASSERT_TRUE(req->getGuidedDecodingParams().has_value()); + EXPECT_EQ(req->getGuidedDecodingParams().value(), guidedParams); + EXPECT_EQ(req->getState(), semanticState); + EXPECT_EQ(req->getLlmRequestType(), semanticType); + + tensorrt_llm::testing::KvCacheManagerTestUtil::simulatePrefillCompletion(*req); + EXPECT_NO_THROW(static_cast(mgr->removeSequence(1, req))); +} + +TEST_F(KVCacheManagerTest, AddSequenceBatchOnboardsOffloadedPrefixForFinalContextToken) +{ + auto const stream = std::make_shared(); + auto mgr = makeBatchTestKVCacheManager(stream, /*blocksInPrimaryPool=*/16, /*blocksInSecondaryPool=*/4); + auto constexpr beamWidth = 1; + auto constexpr promptLen = 9; + auto constexpr reusableLen = promptLen - 1; + auto const inputTokens = std::make_shared(VecTokens{0, 1, 2, 3, 4, 5, 6, 7, 8}); + tr::SamplingConfig const samplingConfig{beamWidth}; + + // Seed two reusable full blocks and remember them before the sequence is + // released. Moving both blocks to secondary memory forces the next batch + // through the real local-onboard path instead of GPU-only radix reuse. + auto seedReq = std::make_shared( + LlmRequest::RequestIdType{0}, SizeType32{0}, inputTokens, samplingConfig, /*isStreaming=*/false); + mgr->addSequenceBatch({{{0, promptLen, beamWidth}}}, {std::ref(*seedReq)}); + auto const windowSize = theOnlyWindowSize(*mgr); + auto const seedBlockIds = mgr->getSequence(0).getCacheBlockIds(windowSize).at(0); + ASSERT_EQ(seedBlockIds.size(), 3); + tensorrt_llm::testing::KvCacheManagerTestUtil::simulatePrefillCompletion(*seedReq); + (void) mgr->removeSequence(0, seedReq); + + // KVCacheManager intentionally exposes BlockManager as read-only. This + // test needs to force a specific reusable block into the secondary tier, + // so keep the mutation local to the fixture instead of widening the + // production interface solely for test setup. + auto& blockManager = const_cast(mgr->getBlockManager()); + for (auto const blockId : {seedBlockIds[0], seedBlockIds[1]}) + { + auto block = blockManager.getBlockById(blockId, windowSize); + ASSERT_TRUE(block->isPrimary()); + blockManager.offloadBlock(block, windowSize); + EXPECT_FALSE(block->isPrimary()); + } + stream->synchronize(); + + // Two requests claiming the same host-resident prefix in one IFB batch + // must both stop at the final prompt token. They may share the immutable + // prefix, but each needs private writable capacity for that final token. + auto req1 = std::make_shared( + LlmRequest::RequestIdType{1}, SizeType32{1}, inputTokens, samplingConfig, /*isStreaming=*/false); + auto req2 = std::make_shared( + LlmRequest::RequestIdType{2}, SizeType32{1}, inputTokens, samplingConfig, /*isStreaming=*/false); + auto const req1State = req1->getState(); + auto const req2State = req2->getState(); + auto const req1Type = req1->getLlmRequestType(); + auto const req2Type = req2->getLlmRequestType(); + + mgr->addSequenceBatch({{{1, promptLen, beamWidth}, {2, promptLen, beamWidth}}}, {std::ref(*req1), std::ref(*req2)}); + // refreshBlocks joins all onboard/copy work to the execution stream. A + // model forward or CUDA graph replay enqueued after this point can consume + // the restored prefix without a host-side synchronization. + mgr->refreshBlocks(); + stream->synchronize(); + + for (auto const& req : {req1, req2}) + { + EXPECT_EQ(req->getPrepopulatedPromptLen(), reusableLen); + EXPECT_EQ(req->getContextCurrentPosition(), reusableLen); + EXPECT_EQ(req->getContextRemainingLength(), 1); + EXPECT_EQ(req->getContextChunkSize(), 1); + EXPECT_TRUE(req->isLastContextChunk()); + } + EXPECT_EQ(req1->getState(), req1State); + EXPECT_EQ(req2->getState(), req2State); + EXPECT_EQ(req1->getLlmRequestType(), req1Type); + EXPECT_EQ(req2->getLlmRequestType(), req2Type); + + auto const req1BlockIds = mgr->getSequence(1).getCacheBlockIds(windowSize).at(0); + auto const req2BlockIds = mgr->getSequence(2).getCacheBlockIds(windowSize).at(0); + ASSERT_EQ(req1BlockIds.size(), 3); + ASSERT_EQ(req2BlockIds.size(), 3); + EXPECT_EQ(req1BlockIds[0], req2BlockIds[0]); + EXPECT_EQ(req1BlockIds[1], req2BlockIds[1]); + EXPECT_NE(req1BlockIds[2], req2BlockIds[2]); + + // Exercise the failure path for one claimant, then complete the other. + // A subsequent request must still reuse the prefix, proving cancellation + // did not leave stale ownership or invalidate the source blocks. + EXPECT_NO_THROW(static_cast(mgr->removeSequence(1, std::nullopt))); + tensorrt_llm::testing::KvCacheManagerTestUtil::simulatePrefillCompletion(*req2); + EXPECT_NO_THROW(static_cast(mgr->removeSequence(2, req2))); + + auto req3 = std::make_shared( + LlmRequest::RequestIdType{3}, SizeType32{1}, inputTokens, samplingConfig, /*isStreaming=*/false); + mgr->addSequenceBatch({{{3, promptLen, beamWidth}}}, {std::ref(*req3)}); + EXPECT_EQ(req3->getContextCurrentPosition(), reusableLen); + EXPECT_EQ(req3->getContextRemainingLength(), 1); + tensorrt_llm::testing::KvCacheManagerTestUtil::simulatePrefillCompletion(*req3); + EXPECT_NO_THROW(static_cast(mgr->removeSequence(3, req3))); + EXPECT_TRUE(blockManager.verifyQueueIntegrity(windowSize)); +} + +TEST_F(KVCacheManagerTest, AddSequenceBatchLeavesOneFinalMultimodalContextTokenAfterReuse) +{ + auto const stream = std::make_shared(); + auto mgr = makeBatchTestKVCacheManager(stream); + auto constexpr beamWidth = 1; + auto constexpr promptLen = 9; + auto constexpr seedRequestId = 0; + auto constexpr requestId = 1; + auto constexpr mropePositionDelta = 7; + auto const inputTokens = std::make_shared(VecTokens{0, 1, 2, 3, 4, 5, 6, 7, 8}); + auto const multimodalHashes = std::make_shared>>( + std::vector>{{1, 2, 3, 4, 5, 6, 7, 8}}); + auto const multimodalPositions = std::make_shared>(std::vector{1}); + auto const multimodalLengths = std::make_shared>(std::vector{4}); + tr::SamplingConfig const samplingConfig{beamWidth}; + auto const makeRequest = [&](LlmRequest::RequestIdType reqId, SizeType32 maxNewTokens) + { + return std::make_shared(reqId, maxNewTokens, inputTokens, samplingConfig, /*isStreaming=*/false, + std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt, + std::nullopt, multimodalHashes, multimodalPositions, multimodalLengths, std::nullopt, std::nullopt, + std::nullopt, mropePositionDelta); + }; + + auto seedReq = makeRequest(seedRequestId, /*maxNewTokens=*/0); + mgr->addSequenceBatch({{{seedRequestId, promptLen, beamWidth}}}, {std::ref(*seedReq)}); + tensorrt_llm::testing::KvCacheManagerTestUtil::simulatePrefillCompletion(*seedReq); + (void) mgr->removeSequence(seedRequestId, seedReq); + + auto req = makeRequest(requestId, /*maxNewTokens=*/2); + auto const initialState = req->getState(); + auto const requestType = req->getLlmRequestType(); + + mgr->addSequenceBatch({{{requestId, promptLen, beamWidth}}}, {std::ref(*req)}); + + EXPECT_EQ(req->getPrepopulatedPromptLen(), promptLen - 1); + EXPECT_EQ(req->getContextCurrentPosition(), promptLen - 1); + EXPECT_EQ(req->getContextRemainingLength(), 1); + EXPECT_EQ(req->getContextChunkSize(), 1); + EXPECT_TRUE(req->isLastContextChunk()); + EXPECT_EQ(req->getState(), initialState); + EXPECT_EQ(req->getLlmRequestType(), requestType); + ASSERT_TRUE(req->getMropePositionDeltas().has_value()); + EXPECT_EQ(req->getMropePositionDeltas().value(), mropePositionDelta); + ASSERT_TRUE(req->getMultimodalHashes().has_value()); + EXPECT_EQ(req->getMultimodalHashes().value(), multimodalHashes); + + tensorrt_llm::testing::KvCacheManagerTestUtil::simulatePrefillCompletion(*req); + EXPECT_NO_THROW(static_cast(mgr->removeSequence(requestId, req))); +} + // Test 1: Two requests in a batch, both partially match the same leaf block. // The tracker should assign reuse to the last request and bump the first to copy. TEST_F(KVCacheManagerTest, BatchAddSequence_LeafPartialThenPartial) diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py index c2a6a5bb4de1..8e0b155f1211 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py @@ -180,9 +180,17 @@ def _create_shared_static_tensors(self): (max_total_tokens, ), device="cuda", dtype=torch.long) def _get_seq_len_mode( - self, - batch: ScheduledRequests, - new_tensors_device: Optional[SampleStateTensors] = None): + self, + batch: ScheduledRequests, + new_tensors_device: Optional[SampleStateTensors] = None, + promoted_context_request_ids: frozenset[int] = frozenset() + ) -> bool: + """Select the sparse-attention graph family for the execution view. + + ``promoted_context_request_ids`` contains semantic final-context rows + that the model engine temporarily placed in the generation list. It is + empty for the existing generation-only path. + """ if (isinstance(self.sparse_config, SeqLenAwareSparseAttentionConfig) and self.sparse_config.needs_separate_short_long_cuda_graphs()): # Some sparse attention algorithms need to use different forward paths for short and long sequences. @@ -202,8 +210,16 @@ def _get_seq_len_mode( is_spec_request = get_draft_token_length( request) > 0 or next_draft_tokens_device is not None num_draft_tokens = self.spec_config.max_draft_len if is_spec_request else 0 + if request.py_request_id in promoted_context_request_ids: + # A promoted context row may retain overlap bookkeeping + # such as py_batch_idx from an earlier context chunk. That + # state describes the previous batch, not the sequence + # length of the final prompt token executed by this graph. + # Use the authoritative context cursor so graph keying + # matches the decode-shaped input prepared for this row. + total_seq_len = request.context_current_position + 1 # First draft - if request.py_is_first_draft: + elif request.py_is_first_draft: # get_num_tokens is O(1); len(get_tokens(0)) marshals the # whole O(seq_len) VecTokens into a Python list just for len. total_seq_len = request.get_num_tokens(0) @@ -229,15 +245,20 @@ def _get_seq_len_mode( return short_seq_len_mode def get_graph_key( - self, - batch: ScheduledRequests, - new_tensors_device: Optional[SampleStateTensors] = None, - spec_resource_manager: Optional[BaseResourceManager] = None, - spec_metadata: Optional[SpecMetadata] = None): + self, + batch: ScheduledRequests, + new_tensors_device: Optional[SampleStateTensors] = None, + spec_resource_manager: Optional[BaseResourceManager] = None, + spec_metadata: Optional[SpecMetadata] = None, + promoted_context_request_ids: frozenset[int] = frozenset() + ) -> KeyType: batch_size = batch.batch_size # Get the sequence length mode. - short_seq_len_mode = self._get_seq_len_mode(batch, new_tensors_device) + # Keep the graph-key tuple unchanged; promoted IDs only correct the + # sequence length observed by sparse short/long graph selection. + short_seq_len_mode = self._get_seq_len_mode( + batch, new_tensors_device, promoted_context_request_ids) # Spec one-engine sampler has two code paths (argmax fast-path vs # advanced sampling kernel). Include this in the key so we capture @@ -305,7 +326,8 @@ def maybe_get_cuda_graph( draft_tokens_cuda: Optional[torch.Tensor] = None, new_tensors_device: Optional[SampleStateTensors] = None, spec_resource_manager: Optional[BaseResourceManager] = None, - ) -> Tuple[Optional[Any], Optional[Any], Optional[Tuple[int, int, bool]]]: + promoted_context_request_ids: frozenset[int] = frozenset(), + ) -> Tuple[Optional[Any], Optional[Any], Optional[KeyType]]: """ Determines if the current batch can be run with a CUDA graph. @@ -313,6 +335,10 @@ def maybe_get_cuda_graph( - The attn_metadata for the graph, if applicable. - The spec_metadata for the graph, if applicable. - The key for the graph, if applicable. + + ``promoted_context_request_ids`` is execution-view metadata. It does + not change request state or type and is used only to build a graph key + consistent with the final-context token that will be executed. """ # disable when doing statistic if ExpertStatistic.should_record(): @@ -342,8 +368,11 @@ def maybe_get_cuda_graph( # do carry a delta must first populate the model-side cache for their # current seq slot before graph replay. return None, None, None + # Propagate the execution-view identity through graph lookup. Existing + # callers pass the empty default and retain generation-only behavior. key = self.get_graph_key(batch, new_tensors_device, - spec_resource_manager, spec_metadata) + spec_resource_manager, spec_metadata, + promoted_context_request_ids) if key in self.graph_metadata: return self.graph_metadata[key][ diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 5afc55544c8c..76b72881771b 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -92,6 +92,54 @@ from .trace_log_utils import log_mem_snapshot +def _make_single_token_context_graph_batch( + scheduled_requests: ScheduledRequests, + is_multimodal_decode_compatible: Optional[Callable[[LlmRequest], + bool]] = None, +) -> tuple[ScheduledRequests, frozenset[int]]: + """Build a decode-shaped graph candidate for final one-token contexts. + + Multimodal rows remain fail-closed unless the engine proves that their one + remaining prompt token is representable by the existing decode provider. + """ + if scheduled_requests.num_context_requests == 0: + return scheduled_requests, frozenset() + + context_requests = scheduled_requests.context_requests_last_chunk + if (scheduled_requests.encoder_requests + or scheduled_requests.context_requests_chunking): + return scheduled_requests, frozenset() + + for request in context_requests: + if (request.context_chunk_size != 1 + or request.context_remaining_length != 1 + or request.context_current_position + 1 != request.py_prompt_len + or request.py_beam_width != 1 + or get_draft_token_length(request) > 0 + or request.py_is_first_draft or request.is_context_only_request + or request.is_generation_only_request() + or request.py_disaggregated_params is not None + or request.py_mm_encoder_event is not None + or (request.py_multimodal_data is not None and + (is_multimodal_decode_compatible is None + or not is_multimodal_decode_compatible(request)))): + return scheduled_requests, frozenset() + + for request in scheduled_requests.generation_requests: + if (request.py_beam_width != 1 or get_draft_token_length(request) > 0 + or request.py_is_first_draft + or request.py_disaggregated_params is not None): + return scheduled_requests, frozenset() + + graph_batch = ScheduledRequests() + graph_batch.generation_requests = list(context_requests) + list( + scheduled_requests.generation_requests) + graph_batch.paused_requests = list(scheduled_requests.paused_requests) + promoted_context_request_ids = frozenset(request.py_request_id + for request in context_requests) + return graph_batch, promoted_context_request_ids + + class ModelEngine(ABC): @abstractmethod @@ -3065,6 +3113,29 @@ def _prepare_multimodal_indices(self, input_ids: list[int]): input_ids, vocab_size=vocab_size, mm_token_ids=mm_token_ids) return text_token_indices, mm_token_indices + def _is_final_multimodal_context_decode_compatible( + self, request: LlmRequest) -> bool: + """Return whether the final prompt token uses the decode input path. + + KV reuse has already materialized every preceding prompt token. A + multimodal final-context row therefore needs its prepared embedding + only when the one remaining token is itself an MM placeholder. Text + tokens can use the existing decode provider; MRoPE deltas are seeded + into the per-sequence cache before graph lookup. An MRoPE request with + real MM payload remains eager until its delta is available. + """ + final_prompt_token = request.get_tokens(0)[ + request.context_current_position] + _, mm_token_indices = self._prepare_multimodal_indices( + [final_prompt_token]) + if mm_token_indices.numel() != 0: + return False + + multimodal_data = request.py_multimodal_data + if not self.use_mrope or not _has_mm_payload_keys(multimodal_data): + return True + return CUDAGraphRunner._get_mrope_position_delta(request) is not None + def _is_encoder_decoder_model(self) -> bool: return bool( getattr(getattr(self.model, "model_config", None), @@ -3861,17 +3932,19 @@ def _apply_steady_gen_fast_prepare( return inputs, None def _prepare_tp_inputs( - self, - scheduled_requests: ScheduledRequests, - kv_cache_manager: Union[KVCacheManager, KVCacheManagerV2], - attn_metadata: AttentionMetadata, - spec_metadata: Optional[SpecMetadata] = None, - new_tensors_device: Optional[SampleStateTensors] = None, - cache_indirection_buffer: Optional[torch.Tensor] = None, - num_accepted_tokens_device: Optional[torch.Tensor] = None, - req_id_to_old_request: Optional[Dict[int, LlmRequest]] = None, - resource_manager: Optional[ResourceManager] = None, - maybe_graph: bool = False): + self, + scheduled_requests: ScheduledRequests, + kv_cache_manager: Union[KVCacheManager, KVCacheManagerV2], + attn_metadata: AttentionMetadata, + spec_metadata: Optional[SpecMetadata] = None, + new_tensors_device: Optional[SampleStateTensors] = None, + cache_indirection_buffer: Optional[torch.Tensor] = None, + num_accepted_tokens_device: Optional[torch.Tensor] = None, + req_id_to_old_request: Optional[Dict[int, LlmRequest]] = None, + resource_manager: Optional[ResourceManager] = None, + maybe_graph: bool = False, + promoted_context_request_ids: frozenset[int] = frozenset() + ) -> Tuple[Dict[str, Any], Optional[torch.Tensor]]: """ Prepare inputs for Pytorch Model. """ @@ -3893,9 +3966,10 @@ def _prepare_tp_inputs( new_tokens=new_tokens_device, runtime_draft_len=self.runtime_draft_len) - if self._can_use_incremental_update(scheduled_requests, - new_tokens_device, - next_draft_tokens_device): + if (not promoted_context_request_ids + and self._can_use_incremental_update(scheduled_requests, + new_tokens_device, + next_draft_tokens_device)): # Spec engines never record the steady-gen cache, but invalidate # defensively so the two fast paths can never interleave if the # gates ever evolve. @@ -3906,10 +3980,10 @@ def _prepare_tp_inputs( num_accepted_tokens_device, req_id_to_old_request, resource_manager) - if self._can_use_steady_gen_fast_prepare(scheduled_requests, - new_tokens_device, - next_draft_tokens_device, - spec_metadata): + if (not promoted_context_request_ids + and self._can_use_steady_gen_fast_prepare( + scheduled_requests, new_tokens_device, + next_draft_tokens_device, spec_metadata)): return self._apply_steady_gen_fast_prepare(kv_cache_manager, attn_metadata, new_tensors_device, @@ -4176,9 +4250,23 @@ def append_cross_attention_state(request: LlmRequest, # a separate iteration over scheduled_requests.generation_requests later. all_gen_request_ids = [] for request in scheduled_requests.generation_requests: - all_gen_request_ids.append(request.py_request_id) - if get_draft_token_length( - request) > 0 or next_draft_tokens_device is not None: + is_promoted_context = (request.py_request_id + in promoted_context_request_ids) + if not is_promoted_context: + all_gen_request_ids.append(request.py_request_id) + # In speculative iterations, keep promoted rows ahead of existing + # generation rows in the extend-request packing order. Although + # their q_len is one, this category provides the "no previous + # speculative tensor" branch needed to source their prompt token + # without disturbing the overlap offsets of ordinary generation + # siblings. Non-speculative promoted rows retain the established + # ordinary generation path below. + if is_promoted_context and self.enable_spec_decode: + extend_requests.append(request) + elif is_promoted_context: + generation_requests.append(request) + elif (get_draft_token_length(request) > 0 + or next_draft_tokens_device is not None): if request.is_dummy: extend_dummy_requests.append(request) else: @@ -4206,6 +4294,8 @@ def append_cross_attention_state(request: LlmRequest, self.runtime_draft_len) runtime_draft_token_buffer_width = runtime_tokens_per_gen_step - 1 for request in extend_requests: + is_promoted_context = (request.py_request_id + in promoted_context_request_ids) if getattr(request, "py_needs_onehot_draft_probs", False): if request.py_seq_slot is not None: padding_gen_slots.append(request.py_seq_slot) @@ -4218,17 +4308,25 @@ def append_cross_attention_state(request: LlmRequest, # (1) next_draft_tokens_device is None, which means overlap scheduler is disabled; or # (2) a dummy request; or # (3) the first step in the generation server of disaggregated serving - if next_draft_tokens_device is None or request.is_dummy or request.py_batch_idx is None: + if (is_promoted_context or next_draft_tokens_device is None + or request.is_dummy or request.py_batch_idx is None): # get token ids, including input token ids and draft token ids. For these dummy requests, # no need to copy the token ids. if not (request.is_attention_dp_dummy or request.is_cuda_graph_dummy): - input_ids.append(request.get_last_tokens(0)) + if is_promoted_context: + input_ids.append( + request.get_tokens(0)[ + request.context_current_position]) + else: + input_ids.append(request.get_last_tokens(0)) input_ids.extend(request.py_draft_tokens) draft_tokens.extend(request.py_draft_tokens) # get other ids and lengths num_draft_tokens = get_draft_token_length(request) - past_seen_token_num = request.max_beam_num_tokens - 1 + past_seen_token_num = (request.context_current_position + if is_promoted_context else + request.max_beam_num_tokens - 1) draft_lens.append(num_draft_tokens) if self.enable_spec_decode and spec_config.spec_dec_mode.extend_ctx( self.attn_backend) and spec_config.is_linear_tree: @@ -4372,11 +4470,17 @@ def append_cross_attention_state(request: LlmRequest, for request in generation_requests: request_ids.append(request.py_request_id) + is_promoted_context = (request.py_request_id + in promoted_context_request_ids) # the request has no previous tensor: # (1) new_tokens_device is None, which means overlap scheduler is disabled; or # (2) a dummy request; or # (3) the first step in the generation server of disaggregated serving - if new_tokens_device is None or request.is_dummy or request.py_batch_idx is None: + if is_promoted_context: + input_ids.append( + request.get_tokens(0)[request.context_current_position]) + past_seen_token_num = request.context_current_position + elif new_tokens_device is None or request.is_dummy or request.py_batch_idx is None: # skip adding input_ids of CUDA graph dummy requests so that new_tokens_device # can be aligned to the correct positions. if not request.is_cuda_graph_dummy: @@ -5646,17 +5750,19 @@ def _get_eager_lora_params_from_requests( @nvtx_range("_prepare_inputs") def _prepare_inputs( - self, - scheduled_requests: ScheduledRequests, - kv_cache_manager: Union[KVCacheManager, KVCacheManagerV2], - attn_metadata: AttentionMetadata, - spec_metadata: Optional[SpecMetadata] = None, - new_tensors_device: Optional[SampleStateTensors] = None, - cache_indirection_buffer: Optional[torch.Tensor] = None, - num_accepted_tokens_device: Optional[torch.Tensor] = None, - req_id_to_old_request: Optional[Dict[int, LlmRequest]] = None, - resource_manager: Optional[ResourceManager] = None, - maybe_graph: bool = False): + self, + scheduled_requests: ScheduledRequests, + kv_cache_manager: Union[KVCacheManager, KVCacheManagerV2], + attn_metadata: AttentionMetadata, + spec_metadata: Optional[SpecMetadata] = None, + new_tensors_device: Optional[SampleStateTensors] = None, + cache_indirection_buffer: Optional[torch.Tensor] = None, + num_accepted_tokens_device: Optional[torch.Tensor] = None, + req_id_to_old_request: Optional[Dict[int, LlmRequest]] = None, + resource_manager: Optional[ResourceManager] = None, + maybe_graph: bool = False, + promoted_context_request_ids: frozenset[int] = frozenset() + ) -> Tuple[Dict[str, Any], Optional[torch.Tensor]]: if self.mapping is not None and 'cp_type' in self.mapping.cp_config: cp_type = self.mapping.cp_config['cp_type'] if CpType.STAR == cp_type: @@ -5697,7 +5803,7 @@ def _prepare_inputs( scheduled_requests, kv_cache_manager, attn_metadata, spec_metadata, new_tensors_device, cache_indirection_buffer, num_accepted_tokens_device, req_id_to_old_request, resource_manager, - maybe_graph) + maybe_graph, promoted_context_request_ids) def _prepare_encoder_inputs( self, @@ -6211,14 +6317,36 @@ def forward(self, inputs, gather_ids=gather_ids, gather_context_logits=gather_context_logits) + + graph_requests = scheduled_requests + promoted_context_request_ids: frozenset[int] = frozenset() + # TODO: Generalize these conservative gates as actual-draft, beam, and + # context-parallel providers for decoder-only LLMs gain support for + # promoted final-context rows. Each relaxation must preserve whole-batch + # fallback on graph miss and prove parity with the provider's native + # q_len=1 path. Encoder-decoder and non-LLM engines remain out of scope. + if (scheduled_requests.num_context_requests > 0 + and self.cuda_graph_runner.enabled + and (not self.enable_spec_decode or + (not self.is_draft_model and self.runtime_draft_len == 0)) + and not self.use_beam_search + and not self._is_encoder_decoder_model() + and not self._is_encode_only + and not self.llm_args.mm_encoder_only + and self.mapping.cp_size == 1): + graph_requests, promoted_context_request_ids = \ + _make_single_token_context_graph_batch( + scheduled_requests, + self._is_final_multimodal_context_decode_compatible) + with self.cuda_graph_runner.pad_batch( - scheduled_requests, resource_manager, - self.runtime_draft_len) as padded_requests: + graph_requests, resource_manager, + self.runtime_draft_len) as padded_graph_requests: # Callee already no-ops when use_mrope=False, but the Python call / # frame setup itself is non-trivial under high concurrency. Gating # at the caller avoids that overhead for non-mrope models. if self.use_mrope: - self._pad_batch_seed_mrope_delta_cache(padded_requests) + self._pad_batch_seed_mrope_delta_cache(padded_graph_requests) # Refresh is_all_greedy_sample for the *current* batch BEFORE the # CUDA graph key is built below. The key includes this flag to pick @@ -6229,11 +6357,11 @@ def forward(self, # unpopulated (greedy) buffers, hanging the run (e.g. MTP nextn>=2). if spec_metadata is not None: spec_metadata.update_is_all_greedy_sample( - padded_requests.all_requests()) + padded_graph_requests.all_requests()) self._sync_group_all_greedy_sample(spec_metadata) maybe_attn_metadata, maybe_spec_metadata, key = self.cuda_graph_runner.maybe_get_cuda_graph( - padded_requests, + padded_graph_requests, enable_spec_decode=self.enable_spec_decode, attn_metadata=attn_metadata, spec_metadata=spec_metadata, @@ -6241,33 +6369,46 @@ def forward(self, if self.is_spec_decode else None, new_tensors_device=new_tensors_device, spec_resource_manager=spec_resource_manager, + promoted_context_request_ids=promoted_context_request_ids, ) can_run_graph = key is not None if can_run_graph: attn_metadata = maybe_attn_metadata spec_metadata = maybe_spec_metadata + execution_requests = padded_graph_requests + execution_promoted_context_ids = promoted_context_request_ids else: attn_metadata = self.attn_metadata if self.enable_spec_decode: spec_metadata = self.spec_metadata else: spec_metadata = None + execution_requests = scheduled_requests + execution_promoted_context_ids = frozenset() # Fill slot-ID buffer for scatter inside draft loop if (self.enable_spec_decode and spec_tree_manager is not None and spec_tree_manager.use_dynamic_tree and not self.is_draft_model): spec_tree_manager.slot_storage.fill_all_slot_ids( - padded_requests.context_requests, - padded_requests.generation_requests, + execution_requests.context_requests, + execution_requests.generation_requests, ) inputs, gather_ids = self._prepare_inputs( - padded_requests, kv_cache_manager, attn_metadata, spec_metadata, - new_tensors_device, cache_indirection_buffer, + execution_requests, kv_cache_manager, attn_metadata, + spec_metadata, new_tensors_device, cache_indirection_buffer, num_accepted_tokens_device, req_id_to_old_request, - resource_manager, can_run_graph) + resource_manager, can_run_graph, execution_promoted_context_ids) + if execution_promoted_context_ids: + self.iter_states[ + 'num_ctx_requests'] = scheduled_requests.num_context_requests + self.iter_states['num_ctx_tokens'] = sum( + request.context_chunk_size + for request in scheduled_requests.context_requests) + self.iter_states[ + 'num_generation_tokens'] = scheduled_requests.num_generation_requests self._prepare_inputs_event = torch.cuda.Event() self._prepare_inputs_event.record() diff --git a/tests/integration/defs/kv_cache/test_final_single_token_context_cuda_graph.py b/tests/integration/defs/kv_cache/test_final_single_token_context_cuda_graph.py new file mode 100644 index 000000000000..6b0643788b43 --- /dev/null +++ b/tests/integration/defs/kv_cache/test_final_single_token_context_cuda_graph.py @@ -0,0 +1,520 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json +import os +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Literal, TypedDict, cast + +import pytest +import torch + +from tensorrt_llm import LLM +from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import CUDAGraphRunner, KeyType +from tensorrt_llm._torch.pyexecutor.model_engine import PyTorchModelEngine +from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor +from tensorrt_llm._torch.pyexecutor.resource_manager import BaseResourceManager +from tensorrt_llm._torch.pyexecutor.sampler import SampleStateTensors +from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests +from tensorrt_llm._torch.speculative.interface import SpecMetadata +from tensorrt_llm._utils import mpi_rank +from tensorrt_llm.executor.executor import GenerationExecutor +from tensorrt_llm.executor.postproc_worker import PostprocWorkerConfig +from tensorrt_llm.executor.proxy import GenerationExecutorProxy +from tensorrt_llm.executor.worker import GenerationExecutorWorker +from tensorrt_llm.llmapi import CudaGraphConfig, Eagle3DecodingConfig, KvCacheConfig, RequestOutput +from tensorrt_llm.llmapi.mpi_session import MpiSession +from tensorrt_llm.sampling_params import GuidedDecodingParams, SamplingParams + +from ..conftest import llm_models_root + +MODEL = f"{llm_models_root()}/llama-models-v2/TinyLlama-1.1B-Chat-v1.0" +SPEC_MODEL = f"{llm_models_root()}/Qwen3/Qwen3-8B" +EAGLE3_MODEL = f"{llm_models_root()}/Qwen3/qwen3_8b_eagle3" +PROMPT_TOKEN_IDS = [1] + [42] * 63 + [43] +CHANGED_FINAL_PROMPT_TOKEN_IDS = PROMPT_TOKEN_IDS[:-1] + [44] +SPEC_PROMPT_TOKEN_IDS = ( + [1] + [42] * 63 + [43], + [1] + [44] * 63 + [45], +) +_TP_GRAPH_PROBE_DIR_ENV = "TLLM_FINAL_CONTEXT_CUDA_GRAPH_PROBE_DIR" + + +class _PromotedContextGraphExecutionReport(TypedDict): + promoted_context_request_count: int + graph_batch_size: int | None + enable_spec_decode: bool + replayed: bool + + +class _RankGraphExecutionReport(TypedDict): + rank: int + executions: list[_PromotedContextGraphExecutionReport] + + +@dataclass +class _PromotedContextGraphExecution: + promoted_context_request_ids: frozenset[int] + key: KeyType | None + enable_spec_decode: bool + replayed: bool = False + + +class _CudaGraphExecutionProbe: + """Observe promoted-context graph selection without replacing execution.""" + + def __init__(self, runner: CUDAGraphRunner) -> None: + self._maybe_get_cuda_graph = runner.maybe_get_cuda_graph + self._replay = runner.replay + self._executions: list[_PromotedContextGraphExecution] = [] + self._pending_execution: _PromotedContextGraphExecution | None = None + + def maybe_get_cuda_graph( + self, + batch: ScheduledRequests, + enable_spec_decode: bool, + attn_metadata: Any, + spec_metadata: SpecMetadata | None = None, + draft_tokens_cuda: torch.Tensor | None = None, + new_tensors_device: SampleStateTensors | None = None, + spec_resource_manager: BaseResourceManager | None = None, + promoted_context_request_ids: frozenset[int] = frozenset(), + ) -> tuple[Any | None, Any | None, KeyType | None]: + # A new decision means the preceding one reached eager execution if it + # did not call replay. Keep that earlier observation unchanged. + self._pending_execution = None + result = self._maybe_get_cuda_graph( + batch, + enable_spec_decode, + attn_metadata, + spec_metadata, + draft_tokens_cuda, + new_tensors_device, + spec_resource_manager, + promoted_context_request_ids, + ) + if promoted_context_request_ids: + execution = _PromotedContextGraphExecution( + promoted_context_request_ids=promoted_context_request_ids, + key=result[2], + enable_spec_decode=enable_spec_decode, + ) + self._executions.append(execution) + self._pending_execution = execution + return result + + def replay( + self, + key: KeyType, + current_inputs: dict[str, Any], + ) -> torch.Tensor | None: + output = self._replay(key, current_inputs) + if self._pending_execution is not None and self._pending_execution.key == key: + self._pending_execution.replayed = True + self._pending_execution = None + return output + + @property + def executions(self) -> tuple[_PromotedContextGraphExecution, ...]: + return tuple(self._executions) + + +def _get_worker_cuda_graph_runner( + worker: GenerationExecutorWorker, +) -> CUDAGraphRunner: + assert isinstance(worker.engine, PyExecutor) + model_engine = worker.engine.model_engine + assert isinstance(model_engine, PyTorchModelEngine) + return model_engine.cuda_graph_runner + + +class _CudaGraphProbeWorker(GenerationExecutorWorker): + """Install the graph probe independently in every TP worker process.""" + + def setup_engine(self) -> None: + super().setup_engine() + runner = _get_worker_cuda_graph_runner(self) + self._cuda_graph_execution_probe = _CudaGraphExecutionProbe(runner) + runner.maybe_get_cuda_graph = self._cuda_graph_execution_probe.maybe_get_cuda_graph + runner.replay = self._cuda_graph_execution_probe.replay + + def shutdown(self) -> None: + # Each TP rank owns a separate Python process and CUDA graph runner. + # Persist one report per rank before the engine is released so the + # parent test can prove that every rank selected and replayed a graph. + if not self.doing_shutdown: + probe_dir = os.getenv(_TP_GRAPH_PROBE_DIR_ENV) + assert probe_dir is not None + report: _RankGraphExecutionReport = { + "rank": mpi_rank(), + "executions": [ + { + "promoted_context_request_count": len( + execution.promoted_context_request_ids + ), + "graph_batch_size": ( + execution.key[0] if execution.key is not None else None + ), + "enable_spec_decode": execution.enable_spec_decode, + "replayed": execution.replayed, + } + for execution in self._cuda_graph_execution_probe.executions + ], + } + report_path = Path(probe_dir) / f"rank-{report['rank']}.json" + report_path.write_text(json.dumps(report), encoding="utf-8") + super().shutdown() + + +def _create_cuda_graph_probe_ipc_executor( + worker_kwargs: dict[str, object], + model_world_size: int, + mpi_session: MpiSession | None, + postproc_worker_config: PostprocWorkerConfig, + is_llm_executor: bool | None, + use_worker: bool = False, +) -> GenerationExecutorProxy: + assert not use_worker, "The TP2 probe requires separate worker processes" + return GenerationExecutorProxy( + worker_kwargs, + model_world_size=model_world_size, + mpi_session=mpi_session, + worker_cls=_CudaGraphProbeWorker, + postproc_worker_config=postproc_worker_config, + is_llm_executor=is_llm_executor, + ) + + +def _get_cuda_graph_runner(llm: LLM) -> CUDAGraphRunner: + assert isinstance(llm._executor, GenerationExecutorWorker) + return _get_worker_cuda_graph_runner(llm._executor) + + +def _assert_reused_context_used_cuda_graph( + executions: tuple[_PromotedContextGraphExecution, ...], +) -> None: + assert len(executions) == 1, ( + "Expected exactly one promoted final-context graph decision for the " + f"reused request, got {len(executions)}" + ) + execution = executions[0] + assert len(execution.promoted_context_request_ids) == 1 + assert execution.key is not None, "The promoted final-context row fell back to eager prefill" + assert execution.key[0] == 1 + assert execution.replayed, ( + "The selected CUDA graph was not replayed for the promoted context row" + ) + + +def _read_rank_graph_execution_reports( + report_dir: Path, + world_size: int, +) -> tuple[_RankGraphExecutionReport, ...]: + reports: list[_RankGraphExecutionReport] = [] + for rank in range(world_size): + report_path = report_dir / f"rank-{rank}.json" + assert report_path.is_file(), f"Missing CUDA graph report for TP rank {rank}" + report = cast( + _RankGraphExecutionReport, + json.loads(report_path.read_text(encoding="utf-8")), + ) + assert report["rank"] == rank + reports.append(report) + return tuple(reports) + + +def _assert_rank_reused_context_used_cuda_graph( + report: _RankGraphExecutionReport, +) -> None: + executions = report["executions"] + assert len(executions) == 1, ( + f"TP rank {report['rank']} observed {len(executions)} promoted " + "final-context graph decisions instead of one" + ) + execution = executions[0] + assert execution["promoted_context_request_count"] == 1 + assert execution["graph_batch_size"] == 1, ( + f"TP rank {report['rank']} fell back to eager prefill" + ) + assert not execution["enable_spec_decode"] + assert execution["replayed"], ( + f"TP rank {report['rank']} selected but did not replay the CUDA graph" + ) + + +def _generate_cold_and_reused( + use_kv_cache_manager_v2: bool, + sampling_params: SamplingParams, + monkeypatch: pytest.MonkeyPatch, + guided_decoding_backend: Literal["xgrammar"] | None = None, +) -> tuple[RequestOutput, RequestOutput, tuple[_PromotedContextGraphExecution, ...]]: + # Two complete 32-token blocks can be reused, leaving exactly the final + # prompt token for the second request. The first and last IDs are distinct + # so an accidental cursor shift is visible in output/logit parity. + kv_cache_config = KvCacheConfig( + enable_block_reuse=True, + use_kv_cache_manager_v2=use_kv_cache_manager_v2, + ) + cuda_graph_config = CudaGraphConfig(batch_sizes=[1], enable_padding=False) + guided_decoding_args: dict[str, str] = {} + if guided_decoding_backend is not None: + guided_decoding_args["guided_decoding_backend"] = guided_decoding_backend + + # Class/instance-level probes do not cross the default TP1 worker process. + # Keep the real PyExecutor in-process so the wrappers below can observe the + # actual graph decision and replay while still calling the original code. + monkeypatch.setenv("TLLM_WORKER_USE_SINGLE_PROCESS", "1") + with LLM( + model=MODEL, + max_batch_size=1, + max_num_tokens=128, + kv_cache_config=kv_cache_config, + cuda_graph_config=cuda_graph_config, + **guided_decoding_args, + ) as llm: + cold = llm.generate([PROMPT_TOKEN_IDS], sampling_params)[0] + runner = _get_cuda_graph_runner(llm) + probe = _CudaGraphExecutionProbe(runner) + monkeypatch.setattr(runner, "maybe_get_cuda_graph", probe.maybe_get_cuda_graph) + monkeypatch.setattr(runner, "replay", probe.replay) + reused = llm.generate([PROMPT_TOKEN_IDS], sampling_params)[0] + + return cold, reused, probe.executions + + +def _generate_changed_final_token_cold_and_reused( + use_kv_cache_manager_v2: bool, + sampling_params: SamplingParams, + monkeypatch: pytest.MonkeyPatch, +) -> tuple[ + RequestOutput, + RequestOutput, + tuple[_PromotedContextGraphExecution, ...], +]: + """Compare a cold changed-tail request with prefix reuse from another tail.""" + monkeypatch.setenv("TLLM_WORKER_USE_SINGLE_PROCESS", "1") + + # Use a separate engine for the cold reference so it cannot populate the + # two shared cache blocks exercised by the promoted request below. + with LLM( + model=MODEL, + max_batch_size=1, + max_num_tokens=128, + kv_cache_config=KvCacheConfig( + enable_block_reuse=True, + use_kv_cache_manager_v2=use_kv_cache_manager_v2, + ), + cuda_graph_config=CudaGraphConfig(batch_sizes=[1], enable_padding=False), + ) as reference_llm: + cold = reference_llm.generate([CHANGED_FINAL_PROMPT_TOKEN_IDS], sampling_params)[0] + + with LLM( + model=MODEL, + max_batch_size=1, + max_num_tokens=128, + kv_cache_config=KvCacheConfig( + enable_block_reuse=True, + use_kv_cache_manager_v2=use_kv_cache_manager_v2, + ), + cuda_graph_config=CudaGraphConfig(batch_sizes=[1], enable_padding=False), + ) as reuse_llm: + reuse_llm.generate([PROMPT_TOKEN_IDS], sampling_params) + runner = _get_cuda_graph_runner(reuse_llm) + probe = _CudaGraphExecutionProbe(runner) + monkeypatch.setattr(runner, "maybe_get_cuda_graph", probe.maybe_get_cuda_graph) + monkeypatch.setattr(runner, "replay", probe.replay) + reused = reuse_llm.generate([CHANGED_FINAL_PROMPT_TOKEN_IDS], sampling_params)[0] + + return cold, reused, probe.executions + + +@pytest.mark.threadleak(enabled=False) +@pytest.mark.parametrize("use_kv_cache_manager_v2", [False, True], ids=["v1", "v2"]) +def test_final_token_reuse_cuda_graph( + use_kv_cache_manager_v2: bool, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Verify the minimal final-token reuse case without optional features.""" + cold, reused, graph_executions = _generate_cold_and_reused( + use_kv_cache_manager_v2, + SamplingParams(max_tokens=4, end_id=-1), + monkeypatch, + ) + + assert cold.outputs[0].token_ids == reused.outputs[0].token_ids + _assert_reused_context_used_cuda_graph(graph_executions) + + +@pytest.mark.threadleak(enabled=False) +@pytest.mark.parametrize("use_kv_cache_manager_v2", [False, True], ids=["v1", "v2"]) +def test_changed_final_token_reuse_cuda_graph( + use_kv_cache_manager_v2: bool, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Verify the promoted row reads a changed final prompt token.""" + cold, reused, graph_executions = _generate_changed_final_token_cold_and_reused( + use_kv_cache_manager_v2, + SamplingParams(max_tokens=4, end_id=-1, temperature=0), + monkeypatch, + ) + + assert PROMPT_TOKEN_IDS[:-1] == CHANGED_FINAL_PROMPT_TOKEN_IDS[:-1] + assert PROMPT_TOKEN_IDS[-1] != CHANGED_FINAL_PROMPT_TOKEN_IDS[-1] + assert cold.outputs[0].token_ids == reused.outputs[0].token_ids + _assert_reused_context_used_cuda_graph(graph_executions) + + +@pytest.mark.threadleak(enabled=False) +@pytest.mark.skip_less_device(2) +@pytest.mark.parametrize("use_kv_cache_manager_v2", [False, True], ids=["v1", "v2"]) +def test_final_token_reuse_cuda_graph_tp2( + use_kv_cache_manager_v2: bool, + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + """Verify every TP2 rank replays the graph for final-token reuse.""" + report_dir = tmp_path / "tp2-cuda-graph-reports" + report_dir.mkdir() + monkeypatch.setenv(_TP_GRAPH_PROBE_DIR_ENV, str(report_dir)) + monkeypatch.setattr( + GenerationExecutor, + "_create_ipc_executor", + staticmethod(_create_cuda_graph_probe_ipc_executor), + ) + + with LLM( + model=MODEL, + tensor_parallel_size=2, + max_batch_size=1, + max_num_tokens=128, + kv_cache_config=KvCacheConfig( + enable_block_reuse=True, + use_kv_cache_manager_v2=use_kv_cache_manager_v2, + ), + cuda_graph_config=CudaGraphConfig(batch_sizes=[1], enable_padding=False), + ) as llm: + sampling_params = SamplingParams(max_tokens=4, end_id=-1) + cold = llm.generate([PROMPT_TOKEN_IDS], sampling_params)[0] + reused = llm.generate([PROMPT_TOKEN_IDS], sampling_params)[0] + + assert cold.outputs[0].token_ids == reused.outputs[0].token_ids + reports = _read_rank_graph_execution_reports(report_dir, world_size=2) + for report in reports: + _assert_rank_reused_context_used_cuda_graph(report) + + +@pytest.mark.threadleak(enabled=False) +@pytest.mark.parametrize("use_kv_cache_manager_v2", [False, True], ids=["v1", "v2"]) +def test_context_logits_after_final_token_reuse( + use_kv_cache_manager_v2: bool, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Verify context logits when reuse leaves one prompt token to compute.""" + cold, reused, graph_executions = _generate_cold_and_reused( + use_kv_cache_manager_v2, + SamplingParams( + max_tokens=4, + end_id=-1, + return_context_logits=True, + ), + monkeypatch, + ) + + assert cold.outputs[0].token_ids == reused.outputs[0].token_ids + assert cold.context_logits is not None + assert reused.context_logits is not None + assert cold.context_logits.shape[0] == len(PROMPT_TOKEN_IDS) + assert reused.context_logits.shape[0] == 1 + _assert_reused_context_used_cuda_graph(graph_executions) + + +@pytest.mark.threadleak(enabled=False) +@pytest.mark.parametrize("use_kv_cache_manager_v2", [False, True], ids=["v1", "v2"]) +def test_guided_decoding_after_final_token_reuse( + use_kv_cache_manager_v2: bool, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Verify guided decoding when reuse leaves one prompt token to compute.""" + cold, reused, graph_executions = _generate_cold_and_reused( + use_kv_cache_manager_v2, + SamplingParams( + max_tokens=4, + end_id=-1, + # Keep the grammar permissive so output equality tests execution- + # path parity rather than narrow-format generation behavior. + guided_decoding=GuidedDecodingParams(regex=r".*"), + ), + monkeypatch, + guided_decoding_backend="xgrammar", + ) + + assert cold.outputs[0].token_ids == reused.outputs[0].token_ids + _assert_reused_context_used_cuda_graph(graph_executions) + + +@pytest.mark.threadleak(enabled=False) +@pytest.mark.parametrize("use_kv_cache_manager_v2", [False, True], ids=["v1", "v2"]) +def test_zero_runtime_draft_speculation_after_final_token_reuse( + use_kv_cache_manager_v2: bool, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Verify a zero-draft speculative iteration replays the decode graph.""" + kv_cache_config = KvCacheConfig( + enable_block_reuse=True, + use_kv_cache_manager_v2=use_kv_cache_manager_v2, + free_gpu_memory_fraction=0.6, + ) + speculative_config = Eagle3DecodingConfig( + max_draft_len=1, + speculative_model=EAGLE3_MODEL, + eagle3_one_model=True, + # Batch size one drafts one token. Larger batches use the implicit + # zero-draft schedule entry and therefore exercise this stage's gate. + draft_len_schedule={1: 1}, + ) + sampling_params = SamplingParams(max_tokens=4, end_id=-1, temperature=0) + + monkeypatch.setenv("TLLM_WORKER_USE_SINGLE_PROCESS", "1") + with LLM( + model=SPEC_MODEL, + max_batch_size=2, + max_num_tokens=256, + kv_cache_config=kv_cache_config, + cuda_graph_config=CudaGraphConfig(batch_sizes=[1, 2], enable_padding=True), + speculative_config=speculative_config, + ) as llm: + cold = [llm.generate([prompt], sampling_params)[0] for prompt in SPEC_PROMPT_TOKEN_IDS] + runner = _get_cuda_graph_runner(llm) + probe = _CudaGraphExecutionProbe(runner) + monkeypatch.setattr(runner, "maybe_get_cuda_graph", probe.maybe_get_cuda_graph) + monkeypatch.setattr(runner, "replay", probe.replay) + reused = llm.generate(list(SPEC_PROMPT_TOKEN_IDS), sampling_params) + + assert [output.outputs[0].token_ids for output in cold] == [ + output.outputs[0].token_ids for output in reused + ] + assert len(probe.executions) == 1 + execution = probe.executions[0] + # Scheduler timing may produce either one promoted final-context row plus + # one generation sibling, or two promoted final-context rows. Both valid + # batch-two shapes must use the zero-runtime-draft graph. + assert 1 <= len(execution.promoted_context_request_ids) <= len(SPEC_PROMPT_TOKEN_IDS) + assert execution.enable_spec_decode + assert execution.key is not None, ( + "The zero-runtime-draft promoted rows fell back to eager prefill" + ) + assert execution.key[:2] == (2, 0) + assert execution.replayed diff --git a/tests/integration/test_lists/test-db/l0_dgx_h100.yml b/tests/integration/test_lists/test-db/l0_dgx_h100.yml index e36416f7e5fa..8b34ef705460 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_h100.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_h100.yml @@ -20,6 +20,7 @@ l0_dgx_h100: - unittest/_torch/multi_gpu -m "not post_merge" TIMEOUT (90) - unittest/_torch/distributed - unittest/_torch/modeling/test_modeling_pixtral.py::test_tensor_parallelism + - kv_cache/test_final_single_token_context_cuda_graph.py::test_final_token_reuse_cuda_graph_tp2 # ------------- Encoder-decoder TP tests --------------- - llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_generate_encoder_decoder_end_to_end[bf16-kv-v1-cuda-graph-on-greedy-tp2-t5-small] - llmapi/test_llm_api_pytorch_bart.py::test_bart_pytorch_generate_encoder_decoder_end_to_end[bf16-kv-v1-cuda-graph-on-greedy-tp2-bart-large-cnn] diff --git a/tests/integration/test_lists/test-db/l0_h100.yml b/tests/integration/test_lists/test-db/l0_h100.yml index f4a3fcefef03..f61cb21d7b69 100644 --- a/tests/integration/test_lists/test-db/l0_h100.yml +++ b/tests/integration/test_lists/test-db/l0_h100.yml @@ -221,6 +221,11 @@ l0_h100: - test_e2e.py::test_openai_chat_harmony_perf_metrics - test_e2e.py::test_openai_responses - test_e2e.py::test_openai_chat_guided_decoding[meta-llama/Llama-3.1-8B-Instruct] + - kv_cache/test_final_single_token_context_cuda_graph.py::test_final_token_reuse_cuda_graph + - kv_cache/test_final_single_token_context_cuda_graph.py::test_changed_final_token_reuse_cuda_graph + - kv_cache/test_final_single_token_context_cuda_graph.py::test_context_logits_after_final_token_reuse + - kv_cache/test_final_single_token_context_cuda_graph.py::test_guided_decoding_after_final_token_reuse + - kv_cache/test_final_single_token_context_cuda_graph.py::test_zero_runtime_draft_speculation_after_final_token_reuse # ------------- Prefix-aware scheduling E2E tests --------------- - kv_cache/test_prefix_aware_scheduling.py::TestServePrefixAwareScheduling::test_multi_round_qa_shared_prefix_smoke - kv_cache/test_prefix_aware_scheduling.py::TestServePrefixAwareScheduling::test_multi_round_qa_shared_prefix[guaranteed-chunked] diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine.py b/tests/unittest/_torch/executor/test_pytorch_model_engine.py index 6347275fd640..4f37479efb3a 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py @@ -2,8 +2,10 @@ # Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. import unittest +from contextlib import nullcontext from dataclasses import dataclass -from unittest.mock import Mock +from types import SimpleNamespace +from unittest.mock import Mock, patch import torch @@ -12,11 +14,15 @@ from tensorrt_llm._torch.pyexecutor.connectors.kv_cache_connector import \ KvCacheConnectorWorker from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import ( - _restore_spec_decode_capture_state, _save_spec_decode_capture_state) + CUDAGraphRunner, _restore_spec_decode_capture_state, + _save_spec_decode_capture_state) from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest from tensorrt_llm._torch.pyexecutor.model_engine import ( - PyTorchModelEngine, _build_request_multimodal_input) -from tensorrt_llm.llmapi.llm_args import TorchLlmArgs + PyTorchModelEngine, _build_request_multimodal_input, + _make_single_token_context_graph_batch) +from tensorrt_llm.llmapi.llm_args import (DecodingBaseConfig, + SeqLenAwareSparseAttentionConfig, + TorchLlmArgs) # isort: off from tensorrt_llm._torch.pyexecutor.resource_manager import (KVCacheManager, @@ -28,8 +34,11 @@ from tensorrt_llm._torch.attention_backend.interface import AttentionMetadata from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests +from tensorrt_llm._torch.speculative.spec_sampler_base import \ + SampleStateTensorsSpec from tensorrt_llm.bindings.executor import KvCacheConfig -from tensorrt_llm.llmapi import CudaGraphConfig, SamplingParams +from tensorrt_llm.llmapi import (CudaGraphConfig, SADecodingConfig, + SamplingParams) from tensorrt_llm.mapping import CpType, Mapping @@ -76,7 +85,12 @@ def forward(self, *args, **kwargs) -> torch.Tensor: class DummyModelEngine(PyTorchModelEngine): - def __init__(self, llm_args: TorchLlmArgs, dtype: torch.dtype) -> None: + def __init__( + self, + llm_args: TorchLlmArgs, + dtype: torch.dtype, + spec_config: DecodingBaseConfig | None = None, + ) -> None: self.dtype = dtype mapping = Mapping(world_size=tensorrt_llm.mpi_world_size(), tp_size=tensorrt_llm.mpi_world_size(), @@ -85,7 +99,8 @@ def __init__(self, llm_args: TorchLlmArgs, dtype: torch.dtype) -> None: super().__init__(model_path="dummy", mapping=mapping, model=model, - llm_args=llm_args) + llm_args=llm_args, + spec_config=spec_config) def _create_request(num_tokens, req_id: int): @@ -107,8 +122,120 @@ def _create_request(num_tokens, req_id: int): return result -def create_model_engine_and_kvcache(llm_args: TorchLlmArgs = None, - execution_stream: torch.cuda.Stream = None): +def _create_request_with_tokens(tokens: list[int], req_id: int) -> LlmRequest: + sampling_params = SamplingParams() + request = LlmRequest( + request_id=req_id, + max_new_tokens=1, + input_tokens=tokens, + sampling_config=tensorrt_llm.bindings.SamplingConfig( + sampling_params._get_sampling_config()), + is_streaming=False, + ) + request.paged_kv_block_ids = [] + return request + + +def _make_request_stub(req_id: int, prompt_len: int = 4) -> SimpleNamespace: + return SimpleNamespace( + py_request_id=req_id, + context_chunk_size=1, + context_remaining_length=1, + context_current_position=prompt_len - 1, + py_prompt_len=prompt_len, + py_beam_width=1, + py_draft_tokens=[], + py_is_first_draft=False, + is_context_only_request=False, + is_generation_only_request=lambda: False, + py_disaggregated_params=None, + py_multimodal_data=None, + py_mm_encoder_event=None, + py_mrope_position_delta=None, + py_return_context_logits=False, + py_batch_idx=None, + is_dummy=False, + max_beam_num_tokens=prompt_len, + state="context", + py_llm_request_type="context_and_generation", + ) + + +def _make_forward_only_engine( + graph_key: tuple[int, int, bool, bool, bool] | None, + runner_enabled: bool = True, +) -> tuple[PyTorchModelEngine, Mock, Mock, Mock, dict[str, object]]: + engine = object.__new__(PyTorchModelEngine) + engine.model = SimpleNamespace( + extra_attrs={}, + model_config=SimpleNamespace(pretrained_config=SimpleNamespace( + rope_scaling=None))) + engine.kv_cache_manager_key = ResourceManagerType.KV_CACHE_MANAGER + engine.enable_spec_decode = False + engine.is_spec_decode = False + engine.is_draft_model = False + engine.guided_decoder = None + engine.max_beam_width = 1 + engine._is_encode_only = False + engine.llm_args = SimpleNamespace(mm_encoder_only=False) + engine.mapping = SimpleNamespace( + cp_size=1, + enable_lm_head_tp_in_adp=False, + ) + engine.runtime_draft_len = 0 + engine.attn_backend = None + engine.model_is_wrapped = False + engine.original_max_draft_len = 0 + engine.original_max_total_draft_tokens = 0 + engine._spec_dec_max_total_draft_tokens = 0 + engine.get_runtime_tokens_per_gen_step = Mock(return_value=1) + engine.iter_states = {} + engine.forward_pass_callable = None + engine._is_encoder_decoder_model = Mock(return_value=False) + engine._get_draft_kv_cache_manager = Mock(return_value=None) + + semantic_attn_metadata = Mock() + graph_attn_metadata = Mock() + engine.attn_metadata = semantic_attn_metadata + engine._set_up_attn_metadata = Mock(return_value=semantic_attn_metadata) + spec_dec_mode = Mock() + spec_dec_mode.attention_need_spec_dec_mode.return_value = False + spec_dec_mode.is_parallel_draft.return_value = False + spec_metadata = Mock( + spec_dec_mode=spec_dec_mode, + is_spec_dec_tree=False, + is_spec_dec_dynamic_tree=False, + ) + engine.spec_metadata = spec_metadata + engine._set_up_spec_metadata = Mock(return_value=spec_metadata) + engine._prepare_inputs = Mock(return_value=({"prepared": True}, None)) + outputs = {"logits": object()} + engine._forward_step = Mock(return_value=outputs) + engine._execute_logit_post_processors = Mock() + + runner = Mock() + runner.enabled = runner_enabled + runner.pad_batch.side_effect = lambda batch, *_args: nullcontext(batch) + runner.maybe_get_cuda_graph.return_value = ((graph_attn_metadata, None, + graph_key) + if graph_key is not None else + (None, None, None)) + runner.get_graph_pool.return_value = None + runner.needs_capture.return_value = False + runner.is_warmup_only = False + runner.replay.return_value = outputs + engine.cuda_graph_runner = runner + + resource_manager = Mock() + resource_manager.get_resource_manager.return_value = object() + return engine, runner, resource_manager, semantic_attn_metadata, outputs + + +def create_model_engine_and_kvcache( + llm_args: TorchLlmArgs | None = None, + execution_stream: torch.cuda.Stream | None = None, + spec_config: DecodingBaseConfig | None = None, +) -> tuple[PyTorchModelEngine, KVCacheManager]: tokens_per_block = 1 max_tokens = 258 # Atleast 1 more than the max seq len num_layers = 1 @@ -127,7 +254,7 @@ def create_model_engine_and_kvcache(llm_args: TorchLlmArgs = None, assert (8 in llm_args.cuda_graph_config.batch_sizes and 16 in llm_args.cuda_graph_config.batch_sizes) - model_engine = DummyModelEngine(llm_args, torch.half) + model_engine = DummyModelEngine(llm_args, torch.half, spec_config) kv_cache_config = KvCacheConfig(max_tokens=max_tokens) mapping = Mapping(world_size=1, tp_size=1, rank=0) @@ -148,6 +275,602 @@ def create_model_engine_and_kvcache(llm_args: TorchLlmArgs = None, return model_engine, kv_cache_manager +class SingleTokenContextGraphBatchTestCase(unittest.TestCase): + + def test_generation_only_is_identity(self) -> None: + generation = _make_request_stub(1) + batch = ScheduledRequests() + batch.generation_requests = [generation] + + graph_batch, promoted_ids = _make_single_token_context_graph_batch( + batch) + + self.assertIs(graph_batch, batch) + self.assertEqual(promoted_ids, frozenset()) + + def test_eligible_batch_has_independent_lists_and_stable_order( + self) -> None: + context_0 = _make_request_stub(10, prompt_len=1) + context_1 = _make_request_stub(11, prompt_len=8) + generation = _make_request_stub(12, prompt_len=16) + paused = object() + batch = ScheduledRequests() + batch.context_requests_last_chunk = [context_0, context_1] + batch.generation_requests = [generation] + batch.paused_requests = [paused] + semantic_lists = ( + batch.encoder_requests, + batch.context_requests_chunking, + batch.context_requests_last_chunk, + batch.generation_requests, + batch.paused_requests, + ) + semantic_snapshot = vars(context_1).copy() + + graph_batch, promoted_ids = _make_single_token_context_graph_batch( + batch) + + self.assertIsNot(graph_batch, batch) + self.assertEqual(graph_batch.context_requests, []) + self.assertEqual(graph_batch.generation_requests, + [context_0, context_1, generation]) + self.assertEqual(graph_batch.paused_requests, [paused]) + graph_lists = ( + graph_batch.encoder_requests, + graph_batch.context_requests_chunking, + graph_batch.context_requests_last_chunk, + graph_batch.generation_requests, + graph_batch.paused_requests, + ) + for semantic_list, graph_list in zip(semantic_lists, graph_lists): + self.assertIsNot(semantic_list, graph_list) + self.assertEqual(promoted_ids, frozenset({10, 11})) + self.assertEqual(vars(context_1), semantic_snapshot) + + graph_batch.generation_requests.append(object()) + self.assertEqual(batch.context_requests_last_chunk, + [context_0, context_1]) + self.assertEqual(batch.generation_requests, [generation]) + + def test_structural_fallbacks_return_semantic_batch(self) -> None: + context = _make_request_stub(1) + + encoder_batch = ScheduledRequests() + encoder_batch.encoder_requests = [object()] + encoder_batch.context_requests_last_chunk = [context] + graph_batch, promoted_ids = _make_single_token_context_graph_batch( + encoder_batch) + self.assertIs(graph_batch, encoder_batch) + self.assertFalse(promoted_ids) + + chunking_batch = ScheduledRequests() + chunking_batch.context_requests_chunking = [context] + graph_batch, promoted_ids = _make_single_token_context_graph_batch( + chunking_batch) + self.assertIs(graph_batch, chunking_batch) + self.assertFalse(promoted_ids) + + def test_context_shape_and_mode_fallback_matrix(self) -> None: + cases = ( + ("multi_token_chunk", "context_chunk_size", 2), + ("more_context_remaining", "context_remaining_length", 2), + ("cursor_prompt_mismatch", "py_prompt_len", 5), + ("beam", "py_beam_width", 2), + ("draft", "py_draft_tokens", [9]), + ("first_draft", "py_is_first_draft", True), + ("context_only", "is_context_only_request", True), + ("disaggregated", "py_disaggregated_params", object()), + ("multimodal", "py_multimodal_data", {}), + ("multimodal_event", "py_mm_encoder_event", object()), + ) + for name, attribute, value in cases: + with self.subTest(name=name): + context = _make_request_stub(1) + setattr(context, attribute, value) + batch = ScheduledRequests() + batch.context_requests_last_chunk = [context] + + graph_batch, promoted_ids = \ + _make_single_token_context_graph_batch(batch) + + self.assertIs(graph_batch, batch) + self.assertFalse(promoted_ids) + + def test_context_logits_use_final_token_graph_candidate(self) -> None: + context = _make_request_stub(1) + context.py_return_context_logits = True + batch = ScheduledRequests() + batch.context_requests_last_chunk = [context] + + graph_batch, promoted_ids = _make_single_token_context_graph_batch( + batch) + + self.assertIsNot(graph_batch, batch) + self.assertEqual(graph_batch.generation_requests, [context]) + self.assertEqual(promoted_ids, frozenset({context.py_request_id})) + + def test_generation_only_request_in_context_list_falls_back(self) -> None: + context = _make_request_stub(1) + context.is_generation_only_request = lambda: True + batch = ScheduledRequests() + batch.context_requests_last_chunk = [context] + graph_batch, promoted_ids = _make_single_token_context_graph_batch( + batch) + self.assertIs(graph_batch, batch) + self.assertFalse(promoted_ids) + + def test_generation_shape_fallback_matrix(self) -> None: + cases = ( + ("beam", "py_beam_width", 2), + ("draft", "py_draft_tokens", [9]), + ("first_draft", "py_is_first_draft", True), + ("disaggregated", "py_disaggregated_params", object()), + ) + for name, attribute, value in cases: + with self.subTest(name=name): + context = _make_request_stub(1) + generation = _make_request_stub(2) + setattr(generation, attribute, value) + batch = ScheduledRequests() + batch.context_requests_last_chunk = [context] + batch.generation_requests = [generation] + + graph_batch, promoted_ids = \ + _make_single_token_context_graph_batch(batch) + + self.assertIs(graph_batch, batch) + self.assertFalse(promoted_ids) + + def test_mixed_one_and_two_token_contexts_fall_back_together(self) -> None: + one_token = _make_request_stub(1) + two_tokens = _make_request_stub(2) + two_tokens.context_current_position -= 1 + two_tokens.context_remaining_length = 2 + two_tokens.context_chunk_size = 2 + batch = ScheduledRequests() + batch.context_requests_last_chunk = [one_token, two_tokens] + + graph_batch, promoted_ids = _make_single_token_context_graph_batch( + batch) + + self.assertIs(graph_batch, batch) + self.assertFalse(promoted_ids) + self.assertEqual(batch.context_requests_last_chunk, + [one_token, two_tokens]) + + def test_mrope_delta_is_supported_by_decode_provider(self) -> None: + context = _make_request_stub(1) + context.py_mrope_position_delta = object() + batch = ScheduledRequests() + batch.context_requests_last_chunk = [context] + + graph_batch, promoted_ids = _make_single_token_context_graph_batch( + batch) + + self.assertIsNot(graph_batch, batch) + self.assertEqual(graph_batch.generation_requests, [context]) + self.assertEqual(promoted_ids, frozenset({context.py_request_id})) + + def test_multimodal_context_requires_compatible_decode_token(self) -> None: + context = _make_request_stub(1) + context.py_multimodal_data = {} + batch = ScheduledRequests() + batch.context_requests_last_chunk = [context] + + graph_batch, promoted_ids = _make_single_token_context_graph_batch( + batch) + self.assertIs(graph_batch, batch) + self.assertFalse(promoted_ids) + + incompatible = Mock(return_value=False) + graph_batch, promoted_ids = _make_single_token_context_graph_batch( + batch, incompatible) + self.assertIs(graph_batch, batch) + self.assertFalse(promoted_ids) + incompatible.assert_called_once_with(context) + + compatible = Mock(return_value=True) + graph_batch, promoted_ids = _make_single_token_context_graph_batch( + batch, compatible) + self.assertIsNot(graph_batch, batch) + self.assertEqual(graph_batch.generation_requests, [context]) + self.assertEqual(promoted_ids, frozenset({context.py_request_id})) + compatible.assert_called_once_with(context) + + def test_multimodal_pending_event_is_rechecked(self) -> None: + context = _make_request_stub(1) + context.py_multimodal_data = {} + context.py_mm_encoder_event = object() + batch = ScheduledRequests() + batch.context_requests_last_chunk = [context] + compatible = Mock(return_value=True) + + graph_batch, promoted_ids = _make_single_token_context_graph_batch( + batch, compatible) + self.assertIs(graph_batch, batch) + self.assertFalse(promoted_ids) + compatible.assert_not_called() + + context.py_mm_encoder_event = None + graph_batch, promoted_ids = _make_single_token_context_graph_batch( + batch, compatible) + self.assertIsNot(graph_batch, batch) + self.assertEqual(promoted_ids, frozenset({context.py_request_id})) + compatible.assert_called_once_with(context) + + def test_multimodal_decode_compatibility_uses_final_prompt_token( + self) -> None: + engine = object.__new__(PyTorchModelEngine) + engine.model = SimpleNamespace( + config=SimpleNamespace(vocab_size=100), + mm_token_ids=torch.tensor([99], dtype=torch.int32), + ) + request = _create_request_with_tokens([11, 99, 22], 1) + + request.context_current_position = 2 + self.assertTrue( + engine._is_final_multimodal_context_decode_compatible(request)) + + request.context_current_position = 1 + self.assertFalse( + engine._is_final_multimodal_context_decode_compatible(request)) + + engine.model.mm_token_ids = None + request = _create_request_with_tokens([11, 100], 2) + request.context_current_position = 1 + self.assertFalse( + engine._is_final_multimodal_context_decode_compatible(request)) + + engine.model.mm_token_ids = torch.tensor([99], dtype=torch.int32) + engine.model.model_config = SimpleNamespace( + pretrained_config=SimpleNamespace(rope_scaling={"type": "mrope"})) + request = _create_request_with_tokens([11, 22], 3) + request.context_current_position = 1 + request.py_multimodal_data = {"mrope_config": {}} + self.assertTrue( + engine._is_final_multimodal_context_decode_compatible(request)) + + request.py_multimodal_data["multimodal_embedding"] = object() + self.assertFalse( + engine._is_final_multimodal_context_decode_compatible(request)) + + request.py_multimodal_data["mrope_config"][ + "mrope_position_deltas"] = object() + self.assertTrue( + engine._is_final_multimodal_context_decode_compatible(request)) + + def test_sparse_sequence_mode_uses_promoted_context_cursor(self) -> None: + sparse_config = Mock(spec=SeqLenAwareSparseAttentionConfig) + sparse_config.needs_separate_short_long_cuda_graphs.return_value = True + sparse_config.seq_len_threshold = 16 + runner = object.__new__(CUDAGraphRunner) + runner.sparse_config = sparse_config + runner.spec_config = None + runner.graphs = {} + runner.graph_outputs = {} + runner.graph_metadata = {} + runner.padding_dummy_requests = {} + runner.memory_pool = None + + request = _make_request_stub(7, prompt_len=8) + request.py_batch_idx = 0 + request.max_beam_num_tokens = 64 + batch = ScheduledRequests() + batch.generation_requests = [request] + overlap_state = SimpleNamespace(new_tokens=object()) + + self.assertTrue( + runner._get_seq_len_mode(batch, overlap_state, + frozenset({request.py_request_id}))) + self.assertFalse( + runner._get_seq_len_mode(batch, overlap_state, frozenset())) + + def test_graph_key_forwards_promoted_context_ids(self) -> None: + runner = Mock() + runner.config = SimpleNamespace(is_draft_model=False) + runner._get_seq_len_mode.return_value = True + request = _make_request_stub(7) + batch = ScheduledRequests() + batch.generation_requests = [request] + promoted_ids = frozenset({request.py_request_id}) + + key = CUDAGraphRunner.get_graph_key( + runner, + batch, + new_tensors_device=None, + promoted_context_request_ids=promoted_ids, + ) + + runner._get_seq_len_mode.assert_called_once_with( + batch, None, promoted_ids) + self.assertEqual(key, (1, 0, False, True, True)) + + def test_graph_lookup_forwards_promoted_context_ids(self) -> None: + runner = Mock() + runner.enabled = True + runner.config = SimpleNamespace( + enable_attention_dp=False, + use_mrope=False, + ) + key = (1, 0, False, True, True) + graph_attn_metadata = object() + graph_spec_metadata = object() + runner.get_graph_key.return_value = key + runner.graphs = {key: object()} + runner.graph_metadata = { + key: { + "attn_metadata": graph_attn_metadata, + "spec_metadata": graph_spec_metadata, + } + } + request = _make_request_stub(7) + batch = ScheduledRequests() + batch.generation_requests = [request] + promoted_ids = frozenset({request.py_request_id}) + + with patch( + "tensorrt_llm._torch.pyexecutor.cuda_graph_runner.ExpertStatistic.should_record", + return_value=False): + result = CUDAGraphRunner.maybe_get_cuda_graph( + runner, + batch, + enable_spec_decode=False, + attn_metadata=object(), + promoted_context_request_ids=promoted_ids, + ) + + runner.get_graph_key.assert_called_once_with(batch, None, None, None, + promoted_ids) + self.assertEqual(result, + (graph_attn_metadata, graph_spec_metadata, key)) + + def test_forward_commits_candidate_only_on_graph_hit(self) -> None: + key = (2, 0, False, False, True) + engine, runner, resource_manager, _, outputs = \ + _make_forward_only_engine(key) + context = _make_request_stub(1) + generation = _make_request_stub(2) + batch = ScheduledRequests() + batch.context_requests_last_chunk = [context] + batch.generation_requests = [generation] + event = Mock() + + with patch( + "tensorrt_llm._torch.pyexecutor.model_engine.torch.cuda.Event", + return_value=event): + actual_outputs = engine.forward(batch, resource_manager) + + self.assertIs(actual_outputs, outputs) + graph_batch = runner.maybe_get_cuda_graph.call_args.args[0] + self.assertIsNot(graph_batch, batch) + self.assertEqual(graph_batch.generation_requests, [context, generation]) + prepare_args = engine._prepare_inputs.call_args.args + self.assertIs(prepare_args[0], graph_batch) + self.assertEqual(prepare_args[-1], frozenset({1})) + runner.replay.assert_called_once_with(key, {"prepared": True}) + engine._forward_step.assert_not_called() + engine._execute_logit_post_processors.assert_called_once_with( + batch, outputs) + self.assertEqual(engine.iter_states['num_ctx_requests'], 1) + self.assertEqual(engine.iter_states['num_ctx_tokens'], 1) + self.assertEqual(engine.iter_states['num_generation_tokens'], 1) + event.record.assert_called_once() + + def test_forward_graph_miss_uses_semantic_eager_batch(self) -> None: + engine, runner, resource_manager, semantic_attn_metadata, outputs = \ + _make_forward_only_engine(None) + context = _make_request_stub(1) + generation = _make_request_stub(2) + batch = ScheduledRequests() + batch.context_requests_last_chunk = [context] + batch.generation_requests = [generation] + + with patch( + "tensorrt_llm._torch.pyexecutor.model_engine.torch.cuda.Event", + return_value=Mock()): + actual_outputs = engine.forward(batch, resource_manager) + + self.assertIs(actual_outputs, outputs) + graph_batch = runner.maybe_get_cuda_graph.call_args.args[0] + self.assertIsNot(graph_batch, batch) + prepare_args = engine._prepare_inputs.call_args.args + self.assertIs(prepare_args[0], batch) + self.assertIs(prepare_args[2], semantic_attn_metadata) + self.assertEqual(prepare_args[-1], frozenset()) + engine._forward_step.assert_called_once() + runner.replay.assert_not_called() + engine._execute_logit_post_processors.assert_called_once_with( + batch, outputs) + + def test_zero_runtime_draft_speculation_commits_graph_candidate( + self) -> None: + key = (2, 0, False, False, True) + engine, runner, resource_manager, semantic_attn_metadata, outputs = \ + _make_forward_only_engine(key) + engine.enable_spec_decode = True + graph_attn_metadata = runner.maybe_get_cuda_graph.return_value[0] + runner.maybe_get_cuda_graph.return_value = ( + graph_attn_metadata, + engine.spec_metadata, + key, + ) + context = _make_request_stub(1) + generation = _make_request_stub(2) + batch = ScheduledRequests() + batch.context_requests_last_chunk = [context] + batch.generation_requests = [generation] + + with patch( + "tensorrt_llm._torch.pyexecutor.model_engine.torch.cuda.Event", + return_value=Mock()): + actual_outputs = engine.forward(batch, resource_manager) + + self.assertIs(actual_outputs, outputs) + graph_batch = runner.maybe_get_cuda_graph.call_args.args[0] + self.assertEqual(graph_batch.generation_requests, [context, generation]) + self.assertTrue( + runner.maybe_get_cuda_graph.call_args.kwargs["enable_spec_decode"]) + engine.spec_metadata.update_is_all_greedy_sample.assert_called_once_with( + graph_batch.all_requests()) + prepare_args = engine._prepare_inputs.call_args.args + self.assertIs(prepare_args[0], graph_batch) + self.assertIs(prepare_args[3], engine.spec_metadata) + self.assertEqual(prepare_args[-1], frozenset({context.py_request_id})) + semantic_attn_metadata.update_spec_dec_param.assert_called_once() + self.assertEqual( + semantic_attn_metadata.update_spec_dec_param.call_args. + kwargs["num_contexts"], 1) + runner.replay.assert_called_once_with(key, {"prepared": True}) + + def test_zero_runtime_draft_speculation_graph_miss_is_semantic_eager( + self) -> None: + engine, runner, resource_manager, semantic_attn_metadata, outputs = \ + _make_forward_only_engine(None) + engine.enable_spec_decode = True + context = _make_request_stub(1) + batch = ScheduledRequests() + batch.context_requests_last_chunk = [context] + + with patch( + "tensorrt_llm._torch.pyexecutor.model_engine.torch.cuda.Event", + return_value=Mock()): + actual_outputs = engine.forward(batch, resource_manager) + + self.assertIs(actual_outputs, outputs) + graph_batch = runner.maybe_get_cuda_graph.call_args.args[0] + self.assertEqual(graph_batch.generation_requests, [context]) + prepare_args = engine._prepare_inputs.call_args.args + self.assertIs(prepare_args[0], batch) + self.assertIs(prepare_args[2], semantic_attn_metadata) + self.assertIs(prepare_args[3], engine.spec_metadata) + self.assertEqual(prepare_args[-1], frozenset()) + engine._forward_step.assert_called_once() + runner.replay.assert_not_called() + + def test_forward_allows_guided_context_logits_on_graph_hit(self) -> None: + key = (1, 0, False, False, True) + engine, runner, resource_manager, _, outputs = \ + _make_forward_only_engine(key) + engine.guided_decoder = Mock() + context = _make_request_stub(1) + context.py_return_context_logits = True + batch = ScheduledRequests() + batch.context_requests_last_chunk = [context] + + with patch( + "tensorrt_llm._torch.pyexecutor.model_engine.torch.cuda.Event", + return_value=Mock()): + actual_outputs = engine.forward(batch, + resource_manager, + gather_context_logits=True) + + self.assertIs(actual_outputs, outputs) + graph_batch = runner.maybe_get_cuda_graph.call_args.args[0] + self.assertEqual(graph_batch.generation_requests, [context]) + prepare_args = engine._prepare_inputs.call_args.args + self.assertIs(prepare_args[0], graph_batch) + self.assertEqual(prepare_args[-1], frozenset({context.py_request_id})) + runner.replay.assert_called_once_with(key, {"prepared": True}) + + def test_multimodal_graph_miss_preserves_semantic_payload(self) -> None: + engine, runner, resource_manager, _, _ = _make_forward_only_engine(None) + engine.model.config = SimpleNamespace(vocab_size=100) + engine.model.mm_token_ids = torch.tensor([99], dtype=torch.int32) + context = _make_request_stub(1, prompt_len=3) + context.get_tokens = Mock(return_value=[99, 11, 22]) + multimodal_data = { + "multimodal_embedding": object(), + "mrope_config": { + "mrope_position_deltas": object() + }, + } + context.py_multimodal_data = multimodal_data + batch = ScheduledRequests() + batch.context_requests_last_chunk = [context] + + with patch( + "tensorrt_llm._torch.pyexecutor.model_engine.torch.cuda.Event", + return_value=Mock()): + engine.forward(batch, resource_manager) + + graph_batch = runner.maybe_get_cuda_graph.call_args.args[0] + self.assertIsNot(graph_batch, batch) + self.assertEqual(graph_batch.generation_requests, [context]) + self.assertIs(engine._prepare_inputs.call_args.args[0], batch) + self.assertIs(context.py_multimodal_data, multimodal_data) + self.assertIn("multimodal_embedding", multimodal_data) + + def test_generation_only_forward_does_not_call_new_selector(self) -> None: + key = (1, 0, False, False, True) + engine, runner, resource_manager, _, _ = _make_forward_only_engine(key) + generation = _make_request_stub(2) + batch = ScheduledRequests() + batch.generation_requests = [generation] + + with patch( + "tensorrt_llm._torch.pyexecutor.model_engine._make_single_token_context_graph_batch" + ) as selector, patch( + "tensorrt_llm._torch.pyexecutor.model_engine.torch.cuda.Event", + return_value=Mock()): + engine.forward(batch, resource_manager) + + selector.assert_not_called() + self.assertIs(runner.maybe_get_cuda_graph.call_args.args[0], batch) + self.assertIs(engine._prepare_inputs.call_args.args[0], batch) + self.assertEqual(engine._prepare_inputs.call_args.args[-1], frozenset()) + + def test_global_incompatibilities_bypass_candidate_selection(self) -> None: + cases = ( + "graphs_disabled", + "speculative_nonzero_runtime_draft", + "speculative_draft_model", + "beam", + "encoder_decoder", + "encode_only", + "mm_encoder_only", + "context_parallel", + ) + for case in cases: + with self.subTest(case=case): + engine, runner, resource_manager, _, _ = \ + _make_forward_only_engine(None) + gather_context_logits = False + if case == "graphs_disabled": + runner.enabled = False + elif case == "speculative_nonzero_runtime_draft": + engine.enable_spec_decode = True + engine.runtime_draft_len = 1 + elif case == "speculative_draft_model": + engine.enable_spec_decode = True + engine.is_draft_model = True + elif case == "beam": + engine.max_beam_width = 2 + elif case == "encoder_decoder": + engine._is_encoder_decoder_model.return_value = True + elif case == "encode_only": + engine._is_encode_only = True + elif case == "mm_encoder_only": + engine.llm_args.mm_encoder_only = True + elif case == "context_parallel": + engine.mapping.cp_size = 2 + + batch = ScheduledRequests() + batch.context_requests_last_chunk = [_make_request_stub(1)] + with patch( + "tensorrt_llm._torch.pyexecutor.model_engine._make_single_token_context_graph_batch" + ) as selector, patch( + "tensorrt_llm._torch.pyexecutor.model_engine.torch.cuda.Event", + return_value=Mock()): + engine.forward( + batch, + resource_manager, + gather_context_logits=gather_context_logits, + ) + + selector.assert_not_called() + self.assertIs(engine._prepare_inputs.call_args.args[0], batch) + + class PyTorchModelEngineTestCase(unittest.TestCase): def test_build_request_multimodal_input_skips_when_cache_disabled( @@ -188,6 +911,131 @@ def test_spec_decode_capture_restores_kv_lens_between_warmups(self) -> None: self.assertEqual(attn_metadata.on_update_kv_lens.call_count, 2) + def test_promoted_context_uses_prompt_token_during_overlap(self) -> None: + model_engine, kv_cache_manager = create_model_engine_and_kvcache() + resource_manager = ResourceManager( + {ResourceManagerType.KV_CACHE_MANAGER: kv_cache_manager}) + attn_metadata = AttentionMetadata(max_num_requests=4, + max_num_tokens=32, + kv_cache_manager=kv_cache_manager) + attn_metadata.is_cuda_graph = False + + context = _create_request_with_tokens([11, 22, 33, 44], 1) + context.context_current_position = 3 + context.context_chunk_size = 1 + context.py_seq_slot = 0 + context.py_batch_idx = 3 + context.py_num_compressed_tokens = 1 + + generation = _create_request_with_tokens([50, 51, 52, 53, 54], 2) + generation.py_seq_slot = 1 + generation.py_batch_idx = 1 + + graph_batch = ScheduledRequests() + graph_batch.generation_requests = [context, generation] + new_tokens = torch.zeros((1, 4, 1), dtype=torch.int32, device="cuda") + new_tokens[0, 0, 0] = 999 + new_tokens[0, 1, 0] = 777 + overlap_state = SimpleNamespace(new_tokens=new_tokens) + model_engine._can_use_incremental_update = Mock(return_value=True) + model_engine._can_use_steady_gen_fast_prepare = Mock(return_value=True) + + inputs, _ = model_engine._prepare_tp_inputs( + scheduled_requests=graph_batch, + kv_cache_manager=kv_cache_manager, + attn_metadata=attn_metadata, + new_tensors_device=overlap_state, + resource_manager=resource_manager, + promoted_context_request_ids=frozenset({context.py_request_id}), + ) + + self.assertEqual(inputs["input_ids"][:2].cpu().tolist(), [44, 777]) + self.assertEqual(inputs["position_ids"][0, :2].cpu().tolist(), [3, 5]) + self.assertEqual( + attn_metadata.kv_cache_params.num_cached_tokens_per_seq, [2, 5]) + self.assertEqual(context.cached_tokens, 3) + model_engine._can_use_incremental_update.assert_not_called() + model_engine._can_use_steady_gen_fast_prepare.assert_not_called() + self.assertEqual( + model_engine.previous_batch_indices_cuda[:1].cpu().tolist(), [1]) + self.assertEqual(attn_metadata.num_contexts, 0) + self.assertEqual(model_engine.previous_request_ids, + [generation.py_request_id]) + kv_cache_manager.shutdown() + + def test_promoted_context_precedes_speculative_overlap_generation( + self) -> None: + spec_config = SADecodingConfig( + max_draft_len=1, + draft_len_schedule={1: 1}, + ) + model_engine, kv_cache_manager = create_model_engine_and_kvcache( + spec_config=spec_config) + model_engine.runtime_draft_len = 0 + resource_manager = ResourceManager( + {ResourceManagerType.KV_CACHE_MANAGER: kv_cache_manager}) + attn_metadata = AttentionMetadata(max_num_requests=4, + max_num_tokens=32, + kv_cache_manager=kv_cache_manager) + attn_metadata.is_cuda_graph = False + spec_metadata = Mock() + + context = _create_request_with_tokens([11, 22, 33, 44], 1) + context.context_current_position = 3 + context.context_chunk_size = 1 + context.py_seq_slot = 0 + # A promoted context row must ignore any stale overlap slot. + context.py_batch_idx = 3 + + generation = _create_request_with_tokens([50, 51, 52, 53, 54], 2) + generation.py_seq_slot = 1 + generation.py_batch_idx = 1 + generation.py_needs_onehot_draft_probs = True + + graph_batch = ScheduledRequests() + graph_batch.generation_requests = [context, generation] + new_tokens = torch.zeros((1, 4, 1), dtype=torch.int32, device="cuda") + new_tokens[0, 0, 0] = 999 + new_tokens[0, 1, 0] = 777 + overlap_state = SampleStateTensorsSpec( + new_tokens=new_tokens, + new_tokens_lens=torch.ones(4, dtype=torch.int32, device="cuda"), + next_draft_tokens=torch.zeros((4, 1), + dtype=torch.int32, + device="cuda"), + ) + + inputs, _ = model_engine._prepare_tp_inputs( + scheduled_requests=graph_batch, + kv_cache_manager=kv_cache_manager, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + new_tensors_device=overlap_state, + resource_manager=resource_manager, + promoted_context_request_ids=frozenset({context.py_request_id}), + ) + + self.assertEqual(inputs["input_ids"][:2].cpu().tolist(), [44, 777]) + self.assertEqual(inputs["position_ids"][0, :2].cpu().tolist(), [3, 4]) + self.assertEqual(attn_metadata.request_ids, + [context.py_request_id, generation.py_request_id]) + self.assertEqual( + attn_metadata.kv_cache_params.num_cached_tokens_per_seq, [3, 5]) + self.assertEqual( + model_engine.previous_batch_indices_cuda[:1].cpu().tolist(), [1]) + self.assertEqual( + model_engine.previous_pos_id_offsets_cuda[:2].cpu().tolist(), + [0, 1]) + self.assertEqual(attn_metadata.num_contexts, 0) + self.assertEqual(model_engine.previous_request_ids, + [generation.py_request_id]) + self.assertEqual(spec_metadata.request_ids, + [context.py_request_id, generation.py_request_id]) + self.assertFalse(generation.py_needs_onehot_draft_probs) + spec_metadata.write_padding_onehot_draft_probs.assert_called_once_with( + [generation.py_seq_slot], 0) + kv_cache_manager.shutdown() + def test_pad_generation_requests(self) -> None: model_engine, kv_cache_manager = create_model_engine_and_kvcache() resource_manager = ResourceManager( @@ -836,7 +1684,67 @@ def test_prepare_tp_inputs_all_text_only_drops_mrope_deltas(self) -> None: torch.testing.assert_close(position_ids, expected, atol=0, rtol=0) kv_cache_manager.shutdown() - def test_kv_cache_manager_with_execution_stream(self): + def test_promoted_mrope_context_uses_decode_state_contract(self) -> None: + model_engine, kv_cache_manager = create_model_engine_and_kvcache() + model_engine.model.model_config.pretrained_config.rope_scaling = { + "type": "mrope" + } + model_engine.mrope_position_ids_cuda = torch.zeros( + (3, 1, model_engine.max_num_tokens), + dtype=torch.int32, + device="cuda", + ) + resource_manager = ResourceManager( + {ResourceManagerType.KV_CACHE_MANAGER: kv_cache_manager}) + attn_metadata = AttentionMetadata(max_num_requests=4, + max_num_tokens=32, + kv_cache_manager=kv_cache_manager) + attn_metadata.is_cuda_graph = False + + context = _create_request_with_tokens([11, 22, 33, 44], 1) + context.context_current_position = 3 + context.context_chunk_size = 1 + context.py_seq_slot = 0 + context.py_batch_idx = 3 + mrope_delta = torch.tensor([[10]], dtype=torch.int32) + context.py_mrope_position_delta = mrope_delta + context.py_mrope_delta_cache_slot = context.py_seq_slot + context.py_multimodal_data = { + "mrope_config": { + "mrope_position_deltas": mrope_delta, + }, + "multimodal_embedding": torch.ones((1, 1), dtype=torch.float16), + } + graph_batch = ScheduledRequests() + graph_batch.generation_requests = [context] + + inputs, _ = model_engine._prepare_tp_inputs( + scheduled_requests=graph_batch, + kv_cache_manager=kv_cache_manager, + attn_metadata=attn_metadata, + resource_manager=resource_manager, + promoted_context_request_ids=frozenset({context.py_request_id}), + ) + + self.assertEqual(inputs["input_ids"][:1].cpu().tolist(), [44]) + expected_positions = torch.full((3, 1, 1), + 13, + dtype=torch.int32, + device="cuda") + torch.testing.assert_close(inputs["position_ids"], + expected_positions, + atol=0, + rtol=0) + self.assertEqual( + attn_metadata.kv_cache_params.num_cached_tokens_per_seq, [3]) + self.assertEqual(inputs["mrope_delta_read_seq_slots"].cpu().tolist(), + [0]) + self.assertNotIn("mrope_delta_write_seq_slots", inputs) + self.assertEqual(attn_metadata.num_contexts, 0) + self.assertEqual(model_engine.previous_request_ids, []) + kv_cache_manager.shutdown() + + def test_kv_cache_manager_with_execution_stream(self) -> None: """Test that KVCacheManager uses the provided execution_stream. """ # Create a dedicated execution stream @@ -870,6 +1778,39 @@ def test_kv_cache_manager_with_execution_stream(self): kv_cache_manager.shutdown() + def test_cuda_graph_replay_observes_execution_stream_dependency( + self) -> None: + """A graph replay on the KV manager stream waits for restored KV data.""" + execution_stream = torch.cuda.Stream() + transfer_stream = torch.cuda.Stream() + _, kv_cache_manager = create_model_engine_and_kvcache( + execution_stream=execution_stream) + + source = torch.zeros(1, dtype=torch.int32, device="cuda") + observed = torch.zeros_like(source) + graph = torch.cuda.CUDAGraph() + torch.cuda.synchronize() + with torch.cuda.graph(graph, stream=execution_stream): + observed.copy_(source) + + ready = torch.cuda.Event() + with torch.cuda.stream(transfer_stream): + # Model the async host-to-device restore completed by the local + # offload manager. Recording the event after the write is the same + # dependency shape that refreshBlocks/resume installs. + source.fill_(7) + ready.record() + + manager_stream = torch.cuda.ExternalStream( + kv_cache_manager._stream.cuda_stream) + with torch.cuda.stream(manager_stream): + manager_stream.wait_event(ready) + graph.replay() + torch.cuda.synchronize() + + self.assertEqual(observed.item(), 7) + kv_cache_manager.shutdown() + if __name__ == "__main__": unittest.main() diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_behavior.py b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_behavior.py index 9c130f12977e..077f4cda20f0 100644 --- a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_behavior.py +++ b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_behavior.py @@ -14,6 +14,7 @@ # limitations under the License. from dataclasses import dataclass, field +from types import SimpleNamespace import pytest import torch @@ -279,6 +280,44 @@ def _assert_request_stats( assert request.missed_blocks == missed +@pytest.mark.parametrize( + ("is_active", "resume_succeeds", "expected_result", "expected_calls"), + [ + (True, False, True, []), + (False, False, False, ["resume"]), + (False, True, True, ["resume", "restore"]), + ], +) +def test_v2_resume_restores_offsets_only_after_execution_stream_ready( + is_active: bool, + resume_succeeds: bool, + expected_result: bool, + expected_calls: list[str], +) -> None: + """V2 restores page tables only after resume joins the execution stream.""" + manager = object.__new__(KVCacheManagerV2) + execution_stream = object() + manager._stream = SimpleNamespace(cuda_stream=execution_stream) + calls: list[str] = [] + + def resume(actual_stream: object) -> bool: + assert actual_stream is execution_stream + calls.append("resume") + return resume_succeeds + + kv_cache = SimpleNamespace(is_active=is_active, resume=resume) + + def restore(request_id: int, actual_cache: object) -> None: + assert request_id == 17 + assert actual_cache is kv_cache + calls.append("restore") + + manager._restore_page_index_bufs = restore + + assert manager._resume_and_restore(17, kv_cache) is expected_result + assert calls == expected_calls + + def _run_v1_context(manager: KVCacheManagerV1, request: LlmRequest): batch = _context_batch(request) manager.prepare_resources(batch) From d7f328721902556cc12e103135627a5cbef1c1ce Mon Sep 17 00:00:00 2001 From: Simeng Liu Date: Tue, 28 Jul 2026 10:50:15 -0700 Subject: [PATCH 2/4] Resolve coderabbit's comment. Signed-off-by: Simeng Liu --- tensorrt_llm/_torch/pyexecutor/model_engine.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 76b72881771b..9680fe43698f 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -4472,14 +4472,14 @@ def append_cross_attention_state(request: LlmRequest, request_ids.append(request.py_request_id) is_promoted_context = (request.py_request_id in promoted_context_request_ids) - # the request has no previous tensor: - # (1) new_tokens_device is None, which means overlap scheduler is disabled; or - # (2) a dummy request; or - # (3) the first step in the generation server of disaggregated serving if is_promoted_context: input_ids.append( request.get_tokens(0)[request.context_current_position]) past_seen_token_num = request.context_current_position + # The request has no previous tensor: + # (1) new_tokens_device is None, which means overlap scheduler is disabled; or + # (2) a dummy request; or + # (3) the first step in the generation server of disaggregated serving. elif new_tokens_device is None or request.is_dummy or request.py_batch_idx is None: # skip adding input_ids of CUDA graph dummy requests so that new_tokens_device # can be aligned to the correct positions. From d013c7d0bf9618feb1cc86ff3eb40838f8681cc5 Mon Sep 17 00:00:00 2001 From: Simeng Liu Date: Tue, 28 Jul 2026 16:10:04 -0700 Subject: [PATCH 3/4] Address Guiju's comments. Signed-off-by: Simeng Liu --- .../_torch/pyexecutor/model_engine.py | 12 +++++-- .../executor/test_pytorch_model_engine.py | 33 +++++++++++++++++++ 2 files changed, 42 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 9680fe43698f..40304ac50ba7 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -6320,15 +6320,21 @@ def forward(self, graph_requests = scheduled_requests promoted_context_request_ids: frozenset[int] = frozenset() + # Non-linear tree input preparation expands runtime_draft_len to the + # total tree width after graph selection. Only linear-tree zero-draft + # iterations can therefore safely reuse a zero-draft graph. + can_promote_spec_decode = (not self.enable_spec_decode + or (not self.is_draft_model + and self.runtime_draft_len == 0 + and self.spec_config is not None + and self.spec_config.is_linear_tree)) # TODO: Generalize these conservative gates as actual-draft, beam, and # context-parallel providers for decoder-only LLMs gain support for # promoted final-context rows. Each relaxation must preserve whole-batch # fallback on graph miss and prove parity with the provider's native # q_len=1 path. Encoder-decoder and non-LLM engines remain out of scope. if (scheduled_requests.num_context_requests > 0 - and self.cuda_graph_runner.enabled - and (not self.enable_spec_decode or - (not self.is_draft_model and self.runtime_draft_len == 0)) + and self.cuda_graph_runner.enabled and can_promote_spec_decode and not self.use_beam_search and not self._is_encoder_decoder_model() and not self._is_encode_only diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine.py b/tests/unittest/_torch/executor/test_pytorch_model_engine.py index 4f37479efb3a..adb3e142c7ac 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py @@ -188,6 +188,7 @@ def _make_forward_only_engine( engine.original_max_draft_len = 0 engine.original_max_total_draft_tokens = 0 engine._spec_dec_max_total_draft_tokens = 0 + engine.spec_config = None engine.get_runtime_tokens_per_gen_step = Mock(return_value=1) engine.iter_states = {} engine.forward_pass_callable = None @@ -688,6 +689,7 @@ def test_zero_runtime_draft_speculation_commits_graph_candidate( engine, runner, resource_manager, semantic_attn_metadata, outputs = \ _make_forward_only_engine(key) engine.enable_spec_decode = True + engine.spec_config = SimpleNamespace(is_linear_tree=True) graph_attn_metadata = runner.maybe_get_cuda_graph.return_value[0] runner.maybe_get_cuda_graph.return_value = ( graph_attn_metadata, @@ -727,6 +729,7 @@ def test_zero_runtime_draft_speculation_graph_miss_is_semantic_eager( engine, runner, resource_manager, semantic_attn_metadata, outputs = \ _make_forward_only_engine(None) engine.enable_spec_decode = True + engine.spec_config = SimpleNamespace(is_linear_tree=True) context = _make_request_stub(1) batch = ScheduledRequests() batch.context_requests_last_chunk = [context] @@ -747,6 +750,36 @@ def test_zero_runtime_draft_speculation_graph_miss_is_semantic_eager( engine._forward_step.assert_called_once() runner.replay.assert_not_called() + def test_zero_runtime_non_linear_tree_speculation_uses_semantic_eager_batch( + self) -> None: + engine, runner, resource_manager, semantic_attn_metadata, outputs = \ + _make_forward_only_engine(None) + engine.enable_spec_decode = True + engine.spec_config = SimpleNamespace(is_linear_tree=False) + context = _make_request_stub(1) + generation = _make_request_stub(2) + batch = ScheduledRequests() + batch.context_requests_last_chunk = [context] + batch.generation_requests = [generation] + + with patch( + "tensorrt_llm._torch.pyexecutor.model_engine._make_single_token_context_graph_batch" + ) as selector, patch( + "tensorrt_llm._torch.pyexecutor.model_engine.torch.cuda.Event", + return_value=Mock()): + actual_outputs = engine.forward(batch, resource_manager) + + self.assertIs(actual_outputs, outputs) + selector.assert_not_called() + self.assertIs(runner.maybe_get_cuda_graph.call_args.args[0], batch) + prepare_args = engine._prepare_inputs.call_args.args + self.assertIs(prepare_args[0], batch) + self.assertIs(prepare_args[2], semantic_attn_metadata) + self.assertIs(prepare_args[3], engine.spec_metadata) + self.assertEqual(prepare_args[-1], frozenset()) + engine._forward_step.assert_called_once() + runner.replay.assert_not_called() + def test_forward_allows_guided_context_logits_on_graph_hit(self) -> None: key = (1, 0, False, False, True) engine, runner, resource_manager, _, outputs = \ From 7a69bd72541bbacf59ff06320596fafaae22875a Mon Sep 17 00:00:00 2001 From: Simeng Liu Date: Wed, 29 Jul 2026 11:02:32 -0700 Subject: [PATCH 4/4] Fix parameterized CUDA graph test list entries Signed-off-by: Simeng Liu --- .../test_lists/test-db/l0_dgx_h100.yml | 3 ++- tests/integration/test_lists/test-db/l0_h100.yml | 15 ++++++++++----- 2 files changed, 12 insertions(+), 6 deletions(-) diff --git a/tests/integration/test_lists/test-db/l0_dgx_h100.yml b/tests/integration/test_lists/test-db/l0_dgx_h100.yml index 8b34ef705460..3cda7d22ca40 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_h100.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_h100.yml @@ -20,7 +20,8 @@ l0_dgx_h100: - unittest/_torch/multi_gpu -m "not post_merge" TIMEOUT (90) - unittest/_torch/distributed - unittest/_torch/modeling/test_modeling_pixtral.py::test_tensor_parallelism - - kv_cache/test_final_single_token_context_cuda_graph.py::test_final_token_reuse_cuda_graph_tp2 + - kv_cache/test_final_single_token_context_cuda_graph.py::test_final_token_reuse_cuda_graph_tp2[v1] + - kv_cache/test_final_single_token_context_cuda_graph.py::test_final_token_reuse_cuda_graph_tp2[v2] # ------------- Encoder-decoder TP tests --------------- - llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_generate_encoder_decoder_end_to_end[bf16-kv-v1-cuda-graph-on-greedy-tp2-t5-small] - llmapi/test_llm_api_pytorch_bart.py::test_bart_pytorch_generate_encoder_decoder_end_to_end[bf16-kv-v1-cuda-graph-on-greedy-tp2-bart-large-cnn] diff --git a/tests/integration/test_lists/test-db/l0_h100.yml b/tests/integration/test_lists/test-db/l0_h100.yml index f61cb21d7b69..a81017fb8566 100644 --- a/tests/integration/test_lists/test-db/l0_h100.yml +++ b/tests/integration/test_lists/test-db/l0_h100.yml @@ -221,11 +221,16 @@ l0_h100: - test_e2e.py::test_openai_chat_harmony_perf_metrics - test_e2e.py::test_openai_responses - test_e2e.py::test_openai_chat_guided_decoding[meta-llama/Llama-3.1-8B-Instruct] - - kv_cache/test_final_single_token_context_cuda_graph.py::test_final_token_reuse_cuda_graph - - kv_cache/test_final_single_token_context_cuda_graph.py::test_changed_final_token_reuse_cuda_graph - - kv_cache/test_final_single_token_context_cuda_graph.py::test_context_logits_after_final_token_reuse - - kv_cache/test_final_single_token_context_cuda_graph.py::test_guided_decoding_after_final_token_reuse - - kv_cache/test_final_single_token_context_cuda_graph.py::test_zero_runtime_draft_speculation_after_final_token_reuse + - kv_cache/test_final_single_token_context_cuda_graph.py::test_final_token_reuse_cuda_graph[v1] + - kv_cache/test_final_single_token_context_cuda_graph.py::test_final_token_reuse_cuda_graph[v2] + - kv_cache/test_final_single_token_context_cuda_graph.py::test_changed_final_token_reuse_cuda_graph[v1] + - kv_cache/test_final_single_token_context_cuda_graph.py::test_changed_final_token_reuse_cuda_graph[v2] + - kv_cache/test_final_single_token_context_cuda_graph.py::test_context_logits_after_final_token_reuse[v1] + - kv_cache/test_final_single_token_context_cuda_graph.py::test_context_logits_after_final_token_reuse[v2] + - kv_cache/test_final_single_token_context_cuda_graph.py::test_guided_decoding_after_final_token_reuse[v1] + - kv_cache/test_final_single_token_context_cuda_graph.py::test_guided_decoding_after_final_token_reuse[v2] + - kv_cache/test_final_single_token_context_cuda_graph.py::test_zero_runtime_draft_speculation_after_final_token_reuse[v1] + - kv_cache/test_final_single_token_context_cuda_graph.py::test_zero_runtime_draft_speculation_after_final_token_reuse[v2] # ------------- Prefix-aware scheduling E2E tests --------------- - kv_cache/test_prefix_aware_scheduling.py::TestServePrefixAwareScheduling::test_multi_round_qa_shared_prefix_smoke - kv_cache/test_prefix_aware_scheduling.py::TestServePrefixAwareScheduling::test_multi_round_qa_shared_prefix[guaranteed-chunked]