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 5cdb2171b0e6..83b230b1ecc9 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/auto_deploy/models/quant_config_reader.py b/tensorrt_llm/_torch/auto_deploy/models/quant_config_reader.py index 6620d3a39bec..c390cc557c2c 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, + parse_modelopt_quant_config, +) + from ..utils.logger import ad_logger @@ -110,14 +115,23 @@ 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}'") + # Accept either modelopt shape: legacy (producer.name == "modelopt" + # with a "quantization" wrapper) or flat (quant_method == "modelopt"). + # A bare producer check would reject flat configs that omit producer. + 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')})" + ) - quant_config = config.get("quantization", {}) + # Downstream auto-deploy transforms read field-by-field from a dict + # using the legacy field names. Parse via the parallel readers and + # re-render as the legacy inner dict for those consumers. + parsed = parse_modelopt_quant_config(config) + quant_config = parsed.to_legacy_inner_dict() - quant_algo = quant_config.get("quant_algo", "").upper() + quant_algo = (quant_config.get("quant_algo") or "").upper() if quant_algo == "MIXED_PRECISION": self._read_mixed_precision_config(quant_config) diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index b96e40d513e3..953f1dc48f14 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -23,6 +23,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 ( + ModelOptQuantConfig, is_modelopt_quant_config, parse_modelopt_quant_config, + warn_if_inline_quant_config_differs) if TYPE_CHECKING: from tensorrt_llm.bindings import ModelConfig as ModelConfigCpp @@ -44,15 +47,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 @@ -295,60 +302,64 @@ def resolve_moe_backend(moe_backend: str, architecture: str) -> 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) + raw = json.load(f) + parsed = parse_modelopt_quant_config(raw) + return ModelConfig._build_quant_config_from_modelopt( + parsed, checkpoint_dir, moe_backend) - json_quant_configs = quant_config_dict['quantization'] + @staticmethod + def _build_quant_config_from_modelopt(parsed: ModelOptQuantConfig, + checkpoint_dir, moe_backend): + quant_config = QuantConfig() + layer_quant_config = None - 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) - quant_config.group_size = json_quant_configs.get('group_size', None) - quant_config.exclude_modules = json_quant_configs.get( - 'exclude_modules', None) + quant_config.quant_algo = parsed.quant_algo + quant_config.kv_cache_quant_algo = parsed.kv_cache_quant_algo + quant_config.group_size = parsed.group_size + quant_config.exclude_modules = list( + parsed.exclude_modules) if parsed.exclude_modules else None if quant_config.quant_algo == QuantAlgo.MIXED_PRECISION: - json_extended_quant_configs: dict = {} + # Optional overlay: quant_cfg.json may carry per-layer overrides + # and/or a kv_cache_quant_algo that must agree with the file. # See tests/unittest/llmapi/test_llm_quant.py + extended: dict = {} try: - mixed_quant_config_file = transformers.utils.hub.cached_file( + ext_file = transformers.utils.hub.cached_file( checkpoint_dir, 'quant_cfg.json') - with open(mixed_quant_config_file) as fm: - json_extended_quant_configs = json.load(fm) + with open(ext_file) as fm: + extended = 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) - mixed_quant_configs = json_quant_configs.get( - 'quantized_layers', None) - if (kv_quant_lhs := json_extended_quant_configs.get( - "kv_cache_quant_algo", None)) is not None and ( - kv_quant_rhs := - quant_config.kv_cache_quant_algo) is not None: - if kv_quant_lhs != kv_quant_rhs: - raise RuntimeError( - f"The kvcache config in 'quant_cfg.json', {kv_quant_lhs}," - f"is different from 'hf_quant_config.json', {kv_quant_rhs}!" - ) + + extended_kv = extended.get("kv_cache_quant_algo") + if (extended_kv is not None + and parsed.kv_cache_quant_algo is not None + and extended_kv != parsed.kv_cache_quant_algo): + raise RuntimeError( + f"The kvcache config in 'quant_cfg.json', {extended_kv}," + f"is different from 'hf_quant_config.json', {parsed.kv_cache_quant_algo}!" + ) + kv_cache_quant_algo = (extended_kv if extended_kv is not None else + parsed.kv_cache_quant_algo) quant_config.kv_cache_quant_algo = kv_cache_quant_algo - for layer in mixed_quant_configs: - 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) - mixed_quant_configs[layer] = config - layer_quant_config = mixed_quant_configs + + # extended.quantized_layers replaces the file's wholesale, matching + # the historical ``dict.update`` overlay semantics. + quantized_layers = extended.get("quantized_layers", + parsed.quantized_layers) + + per_layer = {} + for layer, cfg in quantized_layers.items(): + c = QuantConfig() + c.kv_cache_quant_algo = kv_cache_quant_algo + c.quant_algo = cfg['quant_algo'] + c.group_size = cfg.get('group_size') + per_layer[layer] = c + layer_quant_config = per_layer elif quant_config.quant_algo == QuantAlgo.FP8_BLOCK_SCALES: if quant_config.group_size is None: quant_config.group_size = 128 @@ -374,10 +385,23 @@ 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 + # Inline modelopt config -- legacy ("quantization" wrapper + + # producer.name == "modelopt") or flat (quant_method == "modelopt"). + # The elif chain below only knows fp8 / mxfp4 / compressed-tensors + # / nvfp4, so without this route the modelopt case falls through + # to NotImplementedError. Early-return because the modelopt builder + # returns the full (QuantConfig, layer_quant_config) tuple including + # MIXED_PRECISION layer expansion, which doesn't fit the elif + # "mutate one QuantConfig" pattern. + if is_modelopt_quant_config(hf_quant_config): + parsed = parse_modelopt_quant_config(hf_quant_config) + return ModelConfig._build_quant_config_from_modelopt( + parsed, 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) @@ -673,13 +697,21 @@ 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) + with open(quant_config_file) as f: + hf_quant_config_raw = json.load(f) + parsed = parse_modelopt_quant_config(hf_quant_config_raw) + quant_config, layer_quant_config = cls._build_quant_config_from_modelopt( + parsed, checkpoint_dir, moe_backend) + warn_if_inline_quant_config_differs( + parsed, + getattr(pretrained_config, "quantization_config", None), + source_file="hf_quant_config.json", + ) # 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) + hf_quant_config, moe_backend, 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) 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 afac457a2c2a..a5e78eb7e44c 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,30 @@ 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 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 - 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) @@ -236,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( @@ -1256,15 +1285,27 @@ 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 + # !!! We will not use replay for MTP, to avoid the accuracy drop issue. + enforce_disable_replay = os.environ.get('TRTLLM_USE_MAMBA_REPLAY', + '0') == '0' + if enforce_disable_replay: + logger.info( + "Replay kernel is disabled by TRTLLM_USE_MAMBA_REPLAY=0") + use_replay = False + else: + logger.info( + f"Replay kernel is set as {use_replay} by 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 da682bebc858..bf6868f99b19 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 @@ -60,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): @@ -1000,6 +1026,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 +1086,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 +1175,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 +1195,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 +1255,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 +1339,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 +1368,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 +1427,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 +1527,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/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. " diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index 19ffa08b8a07..04d92ad60ce9 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -336,6 +336,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, @@ -454,13 +455,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) @@ -1202,7 +1209,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}" ) @@ -1218,10 +1232,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, @@ -1532,7 +1552,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. @@ -1555,10 +1580,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/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py b/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py index 0537b3821f9a..5e3a4ab96b23 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 + relax = scheduling_params.attention_dp_relax + return True if relax is None else relax sorted_requests = sorted(new_requests, key=get_relax_value) @@ -576,7 +577,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 + relax = scheduling_params.attention_dp_relax + return True if relax is None else relax sorted_requests = sorted(new_requests, key=get_relax_value) diff --git a/tensorrt_llm/llmapi/llm_utils.py b/tensorrt_llm/llmapi/llm_utils.py index b37cdd03658d..51fe59cacefb 100644 --- a/tensorrt_llm/llmapi/llm_utils.py +++ b/tensorrt_llm/llmapi/llm_utils.py @@ -27,6 +27,10 @@ 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 (ModelOptQuantConfig, + is_modelopt_quant_config, + parse_modelopt_quant_config, + warn_if_inline_quant_config_differs) from .build_cache import (BuildCache, BuildCacheConfig, CachedStage, get_build_cache_config_from_env) # yapf: disable @@ -318,6 +322,64 @@ 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 + @staticmethod + def _apply_modelopt_quant_config( + parsed: ModelOptQuantConfig, quant_config: QuantConfig, + explicit_kv_cache_quant_algo: Optional[QuantAlgo], + kv_cache_dtype: Optional[str]) -> None: + """Apply a parsed modelopt config onto ``quant_config`` in place.""" + hf_quant_algo = parsed.quant_algo + 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 = (QuantAlgo(parsed.kv_cache_quant_algo) + if parsed.kv_cache_quant_algo is not None else + None) + if hf_kv_cache_quant_algo is not None: + 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 consumed separately by the PyTorch model_config + # layer; exclude_modules and group_size are picked up below if set. + if parsed.group_size is not None: + logger.info( + f"Setting group_size={parsed.group_size} from HF quant config.") + quant_config.group_size = parsed.group_size + if parsed.exclude_modules: + logger.info( + f"Setting exclude_modules ({len(parsed.exclude_modules)} entries) from HF quant config." + ) + quant_config.exclude_modules = parsed.exclude_modules + def _update_from_hf_quant_config(self) -> bool: """Update quant_config from the config file of pre-quantized HF checkpoint. @@ -337,75 +399,28 @@ 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) + hf_quant_config_raw = json.load(f) + parsed = parse_modelopt_quant_config(hf_quant_config_raw) - 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) + # Cross-check against inline config.json.quantization_config if any. + hf_config_path = f"{self._model_dir}/config.json" + inline_quant_config = None + try: + with open(hf_config_path, "r") as f: + inline_quant_config = json.load(f).get( + "quantization_config") + except FileNotFoundError: + pass + warn_if_inline_quant_config_differs( + parsed, + inline_quant_config, + source_file="hf_quant_config.json", + ) - # Update the quant_config in llm_args for pytorch + self._apply_modelopt_quant_config(parsed, quant_config, + explicit_kv_cache_quant_algo, + kv_cache_dtype) self.llm_args.quant_config = quant_config - return True hf_config_path = f"{self._model_dir}/config.json" @@ -429,6 +444,14 @@ def _update_from_hf_quant_config(self) -> bool: ) if hf_quant_config is not None: + # Inline modelopt config -- legacy or flat. Reuse the file-path parser. + if is_modelopt_quant_config(hf_quant_config): + parsed = parse_modelopt_quant_config(hf_quant_config) + self._apply_modelopt_quant_config(parsed, quant_config, + explicit_kv_cache_quant_algo, + kv_cache_dtype) + self.llm_args.quant_config = quant_config + return True # DeepSeek V3 FP8 ckpt if hf_quant_config.get("quant_method") == "fp8": if hf_quant_config.get("weight_block_size") 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..7869ef1a3caa --- /dev/null +++ b/tensorrt_llm/quantization/modelopt_config.py @@ -0,0 +1,245 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Unified parser for ModelOpt-produced quantization configs. + +ModelOpt emits ``hf_quant_config.json`` (and the inline +``config.json.quantization_config``) in two on-disk shapes: + +1. **Legacy** (modelopt 0.43.x and earlier) -- nested under a ``quantization`` + wrapper, with ``exclude_modules`` for the deny-list and a string + ``kv_cache_quant_algo``:: + + { + "producer": {...}, + "quantization": { + "quant_algo": ..., + "quantized_layers": {...}, + "exclude_modules": [...], + "kv_cache_quant_algo": ..., + }, + } + +2. **Flat** (modelopt 1.0.x, compressed-tensors-style) -- all fields hoisted + to the top level, with ``ignore`` for the deny-list and a dict + ``kv_cache_scheme`` in place of ``kv_cache_quant_algo``:: + + { + "producer": {...}, + "quant_method": "modelopt", + "quant_algo": ..., + "quantized_layers": {...}, + "ignore": [...], + "config_groups": {...}, + "kv_cache_scheme": {...}, + } + +This module uses parallel readers: one branch per shape, both producing +a common :class:`ModelOptQuantConfig` struct. Each caller in TensorRT-LLM +(PyTorch :class:`ModelConfig`, the LLM-API loader, the AutoDeploy +quant-config reader) consumes the struct in whatever shape it needs. +""" + +from dataclasses import dataclass, field +from typing import Any, Dict, List, Optional + +from ..logger import logger + + +@dataclass +class ModelOptQuantConfig: + """Common struct produced by either legacy or flat shape readers.""" + + quant_algo: Optional[str] = None + kv_cache_quant_algo: Optional[str] = None + group_size: Optional[int] = None + exclude_modules: List[str] = field(default_factory=list) + # Per-layer overrides: {fully_qualified_name: {"quant_algo": ..., "group_size": ...}} + quantized_layers: Dict[str, Dict[str, Any]] = field(default_factory=dict) + # Which shape produced this struct; useful for logging. + source_format: str = "" + + def to_legacy_inner_dict(self) -> Dict[str, Any]: + """Render as the legacy ``quantization`` inner dict. + + Used by AutoDeploy whose downstream consumers operate on a dict + with the legacy field names. + """ + out: Dict[str, Any] = { + "quant_algo": self.quant_algo, + "kv_cache_quant_algo": self.kv_cache_quant_algo, + "group_size": self.group_size, + "exclude_modules": list(self.exclude_modules), + "quantized_layers": dict(self.quantized_layers), + } + # Drop keys whose value is "absent" so downstream `.get(..., default)` + # returns the default rather than an explicit None. + return {k: v for k, v in out.items() if v is not None} + + +def _kv_cache_scheme_to_algo(scheme: Any) -> Optional[str]: + """Translate compressed-tensors ``kv_cache_scheme`` to a string algo.""" + if not isinstance(scheme, dict): + return None + if scheme.get("type") == "float" and scheme.get("num_bits") == 8: + return "FP8" + return None + + +def is_modelopt_quant_config(raw: Dict[str, Any]) -> bool: + """Return True if ``raw`` looks like a ModelOpt-produced quant config. + + Detects either shape via ``producer.name == "modelopt"`` or + ``quant_method`` starting with ``"modelopt"``. + """ + if not isinstance(raw, dict): + return False + if str(raw.get("quant_method", "")).lower().startswith("modelopt"): + return True + if (raw.get("producer") or {}).get("name") == "modelopt": + return True + return False + + +def _read_legacy(raw: Dict[str, Any]) -> ModelOptQuantConfig: + """Reader for the modelopt 0.43.x nested shape.""" + q = raw.get("quantization") + if not isinstance(q, dict): + raise ValueError( + f"'quantization' must be a dict for legacy ModelOpt config, got {type(q).__name__}." + ) + return ModelOptQuantConfig( + quant_algo=q.get("quant_algo"), + kv_cache_quant_algo=q.get("kv_cache_quant_algo"), + group_size=q.get("group_size"), + exclude_modules=list(q.get("exclude_modules") or []), + # Alias the json.load'd quantized_layers (up to ~50k entries); the raw + # dict isn't reused after parse, so a defensive copy would be waste. + quantized_layers=q.get("quantized_layers") or {}, + source_format="legacy", + ) + + +def _read_flat(raw: Dict[str, Any]) -> ModelOptQuantConfig: + """Reader for the modelopt 1.0.x flat / compressed-tensors-style shape.""" + return ModelOptQuantConfig( + quant_algo=raw.get("quant_algo"), + kv_cache_quant_algo=_kv_cache_scheme_to_algo(raw.get("kv_cache_scheme")), + group_size=raw.get("group_size"), + exclude_modules=list(raw.get("ignore") or []), + quantized_layers=raw.get("quantized_layers") or {}, + source_format="flat", + ) + + +def parse_modelopt_quant_config(raw: Dict[str, Any]) -> ModelOptQuantConfig: + """Parse a ModelOpt quant config from either shape. + + Dispatches to :func:`_read_legacy` or :func:`_read_flat` based on the + presence of the ``"quantization"`` wrapper. Both branches read their + native field names directly -- no shape-to-shape conversion occurs. + + Raises ``ValueError`` if ``raw`` is not a recognized ModelOpt config. + """ + if not isinstance(raw, dict): + raise ValueError(f"Expected a dict for ModelOpt quant config, got {type(raw).__name__}") + if "quantization" in raw: + parsed = _read_legacy(raw) + elif is_modelopt_quant_config(raw): + parsed = _read_flat(raw) + else: + producer = (raw.get("producer") or {}).get("name") + raise ValueError( + f"Not a ModelOpt quant config (producer={producer!r}, " + f"quant_method={raw.get('quant_method')!r}); " + "missing 'quantization' wrapper and no modelopt marker found." + ) + + # Canonicalize the ``fp8_pb_wo`` alias once at the parse boundary so + # downstream call sites don't each re-translate it. + if parsed.quant_algo == "fp8_pb_wo": + parsed.quant_algo = "FP8_BLOCK_SCALES" + return parsed + + +# Fields whose mismatch between ``hf_quant_config.json`` and the inline +# ``config.json.quantization_config`` is worth surfacing. +_INVARIANT_FIELDS = ( + "quant_algo", + "kv_cache_quant_algo", + "group_size", +) + + +def _summarize_quantized_layers(layers: Dict[str, Dict[str, Any]]) -> Optional[Dict[str, Any]]: + if not layers: + return None + histogram: Dict[str, int] = {} + for entry in layers.values(): + algo = "UNKNOWN" + if isinstance(entry, dict): + algo = str(entry.get("quant_algo", "UNKNOWN")).upper() + histogram[algo] = histogram.get(algo, 0) + 1 + return {"total": len(layers), "per_algo": histogram} + + +def warn_if_inline_quant_config_differs( + file_parsed: ModelOptQuantConfig, + inline_raw: Optional[Dict[str, Any]], + *, + source_file: str, +) -> None: + """Warn if the inline config diverges from the file-based parsed config. + + Compares ``config.json.quantization_config`` to ``file_parsed`` (which + came from ``hf_quant_config.json``). The file form remains authoritative; + this only logs. If the inline is absent or not a recognizable ModelOpt + config, no warning is emitted. + """ + if not inline_raw or not is_modelopt_quant_config(inline_raw): + return + try: + inline_parsed = parse_modelopt_quant_config(inline_raw) + except ValueError: + return + + diffs = [] + for key in _INVARIANT_FIELDS: + lhs = getattr(file_parsed, key) + rhs = getattr(inline_parsed, key) + if lhs != rhs: + diffs.append(f"{key}: hf_quant_config={lhs!r}, inline={rhs!r}") + + lhs_excl = set(file_parsed.exclude_modules) + rhs_excl = set(inline_parsed.exclude_modules) + if lhs_excl != rhs_excl: + only_file = sorted(lhs_excl - rhs_excl)[:3] + only_inline = sorted(rhs_excl - lhs_excl)[:3] + diffs.append( + f"exclude_modules: |hf_quant_config|={len(lhs_excl)}, " + f"|inline|={len(rhs_excl)}, " + f"only-in-hf-sample={only_file}, " + f"only-in-inline-sample={only_inline}" + ) + + file_ql = _summarize_quantized_layers(file_parsed.quantized_layers) + inline_ql = _summarize_quantized_layers(inline_parsed.quantized_layers) + if file_ql != inline_ql: + diffs.append(f"quantized_layers: hf_quant_config={file_ql}, inline={inline_ql}") + + if diffs: + logger.warning( + "Inline 'config.json.quantization_config' diverges from " + f"'{source_file}'; using the latter. Diffs:\n - " + "\n - ".join(diffs) + ) 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( diff --git a/tests/unittest/_torch/executor/test_adp_router.py b/tests/unittest/_torch/executor/test_adp_router.py index 73af9b8262cb..f47662caece5 100644 --- a/tests/unittest/_torch/executor/test_adp_router.py +++ b/tests/unittest/_torch/executor/test_adp_router.py @@ -203,6 +203,19 @@ def test_target_dp_rank_at_capacity_falls_through(self): assert len(result[0]) == 1 assert len(result[1]) == 0 + def test_attention_dp_relax_none_is_sortable(self): + # Requests created by trtllm-serve carry a default SchedulingParams() + # where attention_dp_relax is None. The sort key in route_requests must + # not produce None or sorted() raises TypeError on None < None. + router = DefaultADPRouter(dist=_mock_dist()) + states = [ + RankState(rank=0, num_active_requests=0, num_active_tokens=0), + RankState(rank=1, num_active_requests=0, num_active_tokens=0), + ] + reqs = [_make_request_item(i, num_tokens=10, attention_dp_relax=None) for i in range(4)] + result, _ = router.route_requests(states, reqs, max_num_active_requests=10) + assert sum(len(v) for v in result.values()) == 4 + def test_favors_less_loaded_rank(self): router = DefaultADPRouter(dist=_mock_dist()) states = [ @@ -926,6 +939,20 @@ def test_cache_affinity_wins(self): assert result[0] == [] assert result[1] == [req] + def test_attention_dp_relax_none_is_sortable(self): + # Mirrors the DefaultADPRouter regression: the KV-aware router also + # sorts new_requests by attention_dp_relax, so a None default value + # must not break sorted(). + router = self._make_router( + tp_size=2, + prefix_matches=[{i: 0 for i in range(4)}, {i: 0 for i in range(4)}], + ) + reqs = [ + _make_request_item(req_id=i, num_tokens=10, attention_dp_relax=None) for i in range(4) + ] + result, _ = router.route_requests(self._rank_states(2), reqs, max_num_active_requests=10) + assert sum(len(v) for v in result.values()) == 4 + def test_match_rate_threshold_gates_cache_affinity(self): """With rank 0 loaded but holding cache, and rank 1 idle with no cache: diff --git a/tests/unittest/_torch/executor/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/test_mamba_cache_manager.py index d959cdfc2547..033dfa9f7949 100644 --- a/tests/unittest/_torch/executor/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/test_mamba_cache_manager.py @@ -3,6 +3,7 @@ """Regression tests for MambaCacheManager padding-slot behavior and CppMambaHybridCacheManager PP-sharding edge cases.""" +import os from types import SimpleNamespace from unittest.mock import MagicMock @@ -17,7 +18,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 +260,188 @@ 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 when running under pytest --run-ray (sets TLLM_DISABLE_MPI=1). In that +# mode ``Mapping`` resolves to ``DeviceMeshTopology``, whose ``pp_rank`` +# requires torch.distributed initialisation that isn't available in this +# single-process unit test. The pp-sharding behaviour exercised here is +# orthogonal to the Ray orchestrator. +_skip_under_ray = pytest.mark.skipif( + os.environ.get("TLLM_DISABLE_MPI") == "1", + reason="pp_size>1 helper builds Mapping with world_size>1 which needs " + "torch.distributed under TLLM_DISABLE_MPI=1 (Ray) sessions", +) + + +@_skip_under_ray +@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 # diff --git a/tests/unittest/llmapi/test_modelopt_quant_config.py b/tests/unittest/llmapi/test_modelopt_quant_config.py new file mode 100644 index 000000000000..9526fd898a6c --- /dev/null +++ b/tests/unittest/llmapi/test_modelopt_quant_config.py @@ -0,0 +1,273 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Unit tests for ``tensorrt_llm.quantization.modelopt_config``.""" + +import logging + +import pytest + +from tensorrt_llm.logger import logger as tllm_logger +from tensorrt_llm.quantization.modelopt_config import ( + ModelOptQuantConfig, + is_modelopt_quant_config, + parse_modelopt_quant_config, + warn_if_inline_quant_config_differs, +) + + +@pytest.fixture +def warn_caplog(caplog): + """Capture warnings emitted via the TensorRT-LLM logger. + + The TRT-LLM Logger wraps the stdlib logger named ``"TRT-LLM"`` and sets + ``propagate = False`` so its messages don't reach pytest's root handler. + Temporarily re-enable propagation, bind caplog to that named logger, and + lower the TRT-LLM severity so warnings actually flow through. + """ + underlying = tllm_logger._logger + prev_propagate = underlying.propagate + prev_severity = tllm_logger._min_severity + underlying.propagate = True + tllm_logger.set_level("warning") + caplog.set_level(logging.WARNING, logger="TRT-LLM") + try: + yield caplog + finally: + underlying.propagate = prev_propagate + tllm_logger.set_level(prev_severity) + + +# ---------- Fixtures --------------------------------------------------------- + +_QUANTIZED_LAYERS = { + "model.layers.0.mlp.up_proj": {"quant_algo": "FP8"}, + "model.layers.0.mlp.down_proj": { + "quant_algo": "NVFP4", + "group_size": 16, + }, +} + +_EXCLUDES = ["lm_head", "model.embed_tokens"] + + +@pytest.fixture +def legacy_config(): + """ModelOpt 0.43.x shape: nested under 'quantization'.""" + return { + "producer": {"name": "modelopt", "version": "0.43.0rc1"}, + "quantization": { + "quant_algo": "MIXED_PRECISION", + "kv_cache_quant_algo": None, + "exclude_modules": list(_EXCLUDES), + "quantized_layers": dict(_QUANTIZED_LAYERS), + }, + } + + +@pytest.fixture +def flat_config(): + """ModelOpt 1.0.x shape: flat at top level with 'ignore' and 'quant_method'.""" + return { + "producer": {"name": "modelopt", "version": "1.0.0"}, + "quant_method": "modelopt", + "quant_algo": "MIXED_PRECISION", + "ignore": list(_EXCLUDES), + "config_groups": { + "group_0": { + "input_activations": {"dynamic": False, "num_bits": 8, "type": "float"}, + "weights": {"dynamic": False, "num_bits": 8, "type": "float"}, + "targets": ["model.layers.0.mlp.up_proj"], + }, + "group_1": { + "input_activations": { + "dynamic": False, + "num_bits": 4, + "type": "float", + "group_size": 16, + }, + "weights": {"dynamic": False, "num_bits": 4, "type": "float", "group_size": 16}, + "targets": ["model.layers.0.mlp.down_proj"], + }, + }, + "quantized_layers": dict(_QUANTIZED_LAYERS), + } + + +# ---------- is_modelopt_quant_config ----------------------------------------- + + +def test_is_modelopt_legacy(legacy_config): + assert is_modelopt_quant_config(legacy_config) + + +def test_is_modelopt_flat(flat_config): + assert is_modelopt_quant_config(flat_config) + + +def test_is_modelopt_compressed_tensors_no_modelopt(): + cfg = {"quant_method": "compressed-tensors", "config_groups": {}, "ignore": []} + assert not is_modelopt_quant_config(cfg) + + +def test_is_modelopt_garbage(): + assert not is_modelopt_quant_config({}) + assert not is_modelopt_quant_config({"producer": {}}) + + +# ---------- parse_modelopt_quant_config: parallel readers -------------------- + + +def test_parse_legacy_branch(legacy_config): + parsed = parse_modelopt_quant_config(legacy_config) + assert isinstance(parsed, ModelOptQuantConfig) + assert parsed.source_format == "legacy" + assert parsed.quant_algo == "MIXED_PRECISION" + assert parsed.kv_cache_quant_algo is None + assert parsed.exclude_modules == _EXCLUDES + assert parsed.quantized_layers == _QUANTIZED_LAYERS + + +def test_parse_flat_branch(flat_config): + parsed = parse_modelopt_quant_config(flat_config) + assert parsed.source_format == "flat" + assert parsed.quant_algo == "MIXED_PRECISION" + assert parsed.exclude_modules == _EXCLUDES + assert parsed.quantized_layers == _QUANTIZED_LAYERS + + +def test_parse_flat_kv_cache_scheme_to_fp8(): + cfg = { + "producer": {"name": "modelopt"}, + "quant_method": "modelopt", + "quant_algo": "FP8", + "kv_cache_scheme": {"type": "float", "num_bits": 8}, + } + assert parse_modelopt_quant_config(cfg).kv_cache_quant_algo == "FP8" + + +def test_parse_flat_kv_cache_scheme_non_fp8_yields_none(): + cfg = { + "producer": {"name": "modelopt"}, + "quant_method": "modelopt", + "quant_algo": "FP8", + "kv_cache_scheme": {"type": "int", "num_bits": 8}, + } + assert parse_modelopt_quant_config(cfg).kv_cache_quant_algo is None + + +def test_parse_legacy_and_flat_produce_equal_structs(legacy_config, flat_config): + """Both branches produce the same struct for equivalent inputs. + + The whole point of parallel readers: both branches produce the same + struct (modulo source_format) for equivalent inputs. + """ + legacy_parsed = parse_modelopt_quant_config(legacy_config) + flat_parsed = parse_modelopt_quant_config(flat_config) + # Drop source_format so the comparison ignores that one provenance field. + for p in (legacy_parsed, flat_parsed): + p.source_format = "" + assert legacy_parsed == flat_parsed + + +def test_parse_rejects_non_modelopt(): + cfg = {"quant_method": "compressed-tensors", "config_groups": {}} + with pytest.raises(ValueError, match="Not a ModelOpt"): + parse_modelopt_quant_config(cfg) + + +def test_parse_rejects_non_dict(): + with pytest.raises(ValueError, match="Expected a dict"): + parse_modelopt_quant_config([]) # type: ignore[arg-type] + + +# ---------- ModelOptQuantConfig.to_legacy_inner_dict ------------------------- + + +def test_to_legacy_inner_dict_roundtrip(legacy_config): + parsed = parse_modelopt_quant_config(legacy_config) + out = parsed.to_legacy_inner_dict() + assert out["quant_algo"] == "MIXED_PRECISION" + assert out["exclude_modules"] == _EXCLUDES + assert out["quantized_layers"] == _QUANTIZED_LAYERS + # kv_cache_quant_algo was None -> not emitted in the dict. + assert "kv_cache_quant_algo" not in out + + +# ---------- warn_if_inline_quant_config_differs ------------------------------ + + +def _captured_warnings(caplog): + return [r.getMessage() for r in caplog.records if r.levelno == logging.WARNING] + + +def test_warn_no_inline_no_warning(legacy_config, warn_caplog): + warn_if_inline_quant_config_differs( + parse_modelopt_quant_config(legacy_config), None, source_file="hf_quant_config.json" + ) + assert _captured_warnings(warn_caplog) == [] + + +def test_warn_legacy_and_flat_inline_match(legacy_config, flat_config, warn_caplog): + """File is legacy, inline is flat with identical content: no warning.""" + warn_if_inline_quant_config_differs( + parse_modelopt_quant_config(legacy_config), flat_config, source_file="hf_quant_config.json" + ) + assert _captured_warnings(warn_caplog) == [] + + +def test_warn_quant_algo_mismatch(legacy_config, warn_caplog): + inline = { + "producer": {"name": "modelopt"}, + "quant_method": "modelopt", + "quant_algo": "FP8", + "ignore": list(_EXCLUDES), + "quantized_layers": dict(_QUANTIZED_LAYERS), + } + warn_if_inline_quant_config_differs( + parse_modelopt_quant_config(legacy_config), inline, source_file="hf_quant_config.json" + ) + msgs = _captured_warnings(warn_caplog) + assert len(msgs) == 1 + assert "quant_algo" in msgs[0] + + +def test_warn_exclude_modules_mismatch(legacy_config, warn_caplog): + inline = { + "producer": {"name": "modelopt"}, + "quant_method": "modelopt", + "quant_algo": "MIXED_PRECISION", + "ignore": _EXCLUDES + ["extra"], + "quantized_layers": dict(_QUANTIZED_LAYERS), + } + warn_if_inline_quant_config_differs( + parse_modelopt_quant_config(legacy_config), inline, source_file="hf_quant_config.json" + ) + msgs = _captured_warnings(warn_caplog) + assert len(msgs) == 1 + assert "exclude_modules" in msgs[0] + + +def test_warn_quantized_layers_count_mismatch(legacy_config, warn_caplog): + inline_layers = dict(_QUANTIZED_LAYERS) + inline_layers["model.layers.1.mlp.up_proj"] = {"quant_algo": "FP8"} + inline = { + "producer": {"name": "modelopt"}, + "quant_method": "modelopt", + "quant_algo": "MIXED_PRECISION", + "ignore": list(_EXCLUDES), + "quantized_layers": inline_layers, + } + warn_if_inline_quant_config_differs( + parse_modelopt_quant_config(legacy_config), inline, source_file="hf_quant_config.json" + ) + msgs = _captured_warnings(warn_caplog) + assert len(msgs) == 1 + assert "quantized_layers" in msgs[0] + + +def test_warn_skips_unparseable_inline(legacy_config, warn_caplog): + inline = {"quant_method": "compressed-tensors", "config_groups": {}} + warn_if_inline_quant_config_differs( + parse_modelopt_quant_config(legacy_config), inline, source_file="hf_quant_config.json" + ) + assert _captured_warnings(warn_caplog) == []