From 435c08e9ee08c6c03b45c418961495e8843e0491 Mon Sep 17 00:00:00 2001 From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> Date: Wed, 20 May 2026 02:00:03 -0700 Subject: [PATCH 1/7] Squashed commit of the following: MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit commit c9933a8f7b67aa7ce14a5f1242a19209c836e05b Author: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Mon May 18 17:58:20 2026 +0800 [None][fix] Scale recurrent-state pool by pp_size for hybrid Mamba models (By Agent) `_calculate_max_num_blocks_for_linear_attention` sizes the recurrent-state slot pool as `max_batch_size + 1`, ignoring `pp_size`. With pipeline parallelism (e.g. `TestNemotronV3Super::test_nvfp4_parallelism[TP4_PP2]`), multiple microbatches are in-flight on the same rank simultaneously, each holding up to `max_batch_size` sequences' Mamba state — so the live-state slot count must be `max_batch_size * pp_size + 1`. The KVCacheManagerV2 path already does this (`max_num_sequences = max_batch_size * pp_size`); this commit aligns the linear-attention sizing path with that invariant. Three call sites are fixed: 1. `KVCacheManager._calculate_max_num_blocks_for_linear_attention` — both `intercept` (memory reservation in the affine model) and `max_snapshots` (slot count) now scale with `pp_size`. 2. The `is_linear_attention` branch in the estimation dry-run code path, which previously used a bare `max_batch_size`. 3. `CppMambaHybridCacheManager.get_cache_size_per_token` — the per-request fixed-cost intercept that drives upstream `_tokens_for_budget`. Without this fix, `PP > 1` runs trip `No free block found. This shouldn't happen!` in `evictionPolicy::getFreeBlock` once requests beyond the first microbatch arrive at the pool. Added a regression unit test that asserts `recurrent_primary >= max_batch_size * pp_size + 1` for `pp_size ∈ {2, 4}`, and verified the fix end-to-end with the TP4_PP2 model run on a TP2_PP2 configuration locally (test passed in 8m53s, no `No free block` error). Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> commit 7a6e90fd9b0172c843f02454418be4eb1bfd312e Author: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Mon May 18 01:45:33 2026 +0800 [None][fix] Cap CppMamba warmup tokens by recurrent-state pool capacity (By Agent) `_create_cuda_graph_warmup_request` constructs the longest dummy request using `available_tokens` returned by `kv_cache_manager.get_num_available_tokens`. For `CppMambaHybridCacheManager`, the parent's implementation only bounds this by the attention KV cache, so the long dummy's `token_num` could balloon to near `max_seq_len` (262K for Qwen3.5-397B-A17B). With block reuse on and `mamba_state_cache_interval=256`, that one dummy then needs ~1024 recurrent-state blocks, exhausting the RS pool when stacked with the (batch_size - 1) short dummies in the same warmup batch and tripping `No free block found` in `evictionPolicy::getFreeBlock`. Override `get_num_available_tokens` in CppMambaHybridCacheManager to also clamp by `(rs_free_blocks - 1) * states_snapshot_interval` whenever the snapshot interval is positive, leaving one block of headroom for the always-allocated trailing snapshot. The override is a no-op when `enable_block_reuse=False` (interval == 0). Verified with `tests/integration/defs/accuracy/test_llm_api_pytorch.py::TestQwen3_5_397B_A17B::test_nvfp4[tep4_block_reuse]`. Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> commit 717edebd1702faa401c330e2ac0cbe5e2bafbd29 Author: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Sat May 16 16:27:02 2026 +0800 fix schedule issue, improve runtime perf, allow change manager Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> commit 181321244e9ba7aedceb3d670e2bbf1e3f7b80e3 Author: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Fri May 15 08:47:09 2026 +0800 [None][fix] Address recurrent-state pool sizing and layer-first transfer bugs (By Agent) Concern 1 (kvCacheTransferManager.cpp): In the layer-first cudaMemcpy2DAsync path, source and destination pitches were both derived from pool.primaryPtr. For host-offload transfers, the secondary pool has mNumSecondaryBlocks which can differ from mNumPrimaryBlocks, giving a different per-layer stride. Fix: compute srcLayerStrideBytes and dstLayerStrideBytes independently from srcPool and dstPool, and use them as separate pitches in cudaMemcpy2DAsync. Concern 2 (resource_manager.py): In _calculate_max_num_blocks_for_linear_attention, the block-reuse branch unconditionally overwrote max_snapshots with max_tokens // interval, discarding the live-state + CUDA-graph-padding floor (max_batch_size + 1 + max_draft_len) set earlier. Fix: use max(max_tokens // interval, max_snapshots) to preserve the floor. Add a regression test (enable_block_reuse=True) that would have caught the dropped floor: with max_batch_size=4 and max_tokens=512/interval=256, the naive branch returns 2 slots (< required floor of 5). Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> commit e00c23093c66f2cbd39463b6995e8a6aa4fe7217 Merge: 69899c5576 be4f57ecef Author: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Wed May 13 18:12:00 2026 +0800 Merge branch 'main' into user/xiweny/mamba-recurrent-cuda-graph-snapshots Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> commit 69899c557625b8e11d98856b16a8078807857c05 Author: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Wed May 13 18:01:04 2026 +0800 fix replay check on sm120 Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> commit 3f0e36c88150ec98833f7bd9428e795724b15151 Author: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Wed May 13 17:57:48 2026 +0800 fix missing is_draft Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> commit 516ea84eaaa772c619c13e85f549245685ceb229 Author: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Tue May 12 23:16:54 2026 +0800 reduce sync overhead Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> commit ace03a2deac89c96c5308a07eb2ffe70cfd294c3 Merge: 5a858312ce 5aad39a889 Author: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Tue May 12 21:27:53 2026 +0800 Merge branch 'user/xiweny/mamba-hybrid-zero-mamba-layers' into mamba-recurrent-cuda-graph-snapshots (By Agent) Brings in the CppMambaHybridCacheManager zero-local-mamba-layers fix and its CodeRabbit follow-ups (PR #13999). Conflict in tests/unittest/_torch/executor/test_mamba_cache_manager.py is resolved by keeping both test sections (recurrent-state pool sizing from this branch, zero-local-mamba early-exit from the merged branch) and unioning the import set. Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> commit 5aad39a8893d2c5e39fa0322019d4535d1f0dd0e Author: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Tue May 12 11:15:21 2026 +0800 [None][fix] CppMambaHybridCacheManager: address CodeRabbit review (By Agent) Two fixes from CodeRabbit's review on #13999: 1. Handle layer_mask=None defensively. The signature has always typed it as Optional[List[bool]], and the new constructor was unconditionally calling layer_mask.copy() / indexing before the zero-mamba guard ran. Treat None as an all-False full-attention mask and validate length against mamba_layer_mask. 2. Seed externally visible mamba fields before the early-return guard. get_mamba_ssm_cache_dtype() reads self.ssm_state_dtype, and the use_replay_state_update property reads self._use_replay_state_update; both were only assigned on the non-early-exit path, so any caller that hit those accessors on a zero-local-mamba rank would have AttributeError'd. Move both assignments above the early-return. Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> commit 5a858312ce70155f98918e5f66d5bb958a683977 Author: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Tue May 12 09:41:25 2026 +0800 [None][fix] KVCacheManager: store spec_config on self in __init__ (By Agent) Address CodeRabbit review on #14003: the linear-attention sizing branch reads self.spec_config, but KVCacheManager.__init__ never assigned the incoming spec_config argument to self. The code worked only by accident, because CppMambaHybridCacheManager's _setup_mtp_intermediate_states runs before super().__init__ and writes self.spec_config on the partially initialized instance. Any other KVCacheManager path that exercises the linear-attention branch with a spec_config would AttributeError. Assign self.spec_config = spec_config at the top of KVCacheManager.__init__ so the attribute is always present regardless of subclass. Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> commit 3310de0a0e9730a494819d8166b1f94c831c96a2 Author: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Mon May 11 23:46:42 2026 +0800 [None][fix] KVCacheManager: fix max_draft_tokens typo and add slot reservation tests (By Agent) Follow-up to the recurrent-state slot reservation: spec_config has no max_draft_tokens attribute (DecodingBaseConfig uses StrictBaseModel with extra="forbid"). The correct field is max_draft_len, which is also the count of additional dummy request ids issued by CUDAGraphRunner._get_padded_batch. Add two unit tests that construct a real CppMambaHybridCacheManager and inspect blocks_per_window[RECURRENT_STATES]: - without spec_config: pool reserves max_batch_size + 1 slots - with spec_config(max_draft_len=N): pool reserves max_batch_size + 1 + N Tests disable block reuse so the override (max_tokens // interval) does not mask the addition under test. Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> commit 992ec1bd08b44257d7d78f9e657c8f3c91f7b584 Author: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Mon May 11 23:34:04 2026 +0800 [None][fix] KVCacheManager: reserve recurrent-state slots for CUDA graph padding and draft-len sentinels (By Agent) The recurrent-state snapshot pool was sized at exactly max_batch_size, which leaves no room for the CUDA graph padding sentinel (CUDA_GRAPH_DUMMY_REQUEST_ID) and the per-draft-len sentinels that CUDAGraphRunner::_get_padded_batch issues during speculative decoding. Reserve one extra slot for the CUDA graph padding dummy, and when a spec_config is present, reserve one additional slot per draft length so all distinct sentinel request IDs can coexist without evicting live recurrent state. The block-reuse path (max_tokens // interval) still overrides this when active. Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> commit ac30ded84936293f65bfa3b9da33e020753cc606 Author: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Mon May 11 20:26:14 2026 +0800 [None][fix] CppMambaHybridCacheManager: handle ranks with zero local mamba layers (By Agent) Under PP sharding, a rank can end up with only full-attention layers and no mamba layers. The previous constructor unconditionally allocated SSM / conv state and set up linear-attention metadata, leading to invalid state on these ranks. Take an early-exit after super().__init__ when local_num_mamba_layers == 0: forward num_layers (not mamba_num_layers + num_layers) and the union of mamba+attention layer masks to the parent KVCacheManager, skip all mamba-only setup. Guard prepare_resources / update_mamba_states / _setup_state_indices with the same condition so they no-op on these ranks. Add a regression test that constructs the manager with a real parent KVCacheManager and verifies: early-exit invariants on mgr state, no mamba-only attributes set, and the three guarded methods no-op without raising. Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> --- .../batch_manager/kvCacheManager.h | 54 +++++- .../batch_manager/capacityScheduler.cpp | 2 +- .../batch_manager/kvCacheManager.cpp | 73 ++++++-- .../batch_manager/kvCacheTransferManager.cpp | 40 ++-- .../nanobind/batch_manager/kvCacheManager.cpp | 4 +- .../_torch/modules/mamba/mamba2_metadata.py | 55 +++++- tensorrt_llm/_torch/pyexecutor/_util.py | 28 ++- .../_torch/pyexecutor/mamba_cache_manager.py | 113 +++++++++--- .../_torch/pyexecutor/resource_manager.py | 53 +++++- .../executor/test_mamba_cache_manager.py | 172 +++++++++++++++++- 10 files changed, 515 insertions(+), 79 deletions(-) diff --git a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h index 11978bc9df7f..1a4c1acee3ff 100644 --- a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h +++ b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h @@ -162,6 +162,46 @@ struct LinearAttentionMetadata return false; } + [[nodiscard]] SizeType32 calcNumBlocksNeededForReq( + SizeType32 promptLen, SizeType32 tokensPerBlock, bool enableReuse) const + { + if (!enableReuse) + { + return 1; + } + SizeType32 count = 0; + if (statesSnapshotInterval > 0) + { + count += promptLen / statesSnapshotInterval; // round down + } + if (saveLastSnapshot + && (promptLen / tokensPerBlock * tokensPerBlock + != promptLen / statesSnapshotInterval * statesSnapshotInterval)) + { + count += 1; + } + if (promptLen % tokensPerBlock == 0) + { + // corner case + count += 1; + } + return count; + } + + [[nodiscard]] SizeType32 calcNumAdditionalBlocksNeededForReq( + SizeType32 numTokens, SizeType32 promptLen, SizeType32 tokensPerBlock, bool enableReuse) const + { + if (!enableReuse) + { + return 0; + } + if (promptLen % tokensPerBlock == 0 && numTokens <= promptLen + 1) + { + return 1; + } + return 0; + } + [[nodiscard]] bool hasRecurrentStatesCache() const { return hasRecurrentStatesCache(cacheType); @@ -872,7 +912,9 @@ class WindowBlockManager //! \brief According to request's current position, copy data from the last full block to the next block (ignoring //! the placeholder block). It should be called after every context chunk is processed. - void copyLinearAttentionBlock(GenerationRequest& sequence, LlmRequest const& llmRequest); + //! \return true iff at least one async block transfer was actually issued. Callers can use this to decide + //! whether a subsequent refreshBlocks()/syncTransfers() is necessary. + bool copyLinearAttentionBlock(GenerationRequest& sequence, LlmRequest const& llmRequest); void replaceSharedBlock(GenerationRequest& sequence, SizeType32 blockIdx); @@ -1387,7 +1429,8 @@ class BlockManager //! \brief According to request's current position, copy data from the last full block to the next block (ignoring //! the placeholder block). It should be called after every context chunk is processed. - void copyLinearAttentionBlock(GenerationRequest& sequence, LlmRequest const& llmRequest); + //! \return true iff at least one async block transfer was actually issued. + bool copyLinearAttentionBlock(GenerationRequest& sequence, LlmRequest const& llmRequest); void replaceSharedBlock(GenerationRequest& sequence, SizeType32 windowSize, SizeType32 blockIdx); @@ -2217,7 +2260,12 @@ class KVCacheManager : public BaseKVCacheManager //! \brief According to request's current position, copy data from the last full block to the next block (ignoring //! the placeholder block). It should be called before every forward step, after adding new tokens. - void copyLinearAttentionBlock(LlmRequest const& llmRequest); + //! \return true iff at least one async block transfer was actually issued for this request. The caller can + //! aggregate this across requests and skip refreshBlocks() (which performs a stream sync) when no copies happened. + bool copyLinearAttentionBlock(LlmRequest const& llmRequest); + + //! \brief Batch variant of copyLinearAttentionBlock. Returns true iff at least one copy was issued. + bool copyLinearAttentionBlockBatch(std::vector> const& llmRequests); void addSequenceBatch( std::vector> const& requestInfos, diff --git a/cpp/tensorrt_llm/batch_manager/capacityScheduler.cpp b/cpp/tensorrt_llm/batch_manager/capacityScheduler.cpp index 3549c1a61f13..60b6cc14ebee 100644 --- a/cpp/tensorrt_llm/batch_manager/capacityScheduler.cpp +++ b/cpp/tensorrt_llm/batch_manager/capacityScheduler.cpp @@ -456,7 +456,7 @@ std::tuple MaxUtilizationScheduler::operator()( // Here we simulate freeing the kvCache blocks associated with that sequence kvCacheManager.schedulingRemoveSequence((*lastStartedReqIt)->mRequestId); pausedRequests.emplace_back(*lastStartedReqIt); - TLLM_LOG_DEBUG("MaxUtilizationScheduler: request ID %lu -> pause", (*lastStartedReqIt)->mRequestId); + TLLM_LOG_INFO("MaxUtilizationScheduler: request ID %lu -> pause", (*lastStartedReqIt)->mRequestId); reqItEnd = std::next(lastStartedReqIt).base(); } else diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp index 11b1d1694eb2..c129a3646214 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp @@ -2222,12 +2222,16 @@ void BlockManager::allocateBlock(GenerationRequest& sequence, SizeType32 windowS mWindowBlockManagers.at(windowSize).allocateBlock(sequence, false); } -void BlockManager::copyLinearAttentionBlock(GenerationRequest& sequence, LlmRequest const& llmRequest) +bool BlockManager::copyLinearAttentionBlock(GenerationRequest& sequence, LlmRequest const& llmRequest) { + bool didCopy = false; for (auto& [windowSize, manager] : mWindowBlockManagers) { - manager.copyLinearAttentionBlock(sequence, llmRequest); + // Use a temp to avoid short-circuiting; every window must run. + bool const windowDidCopy = manager.copyLinearAttentionBlock(sequence, llmRequest); + didCopy = didCopy || windowDidCopy; } + return didCopy; } bool WindowBlockManager::tryAllocatePlaceholderForLinearAttention(GenerationRequest& sequence, bool shareAmongBeams) @@ -2363,11 +2367,11 @@ void WindowBlockManager::allocateBlock(GenerationRequest& sequence, bool shareAm } } -void WindowBlockManager::copyLinearAttentionBlock(GenerationRequest& sequence, LlmRequest const& request) +bool WindowBlockManager::copyLinearAttentionBlock(GenerationRequest& sequence, LlmRequest const& request) { if (!isRecurrentState()) { - return; + return false; } auto const requestId = request.mRequestId; @@ -2376,7 +2380,7 @@ void WindowBlockManager::copyLinearAttentionBlock(GenerationRequest& sequence, L if (mAllocatedBlocksPerSeq.find(requestId) == mAllocatedBlocksPerSeq.end()) { TLLM_LOG_WARNING("%s::copyLinearAttentionBlock - Request %lu not found", mLogPrefix.c_str(), requestId); - return; + return false; } // It points to the next token to be processed/generated @@ -2391,9 +2395,10 @@ void WindowBlockManager::copyLinearAttentionBlock(GenerationRequest& sequence, L if (TLLM_LIKELY(sequence.getBeamWidth() == 1)) { // the block of beam0 is inherited from context phase, no need to copy - return; + return false; } - // copy beam 0 to other beams + // copy beam 0 to other beams: beamWidth >= 2 here, loop runs at least once + // with no skip path, so an onboard is always issued. auto beam0Block = getBlockById(sequence.getCacheBlockIds(mWindowSize).at(0).back()); for (auto beamIdx = 1; beamIdx < sequence.getBeamWidth(); ++beamIdx) { @@ -2406,17 +2411,18 @@ void WindowBlockManager::copyLinearAttentionBlock(GenerationRequest& sequence, L // transfer manager to copy the entire block. sequence.getTransferMode(), sequence.getDirectory()); } - return; + return true; } // copy only happens in context phase or the corner case above if (currentPosition % mTokensPerBlock != 0 || currentPosition > request.getPromptLen() || currentPosition == 0) { - return; + return false; } auto prevBlockIndex = currentPosition / mTokensPerBlock - 1; // signed std::set> onboardedBlocks; + bool didCopy = false; for (auto beamIdx = 0; beamIdx < sequence.getBeamWidth(); ++beamIdx) { auto const& beamBlockIds = sequence.getCacheBlockIds(mWindowSize).at(beamIdx); @@ -2455,7 +2461,9 @@ void WindowBlockManager::copyLinearAttentionBlock(GenerationRequest& sequence, L // manager to copy the entire block. sequence.getTransferMode(), sequence.getDirectory()); onboardedBlocks.insert({prevBlockId, nextBlockId}); + didCopy = true; } + return didCopy; } std::pair> WindowBlockManager::storeBlocks( @@ -3311,6 +3319,11 @@ SizeType32 KVCacheManager::getNeededBlocksOneStep(LlmRequest const& req, bool tw = std::min((isCrossKv() ? req.getEncoderOutputLen() : req.mPromptLen) + maxDraftTokensToAdd, windowSize + mChunkSize) + mSinkBubbleLength; + if (LinearAttentionMetadata::hasLinearCache(windowSize)) + { + return mBlockManager.getLinearAttentionMetadata()->calcNumBlocksNeededForReq( + promptCacheLen, getTokensPerBlock(), mEnableBlockReuse); + } auto const numSharedBlocks = promptCacheLen / getTokensPerBlock(); auto const numUnSharedTokens = promptCacheLen % getTokensPerBlock(); auto const numUnSharedBlocks @@ -3355,6 +3368,12 @@ SizeType32 KVCacheManager::getNeededBlocksOneStep(LlmRequest const& req, bool tw auto const maxTokensToAdd = std::min((twoStepsLookAhead ? 2 : 1) * tokensPerStep, maxTokensToAddToKVCache); auto const numNextTokens = numCurrTokens + maxTokensToAdd; + if (LinearAttentionMetadata::hasLinearCache(windowSize)) + { + return mBlockManager.getLinearAttentionMetadata()->calcNumAdditionalBlocksNeededForReq( + numCurrTokens, req.getPromptLen(), getTokensPerBlock(), mEnableBlockReuse); + } + if (numNextTokens > mBlockManager.getWindowSizeMetadata(windowSize).maxTokenNum) { return 0; @@ -3389,7 +3408,8 @@ SizeType32 KVCacheManager::getRemainingBlocksToCompletion( { if (req.isGenerationInProgressState()) { - return 0; // no need to allocate blocks for recurrent states during generation + return mBlockManager.getLinearAttentionMetadata()->calcNumAdditionalBlocksNeededForReq( + req.getNumTokens(0), req.getPromptLen(), getTokensPerBlock(), mEnableBlockReuse); } else if (!req.isContextFinished()) { @@ -3399,12 +3419,8 @@ SizeType32 KVCacheManager::getRemainingBlocksToCompletion( { return 0; } - if (mEnableBlockReuse) - { - return req.getPromptLen() / mBlockManager.getLinearAttentionMetadata()->statesSnapshotInterval + 1 - + (mBlockManager.getLinearAttentionMetadata()->saveLastSnapshot ? 1 : 0); - } - return 1; + return mBlockManager.getLinearAttentionMetadata()->calcNumBlocksNeededForReq( + req.mPromptLen, getTokensPerBlock(), mEnableBlockReuse); } } @@ -3548,10 +3564,29 @@ void KVCacheManager::addToken(RequestIdType requestId) mBlockManager.adjustBlocksIfNeeded(sequence); } -void KVCacheManager::copyLinearAttentionBlock(LlmRequest const& llmRequest) +bool KVCacheManager::copyLinearAttentionBlock(LlmRequest const& llmRequest) { + if (!mEnableBlockReuse) + { + // When block reuse is disabled, we always use a single slot for a sequence and move it to + // the end of blocks list, so copying is not needed. + return false; + } auto& sequence = getSequence(llmRequest.mRequestId); - mBlockManager.copyLinearAttentionBlock(sequence, llmRequest); + return mBlockManager.copyLinearAttentionBlock(sequence, llmRequest); +} + +bool KVCacheManager::copyLinearAttentionBlockBatch(std::vector> const& llmRequests) +{ + bool copiedAny = false; + for (auto const& req : llmRequests) + { + if (copyLinearAttentionBlock(*req)) + { + copiedAny = true; + } + } + return copiedAny; } void WindowBlockManager::detachFrontBlock(GenerationRequest& sequence) @@ -3741,7 +3776,7 @@ void KVCacheManager::storeContextBlocks(LlmRequest const& llmRequest) } else { - TLLM_LOG_WARNING("[kv cache manager] storeContextBlocks: Can not find sequence for request %lu", requestId); + // TLLM_LOG_WARNING("[kv cache manager] storeContextBlocks: Can not find sequence for request %lu", requestId); } } diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp index 0b6abf361182..019eb45b8ff4 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp @@ -114,22 +114,38 @@ void KVCacheTransferManager::copyBlock(BlockPtr const& src, BlockPtr const& dst, auto const& pool = pools[poolIdx]; // For layer-first layout pools, block data is non-contiguous across layers. - // Copy each layer's block data separately. + // Pool shape: {numLayers, numBlocks, kvFactor, blockSize}. For a fixed block + // index, per-layer slices are contiguous rows of (kvFactor * blockSize) elements, + // separated by a stride of numBlocks rows between layers. Issue this as a single + // pitched cudaMemcpy2DAsync instead of one cudaMemcpyAsync per layer. if (pool.layerFirstLayout) { auto srcPool = src->isPrimary() ? pool.primaryPtr : pool.secondaryPtr; auto dstPool = dst->isPrimary() ? pool.primaryPtr : pool.secondaryPtr; - auto const srcBlockIdx = static_cast(src->getMemoryPoolBlockIndex()); - auto const dstBlockIdx = static_cast(dst->getMemoryPoolBlockIndex()); - - for (SizeType32 layerIdx = 0; layerIdx < pool.numLayers; ++layerIdx) - { - // pool shape: {numLayers, numBlocks, kvFactor, blockSize} - // slice at {layerIdx, blockIdx} gives {1, kvFactor, blockSize} - auto srcBlock = tr::ITensor::slice(srcPool, {layerIdx, srcBlockIdx}, 1); - auto dstBlock = tr::ITensor::slice(dstPool, {layerIdx, dstBlockIdx}, 1); - (isOffload ? mOffloadManager : mOnboardManager).copy(*srcBlock, *dstBlock); - } + auto const srcBlockIdx = static_cast(src->getMemoryPoolBlockIndex()); + auto const dstBlockIdx = static_cast(dst->getMemoryPoolBlockIndex()); + + // Compute pitches from each pool independently: primary and secondary pools + // may have different block counts (mNumPrimaryBlocks vs mNumSecondaryBlocks), + // so their per-layer strides differ. Using the primary shape for both pitches + // would corrupt host-offloaded recurrent state on CPU<->GPU transfers. + auto const& srcShape = srcPool->getShape(); + auto const& dstShape = dstPool->getShape(); + TLLM_CHECK_WITH_INFO(srcShape.nbDims >= 2, + "Expected layer-first KVCache pool to have at least 2 dims, got %d", srcShape.nbDims); + TLLM_CHECK_WITH_INFO(dstShape.nbDims >= 2, + "Expected layer-first KVCache pool to have at least 2 dims, got %d", dstShape.nbDims); + auto const srcLayerStrideBytes = srcPool->getSizeInBytes() / static_cast(pool.numLayers); + auto const dstLayerStrideBytes = dstPool->getSizeInBytes() / static_cast(pool.numLayers); + // rowBytes is the per-block per-layer payload — identical for primary and secondary. + auto const rowBytes = srcLayerStrideBytes / static_cast(srcShape.d[1]); + + auto* srcBase = static_cast(srcPool->data()) + srcBlockIdx * rowBytes; + auto* dstBase = static_cast(dstPool->data()) + dstBlockIdx * rowBytes; + + auto stream = (isOffload ? mOffloadManager : mOnboardManager).getStream().get(); + TLLM_CUDA_CHECK(cudaMemcpy2DAsync(dstBase, dstLayerStrideBytes, srcBase, srcLayerStrideBytes, rowBytes, + static_cast(pool.numLayers), cudaMemcpyDefault, stream)); continue; } diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp index 79d4d5d1ce32..df340f686627 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp @@ -641,7 +641,9 @@ void tb::kv_cache_manager::KVCacheManagerBindings::initBindings(nb::module_& m) .def_prop_ro( "is_variable_window", [](tbk::KVCacheManager& self) { return self.getBlockManager().isVariableWindow(); }) .def("copy_linear_attention_block", &tbk::KVCacheManager::copyLinearAttentionBlock, nb::arg("llm_request"), - nb::call_guard()); + nb::call_guard()) + .def("copy_linear_attention_block_batch", &tbk::KVCacheManager::copyLinearAttentionBlockBatch, + nb::arg("llm_requests"), nb::call_guard()); } void tb::BasePeftCacheManagerBindings::initBindings(nb::module_& m) diff --git a/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py b/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py index beedeccab165..3d1316e40519 100644 --- a/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py +++ b/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py @@ -213,6 +213,10 @@ def __init__(self, max_batch_size: int, chunk_size: int): self.state_indices = torch.zeros(max_batch_size, dtype=torch.int32, device="cuda") + # Stable data_ptr() of the CUDA tensor we alias (if any) — used to + # detect cache-manager buffer reallocation that would silently break + # CUDA graph replays. + self._state_indices_aliased_ptr = None # Pre-allocated buffers. self._arange_buffer = torch.arange(max_batch_size + 1, @@ -242,10 +246,42 @@ def prepare(self, attn_metadata: AttentionMetadata): ] indices = kv_cache_manager.get_state_indices( batch_request_ids, is_padding) - for i, idx in enumerate(indices): - self.state_indices_cpu[i] = idx - self.state_indices[:batch_size].copy_( - self.state_indices_cpu[:batch_size], non_blocking=True) + if isinstance(indices, + torch.Tensor) and indices.device.type == 'cuda': + # Alias the cache manager's CUDA buffer directly instead of + # copying. Iterating a CUDA tensor and assigning each 0-d + # slice to a CPU tensor would trigger one cudaMemcpyAsync + + # cudaStreamSynchronize per element. + # + # Safe under CUDA graphs only when the source buffer has a + # stable data pointer across all calls (currently true for + # CppMambaHybridCacheManager.cuda_state_indices, allocated + # once in __init__). If a future cache manager reallocates + # this buffer between iterations, captured kernels would + # still read from the address seen at capture time, so we + # assert stability here. + if self._state_indices_aliased_ptr is None: + self._state_indices_aliased_ptr = indices.data_ptr() + else: + assert indices.data_ptr( + ) == self._state_indices_aliased_ptr, ( + "kv_cache_manager.get_state_indices() must return a " + "buffer with a stable data pointer when CUDA graphs " + "are used; got a different address than the first " + "call.") + self.state_indices = indices + elif isinstance(indices, torch.Tensor): + # CPU tensor → bulk H2D + self.state_indices_cpu[:batch_size].copy_(indices[:batch_size]) + self.state_indices[:batch_size].copy_( + self.state_indices_cpu[:batch_size], non_blocking=True) + else: + # indices is a Python sequence (e.g. List[int]); data + # already lives on host, CPU staging is fine. + for i, idx in enumerate(indices): + self.state_indices_cpu[i] = idx + self.state_indices[:batch_size].copy_( + self.state_indices_cpu[:batch_size], non_blocking=True) if num_contexts > 0: torch.cumsum(context_lens, @@ -313,3 +349,14 @@ def prepare(self, attn_metadata: AttentionMetadata): self.query_start_loc = None self.query_start_loc_long = self._arange_buffer_long[:batch_size + 1] + + # Complete any deferred recurrent-state block onboards scheduled by + # CppMambaHybridCacheManager.prepare_resources(). prepare_resources + # only enqueues the async cudaMemcpyAsync calls and sets a pending + # flag; we sync the onboard stream here, so the prior CPU-side prep + # work in _prepare_tp_inputs overlaps with the in-flight transfers. + # Cheap no-op on cache managers without this method or when no + # transfers were scheduled this iteration. + flush = getattr(kv_cache_manager, "flush_state_transfers", None) + if flush is not None: + flush() diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 1095ed03c9f6..29bdc335a88c 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -90,9 +90,29 @@ def get_kv_cache_manager_cls(model_config: ModelConfig, logger.info("Hybrid linear model has 0 mamba layers; using " "KVCacheManager without mamba caching") return _non_hybrid_kv_cache_manager_cls(config, kv_cache_config) + if kv_cache_config.enable_block_reuse: + return CppMambaHybridCacheManager if is_disagg or use_cpp_mamba_cache_manager(): return MixedMambaHybridCacheManager - return CppMambaHybridCacheManager + default_cls = CppMambaHybridCacheManager + env_override = os.environ.get('TLLM_MAMBA_MANAGER_PREFERENCE', None) + if env_override is not None: + if env_override.upper() == 'MIXED': + logger.warning( + "Environment variable TLLM_MAMBA_MANAGER_PREFERENCE=MIXED overrides the default Mamba cache manager to MixedMambaHybridCacheManager. This may lead to increased memory usage due to lack of block reuse, but can be necessary for disaggregated setups or to avoid potential issues with the C++ manager. Set TLLM_MAMBA_MANAGER_PREFERENCE=CPP to use the CppMambaHybridCacheManager instead, which is the default for non-disaggregated setups without block reuse explicitly disabled." + ) + return MixedMambaHybridCacheManager + elif env_override.upper() == 'CPP': + logger.warning( + "Environment variable TLLM_MAMBA_MANAGER_PREFERENCE=CPP overrides the default Mamba cache manager to CppMambaHybridCacheManager. This enables block reuse and can reduce memory usage, but may not be compatible with disaggregated setups. Set TLLM_MAMBA_MANAGER_PREFERENCE=MIXED to use the MixedMambaHybridCacheManager instead if you encounter issues with the C++ manager or are running in a disaggregated environment." + ) + return CppMambaHybridCacheManager + else: + logger.warning( + f"Unrecognized value for TLLM_MAMBA_MANAGER_PREFERENCE: {env_override}. " + f"Expected 'CPP' or 'MIXED'. Using default {default_cls.__name__}." + ) + return default_cls else: return _non_hybrid_kv_cache_manager_cls(config, kv_cache_config) @@ -1262,12 +1282,12 @@ def _create_kv_cache_manager( "using legacy MTP path") use_replay = False - # Replay Philox uses PTX cvt.rs.f16x2.f32 which needs sm >= 100. + # Replay Philox uses PTX cvt.rs.f16x2.f32 which needs 100 <= sm < 120. # Flashinfer has a SW fallback at any SM. if (stochastic_rounding and mamba_params.mamba_ssm_cache_dtype == torch.float16 - and sm < 100): - logger.info("Replay kernel Philox requires sm >= 100; " + and sm < 100 or sm in (120, 121)): + logger.info("Replay kernel Philox requires 100 <= sm < 120; " "using legacy MTP path for stochastic rounding support") use_replay = False diff --git a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py index da682bebc858..194bdc7fa87d 100644 --- a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py @@ -32,7 +32,8 @@ from tensorrt_llm._torch.pyexecutor.resource_manager import ( BaseResourceManager, CacheTypeCpp, DataType, KVCacheManager, get_pp_layers) from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests -from tensorrt_llm._utils import nvtx_range, torch_dtype_to_binding +from tensorrt_llm._utils import (nvtx_range, prefer_pinned, + torch_dtype_to_binding) from tensorrt_llm.bindings.internal.batch_manager import ( LinearAttentionMetadata, LinearCacheType) from tensorrt_llm.llmapi.llm_args import KvCacheConfig @@ -1000,6 +1001,7 @@ def __init__( layer_mask: Optional[ List[bool]] = None, # this is the full attention layer mask is_estimating_kv_cache: bool = False, + is_draft: bool = False, use_replay_state_update: bool = False, **kwargs, ) -> None: @@ -1059,6 +1061,7 @@ def __init__( spec_config=spec_config, layer_mask=full_attention_layer_mask, is_estimating_kv_cache=is_estimating_kv_cache, + is_draft=is_draft, ) return @@ -1147,6 +1150,7 @@ def __init__( spec_config=spec_config, layer_mask=layer_mask, is_estimating_kv_cache=is_estimating_kv_cache, + is_draft=is_draft, linear_attention_metadata=self.linear_attention_metadata, ) @@ -1166,6 +1170,12 @@ def __init__( self.cuda_state_indices = torch.zeros([self.max_batch_size], dtype=torch.int32, device="cuda") + self._host_state_indices = torch.zeros([self.max_batch_size], + dtype=torch.int32, + pin_memory=prefer_pinned()) + self._row_indices = torch.arange(self.max_batch_size, + dtype=torch.long, + device="cpu") self.kv_cache_config = kv_cache_config self.is_estimating_kv_cache = is_estimating_kv_cache @@ -1220,9 +1230,13 @@ def get_cache_size_per_token( # Per-request fixed cost. STATIC_SLOTS_PER_REQUEST = 1 today (the # live mamba state); fixed-position snapshots are not yet - # implemented and would simply increment this constant. + # implemented and would simply increment this constant. With + # pipeline parallelism, multiple microbatches are in-flight + # concurrently on the same rank, so each rank holds Mamba state + # for up to ``max_batch_size * pp_size`` concurrent sequences. STATIC_SLOTS_PER_REQUEST = 1 - intercept = (max_batch_size * state_bytes_per_rank * + pp_size = mapping.pp_size if mapping is not None else 1 + intercept = (max_batch_size * pp_size * state_bytes_per_rank * STATIC_SLOTS_PER_REQUEST) # Regular-snapshot bytes per token. None / non-positive intervals @@ -1300,9 +1314,16 @@ def update_resources(self, def _prepare_resources(self, scheduled_batch: ScheduledRequests): self.requests = scheduled_batch.context_requests + \ scheduled_batch.generation_requests - for req in self.requests: - self.impl.copy_linear_attention_block(req) - self.impl.refresh_blocks() + # Issue all async block onboards; defer the syncTransfers() (refresh_blocks) + # until just before forward so the CPU-side prep (attn metadata, input + # tensor assembly, draft model prep, etc.) can overlap with the in-flight + # cudaMemcpyAsync calls instead of being serialized behind them. + # copy_linear_attention_block returns True iff it actually issued a copy; + # we skip refresh_blocks entirely when nothing was scheduled. + self._pending_state_transfers = self.impl.copy_linear_attention_block_batch( + self.requests) + if self._pending_state_transfers: + logger.info(f"Need to transfer mamba state blocks") self._setup_state_indices() # Reset replay double-buffer state for fresh context blocks. A reused # block (prefix-cache hit or block recycled across requests) may carry @@ -1322,6 +1343,16 @@ def prepare_resources(self, scheduled_batch: ScheduledRequests): return self._prepare_resources(scheduled_batch) + @nvtx_range("hybrid_flush_state_transfers") + def flush_state_transfers(self) -> None: + """Complete any deferred state-block onboards scheduled by + prepare_resources(). Must be called before forward() reads recurrent + state blocks. Cheap no-op when nothing was scheduled. + """ + if getattr(self, "_pending_state_transfers", False): + self.impl.refresh_blocks() + self._pending_state_transfers = False + def is_speculative(self) -> bool: return self.spec_config is not None @@ -1371,6 +1402,33 @@ def update_mamba_states(self, num_accepted_draft_tokens] self.all_conv_states[:, state_indices_d, :] = accepted_conv + def get_num_available_tokens(self, + token_num_upper_bound: int, + max_num_draft_tokens: int = 0, + **kwargs) -> int: + # Base bound from attention KV cache pool (the parent's behaviour). + result = super().get_num_available_tokens(token_num_upper_bound, + max_num_draft_tokens, + **kwargs) + # Also bound by the recurrent-state pool capacity: each request needs + # roughly ceil(N / states_snapshot_interval) recurrent-state blocks + # (+1 for the corner case where N is a multiple of tokens_per_block). + # When block reuse is disabled, only one snapshot is needed per + # request, so no additional capping is required here. + interval = (self.linear_attention_metadata.states_snapshot_interval + if self.linear_attention_metadata is not None else 0) + if interval and interval > 0: + stats = self.impl.get_kv_cache_stats() + rs_free = stats.num_free_blocks_per_window_size.get( + LinearCacheType.RECURRENT_STATES.value, 0) + # Reserve 1 block for the always-allocated last block (corner case + # / final live state) so we don't promise more tokens than the + # pool can actually back at allocation time. + usable_rs_blocks = max(0, rs_free - 1) + rs_token_cap = usable_rs_blocks * interval + result = min(result, rs_token_cap) + return max(result, 0) + def get_ssm_states(self, layer_idx: int) -> torch.Tensor: return self.all_ssm_states[self.mamba_layer_offsets[layer_idx]] @@ -1444,34 +1502,41 @@ def _setup_state_indices(self) -> None: self.impl.copy_batch_block_offsets( self.host_block_offsets, [req.py_request_id for req in self.requests], 1, 0) - host_block_offsets = torch.zeros([len(self.requests)], - dtype=torch.int32, - device="cpu") - for i in range(len(self.requests)): - # With layer-first pool layout, setOffsets produces the block index directly - # (no longer multiplied by num_local_mamba_layers) - value = self.host_block_offsets[self.recurrent_states_pool_index, i, - 0, block_indices[i]] - max_blocks = self.blocks_per_window[ - LinearCacheType.RECURRENT_STATES.value][0] - if value < 0 or value >= max_blocks: - req = self.requests[i] + + max_blocks = self.blocks_per_window[ + LinearCacheType.RECURRENT_STATES.value][0] + n = len(self.requests) + self._host_state_indices.zero_() + if n > 0: + # Vectorized gather: replace per-element Python loop with a single + # tensor index op. block_indices is a small Python list so + # torch.tensor() conversion here is O(n) at C level, much cheaper + # than n round-trips through Python indexing. + bi = torch.tensor(block_indices, dtype=torch.long) + rows = self._row_indices[:n] + # host_block_offsets: [num_pools, max_batch_size, 2, max_blocks_per_seq] + values = self.host_block_offsets[self.recurrent_states_pool_index, + rows, 0, bi] + invalid_mask = (values < 0) | (values >= max_blocks) + if invalid_mask.any(): + bad_i = int(invalid_mask.nonzero(as_tuple=False)[0, 0]) + req = self.requests[bad_i] + value = int(values[bad_i]) raise RuntimeError( f"Invalid recurrent state block index {value} " - f"(expected 0 <= index < {max_blocks}) for request {i}, " + f"(expected 0 <= index < {max_blocks}) for request {bad_i}, " f"prompt_len={req.prompt_len}, " f"is_context_finished={req.is_context_finished}, " f"context_current_position={req.context_current_position}, " f"prepopulated_token_num={req.prepopulated_prompt_len}, " f"context_chunk_size={req.context_chunk_size if not req.is_context_finished else 'N/A'}, " - f"block_index for next step is {block_indices[i]}, " + f"block_index for next step is {block_indices[bad_i]}, " f"\nblock_ids={self.impl.get_cache_block_ids(req.py_request_id, LinearCacheType.RECURRENT_STATES.value)}" ) - host_block_offsets[i] = value + self._host_state_indices[:n] = values - torch.fill_(self.cuda_state_indices, 0) - self.cuda_state_indices[:len(self.requests)] = host_block_offsets.cuda() - self._host_state_indices = host_block_offsets.clone() + self.cuda_state_indices.copy_(self._host_state_indices, + non_blocking=True) def get_state_indices( self, diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index 2d02ba2a29aa..c84f21e8d658 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -550,6 +550,7 @@ def __init__( self.mapping = mapping self.dtype = dtype self.kv_cache_type = kv_cache_type + self.spec_config = spec_config self.pp_layers, self.num_layers = get_pp_layers( num_layers, mapping, @@ -668,13 +669,19 @@ def append_to_kv_heads_per_layer(num_kv_heads_per_layer: List[int], # slots live in a separate window: at minimum the live # state per concurrent request, and -- when block reuse is # enabled -- enough room for one regular snapshot per - # snapshot interval over the full token budget. - max_snapshots = self.max_batch_size + # snapshot interval over the full token budget. With + # pipeline parallelism, multiple microbatches can be + # in-flight simultaneously on the same rank, each holding + # up to ``max_batch_size`` sequences' Mamba state, so the + # live-state slot count must scale with ``pp_size``. + pp_size = self.mapping.pp_size if self.mapping is not None else 1 + live_state_slots = self.max_batch_size * pp_size + max_snapshots = live_state_slots if kv_cache_config.enable_block_reuse: max_snapshots = max( kv_cache_config.max_tokens // linear_attention_metadata.states_snapshot_interval, - self.max_batch_size) + live_state_slots) blocks_per_window[LinearCacheType.RECURRENT_STATES.value] = ( int(max_snapshots), 0) @@ -1416,7 +1423,14 @@ def get_batch_cache_indices( return result def get_num_free_blocks(self) -> int: - if self.is_vswa or self.is_linear_attention: + if self.is_linear_attention: + value = self.impl.get_kv_cache_stats( + ).num_free_blocks_per_window_size[self.max_seq_len] + logger.debug( + f"For linear attention case, we return the number of free blocks for the kv cache (not for the recurrent states): {value}" + ) + return value + if self.is_vswa: logger.info( f"For {'linear attention' if self.is_linear_attention else 'VSWA'} case, we return the minimum of the number of free blocks for each window size: {self.impl.get_kv_cache_stats().num_free_blocks_per_window_size}" ) @@ -1432,10 +1446,16 @@ def get_num_available_tokens(self, token_num_upper_bound: int, max_num_draft_tokens: int = 0, **kwargs) -> int: - return min( - token_num_upper_bound, - self.get_num_free_blocks() * self.tokens_per_block - + free_blocks = self.get_num_free_blocks() + result = min( + token_num_upper_bound, free_blocks * self.tokens_per_block - self.num_extra_kv_tokens - max_num_draft_tokens) + logger.debug( + f"[get_num_available_tokens] free_blocks={free_blocks}, " + f"tokens_per_block={self.tokens_per_block}, " + f"num_extra_kv_tokens={self.num_extra_kv_tokens}, " + f"token_num_upper_bound={token_num_upper_bound}, result={result}") + return result def get_buffers(self, layer_idx: int, @@ -1746,7 +1766,12 @@ def _calculate_max_num_blocks_for_linear_attention( slope = attention_slope + mamba_slope # STATIC_SLOTS_PER_REQUEST = 1 (live state); fixed-position # snapshots are not yet implemented. - intercept = self.max_batch_size * state_bytes_local + # With pipeline parallelism, multiple microbatches can be in-flight + # simultaneously on the same rank, so each rank holds Mamba state for + # up to ``max_batch_size * pp_size`` concurrent sequences. Mirror the + # behaviour of KVCacheManagerV2 (see max_num_sequences calculation). + pp_size = self.mapping.pp_size if self.mapping is not None else 1 + intercept = self.max_batch_size * pp_size * state_bytes_local # heuristic: When block reuse is enabled, we assume the mamba snapshots are dominant instead of active states, # otherwise we may run out of kv cache blocks prior to mamba blocks due to the large number of max_batch_size. @@ -1769,10 +1794,18 @@ def _calculate_max_num_blocks_for_linear_attention( # Recurrent state slot count: live state per concurrent request, with # extra room for one regular snapshot per snapshot interval over the # full token budget when block reuse is enabled. - max_snapshots = self.max_batch_size + # With pipeline parallelism, multiple microbatches can be in-flight + # simultaneously on the same rank, each holding up to ``max_batch_size`` + # sequences' Mamba state, so the live-state slot count must scale with + # ``pp_size``. +1 is for the CUDA graph padding dummy. + max_snapshots = self.max_batch_size * pp_size + 1 + if self.spec_config is not None: + # cuda graph has different request ids for different draft len (CUDAGraphRunner::_get_padded_batch) + # TODO: we can use a same slot for all these + max_snapshots += self.spec_config.max_draft_len if (kv_cache_config.enable_block_reuse and interval is not None and interval > 0): - max_snapshots = max_tokens // interval + max_snapshots = max(max_tokens // interval, max_snapshots) secondary_snapshots = int(max_snapshots * (self._secondary_pool_memory_bytes / diff --git a/tests/unittest/_torch/executor/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/test_mamba_cache_manager.py index d959cdfc2547..90950b3e284a 100644 --- a/tests/unittest/_torch/executor/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/test_mamba_cache_manager.py @@ -17,7 +17,8 @@ ) from tensorrt_llm._torch.pyexecutor.resource_manager import CacheTypeCpp from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests -from tensorrt_llm.llmapi.llm_args import KvCacheConfig +from tensorrt_llm.bindings.internal.batch_manager import LinearCacheType +from tensorrt_llm.llmapi.llm_args import KvCacheConfig, MTPDecodingConfig from tensorrt_llm.mapping import Mapping skip_no_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") @@ -258,6 +259,175 @@ def test_cpp_get_state_indices_resolves_sentinel_to_reserved_slot(): assert mgr.get_state_indices(request_ids, is_padding) == indices +# --------------------------------------------------------------------------- +# CppMambaHybridCacheManager: recurrent-state snapshot pool sizing +# +# Sized in KVCacheManager._calculate_max_num_blocks_for_linear_attention. +# Mirrors the MixedMambaCacheManager fix where each kind of padding sentinel +# (CUDA-graph dummy, plus one per draft length under spec decoding) must not +# evict live recurrent state. Wanli's fix made all sentinels share one slot +# in the Python manager (#13489); the C++ hybrid path instead reserves a +# dedicated slot per sentinel kind in the underlying pool — same invariant, +# different mechanism. These tests guard the pool sizing. +# --------------------------------------------------------------------------- + + +def _build_hybrid_with_mamba_layer(spec_config=None, max_batch_size=4, enable_block_reuse=False): + """Construct a real CppMambaHybridCacheManager with one mamba layer + + one full-attention layer so the parent KVCacheManager goes through the + linear-attention pool sizing path.""" + # Layer 0: mamba; Layer 1: full attention. Single rank, no MPI. + mamba_mask = [True, False] + attn_mask = [False, True] + mapping = Mapping(world_size=1, rank=0, tp_size=1, pp_size=1) + # Cap max_tokens to keep the real C++ pool allocation tiny. + kv_cache_config = KvCacheConfig(max_tokens=512, enable_block_reuse=enable_block_reuse) + return CppMambaHybridCacheManager( + mamba_d_state=8, + mamba_d_conv=4, + mamba_num_heads=4, + mamba_n_groups=1, + mamba_head_dim=8, + mamba_num_layers=1, + mamba_layer_mask=mamba_mask, + mamba_cache_dtype=torch.float16, + mamba_ssm_cache_dtype=torch.float16, + kv_cache_config=kv_cache_config, + kv_cache_type=CacheTypeCpp.SELF, + num_layers=1, + num_kv_heads=4, + head_dim=64, + tokens_per_block=32, + max_seq_len=128, + max_batch_size=max_batch_size, + mapping=mapping, + spec_config=spec_config, + layer_mask=attn_mask, + ) + + +@skip_no_cuda +def test_cpp_hybrid_recurrent_pool_reserves_cuda_graph_padding_slot(): + """Without spec decoding, the recurrent-state snapshot pool must + have at least max_batch_size + 1 slots — one extra for the + CUDA-graph padding sentinel (CUDA_GRAPH_DUMMY_REQUEST_ID). Without + it, the padding sentinel evicts live recurrent state under load.""" + max_batch_size = 4 + mgr = _build_hybrid_with_mamba_layer(spec_config=None, max_batch_size=max_batch_size) + recurrent_primary, _ = mgr.blocks_per_window[LinearCacheType.RECURRENT_STATES.value] + assert recurrent_primary >= max_batch_size + 1, ( + f"recurrent-state pool has {recurrent_primary} slots, " + f"need >= max_batch_size + 1 = {max_batch_size + 1} to host the " + f"CUDA-graph padding sentinel without evicting live state" + ) + + +@skip_no_cuda +def test_cpp_hybrid_recurrent_pool_reserves_draft_len_sentinel_slots(): + """With spec decoding, CUDAGraphRunner._get_padded_batch issues a + distinct dummy request id for each runtime_draft_len in + [0, max_draft_len], so the recurrent-state snapshot pool must reserve + one slot per draft length on top of the CUDA-graph padding slot.""" + max_batch_size, max_draft_len = 4, 2 + spec_config = MTPDecodingConfig(max_draft_len=max_draft_len) + mgr = _build_hybrid_with_mamba_layer(spec_config=spec_config, max_batch_size=max_batch_size) + recurrent_primary, _ = mgr.blocks_per_window[LinearCacheType.RECURRENT_STATES.value] + expected_min = max_batch_size + 1 + max_draft_len + assert recurrent_primary >= expected_min, ( + f"recurrent-state pool has {recurrent_primary} slots, " + f"need >= max_batch_size + 1 + max_draft_len = {expected_min} so " + f"per-draft-len sentinels don't collide with live state" + ) + + +def _build_hybrid_with_mamba_layer_pp( + spec_config=None, max_batch_size=4, enable_block_reuse=False, pp_size=2 +): + """Same as ``_build_hybrid_with_mamba_layer`` but with ``pp_size`` >= 1. + + Uses ``world_size = pp_size`` and ``rank = 0`` so the real C++ KVCacheManager + still goes through its single-process path while the Python pool-sizing + code sees ``mapping.pp_size > 1``. Constructs ``pp_size * 2`` total layers + (alternating mamba/attn) so that each PP slice has both a mamba and an + attention layer — otherwise some ranks would hit a slope=0 edge case in + the affine memory model when block reuse is disabled. + """ + pairs = pp_size # one (mamba, attn) pair per PP stage so every rank has both kinds + mamba_mask = [True, False] * pairs + attn_mask = [False, True] * pairs + mamba_num_layers = sum(mamba_mask) + num_layers = sum(attn_mask) + mapping = Mapping(world_size=pp_size, rank=0, tp_size=1, pp_size=pp_size) + kv_cache_config = KvCacheConfig(max_tokens=512, enable_block_reuse=enable_block_reuse) + return CppMambaHybridCacheManager( + mamba_d_state=8, + mamba_d_conv=4, + mamba_num_heads=4, + mamba_n_groups=1, + mamba_head_dim=8, + mamba_num_layers=mamba_num_layers, + mamba_layer_mask=mamba_mask, + mamba_cache_dtype=torch.float16, + mamba_ssm_cache_dtype=torch.float16, + kv_cache_config=kv_cache_config, + kv_cache_type=CacheTypeCpp.SELF, + num_layers=num_layers, + num_kv_heads=4, + head_dim=64, + tokens_per_block=32, + max_seq_len=128, + max_batch_size=max_batch_size, + mapping=mapping, + spec_config=spec_config, + layer_mask=attn_mask, + ) + + +@skip_no_cuda +@pytest.mark.parametrize("pp_size", [2, 4]) +def test_cpp_hybrid_recurrent_pool_scales_with_pp_size(pp_size): + """With pipeline parallelism, multiple microbatches are in-flight on the + same rank concurrently, each holding up to ``max_batch_size`` sequences' + Mamba state. The recurrent-state pool must therefore size for + ``max_batch_size * pp_size`` live slots (plus the CUDA-graph padding + sentinel). Without this scaling, the first inference batch under PP>1 + trips ``No free block found`` once requests beyond the first microbatch + enter the pool (cf. TestNemotronV3Super::test_nvfp4_parallelism[TP4_PP2]). + """ + max_batch_size = 4 + mgr = _build_hybrid_with_mamba_layer_pp( + spec_config=None, max_batch_size=max_batch_size, pp_size=pp_size + ) + recurrent_primary, _ = mgr.blocks_per_window[LinearCacheType.RECURRENT_STATES.value] + expected_min = max_batch_size * pp_size + 1 + assert recurrent_primary >= expected_min, ( + f"recurrent-state pool has {recurrent_primary} slots with pp_size={pp_size}, " + f"need >= max_batch_size * pp_size + 1 = {expected_min} so concurrent " + f"in-flight microbatches don't exhaust live-state slots" + ) + + +@skip_no_cuda +def test_cpp_hybrid_recurrent_pool_floor_with_block_reuse(): + """With block reuse enabled, the block-reuse branch must not drop the + live-state + CUDA-graph-padding floor. + + With max_batch_size=4, mamba_state_cache_interval=256, max_tokens=512: + naive: max_snapshots = 512 // 256 = 2 (drops live-state floor!) + fixed: max_snapshots = max(2, 4 + 1) = 5 + """ + max_batch_size = 4 + mgr = _build_hybrid_with_mamba_layer( + spec_config=None, max_batch_size=max_batch_size, enable_block_reuse=True + ) + recurrent_primary, _ = mgr.blocks_per_window[LinearCacheType.RECURRENT_STATES.value] + assert recurrent_primary >= max_batch_size + 1, ( + f"recurrent-state pool has {recurrent_primary} slots with block reuse enabled, " + f"need >= max_batch_size + 1 = {max_batch_size + 1} to prevent the padding " + f"sentinel from evicting live recurrent state" + ) + + # --------------------------------------------------------------------------- # CppMambaHybridCacheManager: rank with zero local mamba layers # From da9f86eb92d71719c3a89f3a36b27db7f4134368 Mon Sep 17 00:00:00 2001 From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> Date: Wed, 20 May 2026 02:01:41 -0700 Subject: [PATCH 2/7] Squashed commit of the following: commit 1d76cf70a9e4acf965f6a3fbda8b6d6de921629a Author: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> Date: Wed May 13 01:57:44 2026 -0700 [None][feat] Refactor to support legacy and 1.x modelopt quant config format Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> --- .../auto_deploy/models/quant_config_reader.py | 21 +- tensorrt_llm/_torch/model_config.py | 91 +++- tensorrt_llm/llmapi/llm_utils.py | 166 +++--- tensorrt_llm/quantization/modelopt_config.py | 140 +++++ tests/unittest/llmapi/test_llm_quant.py | 489 +++++++++++++++++- 5 files changed, 804 insertions(+), 103 deletions(-) create mode 100644 tensorrt_llm/quantization/modelopt_config.py diff --git a/tensorrt_llm/_torch/auto_deploy/models/quant_config_reader.py b/tensorrt_llm/_torch/auto_deploy/models/quant_config_reader.py index 6620d3a39bec..92ca23fd272b 100644 --- a/tensorrt_llm/_torch/auto_deploy/models/quant_config_reader.py +++ b/tensorrt_llm/_torch/auto_deploy/models/quant_config_reader.py @@ -25,6 +25,11 @@ from abc import ABC, abstractmethod from typing import Any, Callable, Dict, Optional, Tuple, Type +from tensorrt_llm.quantization.modelopt_config import ( + is_modelopt_quant_config, + read_modelopt_quant_config, +) + from ..utils.logger import ad_logger @@ -110,12 +115,16 @@ class ModelOPTQuantConfigReader(QuantConfigReader): DEFAULT_KV_CACHE_DTYPE = "fp8" def read_config(self, config: Dict) -> Dict: - producer = config.get("producer", {}).get("name") - # sanity check - if producer != "modelopt": - raise ValueError(f"Expected producer 'modelopt', got '{producer}'") - - quant_config = config.get("quantization", {}) + # Accept either modelopt shape: legacy (producer.name == "modelopt" + # with a "quantization" wrapper) or flat (quant_method == "modelopt"). + if not is_modelopt_quant_config(config): + raise ValueError( + "Expected a modelopt quant config " + f"(producer={config.get('producer')}, " + f"quant_method={config.get('quant_method')})" + ) + # Downstream auto-deploy transforms operate on the legacy field names. + quant_config = read_modelopt_quant_config(config) quant_algo = quant_config.get("quant_algo", "").upper() diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index cba1fd0996de..4e8ad3e73652 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -24,6 +24,9 @@ from tensorrt_llm.mapping import Mapping from tensorrt_llm.models.modeling_utils import QuantConfig from tensorrt_llm.quantization.mode import QuantAlgo +from tensorrt_llm.quantization.modelopt_config import ( + is_modelopt_quant_config, read_modelopt_quant_config, + warn_if_inline_diverges) if TYPE_CHECKING: from tensorrt_llm.bindings import ModelConfig as ModelConfigCpp @@ -306,23 +309,38 @@ def resolve_moe_backend(moe_backend: str, @staticmethod def load_modelopt_quant_config(quant_config_file, checkpoint_dir, moe_backend): - quant_config = QuantConfig() - layer_quant_config = None - with open(quant_config_file) as f: quant_config_dict = json.load(f) + return ModelConfig._build_modelopt_quant_config( + read_modelopt_quant_config(quant_config_dict), checkpoint_dir, + moe_backend) - json_quant_configs = quant_config_dict['quantization'] + @staticmethod + def _build_modelopt_quant_config(json_quant_configs, checkpoint_dir, + moe_backend): + """Build (quant_config, layer_quant_config) from a normalized modelopt 'quantization' inner dict. - quant_config.quant_algo = json_quant_configs.get('quant_algo', None) - # fp8_pb_wo from modelopt is the same as FP8_BLOCK_SCALES - if quant_config.quant_algo == "fp8_pb_wo": - quant_config.quant_algo = 'FP8_BLOCK_SCALES' - quant_config.kv_cache_quant_algo = json_quant_configs.get( - 'kv_cache_quant_algo', None) + ``json_quant_configs`` should be a dict as produced by + :func:`read_modelopt_quant_config`. May be mutated in place via + ``.update()`` when overlaying ``quant_cfg.json``. + """ + quant_config = QuantConfig() + layer_quant_config = None + + quant_config.quant_algo = (QuantAlgo(json_quant_configs['quant_algo']) + if json_quant_configs.get('quant_algo') + is not None else None) + quant_config.kv_cache_quant_algo = ( + QuantAlgo(json_quant_configs['kv_cache_quant_algo']) if + json_quant_configs.get('kv_cache_quant_algo') is not None else None) quant_config.group_size = json_quant_configs.get('group_size', None) quant_config.exclude_modules = json_quant_configs.get( 'exclude_modules', None) + # AWQ-specific extras; only override defaults when present in JSON. + if 'has_zero_point' in json_quant_configs: + quant_config.has_zero_point = json_quant_configs['has_zero_point'] + if 'pre_quant_scale' in json_quant_configs: + quant_config.pre_quant_scale = json_quant_configs['pre_quant_scale'] if quant_config.quant_algo == QuantAlgo.MIXED_PRECISION: json_extended_quant_configs: dict = {} @@ -334,12 +352,14 @@ def load_modelopt_quant_config(quant_config_file, checkpoint_dir, json_extended_quant_configs = json.load(fm) except Exception: logger.info( - f"No quant_cfg.json found for layer quant info, using hf_quant_config.json." + "No quant_cfg.json found for layer quant info, using hf_quant_config.json." ) json_quant_configs.update(json_extended_quant_configs) # kv_cache_quant_algo is global regardless of MIXED_PRECISION - kv_cache_quant_algo = json_quant_configs.get( - 'kv_cache_quant_algo', None) + kv_cache_quant_algo = (QuantAlgo( + json_quant_configs['kv_cache_quant_algo']) if + json_quant_configs.get('kv_cache_quant_algo') + is not None else None) mixed_quant_configs = json_quant_configs.get( 'quantized_layers', None) if (kv_quant_lhs := json_extended_quant_configs.get( @@ -352,19 +372,30 @@ def load_modelopt_quant_config(quant_config_file, checkpoint_dir, f"is different from 'hf_quant_config.json', {kv_quant_rhs}!" ) quant_config.kv_cache_quant_algo = kv_cache_quant_algo + quant_config.group_size = json_quant_configs.get( + 'group_size', quant_config.group_size) + quant_config.exclude_modules = json_quant_configs.get( + 'exclude_modules', quant_config.exclude_modules) for layer in mixed_quant_configs: + layer_cfg = mixed_quant_configs[layer] config = QuantConfig() config.kv_cache_quant_algo = kv_cache_quant_algo - config.quant_algo = mixed_quant_configs[layer]['quant_algo'] - config.group_size = mixed_quant_configs[layer].get( - 'group_size', None) + config.quant_algo = QuantAlgo(layer_cfg['quant_algo']) + config.group_size = layer_cfg.get('group_size', None) + # AWQ-specific extras emitted by modelopt per-layer. + if 'has_zero_point' in layer_cfg: + config.has_zero_point = layer_cfg['has_zero_point'] + if 'pre_quant_scale' in layer_cfg: + config.pre_quant_scale = layer_cfg['pre_quant_scale'] mixed_quant_configs[layer] = config layer_quant_config = mixed_quant_configs elif quant_config.quant_algo == QuantAlgo.FP8_BLOCK_SCALES: if quant_config.group_size is None: quant_config.group_size = 128 - if moe_backend == 'TRTLLM' and quant_config.quant_algo == "FP8_BLOCK_SCALES" and quant_config.exclude_modules is None: + if (moe_backend == 'TRTLLM' + and quant_config.quant_algo == QuantAlgo.FP8_BLOCK_SCALES + and quant_config.exclude_modules is None): quant_config.exclude_modules = [ "*kv_b_proj*", "*k_b_proj*", "*eh_proj" ] @@ -385,10 +416,16 @@ def get_mxfp4_quant_algo(moe_backend, is_dynamic_quant=False): return quant_algo @staticmethod - def load_hf_quant_config(hf_quant_config, moe_backend): + def load_hf_quant_config(hf_quant_config, moe_backend, checkpoint_dir=None): quant_config = QuantConfig() layer_quant_config = None + # Route inline modelopt configs (legacy or flat) to the modelopt builder. + if is_modelopt_quant_config(hf_quant_config): + return ModelConfig._build_modelopt_quant_config( + read_modelopt_quant_config(hf_quant_config), checkpoint_dir, + moe_backend) + # Read exclude_modules from HF config if present (HF format module names) hf_exclude_modules = hf_quant_config.get('modules_to_not_convert', None) @@ -686,13 +723,25 @@ def _recursive_update_config(config: transformers.PretrainedConfig, # quantized ckpt in modelopt format if quant_config_file := cached_file(checkpoint_dir, 'hf_quant_config.json'): - quant_config, layer_quant_config = cls.load_modelopt_quant_config( - quant_config_file, checkpoint_dir, moe_backend_hint) + with open(quant_config_file) as f: + normalized = read_modelopt_quant_config(json.load(f)) + # The file is authoritative; warn if the inline copy disagrees. + # Done before _build_modelopt_quant_config since the builder may + # mutate ``normalized`` via ``.update`` from quant_cfg.json. + warn_if_inline_diverges( + normalized, + getattr(pretrained_config, "quantization_config", None), + source_file="hf_quant_config.json", + ) + quant_config, layer_quant_config = cls._build_modelopt_quant_config( + normalized, checkpoint_dir, moe_backend_hint) # quantized ckpt in other formats elif hasattr(pretrained_config, "quantization_config"): hf_quant_config = pretrained_config.quantization_config quant_config, layer_quant_config = cls.load_hf_quant_config( - hf_quant_config, moe_backend_hint) + hf_quant_config, + moe_backend_hint, + checkpoint_dir=checkpoint_dir) elif quant_config_file := cached_file(checkpoint_dir, 'dtypes.json'): quant_config, layer_quant_config = cls.load_quant_config_from_dtypes_json( quant_config_file, moe_backend_hint) diff --git a/tensorrt_llm/llmapi/llm_utils.py b/tensorrt_llm/llmapi/llm_utils.py index b37cdd03658d..1d362a96d63c 100644 --- a/tensorrt_llm/llmapi/llm_utils.py +++ b/tensorrt_llm/llmapi/llm_utils.py @@ -27,6 +27,9 @@ from ..models.automodel import MODEL_MAP, AutoConfig, AutoModelForCausalLM from ..models.modeling_utils import PretrainedConfig, QuantAlgo, QuantConfig from ..module import Module +from ..quantization.modelopt_config import (is_modelopt_quant_config, + read_modelopt_quant_config, + warn_if_inline_diverges) from .build_cache import (BuildCache, BuildCacheConfig, CachedStage, get_build_cache_config_from_env) # yapf: disable @@ -81,6 +84,7 @@ def from_module(cls, module: Module): @dataclass class _ModelRuntimeContext: ''' _ModelRuntimeContext holds the minimum runtime resources for running a model. + It could be a runtime cache in MPI nodes. ''' engine: Optional[Engine] = None @@ -98,6 +102,7 @@ def model_arch(self) -> str: class ModelLoader: ''' The ModelLoader is used to build an end-to-end model for a single-gpu. + It accepts model name or a local model dir, and will download the model if necessary. ''' @@ -318,13 +323,78 @@ def _download_hf_model(self): self.model_obj.model_dir = self._model_dir # mark as a local model assert self.model_obj.is_local_model + def _apply_modelopt_quant_config(self, hf_quant_config: Dict[str, Any], + explicit_kv_cache_quant_algo) -> None: + """Apply a normalized modelopt ``quantization`` inner dict onto ``self.llm_args.quant_config``. + + Pops the well-known fields, validates them, then forwards any + remaining ``QuantConfig`` fields (e.g. AWQ ``has_zero_point`` / + ``pre_quant_scale``) via setattr. + """ + quant_config = self.llm_args.quant_config + hf_quant_algo = hf_quant_config.pop("quant_algo", None) + if hf_quant_algo is None: + raise ValueError("Pre-quantized checkpoint must have quant_algo.") + hf_quant_algo = QuantAlgo(hf_quant_algo) + if quant_config.quant_algo is None: + logger.info( + f"Setting quant_algo={hf_quant_algo} from HF quant config.") + quant_config.quant_algo = hf_quant_algo + elif quant_config.quant_algo != hf_quant_algo: + raise ValueError( + f"Specified quant_algo={quant_config.quant_algo}, conflicting with quant_algo={hf_quant_algo} from HF quant config." + ) + + hf_kv_cache_quant_algo = hf_quant_config.pop("kv_cache_quant_algo", + None) + if hf_kv_cache_quant_algo is not None: + hf_kv_cache_quant_algo = QuantAlgo(hf_kv_cache_quant_algo) + if explicit_kv_cache_quant_algo is not None: + if explicit_kv_cache_quant_algo != hf_kv_cache_quant_algo: + logger.warning( + f"Overriding checkpoint kv_cache_quant_algo={hf_kv_cache_quant_algo} with explicit kv_cache_config.dtype={explicit_kv_cache_quant_algo}." + ) + quant_config.kv_cache_quant_algo = explicit_kv_cache_quant_algo + elif quant_config.kv_cache_quant_algo is None: + logger.info( + f"Setting kv_cache_quant_algo={hf_kv_cache_quant_algo} from HF quant config." + ) + quant_config.kv_cache_quant_algo = hf_kv_cache_quant_algo + elif quant_config.kv_cache_quant_algo != hf_kv_cache_quant_algo: + raise ValueError( + f"Specified kv_cache_quant_algo={quant_config.kv_cache_quant_algo}, conflicting with kv_cache_quant_algo={hf_kv_cache_quant_algo} from HF quant config." + ) + else: + if quant_config.kv_cache_quant_algo not in [ + None, QuantAlgo.FP8, QuantAlgo.NVFP4 + ]: + raise ValueError( + f"Only kv_cache_quant_algo={QuantAlgo.FP8} or {QuantAlgo.NVFP4} is allowed for pre-quantized checkpoint, got {quant_config.kv_cache_quant_algo}." + ) + + # quantized_layers is handled separately (e.g. via LayerQuantConfig + # in PretrainedConfig for TRT, or _torch/model_config.py for PyTorch) + hf_quant_config.pop("quantized_layers", None) + + quant_config_fields = set(quant_config.model_fields.keys()) + for key, value in hf_quant_config.items(): + if key not in quant_config_fields: + logger.warning( + f"Ignoring unknown field '{key}' from HF quant config (not a QuantConfig field)." + ) + continue + logger.info( + f"Setting {key}={str(value)[:100]}{'...' if len(str(value)) > 100 else ''} from HF quant config." + ) + setattr(quant_config, key, value) + self.llm_args.quant_config = quant_config + def _update_from_hf_quant_config(self) -> bool: """Update quant_config from the config file of pre-quantized HF checkpoint. Returns: prequantized (bool): Whether the checkpoint is pre-quantized. """ - quant_config = self.llm_args.quant_config kv_cache_dtype = self.llm_args.kv_cache_config.dtype explicit_kv_cache_quant_algo = { "fp8": QuantAlgo.FP8, @@ -337,75 +407,21 @@ def _update_from_hf_quant_config(self) -> bool: f"Found {hf_quant_config_path}, pre-quantized checkpoint is used." ) with open(hf_quant_config_path, "r") as f: - hf_quant_config = json.load(f) - hf_quant_config = hf_quant_config["quantization"] - - hf_quant_algo = hf_quant_config.pop("quant_algo", None) - if hf_quant_algo is not None: - # fp8_pb_wo from modelopt is the same as fp8_block_scales - if hf_quant_algo == "fp8_pb_wo": - hf_quant_algo = QuantAlgo.FP8_BLOCK_SCALES - else: - hf_quant_algo = QuantAlgo(hf_quant_algo) - if quant_config.quant_algo is None: - logger.info( - f"Setting quant_algo={hf_quant_algo} from HF quant config." - ) - quant_config.quant_algo = hf_quant_algo - elif quant_config.quant_algo != hf_quant_algo: - raise ValueError( - f"Specified quant_algo={quant_config.quant_algo}, conflicting with quant_algo={hf_quant_algo} from HF quant config." - ) - else: - raise ValueError( - "Pre-quantized checkpoint must have quant_algo.") - - hf_kv_cache_quant_algo = hf_quant_config.pop( - "kv_cache_quant_algo", None) - if hf_kv_cache_quant_algo is not None: - hf_kv_cache_quant_algo = QuantAlgo(hf_kv_cache_quant_algo) - if explicit_kv_cache_quant_algo is not None: - if explicit_kv_cache_quant_algo != hf_kv_cache_quant_algo: - logger.warning( - f"Overriding checkpoint kv_cache_quant_algo={hf_kv_cache_quant_algo} with explicit kv_cache_config.dtype={kv_cache_dtype}." - ) - quant_config.kv_cache_quant_algo = explicit_kv_cache_quant_algo - elif quant_config.kv_cache_quant_algo is None: - logger.info( - f"Setting kv_cache_quant_algo={hf_kv_cache_quant_algo} from HF quant config." - ) - quant_config.kv_cache_quant_algo = hf_kv_cache_quant_algo - elif quant_config.kv_cache_quant_algo != hf_kv_cache_quant_algo: - raise ValueError( - f"Specified kv_cache_quant_algo={quant_config.kv_cache_quant_algo}, conflicting with kv_cache_quant_algo={hf_kv_cache_quant_algo} from HF quant config." - ) - else: - if quant_config.kv_cache_quant_algo not in [ - None, QuantAlgo.FP8, QuantAlgo.NVFP4 - ]: - raise ValueError( - f"Only kv_cache_quant_algo={QuantAlgo.FP8} or {QuantAlgo.NVFP4} is allowed for pre-quantized checkpoint, got {quant_config.kv_cache_quant_algo}." - ) - - # quantized_layers is handled separately (e.g. via LayerQuantConfig - # in PretrainedConfig for TRT, or _torch/model_config.py for PyTorch) - hf_quant_config.pop("quantized_layers", None) - - quant_config_fields = set(quant_config.model_fields.keys()) - for key, value in hf_quant_config.items(): - if key not in quant_config_fields: - logger.warning( - f"Ignoring unknown field '{key}' from HF quant config (not a QuantConfig field)." + normalized = read_modelopt_quant_config(json.load(f)) + # Cross-check against inline config.json.quantization_config if any. + # Done before _apply_modelopt_quant_config since the apply step + # mutates ``normalized`` via ``.pop()``. + try: + with open(f"{self._model_dir}/config.json", "r") as f: + warn_if_inline_diverges( + normalized, + json.load(f).get("quantization_config"), + source_file="hf_quant_config.json", ) - continue - logger.info( - f"Setting {key}={str(value)[:100]}{'...' if len(str(value)) > 100 else ''} from HF quant config." - ) - setattr(quant_config, key, value) - - # Update the quant_config in llm_args for pytorch - self.llm_args.quant_config = quant_config - + except FileNotFoundError: + pass + self._apply_modelopt_quant_config(normalized, + explicit_kv_cache_quant_algo) return True hf_config_path = f"{self._model_dir}/config.json" @@ -429,6 +445,12 @@ def _update_from_hf_quant_config(self) -> bool: ) if hf_quant_config is not None: + if is_modelopt_quant_config(hf_quant_config): + self._apply_modelopt_quant_config( + read_modelopt_quant_config(hf_quant_config), + explicit_kv_cache_quant_algo) + return True + quant_config = self.llm_args.quant_config # DeepSeek V3 FP8 ckpt if hf_quant_config.get("quant_method") == "fp8": if hf_quant_config.get("weight_block_size") is not None: @@ -626,8 +648,9 @@ def _build_engine(self): logger_debug(f"rank{mpi_rank()} build engine done\n", "green") def _save_engine_for_runtime(self): - ''' - Persist the engine to disk for the cpp runtime. Currently, the cpp runtime can accept an engine path, + '''Persist the engine to disk for the cpp runtime. + + Currently, the cpp runtime can accept an engine path, that requires the engine should always be saved to disk. This explicit saving will be removed in the future when the cpp runtime can accept the engine buffer directly. @@ -855,6 +878,7 @@ def serialize(d) -> str: def get_pretrained_config(self) -> PretrainedConfig: ''' Get the PretrainedConfig for cache key. + NOTE, this is not the HF model's config, but the TRT-LLM's config. We use this as a generic information for HF and other models. ''' assert self._hf_model_dir is not None diff --git a/tensorrt_llm/quantization/modelopt_config.py b/tensorrt_llm/quantization/modelopt_config.py new file mode 100644 index 000000000000..afd482f0ea95 --- /dev/null +++ b/tensorrt_llm/quantization/modelopt_config.py @@ -0,0 +1,140 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Unified normalizer for ModelOpt-produced quantization configs. + +ModelOpt emits ``hf_quant_config.json`` (and the inline +``config.json.quantization_config``) in two on-disk shapes: + +- **Legacy** (modelopt 0.x): ``{producer, quantization: {quant_algo, + kv_cache_quant_algo, exclude_modules, quantized_layers, ...}}`` +- **Flat 1.x** (compressed-tensors-style): ``{producer, quant_method="modelopt", + quant_algo, kv_cache_scheme, ignore, config_groups, ...}`` + +:func:`read_modelopt_quant_config` collapses either shape into the legacy +"quantization" inner-dict shape, which is what every TensorRT-LLM call site +consumes. +""" + +from typing import Any, Dict, Optional + +from ..logger import logger + + +def is_modelopt_quant_config(raw: Any) -> bool: + """Return True if ``raw`` looks like a modelopt config (either shape). + + Detects either ``producer.name == "modelopt"`` or a ``quant_method`` + starting with ``"modelopt"``. + """ + if not isinstance(raw, dict): + return False + if str(raw.get("quant_method", "")).lower().startswith("modelopt"): + return True + return (raw.get("producer") or {}).get("name") == "modelopt" + + +# Modelopt 1.x supports the same KV-cache algos as 0.x (FP8/NVFP4/INT8), +# but encodes ``kv_cache_scheme`` in two different shapes depending on the +# algo: +# - FP8 -> dict {"type": "float", "num_bits": 8} +# - NVFP4, INT8 -> string "NVFP4" / "INT8" +# Both maps mirror `KV_CACHE_QUANT_ALGO_LIST = [FP8, NVFP4, INT8]`. +_KV_SCHEME_DICT_MAP = { + ("float", 8): "FP8", + ("float", 4): "NVFP4", + ("int", 8): "INT8", +} +_KV_SCHEME_STRING_ALGOS = {"FP8", "NVFP4", "INT8"} + + +def _kv_cache_scheme_to_algo(scheme: Any) -> Optional[str]: + """Translate modelopt 1.x ``kv_cache_scheme`` to a legacy algo name. + + Returns ``None`` (and warns) for any unrecognized shape so a silently + dropped KV-cache quant setting surfaces in the logs. A missing scheme + (``None``) is the documented "no KV-cache quant" case and does not warn. + """ + if scheme is None: + return None + if isinstance(scheme, str): + algo = scheme.upper() + if algo in _KV_SCHEME_STRING_ALGOS: + return algo + elif isinstance(scheme, dict): + mapped = _KV_SCHEME_DICT_MAP.get((scheme.get("type"), scheme.get("num_bits"))) + if mapped is not None: + return mapped + logger.warning(f"Unrecognized 'kv_cache_scheme' {scheme!r}; KV-cache quant disabled.") + return None + + +def read_modelopt_quant_config(raw: Dict[str, Any]) -> Dict[str, Any]: + """Normalize either modelopt shape into the legacy ``quantization`` inner dict. + + Raises ``ValueError`` if ``raw`` is not a recognized modelopt config. + """ + if not isinstance(raw, dict): + raise ValueError(f"Expected dict for modelopt quant config, got {type(raw).__name__}") + if "quantization" in raw: + # Legacy shape; the inner value must be a dict. + q = raw["quantization"] + if not isinstance(q, dict): + raise ValueError(f"'quantization' must be a dict, got {type(q).__name__}") + result = dict(q) + elif is_modelopt_quant_config(raw): + # Flat → legacy: rename `ignore` → `exclude_modules`, `kv_cache_scheme` + # → `kv_cache_quant_algo`, drop modelopt-1.x-only metadata. + _SKIP = {"producer", "quant_method", "ignore", "kv_cache_scheme", "config_groups"} + result = {k: v for k, v in raw.items() if k not in _SKIP} + if "ignore" in raw: + result["exclude_modules"] = raw["ignore"] + if "kv_cache_scheme" in raw: + algo = _kv_cache_scheme_to_algo(raw["kv_cache_scheme"]) + if algo is not None: + result["kv_cache_quant_algo"] = algo + else: + raise ValueError( + f"Not a modelopt quant config (producer={raw.get('producer')!r}, " + f"quant_method={raw.get('quant_method')!r})" + ) + # Canonicalize the fp8_pb_wo legacy alias. + if result.get("quant_algo") == "fp8_pb_wo": + result["quant_algo"] = "FP8_BLOCK_SCALES" + return result + + +def warn_if_inline_diverges( + file_quant: Dict[str, Any], inline_raw: Any, *, source_file: str +) -> None: + """Warn if ``config.json.quantization_config`` diverges from the file-based config. + + The file remains authoritative; this only logs. Skips if the inline is + absent or not recognizably modelopt. + """ + if not inline_raw or not is_modelopt_quant_config(inline_raw): + return + try: + inline_quant = read_modelopt_quant_config(inline_raw) + except ValueError as e: + logger.warning( + f"Inline modelopt config failed to parse ({e}); " + f"skipping divergence check against '{source_file}'." + ) + return + diffs = [] + for key in ("quant_algo", "kv_cache_quant_algo", "group_size"): + if file_quant.get(key) != inline_quant.get(key): + diffs.append(f"{key}: file={file_quant.get(key)!r}, inline={inline_quant.get(key)!r}") + f_excl = set(file_quant.get("exclude_modules") or []) + i_excl = set(inline_quant.get("exclude_modules") or []) + if f_excl != i_excl: + diffs.append(f"exclude_modules: |file|={len(f_excl)}, |inline|={len(i_excl)}") + f_ql = len(file_quant.get("quantized_layers") or {}) + i_ql = len(inline_quant.get("quantized_layers") or {}) + if f_ql != i_ql: + diffs.append(f"quantized_layers count: file={f_ql}, inline={i_ql}") + if diffs: + logger.warning( + f"Inline 'config.json.quantization_config' diverges from " + f"'{source_file}':\n - " + "\n - ".join(diffs) + ) diff --git a/tests/unittest/llmapi/test_llm_quant.py b/tests/unittest/llmapi/test_llm_quant.py index 573a8bf0ef9f..642571814267 100644 --- a/tests/unittest/llmapi/test_llm_quant.py +++ b/tests/unittest/llmapi/test_llm_quant.py @@ -97,7 +97,9 @@ def test_quant_cfg_from_quant_cfg_json(): }, "model.layers.1.mlp.gate_proj": { "quant_algo": "W4A8_AWQ", - "group_size": 128 + "group_size": 128, + "has_zero_point": False, + "pre_quant_scale": True, } } } @@ -130,10 +132,51 @@ def test_quant_cfg_from_quant_cfg_json(): "model.layers.0.self_attn.q_proj"].quant_algo == "FP8" assert layer_quant_config[ "model.layers.0.self_attn.q_proj"].kv_cache_quant_algo == "FP8" - assert layer_quant_config[ - "model.layers.1.mlp.gate_proj"].quant_algo == "W4A8_AWQ" - assert layer_quant_config[ - "model.layers.1.mlp.gate_proj"].group_size == 128 + awq_layer = layer_quant_config["model.layers.1.mlp.gate_proj"] + assert awq_layer.quant_algo == "W4A8_AWQ" + assert awq_layer.group_size == 128 + assert awq_layer.has_zero_point is False + assert awq_layer.pre_quant_scale is True + + +def test_quant_cfg_top_level_overlay(): + """quant_cfg.json's top-level group_size/exclude_modules override hf_quant_config.json.""" + with tempfile.TemporaryDirectory() as tmp_dir: + model_dir = Path(tmp_dir) + + # quant_cfg.json overrides top-level group_size and exclude_modules. + quant_cfg_content = { + "quant_algo": "MIXED_PRECISION", + "kv_cache_quant_algo": "FP8", + "group_size": 64, + "exclude_modules": ["lm_head", "model.embed_tokens"], + "quantized_layers": { + "model.layers.0.self_attn.q_proj": { + "quant_algo": "FP8" + } + }, + } + quant_cfg_file = model_dir / "quant_cfg.json" + with open(quant_cfg_file, 'w') as f: + json.dump(quant_cfg_content, f) + + hf_quant_config_content = { + "quantization": { + "quant_algo": "MIXED_PRECISION", + "kv_cache_quant_algo": "FP8", + "group_size": 128, + "exclude_modules": ["foo"], + } + } + hf_quant_config_file = model_dir / "hf_quant_config.json" + with open(hf_quant_config_file, 'w') as f: + json.dump(hf_quant_config_content, f) + + quant_config, _ = ModelConfig.load_modelopt_quant_config( + hf_quant_config_file, model_dir, None) + + assert quant_config.group_size == 64 + assert quant_config.exclude_modules == ["lm_head", "model.embed_tokens"] def test_quant_cfg_from_hf_quant_config(): @@ -173,6 +216,442 @@ def test_quant_cfg_from_hf_quant_config(): assert layer_quant_config["model.layers.0.mlp.up_proj"].group_size == 64 +def _write_hf_quant_config(model_dir: Path, content: dict) -> Path: + """Write a ``hf_quant_config.json`` under ``model_dir`` and return its path.""" + path = model_dir / "hf_quant_config.json" + with open(path, 'w') as f: + json.dump(content, f) + return path + + +def test_quant_cfg_fp8_legacy_shape(): + """Plain FP8 modelopt 0.x checkpoint: legacy 'quantization' wrapper.""" + with tempfile.TemporaryDirectory() as tmp_dir: + model_dir = Path(tmp_dir) + hf_quant_config_file = _write_hf_quant_config( + model_dir, { + "producer": { + "name": "modelopt", + "version": "0.43.0" + }, + "quantization": { + "quant_algo": "FP8", + "kv_cache_quant_algo": "FP8", + "exclude_modules": ["lm_head"], + }, + }) + quant_config, layer_quant_config = ModelConfig.load_modelopt_quant_config( + hf_quant_config_file, model_dir, None) + assert quant_config.quant_algo == QuantAlgo.FP8 + assert quant_config.kv_cache_quant_algo == QuantAlgo.FP8 + assert quant_config.exclude_modules == ["lm_head"] + assert layer_quant_config is None + + +def test_quant_cfg_flat_shape_with_ignore_rename(): + """Modelopt 1.x flat shape: ``ignore`` is renamed to ``exclude_modules``.""" + with tempfile.TemporaryDirectory() as tmp_dir: + model_dir = Path(tmp_dir) + hf_quant_config_file = _write_hf_quant_config( + model_dir, { + "producer": { + "name": "modelopt", + "version": "1.0.0" + }, + "quant_method": "modelopt", + "quant_algo": "FP8", + "ignore": ["lm_head", "model.embed_tokens"], + }) + quant_config, _ = ModelConfig.load_modelopt_quant_config( + hf_quant_config_file, model_dir, None) + assert quant_config.quant_algo == QuantAlgo.FP8 + assert quant_config.exclude_modules == ["lm_head", "model.embed_tokens"] + + +def test_quant_cfg_flat_shape_kv_cache_scheme_dict(): + """Flat shape with compressed-tensors-style kv_cache_scheme dict (FP8).""" + with tempfile.TemporaryDirectory() as tmp_dir: + model_dir = Path(tmp_dir) + hf_quant_config_file = _write_hf_quant_config( + model_dir, { + "producer": { + "name": "modelopt", + "version": "1.0.0" + }, + "quant_method": "modelopt", + "quant_algo": "FP8", + "kv_cache_scheme": { + "dynamic": False, + "num_bits": 8, + "type": "float", + }, + "ignore": [], + }) + quant_config, _ = ModelConfig.load_modelopt_quant_config( + hf_quant_config_file, model_dir, None) + assert quant_config.kv_cache_quant_algo == QuantAlgo.FP8 + + +def test_quant_cfg_flat_shape_kv_cache_scheme_string_nvfp4(): + """Flat shape with bare-string kv_cache_scheme fallback (NVFP4).""" + with tempfile.TemporaryDirectory() as tmp_dir: + model_dir = Path(tmp_dir) + hf_quant_config_file = _write_hf_quant_config( + model_dir, { + "producer": { + "name": "modelopt", + "version": "1.0.0" + }, + "quant_method": "modelopt", + "quant_algo": "NVFP4", + "kv_cache_scheme": "NVFP4", + "ignore": [], + }) + quant_config, _ = ModelConfig.load_modelopt_quant_config( + hf_quant_config_file, model_dir, None) + assert quant_config.quant_algo == QuantAlgo.NVFP4 + assert quant_config.kv_cache_quant_algo == QuantAlgo.NVFP4 + + +def test_quant_cfg_fp8_pb_wo_alias_canonicalized(): + """Legacy ``fp8_pb_wo`` alias is canonicalized to FP8_BLOCK_SCALES.""" + with tempfile.TemporaryDirectory() as tmp_dir: + model_dir = Path(tmp_dir) + hf_quant_config_file = _write_hf_quant_config( + model_dir, { + "producer": { + "name": "modelopt" + }, + "quantization": { + "quant_algo": "fp8_pb_wo" + }, + }) + quant_config, _ = ModelConfig.load_modelopt_quant_config( + hf_quant_config_file, model_dir, None) + assert quant_config.quant_algo == QuantAlgo.FP8_BLOCK_SCALES + # FP8_BLOCK_SCALES default group_size. + assert quant_config.group_size == 128 + + +def test_quant_cfg_fp8_block_scales_trtllm_default_excludes(): + """TRTLLM moe_backend + FP8_BLOCK_SCALES + no excludes → defaults applied.""" + with tempfile.TemporaryDirectory() as tmp_dir: + model_dir = Path(tmp_dir) + hf_quant_config_file = _write_hf_quant_config( + model_dir, { + "producer": { + "name": "modelopt" + }, + "quantization": { + "quant_algo": "FP8_BLOCK_SCALES" + }, + }) + quant_config, _ = ModelConfig.load_modelopt_quant_config( + hf_quant_config_file, model_dir, "TRTLLM") + assert quant_config.exclude_modules == [ + "*kv_b_proj*", "*k_b_proj*", "*eh_proj" + ] + + +def test_quant_cfg_explicit_empty_excludes_preserved(): + """Explicit ``exclude_modules: []`` is preserved (no defaults applied).""" + with tempfile.TemporaryDirectory() as tmp_dir: + model_dir = Path(tmp_dir) + hf_quant_config_file = _write_hf_quant_config( + model_dir, { + "producer": { + "name": "modelopt" + }, + "quantization": { + "quant_algo": "FP8_BLOCK_SCALES", + "exclude_modules": [], + }, + }) + quant_config, _ = ModelConfig.load_modelopt_quant_config( + hf_quant_config_file, model_dir, "TRTLLM") + # Explicit [] must NOT trigger the TRTLLM default-excludes branch. + assert quant_config.exclude_modules == [] + + +def test_quant_cfg_mixed_precision_kv_cache_conflict_raises(): + """quant_cfg.json kv_cache_quant_algo conflicting with hf_quant_config.json raises.""" + with tempfile.TemporaryDirectory() as tmp_dir: + model_dir = Path(tmp_dir) + with open(model_dir / "quant_cfg.json", 'w') as f: + json.dump( + { + "quant_algo": "MIXED_PRECISION", + "kv_cache_quant_algo": "NVFP4", + "quantized_layers": { + "l0": { + "quant_algo": "FP8" + } + }, + }, f) + hf_quant_config_file = _write_hf_quant_config( + model_dir, { + "producer": { + "name": "modelopt" + }, + "quantization": { + "quant_algo": "MIXED_PRECISION", + "kv_cache_quant_algo": "FP8", + }, + }) + with pytest.raises(RuntimeError, match="kvcache config"): + ModelConfig.load_modelopt_quant_config(hf_quant_config_file, + model_dir, None) + + +def test_quant_cfg_awq_extra_fields_preserved_via_load_hf_quant_config(): + """AWQ extras (``has_zero_point``, ``pre_quant_scale``) flow through ``load_hf_quant_config``.""" + inline_modelopt_awq = { + "producer": { + "name": "modelopt" + }, + "quantization": { + "quant_algo": "W4A16_AWQ", + "group_size": 128, + "has_zero_point": False, + "pre_quant_scale": True, + }, + } + quant_config, _ = ModelConfig.load_hf_quant_config(inline_modelopt_awq, + moe_backend=None, + checkpoint_dir=None) + assert quant_config.quant_algo == QuantAlgo.W4A16_AWQ + assert quant_config.group_size == 128 + assert quant_config.has_zero_point is False + assert quant_config.pre_quant_scale is True + + +@pytest.mark.parametrize("config,expected", [ + ({ + "producer": { + "name": "modelopt" + } + }, True), + ({ + "quant_method": "modelopt" + }, True), + ({ + "quant_method": "modelopt-flat" + }, True), + ({ + "producer": { + "name": "other" + } + }, False), + ({ + "quant_method": "fp8" + }, False), + ({}, False), + ("not a dict", False), + (None, False), +]) +def test_is_modelopt_quant_config(config, expected): + """Producer name or quant_method prefix must signal modelopt.""" + from tensorrt_llm.quantization.modelopt_config import \ + is_modelopt_quant_config + assert is_modelopt_quant_config(config) is expected + + +@pytest.mark.parametrize( + "scheme,expected", + [ + (None, None), + # Bare-string fallback. + ("FP8", "FP8"), + ("NVFP4", "NVFP4"), + ("INT8", "INT8"), + ("int8", "INT8"), + # Compressed-tensors dict form. + ({ + "type": "float", + "num_bits": 8 + }, "FP8"), + ({ + "type": "float", + "num_bits": 4 + }, "NVFP4"), + ({ + "type": "int", + "num_bits": 8 + }, "INT8"), + # Unrecognized -> None. + ("UNKNOWN_ALGO", None), + ({ + "type": "float", + "num_bits": 16 + }, None), + (123, None), + ]) +def test_kv_cache_scheme_to_algo(scheme, expected): + """``_kv_cache_scheme_to_algo`` covers string + dict + None inputs.""" + from tensorrt_llm.quantization.modelopt_config import \ + _kv_cache_scheme_to_algo + assert _kv_cache_scheme_to_algo(scheme) == expected + + +@pytest.mark.parametrize("raw,match", [ + ("not a dict", "Expected dict"), + ({ + "producer": { + "name": "other" + } + }, "Not a modelopt quant config"), + ({ + "producer": { + "name": "modelopt" + }, + "quantization": "not a dict", + }, "'quantization' must be a dict"), +]) +def test_read_modelopt_quant_config_invalid_raises(raw, match): + """Non-dict / non-modelopt / malformed configs raise ValueError.""" + from tensorrt_llm.quantization.modelopt_config import \ + read_modelopt_quant_config + with pytest.raises(ValueError, match=match): + read_modelopt_quant_config(raw) + + +def test_quant_cfg_quant_algo_fields_are_enum_typed(): + """Top-level and per-layer ``quant_algo``/``kv_cache_quant_algo`` are QuantAlgo enums.""" + with tempfile.TemporaryDirectory() as tmp_dir: + model_dir = Path(tmp_dir) + hf_quant_config_file = _write_hf_quant_config( + model_dir, { + "producer": { + "name": "modelopt" + }, + "quantization": { + "quant_algo": "MIXED_PRECISION", + "kv_cache_quant_algo": "FP8", + "quantized_layers": { + "layer.0": { + "quant_algo": "NVFP4" + }, + }, + }, + }) + quant_config, layer_quant_config = ModelConfig.load_modelopt_quant_config( + hf_quant_config_file, model_dir, None) + # Top-level. + assert isinstance(quant_config.quant_algo, QuantAlgo) + assert isinstance(quant_config.kv_cache_quant_algo, QuantAlgo) + assert quant_config.kv_cache_quant_algo is QuantAlgo.FP8 + # Per-layer: quant_algo from the layer dict; kv_cache_quant_algo inherited. + layer = layer_quant_config["layer.0"] + assert isinstance(layer.quant_algo, QuantAlgo) + assert isinstance(layer.kv_cache_quant_algo, QuantAlgo) + assert layer.quant_algo is QuantAlgo.NVFP4 + assert layer.kv_cache_quant_algo is QuantAlgo.FP8 + + +@pytest.mark.parametrize("scheme", ["INT8", {"type": "int", "num_bits": 8}]) +def test_quant_cfg_flat_shape_kv_cache_scheme_int8(scheme): + """Flat shape: INT8 ``kv_cache_scheme`` honored via both string and dict forms.""" + with tempfile.TemporaryDirectory() as tmp_dir: + model_dir = Path(tmp_dir) + hf_quant_config_file = _write_hf_quant_config( + model_dir, { + "quant_method": "modelopt", + "quant_algo": "FP8", + "kv_cache_scheme": scheme, + "ignore": [], + }) + quant_config, _ = ModelConfig.load_modelopt_quant_config( + hf_quant_config_file, model_dir, None) + assert quant_config.kv_cache_quant_algo is QuantAlgo.INT8 + + +def test_quant_cfg_awq_extras_default_when_absent(): + """When AWQ extras are absent from the JSON, ``QuantConfig`` defaults are preserved.""" + with tempfile.TemporaryDirectory() as tmp_dir: + model_dir = Path(tmp_dir) + hf_quant_config_file = _write_hf_quant_config( + model_dir, { + "producer": { + "name": "modelopt" + }, + "quantization": { + "quant_algo": "FP8" + }, + }) + quant_config, _ = ModelConfig.load_modelopt_quant_config( + hf_quant_config_file, model_dir, None) + # QuantConfig defaults: has_zero_point=False, pre_quant_scale=False. + assert quant_config.has_zero_point is False + assert quant_config.pre_quant_scale is False + + +def test_load_hf_quant_config_fp8_block_scales_deepseek_v3(): + """DeepSeek V3 ``quant_method=fp8`` with weight_block_size=(128,128).""" + quant_config, _ = ModelConfig.load_hf_quant_config( + { + "quant_method": "fp8", + "weight_block_size": [128, 128] + }, + moe_backend=None) + assert quant_config.quant_algo == QuantAlgo.FP8_BLOCK_SCALES + assert quant_config.group_size == 128 + # Default excludes for FP8_BLOCK_SCALES include kv/eh-proj patterns. + assert "*kv_b_proj*" in quant_config.exclude_modules + + +@pytest.mark.parametrize("weights_strategy,inputs_strategy,expected_algo", [ + ("channel", "token", QuantAlgo.FP8_PER_CHANNEL_PER_TOKEN), + ("block", "group", QuantAlgo.FP8_BLOCK_SCALES), +]) +def test_load_hf_quant_config_compressed_tensors(weights_strategy, + inputs_strategy, + expected_algo): + """LLM-compressor ``compressed-tensors``: strategy combinations map to TRT-LLM algos.""" + inputs_cfg = {"num_bits": 8, "strategy": inputs_strategy} + if inputs_strategy == "group": + inputs_cfg["group_size"] = 128 + quant_config, _ = ModelConfig.load_hf_quant_config( + { + "quant_method": "compressed-tensors", + "config_groups": { + "group_0": { + "weights": { + "num_bits": 8, + "strategy": weights_strategy + }, + "input_activations": inputs_cfg, + }, + }, + "ignore": ["lm_head"], + }, + moe_backend=None) + assert quant_config.quant_algo == expected_algo + assert quant_config.exclude_modules == ["lm_head"] + + +def test_load_hf_quant_config_nvfp4_native_with_modules_to_not_convert(): + """HF nvfp4 schema: ``modules_to_not_convert`` is merged into ``exclude_modules``.""" + quant_config, _ = ModelConfig.load_hf_quant_config( + { + "quant_method": "nvfp4", + "group_size": 16, + "modules_to_not_convert": ["custom_layer"], + }, + moe_backend=None) + assert quant_config.quant_algo == QuantAlgo.NVFP4 + assert quant_config.group_size == 16 + assert "custom_layer" in quant_config.exclude_modules + assert "*.mlp.gate" in quant_config.exclude_modules # default + assert "lm_head" in quant_config.exclude_modules # default + + +def test_load_hf_quant_config_no_match_returns_empty_quant_config(): + """An unrecognized ``quant_method`` returns an empty QuantConfig (no algo set).""" + quant_config, _ = ModelConfig.load_hf_quant_config( + {"quant_method": "unknown_format"}, moe_backend=None) + assert quant_config.quant_algo is None + + if __name__ == "__main__": test_llm_int4_awq_quantization() test_llm_fp8_quantization_modelOpt_ckpt() From ddb8a79b6371a82d3adb01217b29ecf66b3e1324 Mon Sep 17 00:00:00 2001 From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> Date: Tue, 19 May 2026 23:41:44 -0700 Subject: [PATCH 3/7] [None][feature] Add env variables to help debugging mamba modules Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> --- tensorrt_llm/_torch/model_config.py | 8 +++- tensorrt_llm/_torch/pyexecutor/_util.py | 42 ++++++++++++++----- .../_torch/pyexecutor/mamba_cache_manager.py | 27 +++++++++++- .../_torch/pyexecutor/py_executor_creator.py | 11 +++++ 4 files changed, 74 insertions(+), 14 deletions(-) diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index 4e8ad3e73652..d9e3fb8cf17a 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -48,15 +48,19 @@ def _unified_kv_pool_includes_mamba( * disaggregated serving forces the C++ mamba manager (``TRTLLM_USE_CPP_MAMBA=1`` enables the same path locally), or + * ``TRTLLM_USE_PY_MAMBA=1`` forces the Python mamba manager locally + (agg-mode override), or * one-model speculative decoding splits mamba and attention into separate caches. Single source of truth for the binding-side layer-counting decision; do not duplicate the predicate at call sites. """ - use_disagg = is_disagg or os.environ.get('TRTLLM_USE_CPP_MAMBA', '0') == '1' + use_split_pool = is_disagg \ + or os.environ.get('TRTLLM_USE_CPP_MAMBA', '0') == '1' \ + or os.environ.get('TRTLLM_USE_PY_MAMBA', '0') == '1' use_spec = spec_config is not None - return not (use_disagg or use_spec) + return not (use_split_pool or use_spec) @contextlib.contextmanager diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 29bdc335a88c..f8ff19ff519c 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -41,7 +41,8 @@ from .mamba_cache_manager import (BaseMambaCacheManager, CppMambaHybridCacheManager, MixedMambaHybridCacheManager, - use_cpp_mamba_cache_manager) + use_cpp_mamba_cache_manager, + use_py_mamba_cache_manager) from .model_engine import PyTorchModelEngine from .py_executor import PyExecutor from .resource_manager import (KVCacheManager, KVCacheManagerV2, @@ -75,9 +76,14 @@ def get_kv_cache_manager_cls(model_config: ModelConfig, """Resolve the concrete KV cache manager class for ``model_config``. For hybrid mamba models the choice between ``Mixed`` (separate pools, - needed for disagg / TRTLLM_USE_CPP_MAMBA) and ``Cpp`` (unified pool with - block reuse) is made here. Callers that don't care about disagg can omit - ``is_disagg`` and get the unified-pool default. + needed for disagg / TRTLLM_USE_CPP_MAMBA / TRTLLM_USE_PY_MAMBA) and + ``Cpp`` (unified pool with block reuse) is made here. Callers that don't + care about disagg can omit ``is_disagg`` and get the unified-pool default. + + Env-var overrides (agg mode only — disagg picks its inner impl via + ``cache_transceiver_config.transceiver_runtime``): + * ``TRTLLM_USE_CPP_MAMBA=1`` — Mixed manager with CppMambaCacheManager. + * ``TRTLLM_USE_PY_MAMBA=1`` — Mixed manager with PythonMambaCacheManager. """ config = model_config.pretrained_config sparse_attn_config = model_config.sparse_attention_config @@ -90,9 +96,10 @@ def get_kv_cache_manager_cls(model_config: ModelConfig, logger.info("Hybrid linear model has 0 mamba layers; using " "KVCacheManager without mamba caching") return _non_hybrid_kv_cache_manager_cls(config, kv_cache_config) - if kv_cache_config.enable_block_reuse: - return CppMambaHybridCacheManager - if is_disagg or use_cpp_mamba_cache_manager(): + # if kv_cache_config.enable_block_reuse: + # return CppMambaHybridCacheManager + if is_disagg or use_cpp_mamba_cache_manager( + ) or use_py_mamba_cache_manager(): return MixedMambaHybridCacheManager default_cls = CppMambaHybridCacheManager env_override = os.environ.get('TLLM_MAMBA_MANAGER_PREFERENCE', None) @@ -256,14 +263,16 @@ def _get_model_kv_cache_manager_cls(self, model_engine: PyTorchModelEngine): "event buffer max size > 0, or cache transceiver. Falling back to KVCacheManager." ) cls = KVCacheManager - # The V1-route hybrid mamba managers (disagg via TRTLLM_USE_CPP_MAMBA, - # or one-model speculative decoding) keep mamba state in a separate - # cache that doesn't honor block reuse. Warn at the routing site so - # users see the warning where the decision is actually made. + # The V1-route hybrid mamba managers (disagg, TRTLLM_USE_CPP_MAMBA, + # TRTLLM_USE_PY_MAMBA, or one-model speculative decoding) keep mamba + # state in a separate cache that doesn't honor block reuse. Warn at + # the routing site so users see the warning where the decision is + # actually made. if is_hybrid_linear(model_engine.model.model_config.pretrained_config) \ and self._kv_cache_config.enable_block_reuse: uses_v1_mamba_route = self._is_disagg \ or os.environ.get('TRTLLM_USE_CPP_MAMBA', '0') == '1' \ + or os.environ.get('TRTLLM_USE_PY_MAMBA', '0') == '1' \ or self._speculative_config is not None if uses_v1_mamba_route: logger.warning( @@ -1291,6 +1300,17 @@ def _create_kv_cache_manager( "using legacy MTP path for stochastic rounding support") use_replay = False + # Use replay algorithm for mamba (default is on). + enforce_disable_replay = os.environ.get('TRTLLM_USE_MAMBA_REPLAY', + '1') == '0' + if enforce_disable_replay: + logger.info( + "Replay kernel is disabled by TRTLLM_USE_MAMBA_REPLAY=0") + use_replay = False + else: + logger.info( + "Replay kernel is not changed since TRTLLM_USE_MAMBA_REPLAY=1") + kv_cache_manager = kv_cache_manager_cls( # mamba cache parameters mamba_params.state_size, diff --git a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py index 194bdc7fa87d..bf6868f99b19 100644 --- a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py @@ -61,7 +61,32 @@ def use_cpp_mamba_cache_manager() -> bool: Returns True if TRTLLM_USE_CPP_MAMBA='1' is set, False otherwise. By default, PythonMambaCacheManager is used. """ - return os.environ.get('TRTLLM_USE_CPP_MAMBA', '0') == '1' + cpp = os.environ.get('TRTLLM_USE_CPP_MAMBA', '0') == '1' + py = os.environ.get('TRTLLM_USE_PY_MAMBA', '0') == '1' + if cpp and py: + raise ValueError( + "TRTLLM_USE_CPP_MAMBA=1 and TRTLLM_USE_PY_MAMBA=1 are mutually " + "exclusive; unset one of them.") + return cpp + + +def use_py_mamba_cache_manager() -> bool: + """Check if PythonMambaCacheManager should be forced (agg mode override). + + Returns True if TRTLLM_USE_PY_MAMBA='1' is set, False otherwise. + + Agg-mode-only override: forces the V1-route MixedMambaHybridCacheManager + with PythonMambaCacheManager inside instead of the default unified-pool + CppMambaHybridCacheManager. Disagg mode is unaffected — it already picks + PythonMambaCacheManager when transceiver_runtime='PYTHON'. + """ + cpp = os.environ.get('TRTLLM_USE_CPP_MAMBA', '0') == '1' + py = os.environ.get('TRTLLM_USE_PY_MAMBA', '0') == '1' + if cpp and py: + raise ValueError( + "TRTLLM_USE_CPP_MAMBA=1 and TRTLLM_USE_PY_MAMBA=1 are mutually " + "exclusive; unset one of them.") + return py class BaseMambaCacheManager(ABC): diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index fa33e0ce8643..de2fac2b6330 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -829,6 +829,17 @@ def drafting_loop_wrapper(model): is_hybrid = is_hybrid_linear(config) if is_disagg and is_hybrid: + # NOTE: TRTLLM_USE_PY_MAMBA is an agg-mode-only override and has + # no effect in disagg. The disagg manager choice is driven solely + # by transceiver_runtime: PYTHON => PythonMambaCacheManager, + # otherwise CppMambaCacheManager. Clear the var here so the + # mutual-exclusion check in use_cpp_mamba_cache_manager() does + # not fire after we force TRTLLM_USE_CPP_MAMBA=1 below. + if os.environ.pop("TRTLLM_USE_PY_MAMBA", "0") == "1": + logger.warning( + "TRTLLM_USE_PY_MAMBA is ignored in disaggregated serving; " + "use cache_transceiver_config.transceiver_runtime='PYTHON' " + "to select PythonMambaCacheManager.") if cache_transceiver_config.transceiver_runtime != "PYTHON" or os.environ.get( "TRTLLM_USE_CPP_MAMBA") == "1": logger.info("Disaggregated serving with hybrid model detected. " From 521c8c67b52eb4b87101943cdc2b3f47194c1ada Mon Sep 17 00:00:00 2001 From: nv-guomingz <137257613+nv-guomingz@users.noreply.github.com> Date: Mon, 18 May 2026 13:50:57 +0000 Subject: [PATCH 4/7] [None][fix] ADP router crashes on serve when scheduling_params.attention_dp_relax is None Signed-off-by: nv-guomingz <137257613+nv-guomingz@users.noreply.github.com> --- .../_torch/pyexecutor/scheduler/adp_router.py | 6 ++- .../_torch/executor/test_adp_router.py | 47 +++++++++++++++++++ 2 files changed, 51 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py b/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py index 9a634ddb98f0..07eb9e13d4c2 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py @@ -304,7 +304,8 @@ def get_relax_value(req_item): scheduling_params = getattr(req_item.request, "py_scheduling_params", None) if scheduling_params is None: return True - return scheduling_params.attention_dp_relax + val = scheduling_params.attention_dp_relax + return True if val is None else val sorted_requests = sorted(new_requests, key=get_relax_value) @@ -581,7 +582,8 @@ def get_relax_value(req_item): scheduling_params = getattr(req_item.request, "py_scheduling_params", None) if scheduling_params is None: return True - return scheduling_params.attention_dp_relax + val = scheduling_params.attention_dp_relax + return True if val is None else val sorted_requests = sorted(new_requests, key=get_relax_value) diff --git a/tests/unittest/_torch/executor/test_adp_router.py b/tests/unittest/_torch/executor/test_adp_router.py index 9cf7a149ad0e..4ead152409b7 100644 --- a/tests/unittest/_torch/executor/test_adp_router.py +++ b/tests/unittest/_torch/executor/test_adp_router.py @@ -463,6 +463,53 @@ def test_schedule_attention_dp_requests_empty_lists( assert len(result) == 0 +def test_schedule_attention_dp_requests_serve_default_relax_None_does_not_crash( + attention_dp_config, all_ranks_num_active_requests, all_ranks_num_active_tokens +): + """Regression: trtllm-serve produces SchedulingParams with attention_dp_relax=None. + + openai_server.py builds ``SchedulingParams(agent_hierarchy=...)`` on every + chat request, leaving ``attention_dp_rank`` and ``attention_dp_relax`` at + their dataclass default of None. With >=2 such requests in one scheduling + iteration, ``sorted(new_requests, key=get_relax_value)`` previously raised + ``TypeError: '<' not supported between instances of 'NoneType' and + 'NoneType'`` because the key returned None for every item. + + Why prior coverage missed it: + * Every ADP CI E2E test calls ``LLM(...).generate`` where + ``request.scheduling_params`` is None, so ``base_worker`` leaves + ``py_scheduling_params=None`` and ``get_relax_value`` short-circuits + on the ``if scheduling_params is None`` branch. + * The mock factory in this file defaults ``attention_dp_relax=False``, + not None, so existing unit tests never construct the buggy shape. + + Use the real ``SchedulingParams`` dataclass so the test tracks the actual + production default — if it ever flips back to None this stays accurate. + """ + from tensorrt_llm.scheduling_params import SchedulingParams + + def _make_serve_shaped_request(req_id): + mock_request = Mock() + mock_request.py_scheduling_params = SchedulingParams() + mock_request.input_token_ids = [1, 2, 3] + return RequestQueueItem(req_id, mock_request) + + new_requests = [_make_serve_shaped_request(i) for i in range(4)] + + # No raise = sort key is None-safe. All four items have + # attention_dp_rank=None so they fall into remaining_unscheduled and + # then balance across ranks; with capacity 8 per rank they all land. + all_ranks_new_requests, _ = _assign( + all_ranks_num_active_requests, + all_ranks_num_active_tokens, + new_requests, + attention_dp_config["max_num_active_requests"], + ) + + total_assigned = sum(len(reqs) for reqs in all_ranks_new_requests.values()) + assert total_assigned == 4 + + def test_schedule_attention_dp_requests_expected_num_active_calculation( attention_dp_config, all_ranks_num_active_requests, all_ranks_num_active_tokens ): From c5b2ac795dbd470534c1b0c326a20331dd520702 Mon Sep 17 00:00:00 2001 From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> Date: Tue, 19 May 2026 22:40:22 -0700 Subject: [PATCH 5/7] [None][fix] Hide trailing EOS from generated text Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> --- tensorrt_llm/executor/result.py | 22 +++++++++++++++++++--- 1 file changed, 19 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/executor/result.py b/tensorrt_llm/executor/result.py index 420fb16ab3b1..9751e586f868 100644 --- a/tensorrt_llm/executor/result.py +++ b/tensorrt_llm/executor/result.py @@ -378,6 +378,8 @@ def _handle_sequence(self, if sequence_is_finished: if finish_reasons[src_idx] == tllm.FinishReason.END_ID: + # Keep the trailing EOS in token_ids; the detokenizer hides + # it from text. stop_reason stays None to mark "natural EOS". output.finish_reason = 'stop' elif finish_reasons[src_idx] == tllm.FinishReason.STOP_WORDS: output.finish_reason = 'stop' @@ -748,22 +750,36 @@ def _handle_response(self, response: "GenerationExecutor.Response"): 'spaces_between_special_tokens': self.sampling_params.spaces_between_special_tokens } + end_id = self.sampling_params.end_id + include_stop = self.sampling_params.include_stop_str_in_output + if self.sampling_params.detokenize and self.tokenizer is not None: for beam_output in self.outputs: beam_output._last_text_len = len(beam_output.text) + # Hide trailing EOS from decoded text without mutating + # token_ids (matches vLLM). stop_reason is None gates on + # END_ID; STOP_WORDS sets it and was stripped upstream. + trim_eos = (beam_output.finish_reason == 'stop' + and beam_output.stop_reason is None + and not include_stop and end_id is not None) if hasattr( self.tokenizer, 'decode_incrementally' ) and self._streaming and not self.sampling_params.use_beam_search: + diff = beam_output.token_ids_diff + if trim_eos and diff and diff[-1] == end_id: + diff = diff[:-1] beam_output.text, beam_output._incremental_states = self.tokenizer.decode_incrementally( - beam_output.token_ids_diff, + diff, prev_text=beam_output.text, states=beam_output._incremental_states, flush=self._done, stream_interval=self.sampling_params._stream_interval, **kwargs) else: - beam_output.text = self.tokenizer.decode( - beam_output.token_ids, **kwargs) + tids = beam_output.token_ids + if trim_eos and tids and tids[-1] == end_id: + tids = tids[:-1] + beam_output.text = self.tokenizer.decode(tids, **kwargs) is_generating = not self._done is_finished_with_stop_or_length = ( From 2b44d0d48bf039220c2494d42315d2447efc4cac Mon Sep 17 00:00:00 2001 From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> Date: Mon, 18 May 2026 22:27:28 -0700 Subject: [PATCH 6/7] Fix chat-template missing bug * only with issue with transformers 4.x Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> --- tensorrt_llm/serve/openai_server.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index 1b0b113f0a36..8f07d070f32c 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -374,7 +374,22 @@ def _init_llm(self, chat_template: Optional[str] = None): logger.debug("Failed to load AutoConfig for %s", hf_tokenizer_path) self.model_config = None + # Resolution order: + # 1. Explicit --chat_template (file path or literal Jinja string). + # 2. Sidecar /chat_template.jinja, if present. + # Step 2 is required because transformers < 5.0 does not auto-load + # the sidecar into tokenizer.chat_template, and the model's + # tokenizer_class may not be recognised either. Without this fallback, + # any /v1/chat/completions request would raise + # "No chat template found for the given tokenizer and tools." from + # tensorrt_llm/inputs/utils.py. self.chat_template = load_chat_template(chat_template) + if self.chat_template is None and hf_tokenizer_path: + sidecar = Path(hf_tokenizer_path) / "chat_template.jinja" + if sidecar.is_file(): + self.chat_template = load_chat_template(sidecar) + logger.info("Loaded chat template from sidecar file: %s", + sidecar) # Enable response storage for Responses API self.enable_store = (len( From c6c611d7a5e64c4c374954ef1436fb4ed9ef180a Mon Sep 17 00:00:00 2001 From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> Date: Wed, 20 May 2026 21:14:09 -0700 Subject: [PATCH 7/7] Disable replay by default Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/_util.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index f8ff19ff519c..5bd3c45b12e9 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -1300,9 +1300,9 @@ def _create_kv_cache_manager( "using legacy MTP path for stochastic rounding support") use_replay = False - # Use replay algorithm for mamba (default is on). + # Use replay algorithm for mamba (default is off). enforce_disable_replay = os.environ.get('TRTLLM_USE_MAMBA_REPLAY', - '1') == '0' + '0') == '0' if enforce_disable_replay: logger.info( "Replay kernel is disabled by TRTLLM_USE_MAMBA_REPLAY=0")