diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 645e30749ac1..014e4dae8b86 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -1,82 +1,59 @@ -# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - import copy import dataclasses import os -from typing import Any, Dict, List, Optional, Sequence, Union +from typing import Dict, List, Optional, Union import torch import tensorrt_llm import tensorrt_llm.bindings.executor as trtllm +from tensorrt_llm._torch.models.modeling_multimodal_utils import _is_mm_disagg +from tensorrt_llm._torch.models.modeling_utils import \ + MODEL_CLASS_VISION_ENCODER_MAPPING from tensorrt_llm._utils import (confidential_compute_enabled, get_sm_version, - is_sm_100f, prefer_pinned, - str_dtype_to_binding, torch_dtype_to_str) -from tensorrt_llm.inputs.multimodal import MultimodalParams + is_device_integrated, str_dtype_to_binding, + torch_dtype_to_str) +from tensorrt_llm.bindings.executor import DecodingMode # isort: off from tensorrt_llm.llmapi.llm_args import ( CacheTransceiverConfig, CapacitySchedulerPolicy, EagleDecodingConfig, - KVEventsConfig, KvCacheCompressionConfig, KvCacheConfig, MTPDecodingConfig, - MultimodalEncoderSchedulingPolicy, PeftCacheConfig, SchedulerConfig, - SparseAttentionConfig, SpeculativeConfig, TorchLlmArgs, WaitingQueuePolicy) + KvCacheConfig, MTPDecodingConfig, PeftCacheConfig, SamplerType, + SchedulerConfig, SparseAttentionConfig, SpeculativeConfig, TorchLlmArgs, + WaitingQueuePolicy) # isort: on -from tensorrt_llm._torch.peft.lora.config import ( - LoraConfig, get_default_trtllm_modules_to_hf_modules) -from tensorrt_llm._torch.peft.lora.manager import (load_torch_lora, - supports_native_fp8_lora) from tensorrt_llm.logger import logger +from tensorrt_llm.lora_helper import (LoraConfig, + get_default_trtllm_modules_to_hf_modules) +from tensorrt_llm.lora_manager import load_torch_lora from tensorrt_llm.mapping import CpType, Mapping -from tensorrt_llm.quantization import QuantAlgo -from ..attention.backends import get_sparse_attn_kv_cache_manager -from ..disaggregation.kv_cache_transceiver import ( - AttentionTypeCpp, create_kv_cache_transceiver, - maybe_enable_fabric_memory_for_python_transceiver) -from ..hostfunc import set_low_latency_dispatch +from ..attention_backend import get_sparse_attn_kv_cache_manager from ..model_config import ModelConfig -from ..models.modeling_multimodal_mixin import MultimodalModelMixin -from ..speculative import (draft_prompt_lookahead, get_num_extra_kv_tokens, - get_num_spec_layers, get_spec_decoder, - should_use_separate_draft_kv_cache) -from ..utils import is_gdn_replay_enabled -from .config_utils import (MambaKVCacheParams, extract_mamba_kv_cache_params, - extract_qwen4_exp_ple_cache_params, - get_layer_attention_window, is_gemma4_hybrid, - is_hybrid_linear, is_kimi_linear, is_mla, - is_nemotron_hybrid, is_qwen3_hybrid, is_qwen4_exp, - uses_vswa_kv_cache_layout) +from ..speculative import (get_num_extra_kv_tokens, get_num_spec_layers, + get_spec_decoder, should_use_separate_draft_kv_cache) +from .config_utils import (extract_mamba_kv_cache_params, is_gemma4_hybrid, + is_hybrid_linear, is_mla, is_nemotron_hybrid, + is_qwen3_hybrid) from .connectors.kv_cache_connector import KvCacheConnectorManager from .dwdp import DwdpManager from .guided_decoder import GuidedDecoder -from .kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 -from .kv_cache.mamba_cache_manager import (BaseMambaCacheManager, - CppMambaHybridCacheManager, - MambaHybridCacheManagerV2, - MixedMambaHybridCacheManager, - use_py_mamba_cache_manager) +from .kv_cache_manager_v2 import KVCacheManagerV2 +from .kv_cache_transceiver import AttentionTypeCpp, create_kv_cache_transceiver from .llm_request import ExecutorResponse, LlmRequestState +from .mamba_cache_manager import (BaseMambaCacheManager, + CppMambaHybridCacheManager, + MixedMambaHybridCacheManager, + use_cpp_mamba_cache_manager, + use_py_mamba_cache_manager) from .model_engine import PyTorchModelEngine from .py_executor import PyExecutor -from .resource_manager import (KVCacheCompressionManager, KVCacheManager, - PeftCacheManager, ResourceManager, - ResourceManagerType) -from .sampler import EarlyStopSampler, EarlyStopWithMMResult, TorchSampler +from .resource_manager import (KVCacheManager, PeftCacheManager, + ResourceManager, ResourceManagerType) +from .sampler import (EarlyStopSampler, EarlyStopWithMMResult, TorchSampler, + TRTLLMSampler) from .scheduler import (BindCapacityScheduler, BindMicroBatchScheduler, - KVCacheV2Scheduler, MultimodalEagerEncoderScheduler, - MultimodalScheduler, SimpleScheduler, + KVCacheV2Scheduler, SimpleScheduler, SimpleUnifiedScheduler) from .seq_slot_manager import SeqSlotManager @@ -87,40 +64,14 @@ def ceil_div(a: int, b: int) -> int: return (a + b - 1) // b -def _get_initial_lora_data_type( - configured_lora_data_type: Optional[torch.dtype], -) -> Optional[torch.dtype]: - if configured_lora_data_type != torch.float8_e4m3fn: - return None - if supports_native_fp8_lora(torch.cuda.get_device_capability()): - return configured_lora_data_type - return None - - def _non_hybrid_kv_cache_manager_cls(config, kv_cache_config: KvCacheConfig): # Models with per-layer head_dim (e.g., Gemma4 hybrid attention) # require KVCacheManagerV2 for per-layer buffer sizes. - needs_v2 = (kv_cache_config.use_kv_cache_manager_v2 is True + needs_v2 = (kv_cache_config.use_kv_cache_manager_v2 or is_gemma4_hybrid(config)) return KVCacheManagerV2 if needs_v2 else KVCacheManager -def _resolve_disagg_transceiver_route( - cache_transceiver_config: Optional[CacheTransceiverConfig], -) -> tuple[Optional[str], Optional[str]]: - """Return the effective backend and runtime used for manager routing.""" - if cache_transceiver_config is None: - return None, None - - backend, _ = cache_transceiver_config._resolve_default_backend() - runtime = cache_transceiver_config.transceiver_runtime - if runtime == "auto": - # Model loading normally resolves ``auto``. Paths that skip model - # defaults use the global C++ fallback, matching transceiver creation. - runtime = None - return backend, runtime - - def get_kv_cache_manager_cls( model_config: ModelConfig, kv_cache_config: KvCacheConfig, @@ -128,188 +79,66 @@ def get_kv_cache_manager_cls( cache_transceiver_config: Optional[CacheTransceiverConfig] = None): """Resolve the concrete KV cache manager class for ``model_config``. - For hybrid mamba models the choice between - ``MambaHybridCacheManagerV2`` and compatibility managers is made here. - Callers that don't care about disagg can omit ``is_disagg`` and get the - unified-pool default. - - Model loading resolves ``use_kv_cache_manager_v2="auto"`` to V2 for - supported hybrid Mamba models. An explicit ``False`` selects a - compatibility manager. In disaggregated serving, V2 additionally requires - the Python transceiver with the NIXL backend. Unsupported V2 routes fail - rather than falling back to a different manager. + For hybrid mamba models the choice between ``Mixed`` ( 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: - * ``TRTLLM_USE_PY_MAMBA=1`` — Mixed manager in aggregated serving. - * ``TLLM_MAMBA_MANAGER_PREFERENCE`` — explicit manager preference. + 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 - sparse_attn_algorithm = getattr(sparse_attn_config, "algorithm", None) - use_v2 = kv_cache_config.use_kv_cache_manager_v2 is True - if is_hybrid_linear(config): + sparse_attention_config = model_config.sparse_attention_config + if sparse_attention_config is not None: + return get_sparse_attn_kv_cache_manager(sparse_attention_config) + elif is_hybrid_linear(config): # Degenerate case: model is flagged as hybrid but the config has zero - # mamba layers. Fall through to the standard non-hybrid routes. + # mamba layers. Fall through to the standard non-hybrid manager. if model_config.get_num_mamba_layers() == 0: logger.info("Hybrid linear model has 0 mamba layers; using " - "KV cache manager without mamba caching") - if sparse_attn_config is not None: - return get_sparse_attn_kv_cache_manager( - sparse_attn_config, use_kv_cache_manager_v2=use_v2) + "KVCacheManager without mamba caching") return _non_hybrid_kv_cache_manager_cls(config, kv_cache_config) - - if sparse_attn_algorithm == "qsa" and not use_v2: - raise ValueError( - "QSA with hybrid Mamba / linear-attention models requires " - "use_kv_cache_manager_v2=True.") - if (sparse_attn_config is not None - and sparse_attn_algorithm not in ("qsa", "skip_softmax")): - raise ValueError( - f"Sparse attention algorithm {sparse_attn_algorithm!r} is not " - "supported with hybrid Mamba / linear-attention models.") - - state_config = kv_cache_config.mamba_state_config - has_additional_snapshots = bool( - state_config.additional_snapshot_offsets_from_start - or state_config.additional_snapshot_offsets_from_end) - if has_additional_snapshots and not use_v2: - raise ValueError("Mamba additional snapshot offsets require " - "use_kv_cache_manager_v2=True; V1 supports only " - "periodic_snapshot_interval.") - - # Kimi K3 (KDA + MLA hybrid): block reuse uses the unified C++ pool - # (CppMambaHybridCacheManager) like the other hybrid linear models — - # per-block KDA state snapshots every mamba_state_cache_interval - # tokens with FORCE_CHUNK context chunking. Without block reuse the - # Mixed manager (separate KV / recurrent-state pools) stays the - # default. SA speculative decoding is validated on the Mixed - # manager's SpeculativeState scratch path only; reuse + SA is - # unvalidated. Disaggregated serving (TRTLLM-14815) routes through - # the shared hybrid transceiver validation below: the Python NIXL - # transceiver selects the Mixed manager, whose KDA recurrent/conv - # states transfer through the bounce buffer. - if is_kimi_linear(config) and not use_v2 and not is_disagg: - if kv_cache_config.enable_block_reuse: - logger.info( - "Using CppMambaHybridCacheManager for Kimi K3 hybrid " - "model (block reuse enabled)") - return CppMambaHybridCacheManager - logger.info( - "Using MixedMambaHybridCacheManager for Kimi K3 hybrid model") - return MixedMambaHybridCacheManager - - # Skip Softmax only changes attention kernels. Hybrid models still - # need a Mamba-capable cache manager for recurrent state. - if is_disagg: - backend, runtime = _resolve_disagg_transceiver_route( - cache_transceiver_config) - if is_kimi_linear(config) and (runtime != "PYTHON" - or backend != "NIXL"): - # Only the Python NIXL transceiver can move KDA recurrent - # state; the C++ transceiver would silently serve wrong - # results. Model loading resolves ``auto`` to PYTHON via - # KimiLinearForCausalLM.get_preferred_transceiver_runtime - # (NIXL-gated); this rejects explicit non-Python routes and - # paths that skip model defaults (e.g. AutoDeploy). - raise ValueError( - "Kimi K3 disaggregated serving requires the Python " - "transceiver: set cache_transceiver_config " - "backend='NIXL' with transceiver_runtime='PYTHON' (or " - "leave transceiver_runtime='auto' with the NIXL " - "backend). The C++ transceiver cannot transfer KDA " - f"recurrent state (got backend={backend!r}, " - f"transceiver_runtime={runtime!r}).") - if use_v2: - if runtime != "PYTHON" or backend != "NIXL": - raise ValueError( - "KV cache manager V2 for hybrid Mamba disaggregated " - "serving requires transceiver_runtime='PYTHON' with " - "backend='NIXL'.") - else: - if (kv_cache_config.enable_block_reuse and runtime == "PYTHON"): - raise ValueError( - "Hybrid Mamba disaggregated serving with block reuse " - "and transceiver_runtime='PYTHON' requires " - "use_kv_cache_manager_v2=True.") - if kv_cache_config.enable_block_reuse: - return CppMambaHybridCacheManager - if runtime == "PYTHON" and backend == "NIXL": - logger.info("Python transceiver detected; using " - "MixedMambaHybridCacheManager for hybrid model") - return MixedMambaHybridCacheManager - return CppMambaHybridCacheManager - - if use_py_mamba_cache_manager() and not is_disagg: - if use_v2: - raise ValueError( - "TRTLLM_USE_PY_MAMBA=1 conflicts with explicit " - "use_kv_cache_manager_v2=True.") + if use_py_mamba_cache_manager(): if kv_cache_config.enable_block_reuse: raise ValueError( "TRTLLM_USE_PY_MAMBA=1 forces " "MixedMambaHybridCacheManager, which does not support " "block reuse. Disable block reuse or unset " - "TRTLLM_USE_PY_MAMBA to use the configured cache manager.") + "TRTLLM_USE_PY_MAMBA to use CppMambaHybridCacheManager.") + logger.info( + "Using MixedMambaHybridCacheManager for hybrid mamba model") + return MixedMambaHybridCacheManager + if kv_cache_config.enable_block_reuse: + return CppMambaHybridCacheManager + if use_cpp_mamba_cache_manager(): logger.info( "Using MixedMambaHybridCacheManager for hybrid mamba model") return MixedMambaHybridCacheManager + if (cache_transceiver_config is not None + and cache_transceiver_config.transceiver_runtime == "PYTHON"): + logger.info("Python transceiver detected; using " + "MixedMambaHybridCacheManager for hybrid mamba model") + return MixedMambaHybridCacheManager + default_cls = CppMambaHybridCacheManager env_override = os.environ.get('TLLM_MAMBA_MANAGER_PREFERENCE', None) if env_override is not None: - env_override = env_override.upper() - if env_override == 'MIXED': - if use_v2: - raise ValueError( - "TLLM_MAMBA_MANAGER_PREFERENCE=MIXED conflicts with " - "explicit use_kv_cache_manager_v2=True.") - if kv_cache_config.enable_block_reuse: - raise ValueError( - "TLLM_MAMBA_MANAGER_PREFERENCE=MIXED forces " - "MixedMambaHybridCacheManager, which does not support " - "block reuse. Disable block reuse, use the CPP " - "preference, or explicitly enable KV cache manager " - "V2.") + if env_override.upper() == 'MIXED': logger.warning( - "Environment variable TLLM_MAMBA_MANAGER_PREFERENCE=MIXED " - "overrides the default Mamba cache manager to " - "MixedMambaHybridCacheManager.") + "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 - if env_override == 'CPP': - if use_v2: - raise ValueError( - "TLLM_MAMBA_MANAGER_PREFERENCE=CPP conflicts with " - "explicit use_kv_cache_manager_v2=True.") + elif env_override.upper() == 'CPP': logger.warning( - "Environment variable TLLM_MAMBA_MANAGER_PREFERENCE=CPP " - "overrides the default Mamba cache manager to " - "CppMambaHybridCacheManager.") + "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 - logger.warning( - f"Unrecognized value for TLLM_MAMBA_MANAGER_PREFERENCE: {env_override}. " - "Expected 'CPP' or 'MIXED'. Using the configured " - "KV cache manager default.") - - if not use_v2: - return CppMambaHybridCacheManager - - if (kv_cache_config.enable_block_reuse - and kv_cache_config.enable_kv_pool_rebalance): - raise ValueError( - "V2 Mamba block reuse is not compatible with " - "enable_kv_pool_rebalance because the rebalancer does not " - "yet model retained recurrent-state snapshots.") - if sparse_attn_algorithm == "qsa": - return get_sparse_attn_kv_cache_manager( - sparse_attn_config, use_kv_cache_manager_v2=True) - return MambaHybridCacheManagerV2 - elif sparse_attn_config is not None: - if sparse_attn_algorithm == "qsa": - # The QSA manager extends the hybrid manager because it must retain - # both paged attention data and recurrent GDN state. - raise ValueError( - "QSA sparse attention currently requires a hybrid Mamba / " - "linear-attention model.") - return get_sparse_attn_kv_cache_manager(sparse_attn_config, - use_kv_cache_manager_v2=use_v2) + 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) @@ -319,15 +148,14 @@ def get_kv_cache_manager_cls( # KVCacheManager.get_cache_size_per_token may return either an ``int`` # (legacy proportional model ``bytes = slope * tokens``) or an affine # ``(slope, intercept)`` tuple (CppMambaHybridCacheManager, where mamba -# state introduces a per-batch fixed cost). CacheCost normalizes the combined -# shape so the rest of the file does plain attribute access and method calls -# instead of branching on type. Managers are responsible for including any -# pool-specific alignment or capacity headroom in the returned cost. +# state introduces a per-batch fixed cost). CacheCost normalizes both +# shapes so the rest of the file does plain attribute access and method +# calls instead of branching on type. @dataclasses.dataclass(frozen=True) class CacheCost: - """Affine KV cache budget: ``bytes = slope * tokens + intercept``. + """Affine KV cache cost: ``bytes = slope * tokens + intercept``. The legacy proportional case is just ``intercept = 0``. """ @@ -368,235 +196,15 @@ def bytes_for_tokens(self, tokens: int) -> int: return self.slope * tokens + self.intercept -def get_attention_workspace_bytes_per_token(model_config, mapping) -> int: - """Per-token workspace headroom the model's selected attention backend declares. - - The KV-cache profiling forward under-measures any attention workspace sized by a runtime quantity it - does not drive to its serving maximum (e.g. ``total_kv_len``, inflated by KV reuse). Backends declare - such a buffer via ``AttentionBackend.runtime_workspace_bytes_per_token``; this resolves the model's - backend and returns its rate. A backend that stages no such buffer inherits the default 0, so no - workspace is reserved and no admission cap is installed for it. See ``ATTENTION_DEVELOPER_GUIDE.md`` - §2.3. - """ - from ..attention.backends.utils import get_attention_backend - - # Resolved without ``sparse_params``: those are per-layer, while this workspace is one buffer shared - # across every attention layer, so the declaration is a whole-model question. The sparse backends all - # derive from the dense class and inherit its declaration, which reads the sparse gate off model_config. - return get_attention_backend( - model_config.attn_backend).runtime_workspace_bytes_per_token( - model_config, mapping) - - -def get_attention_workspace_is_chunked_prefill_bounded(model_config) -> bool: - """Whether chunked prefill bounds the selected backend's runtime workspace.""" - from ..attention.backends.utils import get_attention_backend - - return get_attention_backend( - model_config.attn_backend).runtime_workspace_is_chunked_prefill_bounded( - model_config) - - -def get_mla_context_workspace_kv_len_cap( - kv_cache_config, - max_batch_size, - max_num_tokens, - max_seq_len, - enable_chunked_prefill, - workspace_is_chunked_prefill_bounded=True): - """Max summed attended-KV length covered by the context-MLA workspace reserve. - - KV-cache reuse can grow this workspace beyond the fresh-prefill profiling - floor. Chunked prefill normally prevents that by staging one bounded KV - chunk per launch. An implementation that consumes the complete attended - prefix sets ``workspace_is_chunked_prefill_bounded=False`` and receives the - same reservation and scheduler admission protection as cache reuse. - - Otherwise the default (no override) is the never-stall worst case ``min(max_batch_size, max_num_tokens) - * max_seq_len``: at most that many context requests run in a step, each attending at most ``max_seq_len`` - KV, so reserving for it never defers a request. An explicit ``fp8_context_mla_kv_len_cap`` override - reserves less workspace (freeing KV pool) and lets the scheduler defer over-cap requests; it is floored - at ``max_seq_len`` (one request must always fit) and capped at the worst case. - """ - workspace_can_exceed_profile = ( - kv_cache_config.enable_block_reuse and not enable_chunked_prefill) or ( - enable_chunked_prefill and not workspace_is_chunked_prefill_bounded) - if not workspace_can_exceed_profile: - return None - worst_case = min(max_batch_size, max_num_tokens) * max_seq_len - override = kv_cache_config.fp8_context_mla_kv_len_cap - if override is None: - return worst_case - return min(max(int(override), max_seq_len), worst_case) - - -def get_mla_context_workspace_reserve(budget_bytes, k_bytes_per_token, - w_bytes_per_token, kv_len_cap): - """Bytes to reserve for the fp8 context-MLA workspace, and the token admission cap that reserve covers. - - Reserve ``w * kv_len_cap`` (the worst-case summed attended KV), clamped to the per-token split - ``budget * w / (k + w)`` so a memory-constrained node shares the budget at a common token count rather - than starving the KV pool. The admission cap is ``reserve / w == min(kv_len_cap, budget / (k + w))`` - tokens; the scheduler admits at most that much summed attended KV, so the fp8 dequant staging buffer - this reserve covers stays within it. This accounts for the fp8 staging term only -- the separate BF16 - full-gather buffers on the reuse path are not yet charged here (tracked as a follow-up), so this bounds - but does not by itself guarantee the reuse-path peak. Returns ``(reserve_bytes, cap_tokens)``, or - ``(0, None)`` when any input is non-positive. - """ - if not (budget_bytes > 0 and k_bytes_per_token > 0 and w_bytes_per_token > 0 - and kv_len_cap and kv_len_cap > 0): - return 0, None - reserve = min( - w_bytes_per_token * kv_len_cap, budget_bytes * w_bytes_per_token / - (k_bytes_per_token + w_bytes_per_token)) - return reserve, int(reserve / w_bytes_per_token) - - -def _normalize_attention_windows( - max_attention_window: List[Optional[int]], - max_seq_len: int, -) -> Optional[List[int]]: - normalized = [ - max_seq_len if window is None else min(max_seq_len, window) - for window in max_attention_window - ] - if all(window == max_seq_len for window in normalized): - return None - if len(set(normalized)) == 1: - return [normalized[0]] - return normalized - - -def _get_num_pool_groups_for_estimation( - model_config: object, - max_seq_len: int, - fallback_attention_windows: Optional[List[Optional[int]]], -) -> int: - """Infer the number of V2 KV-cache pools needed during estimation. - - Sliding/full-attention hybrids are best distinguished by their effective - windows. Hybrid linear-attention models with mixed layer types fall back to - their distinct layer types. Unsupported target window metadata must not - make estimation fail; in that - case preserve the legacy layer-type/window heuristic. - """ - num_layers = getattr(model_config, "num_hidden_layers", None) - layer_types = getattr(model_config, "layer_types", None) - attention_windows = None - if isinstance(num_layers, int) and num_layers > 0: - try: - inferred_windows = [ - get_layer_attention_window(model_config, layer_idx) - for layer_idx in range(num_layers) - ] - except (NotImplementedError, ValueError) as error: - logger.warning( - "Unable to infer target attention windows for KV-cache " - f"estimation ({error}); falling back to layer metadata.") - else: - if any(window is not None for window in inferred_windows): - attention_windows = [ - max_seq_len if window is None else window - for window in inferred_windows - ] - - if attention_windows is not None: - normalized_windows = _normalize_attention_windows( - attention_windows, max_seq_len) - if normalized_windows is None: - return 1 - return len(set(normalized_windows)) - - if isinstance(layer_types, (list, tuple)): - num_layer_types = len(set(layer_types)) - if num_layer_types > 1: - return num_layer_types - - if fallback_attention_windows is not None: - normalized_windows = _normalize_attention_windows( - fallback_attention_windows, max_seq_len) - if normalized_windows is not None: - return len(set(normalized_windows)) - - return 1 - - -def draft_config_defines_attention_layout( - draft_pretrained_config: object, ) -> bool: - """Return whether the draft HF config explicitly defines its attention layout. - - A ``True`` result makes the draft settings authoritative, including an - explicit full-attention layout. For example, a config with - ``use_sliding_window=False`` and ``sliding_window=4096`` returns ``True``: - its layers should attend to ``max_seq_len`` instead of inheriting the - target model's window. A config that provides none of - ``use_sliding_window``, ``sliding_window``, or ``layer_types`` returns - ``False`` so the legacy uniform-target fallback can be used. - """ - return ( - getattr(draft_pretrained_config, "use_sliding_window", None) is not None - or getattr(draft_pretrained_config, "sliding_window", None) is not None - or bool(getattr(draft_pretrained_config, "layer_types", None))) - - -def _expand_attention_window_pattern_to_global_layers( - max_attention_window: Optional[Sequence[int]], - layer_mask: Sequence[bool], -) -> Optional[List[int]]: - """Expand an enabled-layer pattern into physical global-layer order.""" - if max_attention_window is None: - return None - - pattern = list(max_attention_window) - global_windows = [pattern[0]] * len(layer_mask) - enabled_layer_offset = 0 - for layer_idx, enabled in enumerate(layer_mask): - if enabled: - global_windows[layer_idx] = pattern[enabled_layer_offset % - len(pattern)] - enabled_layer_offset += 1 - return global_windows - - -def _derive_draft_max_attention_window( - kv_cache_config: KvCacheConfig, - draft_pretrained_config: object, - max_seq_len: int, - num_draft_layers: int, -) -> Optional[List[int]]: - layer_windows = [ - get_layer_attention_window(draft_pretrained_config, layer_idx) - for layer_idx in range(num_draft_layers) - ] - if draft_config_defines_attention_layout(draft_pretrained_config): - draft_windows = [ - max_seq_len if window is None else window - for window in layer_windows - ] - return _normalize_attention_windows(draft_windows, max_seq_len) - - if not uses_vswa_kv_cache_layout(kv_cache_config.max_attention_window): - max_attention_window = kv_cache_config.max_attention_window - if max_attention_window is None: - return None - return _normalize_attention_windows(max_attention_window, max_seq_len) - - return None +def is_vswa_enabled(kv_cache_config): + max_attention_window = kv_cache_config.max_attention_window + return max_attention_window is not None and len( + set(max_attention_window)) > 1 class KvCacheCreator: """Groups together logic related to KV cache construction.""" - # Byte budgets that back an offload tier reserved in full at manager - # construction: the host tier is prefaulted and page-locked, the disk tier - # is preallocated. Every live manager reserves its own, so these budgets - # must be divided rather than handed out whole. - _OFFLOAD_TIER_BUDGET_ATTRS = ("host_cache_size", "disk_cache_size") - - # Paired-reuse protocol flag. ``build_managers`` resolves it once and hands - # it to both constructors; the managers must not re-derive it. - _joint_kv_cache_reuse = False - def __init__( self, *, @@ -625,13 +233,12 @@ def __init__( self._mapping = mapping self._kv_cache_config = kv_cache_config self._max_kv_tokens_in = self._kv_cache_config.max_tokens - self._max_gpu_total_bytes_in = self._kv_cache_config.max_gpu_total_bytes - self._pool_ratio_in = self._kv_cache_config.pool_ratio - self._avg_seq_len_in = self._kv_cache_config.avg_seq_len self._max_num_tokens = max_num_tokens self._max_beam_width = max_beam_width self._kv_connector_manager = kv_connector_manager self._llm_args = llm_args + # For V2 fallback use only, will be removed after V2 is stable + self._cache_transceiver_config = llm_args.cache_transceiver_config self._speculative_config = speculative_config self._sparse_attention_config = sparse_attention_config self._tokens_per_block = tokens_per_block @@ -639,26 +246,13 @@ def __init__( self._max_batch_size = max_batch_size self._net_max_seq_len = net_max_seq_len self._dummy_reqs = None - self._dummy_encoder_inputs: List[MultimodalParams] = [] self._profiling_stage_data = profiling_stage_data self._is_disagg = is_disagg - self._cache_transceiver_config = llm_args.cache_transceiver_config self._execution_stream = execution_stream self._kv_cache_manager_cls = self._get_model_kv_cache_manager_cls( model_engine) - self._is_kv_cache_manager_v2 = issubclass(self._kv_cache_manager_cls, - KVCacheManagerV2) self._draft_config = draft_config self._skip_est = skip_est - # Admission cap (tokens of summed context attended-KV) that the fp8 context-MLA workspace reservation - # covers, computed in configure_kv_cache_capacity and carried to the KV manager so the scheduler - # reads it directly instead of re-deriving it from pool layout. None until reserved (or w == 0). - self._fp8_ctx_mla_kv_len_cap = None - self._maybe_enable_fabric_memory_for_python_transceiver() - - def _maybe_enable_fabric_memory_for_python_transceiver(self) -> None: - maybe_enable_fabric_memory_for_python_transceiver( - self._cache_transceiver_config, self._kv_cache_manager_cls) def _get_model_kv_cache_manager_cls( self, @@ -668,90 +262,83 @@ def _get_model_kv_cache_manager_cls( kv_cache_config = (kv_cache_config_override if kv_cache_config_override is not None else self._kv_cache_config) model_config = model_engine.model.model_config + config = model_config.pretrained_config cls = get_kv_cache_manager_cls( model_config, kv_cache_config, is_disagg=self._is_disagg, cache_transceiver_config=self._cache_transceiver_config) - cls = self._validate_or_fallback_kv_cache_manager_v2( - cls, model_config, kv_cache_config) - if is_hybrid_linear(model_config.pretrained_config): - logger.info_once( - f"Selected hybrid KV cache manager: {cls.__name__}", - key=f"hybrid_kv_cache_manager_{cls.__name__}") - # Compatibility managers do not support MTP block reuse. Warn at the - # routing site so users see the concrete manager selected for the - # incompatible combination. - if is_hybrid_linear(model_engine.model.model_config.pretrained_config) \ - and kv_cache_config.enable_block_reuse \ - and self._speculative_config is not None: - if not issubclass(cls, MambaHybridCacheManagerV2): - logger.warning( - "Block reuse does not work with MTP for hybrid linear models " - f"when using non-V2 Mamba cache manager {cls.__name__}") - return cls - - def _validate_or_fallback_kv_cache_manager_v2( - self, - kv_cache_manager_cls, - model_config: ModelConfig, - kv_cache_config: Optional[KvCacheConfig] = None): - config = model_config.pretrained_config # Use ``issubclass`` rather than identity equality so V2 subclasses # (e.g. ``MiniMaxM3KVCacheManagerV2`` from the sparse-attention path) - # also go through the V2-incompatible-feature gate below. - if issubclass(kv_cache_manager_cls, KVCacheManagerV2): - sparse_attn_config = model_config.sparse_attention_config + # also go through the V2-incompatible-feature gate below. The earlier + # ``cls == KVCacheManagerV2`` check silently bypassed all V2 subclasses + # and let the bare ``assert event_buffer_max_size == 0`` in + # ``KVCacheManagerV2.__init__`` trip at executor construction with + # a useless error message. + if issubclass(cls, KVCacheManagerV2): incompat: List[str] = [] if self._kv_connector_manager is not None: incompat.append("kv_connector_manager") if self._max_beam_width is not None and self._max_beam_width > 1: - incompat.append("max_beam_width > 1") + incompat.append("beam_width > 1") + if kv_cache_config.event_buffer_max_size > 0: + incompat.append("event_buffer_max_size > 0") + if (self._cache_transceiver_config is not None + and self._cache_transceiver_config.backend is not None): + incompat.append("cache_transceiver") if incompat: incompat_str = ", ".join(incompat) - # Never silently replace a sparse V2 manager with V1. Some - # sparse models require V2 structurally; for models such as DSA - # that support both managers, fallback would ignore the user's - # explicit manager selection. + # Some models are structurally bound to V2 and cannot fall + # back to V1 without producing wrong outputs: + # * Sparse-attention models (e.g. MiniMax-M3) need V2's + # per-layer split-pool to allocate the per-sparse-layer + # INDEX_KEY pool with a different stride than the main + # K/V pool. V1's unified pool cannot represent that. + # * Gemma4 hybrid uses per-layer head_dim that V1 would + # coerce to ``max(head_dim)``, changing per-layer KV + # byte sizes — correctness bug, not just efficiency. + sparse_attn_config = model_config.sparse_attention_config if sparse_attn_config is not None: raise NotImplementedError( - f"KVCacheManagerV2 for sparse-attention models " - f"(algorithm={sparse_attn_config.algorithm!r}) is not " - f"supported with " - f"{incompat_str}. Disable the incompatible features to " - f"run sparse-attention models.") - # Gemma4 hybrid uses per-layer head_dim that V1 would coerce to - # ``max(head_dim)``, changing per-layer KV byte sizes. + f"Sparse-attention models " + f"(algorithm={sparse_attn_config.algorithm!r}) require " + f"KVCacheManagerV2, which is not yet supported with " + f"{incompat_str}. Disable these KvCacheConfig features " + f"to run sparse-attention models.") if is_gemma4_hybrid(config): raise NotImplementedError( f"Gemma4 hybrid attention requires KVCacheManagerV2, " f"which is not yet supported with {incompat_str}. " f"Disable these features to run Gemma4 hybrid models.") - if is_hybrid_linear(config): - raise NotImplementedError( - "Hybrid Mamba cache managers do not support " - f"{incompat_str}; CppMambaHybridCacheManager does not " - "provide a compatible fallback. Use max_beam_width=1 " - "and disable the KV connector.") - # Plain V2 (explicitly enabled or selected by a model preference): - # V2 was a preference, not a structural requirement, so we can - # safely fall back to V1. + # Plain V2 (user opt-in via ``use_kv_cache_manager_v2=True``): + # V2 was a preference, not a structural requirement, so we + # can safely fall back to V1. logger.warning( "KVCacheManagerV2 is not supported with %s. " "Falling back to KVCacheManager.", incompat_str) - return KVCacheManager - return kv_cache_manager_cls - - def _enable_kv_cache_stats(self) -> bool: - return (self._llm_args.enable_iter_perf_stats - or getattr(self._llm_args, "return_perf_metrics", False)) + cls = KVCacheManager + # 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 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( + "Block reuse does not work with MTP for hybrid linear models " + "when using the legacy MambaCacheManager (TRTLLM_USE_CPP_MAMBA=1)" + ) + return cls def _per_manager_cache_cost(self, manager_cls, model_config, kv_cache_config: Optional[KvCacheConfig] = None, - *, - is_draft: bool = False, **extra_kwargs) -> CacheCost: kv_cache_config = (kv_cache_config if kv_cache_config is not None else self._kv_cache_config) @@ -760,23 +347,10 @@ def _per_manager_cache_cost(self, model_config, self._mapping, tokens_per_block=self._tokens_per_block, - max_seq_len=self._max_seq_len, max_batch_size=self._max_batch_size, - max_num_tokens=self._max_num_tokens if is_draft else 0, kv_cache_config=kv_cache_config, - spec_config=self._speculative_config, - is_draft=is_draft, **extra_kwargs)) - def _get_one_model_draft_layer_mask(self) -> List[bool]: - """Return the same draft-only mask used by runtime construction.""" - num_draft_layers = self._get_num_draft_layers() - if self._speculative_config.spec_dec_mode.is_external_drafter(): - return [True] * num_draft_layers - target_num_layers = (self._model_engine.model.model_config. - pretrained_config.num_hidden_layers) - return [False] * target_num_layers + [True] * num_draft_layers - def _get_kv_size_per_token(self, kv_cache_config: Optional[KvCacheConfig] = None ) -> CacheCost: @@ -788,13 +362,8 @@ def _get_kv_size_per_token(self, kv_cache_config = (kv_cache_config if kv_cache_config is not None else self._kv_cache_config) model_config = self._model_engine.model.model_config - use_separate_draft_kv_cache = ( - self._should_create_separate_draft_kv_cache()) - total = self._per_manager_cache_cost( - self._kv_cache_manager_cls, - model_config, - kv_cache_config, - use_separate_draft_kv_cache=use_separate_draft_kv_cache) + total = self._per_manager_cache_cost(self._kv_cache_manager_cls, + model_config, kv_cache_config) if self._is_encoder_decoder(): total += CacheCost.from_raw(self._get_cross_kv_size_per_token()) if self._draft_model_engine is not None: @@ -804,7 +373,7 @@ def _get_kv_size_per_token(self, total += self._per_manager_cache_cost(draft_kv_cache_manager_cls, draft_model_config, kv_cache_config) - elif use_separate_draft_kv_cache: + elif self._should_create_separate_draft_kv_cache(): # One-model draft with separate KV cache layout. # Pass num_layers explicitly since the HF config may report a # different layer count than what is actually used at runtime @@ -812,29 +381,24 @@ def _get_kv_size_per_token(self, # For PP, draft layers are only on the last rank (see # get_pp_layers), so only that rank should include draft cost. effective_draft_config = self._get_effective_draft_config() - draft_kv_cache_config = self._get_one_model_draft_kv_cache_config( - kv_cache_config, self._max_seq_len) if self._speculative_config.spec_dec_mode.is_external_drafter(): # External drafter: layers start from 0, normal PP distribution # Resolve draft manager class from draft config — may differ # from target (e.g. hybrid target + plain transformer draft). draft_kv_cache_manager_cls = get_kv_cache_manager_cls( effective_draft_config, - draft_kv_cache_config, + kv_cache_config, is_disagg=self._is_disagg) total += self._per_manager_cache_cost( - draft_kv_cache_manager_cls, - effective_draft_config, - draft_kv_cache_config, - is_draft=True) + draft_kv_cache_manager_cls, effective_draft_config, + kv_cache_config) elif self._mapping.is_last_pp_rank(): # EAGLE3/MTP: draft layers only on last PP rank total += self._per_manager_cache_cost( self._kv_cache_manager_cls, effective_draft_config, - draft_kv_cache_config, - num_layers=self._get_num_draft_layers(), - is_draft=True) + kv_cache_config, + num_layers=self._get_num_draft_layers()) return total def _cal_max_memory(self, peak_memory, total_gpu_memory, fraction, @@ -856,11 +420,92 @@ def _cal_max_memory(self, peak_memory, total_gpu_memory, fraction, ) return int(available_kv_mem) + def _create_dummy_mm_context_request( + self, input_seq_len: int) -> List[trtllm.Request]: + requests = [] + if isinstance( + self._profiling_stage_data, + dict) and not self._profiling_stage_data.get("enable_mm_reqs"): + return requests + + input_processor = self._model_engine.input_processor + if not (hasattr(input_processor, "get_dummy_prompt")): + logger.warning("The input processor of the model does not have the method [get_dummy_prompt] implemented." \ + "Profiling with the default input dummy context request. This may not take into account the memory consumption of " \ + "the image encoder") + return requests + + max_num_tokens = self._max_num_tokens + max_beam_width = self._max_beam_width + vocab_size = self._model_engine.model.model_config.pretrained_config.vocab_size + + input_seq_len = min(max_num_tokens, input_seq_len) + remaining_tokens = max_num_tokens + while remaining_tokens > 0: + input_seq_len = min(input_seq_len, remaining_tokens) + dummy_mm_prompt = input_processor.get_dummy_prompt(input_seq_len) + + if dummy_mm_prompt is not None: + prompt_token_ids, extra_processed_inputs = self._model_engine.input_processor_with_hash( + dummy_mm_prompt, sampling_params=None) + + multimodal_input = extra_processed_inputs.get( + 'multimodal_input') + multimodal_data = extra_processed_inputs.get('multimodal_data') + req_mm_input = multimodal_input.to_binding( + trtllm) if multimodal_input else None + + request = trtllm.Request(prompt_token_ids, + max_tokens=1, + streaming=False, + sampling_config=trtllm.SamplingConfig( + beam_width=max_beam_width, ), + output_config=trtllm.OutputConfig(), + end_id=-1, + multimodal_input=req_mm_input) + request.py_multimodal_data = multimodal_data + else: + # Fall back to text-only prompt when we could not find the small image size. + prompt_token_ids = torch.randint( + low=0, high=vocab_size, size=(input_seq_len, )).tolist() + request = trtllm.Request(prompt_token_ids, + max_tokens=1, + streaming=False, + sampling_config=trtllm.SamplingConfig( + beam_width=max_beam_width, ), + output_config=trtllm.OutputConfig(), + end_id=-1) + if self._model_engine.use_mrope: + request.py_multimodal_data = { + "mrope_config": { + "mrope_position_ids": + torch.zeros(3, 1, input_seq_len, dtype=torch.int32), + "mrope_position_deltas": + torch.zeros(1, 1, dtype=torch.int32) + } + } + remaining_tokens -= len(prompt_token_ids) + requests.append(request) + + if self._mapping.enable_attention_dp: + requests = requests * self._mapping.tp_size + + return requests + def _create_dummy_context_requests( self, input_seq_len: int) -> List[trtllm.Request]: - # Keep the LLM dummy text-only so it can always fill max_num_tokens. - # The MM encoder is profiled independently at its own token budget. requests = [] + # Disaggregated workers receive multimodal embeddings instead of raw + # pixel inputs, so capacity probing must use the text-only fallback. + if (not _is_mm_disagg() + and hasattr(self._model_engine.model, "original_arch") + and MODEL_CLASS_VISION_ENCODER_MAPPING.get( + self._model_engine.model.original_arch, None)): + requests = self._create_dummy_mm_context_request(input_seq_len) + # if succeed profiling with multimodal requests then return, otherwise profile + # with default case + if requests: + return requests vocab_size = self._model_engine.model.model_config.pretrained_config.vocab_size max_num_tokens = self._max_num_tokens max_beam_width = self._max_beam_width @@ -888,122 +533,12 @@ def _create_dummy_context_requests( torch.zeros(1, 1, dtype=torch.int32) } } - request.py_conversation_params = None requests.append(request) remaining_tokens -= input_seq_len if self._mapping.enable_attention_dp: requests = requests * self._mapping.tp_size return requests - def _create_dummy_encoder_inputs(self) -> List[MultimodalParams]: - """Build one processed MM encoder batch at its scheduling limits.""" - if not isinstance(self._model_engine.model, MultimodalModelMixin): - return [] - if isinstance( - self._profiling_stage_data, - dict) and not self._profiling_stage_data.get("enable_mm_reqs"): - return [] - if self._llm_args.disable_mm_encoder: - return [] - # MM E/P disaggregation may remove an otherwise exposed encoder. - if (hasattr(self._model_engine.model, "mm_encoder") - and self._model_engine.model.mm_encoder is None): - return [] - - input_processor = self._model_engine.input_processor - encoder_max_num_tokens = self._model_engine.encoder_max_num_tokens - if encoder_max_num_tokens is None or encoder_max_num_tokens <= 0: - return [] - - try: - max_tokens_per_item = input_processor.get_mm_max_tokens_per_item( - max_num_encoder_tokens=encoder_max_num_tokens) - for modality, num_tokens in max_tokens_per_item.items(): - if not modality: - raise ValueError("Multimodal modality name cannot be empty") - if num_tokens <= 0: - raise ValueError( - "Multimodal encoder token counts must be positive; " - f"got {num_tokens} for {modality}") - max_tokens_per_item = { - modality: num_tokens - for modality, num_tokens in max_tokens_per_item.items() - if num_tokens <= encoder_max_num_tokens - } - if not max_tokens_per_item: - return [] - - modality, num_tokens_per_item = max( - max_tokens_per_item.items(), - key=lambda item: (item[1], item[0]), - ) - num_items = min( - self._model_engine.encoder_batch_size, - encoder_max_num_tokens // num_tokens_per_item, - ) - mm_data = input_processor.get_dummy_mm_data( - max_num_encoder_tokens=encoder_max_num_tokens, - mm_counts={modality: num_items}, - dtype=self._model_engine.model.dtype, - ) - except NotImplementedError: - logger.info("Multimodal memory profiling skipped: " - f"{type(input_processor).__name__} does not implement " - "get_dummy_mm_data().") - return [] - if not mm_data: - return [] - if not isinstance(mm_data, dict): - raise ValueError( - "get_dummy_mm_data() must return a multimodal_data " - "dictionary") - return [MultimodalParams(multimodal_data=mm_data)] - - def _encode_dummy_inputs(self) -> Optional[torch.Tensor]: - """Run the full-budget MM encoder and retain request-owned output storage.""" - if not self._dummy_encoder_inputs: - return None - - encoder_inputs = self._dummy_encoder_inputs - try: - with torch.inference_mode(): - for encoder_input in encoder_inputs: - encoder_input.to_device( - "multimodal_data", - "cuda", - pin_memory=prefer_pinned(), - target_keywords=getattr( - self._model_engine.model, - "multimodal_data_device_paths", - None, - ), - ) - output = self._model_engine.model.encode_multimodal_inputs( - encoder_inputs) - # Runtime item state owns detached copies rather than views of - # an encoder batch. Reproduce that allocation boundary here. - return output.detach().clone() - finally: - self._dummy_encoder_inputs = [] - - def _get_multimodal_encoder_memory_reserve(self, - profiled_output_bytes: int = 0 - ) -> int: - """Return output and cache capacity absent from the measured peak.""" - output_budget = getattr(self._model_engine, - "mm_encoder_output_budget_bytes", None) - unprofiled_output_bytes = max(0, (output_budget or 0) - - profiled_output_bytes) - - model = self._model_engine.model - cache_bytes = 0 - if (isinstance(model, MultimodalModelMixin) - and model.encoder_cache_active - and model.model_config.multimodal_config is not None): - cache_bytes = ( - model.model_config.multimodal_config.encoder_cache_max_bytes) - return unprofiled_output_bytes + cache_bytes - def _get_token_num_for_estimation(self) -> int: """Compute KV cache capacity required for estimate_max_kv_cache_tokens to succeed.""" if 'cp_type' in self._mapping.cp_config: @@ -1024,7 +559,6 @@ def _get_token_num_for_estimation(self) -> int: if self._dummy_reqs is None: self._dummy_reqs = self._create_dummy_context_requests( max(1, self._net_max_seq_len - 1)) - self._dummy_encoder_inputs = self._create_dummy_encoder_inputs() for req in self._dummy_reqs: num_req_tokens = len(req.input_token_ids) + num_extra_tokens_per_seq # Requests cannot share KV cache blocks. Round up to nearest integer multiple of block size. @@ -1051,40 +585,56 @@ def _get_token_num_for_estimation(self) -> int: # If not able to allocate self._model_engine.batch_size blocks, the max batch size should be adjusted. num_cache_blocks = max(num_cache_blocks, self._model_engine.batch_size) - # KVCacheManagerV2 divides the quota derived from max_tokens across its - # pool groups. Scale the dummy workload by the inferred group count so - # each pool can hold a max-length request. This covers both VSWA pools - # (distinct attention windows) and hybrid recurrent/attention pools - # (distinct layer types without sliding windows). + # For VSWA (variable sliding window attention) models such as Gemma4 + # hybrid, KVCacheManagerV2 creates a separate pool group per distinct + # attention window size. The quota passed via max_tokens is split + # across pool groups proportionally, so each pool ends up with roughly + # num_cache_blocks/num_pool_groups blocks in the worst case. A single + # context request of max_seq_len tokens then exceeds the full-attention + # pool's block budget and resize_context livelocks on suspend/retry + # (observed for Gemma4 multimodal at max_seq_len>=8K, e.g. MMMU Pro). + # Scale num_cache_blocks by the number of distinct pool groups so that + # each pool has enough blocks for the dummy request even after the + # proportional split. Inferred from the model config since the hybrid + # max_attention_window hasn't been populated in kv_cache_config yet at + # this stage (it's filled in later by _create_kv_cache_manager). + # Only V2 has split-pool semantics — Mamba hybrid (which also has + # heterogeneous layer_types) uses MambaHybridCacheManager and would + # have its max_tokens estimate inflated incorrectly otherwise. num_pool_groups = 1 - if self._is_kv_cache_manager_v2: + if self._kv_cache_manager_cls == KVCacheManagerV2: model_cfg = self._model_engine.model.model_config.pretrained_config - num_pool_groups = _get_num_pool_groups_for_estimation( - model_cfg, - self._model_engine.max_seq_len, - self._kv_cache_config.max_attention_window, - ) + layer_types = getattr(model_cfg, "layer_types", None) + if isinstance(layer_types, (list, tuple)): + distinct = len(set(layer_types)) + if distinct > 1: + num_pool_groups = distinct + elif (self._kv_cache_config.max_attention_window is not None + and len(set(self._kv_cache_config.max_attention_window)) > 1): + num_pool_groups = len( + set(self._kv_cache_config.max_attention_window)) num_cache_blocks *= num_pool_groups - # Dummy context requests use the configured maximum beam width. Scale - # their block budget by the same value so the temporary KV cache used - # during warm-up can accommodate those requests. - num_cache_blocks *= self._max_beam_width - - max_num_tokens_for_estimation = (num_cache_blocks * - self._tokens_per_block) - # V2 capacity is controlled by max_gpu_total_bytes; max_tokens only - # describes the dummy workload needed for estimation. - if self._is_kv_cache_manager_v2: - return max_num_tokens_for_estimation - - free_mem, _ = torch.cuda.mem_get_info() + free_mem, total_mem = torch.cuda.mem_get_info() max_memory = self._kv_cache_config.free_gpu_memory_fraction * free_mem kv_size_per_token = self._get_kv_size_per_token() + if kv_size_per_token.intercept > 0 and is_device_integrated(): + # On unified-memory GPUs the free counter can already be depressed + # by mmap-backed weights sharing the same physical pool. During the + # estimation dry run, subtracting the recurrent-state fixed cost + # here can collapse the provisional token cap to zero even though + # post-profiling affine sizing succeeds. Keep the final affine + # sizing unchanged; relax only this provisional cap. + kv_size_per_token = CacheCost(slope=kv_size_per_token.slope) max_num_tokens_in_memory = ( kv_size_per_token.tokens_for_budget(max_memory) // self._tokens_per_block * self._tokens_per_block) - return min(max_num_tokens_for_estimation, max_num_tokens_in_memory) + + # Multiply by beam width, to prevent rescaling of the max_seq_len caused by the influence of beam width during the preparation for kv_cache_estimation + return min( + num_cache_blocks * self._tokens_per_block * + self._dummy_reqs[0].sampling_config.beam_width, + max_num_tokens_in_memory) def try_prepare_estimation(self) -> bool: """Prepare for possible KV cache capacity estimation. @@ -1094,98 +644,21 @@ def try_prepare_estimation(self) -> bool: """ if self._skip_est: return False - - estimating_kv_cache = True - if 'cp_type' in self._mapping.cp_config: - estimating_kv_cache = False - logger.info( - "KV cache size estimation is not supported for context parallelism, disable it." - ) - if (self._is_kv_cache_manager_v2 - and self._mapping.cp_config.get('cp_type') == CpType.HELIX): - # Promote like the encoder-decoder case so build_managers - # runs configure_kv_cache_capacity(), which sets the quota - # KVCacheManagerV2 requires at construction (V1 stays local). - # HELIX only: configure_kv_cache_capacity has no sizing path - # for other CP types and would hit its assertion. - self._skip_est = True + estimating_kv_cache = False + if 'cp_type' not in self._mapping.cp_config: + estimating_kv_cache = True + estimate_max_tokens = self._get_token_num_for_estimation() + self._kv_cache_config.max_tokens = min( + estimate_max_tokens, self._kv_cache_config.max_tokens + ) if self._kv_cache_config.max_tokens is not None else estimate_max_tokens model_config = self._model_engine.model.model_config if model_config.attn_backend == "VANILLA": - estimating_kv_cache = False logger.info( "KV cache size estimation is not supported for Vanilla attention backend, disable it." ) - if getattr(model_config, "is_encoder_decoder", False): - # The estimation dummies are text-only, and the cross-KV block - # accounting needs an encoder length (getEncoderOutputLen throws). - # _skip_est (not just the local flag) so build_managers runs - # configure_kv_cache_capacity(), which KVCacheManagerV2 needs for - # its memory quota — the TRTLLM_SKIP_KV_CACHE_ESTIMATION=1 path. - self._skip_est = True estimating_kv_cache = False - logger.info( - "KV cache size estimation is not supported for encoder-decoder " - "models, disable it.") - - if estimating_kv_cache: - estimate_max_tokens = self._get_token_num_for_estimation() - max_tokens = min( - estimate_max_tokens, self._kv_cache_config.max_tokens - ) if self._kv_cache_config.max_tokens is not None else estimate_max_tokens - # User-provided pool sizing can underprovision the temporary - # estimation cache and cause warmup to hang or fail. Override it - # for estimation, then restore it in configure_kv_cache_capacity(). - self._kv_cache_config.pool_ratio = None - self._kv_cache_config.avg_seq_len = self._max_seq_len - if self._is_kv_cache_manager_v2: - free_mem, _ = torch.cuda.mem_get_info() - max_gpu_total_bytes = int( - self._kv_cache_config.free_gpu_memory_fraction * free_mem) - if (self._max_gpu_total_bytes_in is not None - and self._max_gpu_total_bytes_in > 0): - max_gpu_total_bytes = min(max_gpu_total_bytes, - self._max_gpu_total_bytes_in) - self._kv_cache_config.max_gpu_total_bytes = max_gpu_total_bytes - self._kv_cache_config.max_tokens = max_tokens - else: - self._kv_cache_config.max_tokens = max_tokens return estimating_kv_cache - def _configure_helix_kv_cache_capacity(self) -> None: - """Set the helix KV quota without profiling (not CP-aware). - - Explicit quotas pass through; otherwise fraction sizing sets - ``max_gpu_total_bytes`` (a rank-local byte cap the manager consumes - as-is). Setting ``max_tokens`` here would overshoot the fraction: - the manager inflates that knob by 1 / max_util_for_resume. - """ - if (self._kv_cache_config.max_tokens is not None - and self._kv_cache_config.max_tokens <= 0): - raise ValueError( - "Helix CP: kv_cache_config.max_tokens must be positive when " - f"set, got {self._kv_cache_config.max_tokens}.") - if (self._kv_cache_config.max_gpu_total_bytes or 0) > 0 or \ - (self._kv_cache_config.max_tokens or 0) > 0: - logger.info("Helix CP: skipping KV cache capacity profiling; using " - "the explicitly configured quota.") - return - fraction = self._kv_cache_config.free_gpu_memory_fraction - free_mem, _total = torch.cuda.mem_get_info() - budget_bytes = int(free_mem * fraction) - if budget_bytes <= 0: - raise ValueError( - "Helix CP: fraction-based KV sizing found no usable free " - "memory; set kv_cache_config.max_tokens or " - "max_gpu_total_bytes.") - logger.warning( - "Helix CP: capacity profiling is unsupported; sizing the KV " - f"cache as fraction {fraction} of free memory -> " - f"max_gpu_total_bytes={budget_bytes} (rank-local byte cap; the " - "manager min-syncs across ranks and converts to global tokens). " - "Set kv_cache_config.max_tokens or max_gpu_total_bytes to " - "override.") - self._kv_cache_config.max_gpu_total_bytes = budget_bytes - def configure_kv_cache_capacity(self, py_executor: PyExecutor = None) -> None: """Perform KV cache capacity estimation. @@ -1196,16 +669,6 @@ def configure_kv_cache_capacity(self, mapping = self._mapping # TODO: support CP by generating dummy requests for it. - if mapping.cp_config.get('cp_type') == CpType.HELIX: - if not self._is_kv_cache_manager_v2: - # The helix sizing below emits V2 ledger (global) quotas; - # V1 reads max_tokens as rank-local. Reject explicitly. - raise NotImplementedError( - "TRTLLM_SKIP_KV_CACHE_ESTIMATION with helix CP requires " - "the V2 KV cache manager " - "(kv_cache_config.use_kv_cache_manager_v2=True).") - self._configure_helix_kv_cache_capacity() - return assert 'cp_type' not in mapping.cp_config fraction = self._kv_cache_config.free_gpu_memory_fraction @@ -1222,16 +685,7 @@ def configure_kv_cache_capacity(self, f"Memory used after loading model weights (outside torch) in memory usage profiling: {((total_used_bytes - model_bytes) if total_used_bytes > model_bytes else 0) / (GB):.2f} GiB" ) - profiled_output_bytes = 0 - if py_executor is not None and not self._skip_est: - # Run the MM encoder at its independent token budget, then keep the - # resulting request-owned embeddings resident while the text-only - # LLM dummy fills max_num_tokens. - encoder_profile_output = self._encode_dummy_inputs() - if encoder_profile_output is not None: - profiled_output_bytes = (encoder_profile_output.numel() * - encoder_profile_output.element_size()) py_executor.set_gather_responses(True) origin_iter_stats = py_executor.enable_iter_perf_stats py_executor.enable_iter_perf_stats = False @@ -1254,20 +708,12 @@ def configure_kv_cache_capacity(self, torch_peak_memory = torch.cuda.memory_stats( )["allocated_bytes.all.peak"] - # Release before measuring current usage so the retained - # embeddings count toward the peak but not the steady state. - encoder_profile_output = None - # Clear the caching allocator before measuring the current memory usage torch.cuda.empty_cache() end, total_gpu_memory = torch.cuda.mem_get_info() torch_used_bytes = torch.cuda.memory_stats( )["allocated_bytes.all.current"] finally: - # Redundant on the success path, but a failed dummy run would - # otherwise keep the profiling embeddings alive through - # teardown -- exactly when memory is already scarce. - encoder_profile_output = None # get kv cache stats for both model and draft model kv_stats = py_executor.resource_manager.resource_managers.get( ResourceManagerType.KV_CACHE_MANAGER).get_kv_cache_stats() @@ -1303,65 +749,11 @@ def configure_kv_cache_capacity(self, allocated_bytes = 0 activation_bytes = 0 - multimodal_encoder_memory_reserve = ( - self._get_multimodal_encoder_memory_reserve( - profiled_output_bytes=profiled_output_bytes)) - peak_memory += multimodal_encoder_memory_reserve - if multimodal_encoder_memory_reserve > 0: - mem_gb = multimodal_encoder_memory_reserve / GB - logger.info( - f"Reserving {mem_gb:.2f} GiB for multimodal encoder memory " - "not materialized by the profiling run.") - # calculate max memory from peak memory and free gpu memory fraction kv_cache_max_memory = self._cal_max_memory(peak_memory, total_gpu_memory, fraction, allocated_bytes) - # Estimation uses inferred pool sizing; the final manager uses the - # user-provided configuration. - self._kv_cache_config.pool_ratio = self._pool_ratio_in - self._kv_cache_config.avg_seq_len = self._avg_seq_len_in - - # Reserve headroom for attention workspace the selected backend declares and the profiling forward - # under-measures. KV-cache reuse can push summed attended KV past the profiled floor. Chunked prefill - # usually bounds each attention launch by its chunk buffer, except for implementations such as the - # NVFP4 DSA context gather that consume the complete attended prefix. The backend declares that - # distinction through runtime_workspace_is_chunked_prefill_bounded. When a reserve applies, - # reserve w * L_cap bytes -- covering the worst-case summed attended KV the scheduler admits - # (get_mla_context_workspace_kv_len_cap) -- but clamp it to the per-token split budget * w / (k + w) - # so a memory-constrained node shares the budget at a common token count instead of starving the - # pool. Equivalently the pool keeps max((budget - w*L_cap)/k, budget/(k+w)) tokens. The reserve - # covers exactly reserve/w tokens of summed attended KV; that count is carried to the KV manager as - # the scheduler's admission cap so it never re-derives the cap from pool layout (which V2 - # overstates). No cap or w == 0 -> no-op. - w_bytes_per_token = get_attention_workspace_bytes_per_token( - self._model_engine.model.model_config, self._mapping) - workspace_is_chunked_prefill_bounded = True - if w_bytes_per_token > 0: - workspace_is_chunked_prefill_bounded = ( - get_attention_workspace_is_chunked_prefill_bounded( - self._model_engine.model.model_config)) - kv_len_cap = get_mla_context_workspace_kv_len_cap( - self._kv_cache_config, self._max_batch_size, self._max_num_tokens, - self._max_seq_len, self._llm_args.enable_chunked_prefill, - workspace_is_chunked_prefill_bounded) - if w_bytes_per_token > 0 and kv_len_cap: - budget_before = kv_cache_max_memory - workspace_reserve, self._fp8_ctx_mla_kv_len_cap = ( - get_mla_context_workspace_reserve( - budget_before, - self._get_kv_size_per_token().slope, w_bytes_per_token, - kv_len_cap)) - if workspace_reserve > 0: - kv_cache_max_memory = int(budget_before - workspace_reserve) - logger.info( - f"Reserving {workspace_reserve / (GB):.2f} GiB for the context-MLA attention " - f"workspace (w={w_bytes_per_token} B/token, admitting up to " - f"{self._fp8_ctx_mla_kv_len_cap} tokens of summed attended KV): KV cache budget " - f"{budget_before / (GB):.2f} -> {kv_cache_max_memory / (GB):.2f} GiB." - ) - # NOTE: # For KVCacheManager, KvCacheCreator currently controls capacity using two parameters in KVCacheConfig: # • max_tokens @@ -1370,15 +762,14 @@ def configure_kv_cache_capacity(self, # This leaves max_tokens as a user-defined constraint. # ---------------------------handle max_tokens--------------------------------- - if self._is_kv_cache_manager_v2: + if issubclass(self._kv_cache_manager_cls, KVCacheManagerV2): # KVCacheManagerV2 doesn't rely on max_tokens to control capacity, so restore user provided value self._kv_cache_config.max_tokens = self._max_kv_tokens_in else: # handle user provided max_tokens if self._max_kv_tokens_in is not None: # raise error if it is VSWA case - is_vswa = uses_vswa_kv_cache_layout( - self._kv_cache_config.max_attention_window) + is_vswa = is_vswa_enabled(self._kv_cache_config) # raise error if it is VSWA case if is_vswa: @@ -1402,12 +793,11 @@ def configure_kv_cache_capacity(self, # ---------------------------handle max_gpu_total_bytes--------------------------------- # if user provided max_gpu_total_bytes, set max memory from max_gpu_total_bytes - if (self._max_gpu_total_bytes_in is not None - and self._max_gpu_total_bytes_in > 0): + if self._kv_cache_config.max_gpu_total_bytes > 0: kv_cache_max_memory = min(kv_cache_max_memory, - self._max_gpu_total_bytes_in) + self._kv_cache_config.max_gpu_total_bytes) logger.info( - f"max_gpu_total_bytes={self._max_gpu_total_bytes_in / (GB):.2f} GiB is provided. New max memory is {kv_cache_max_memory / (GB):.2f} GiB" + f"max_gpu_total_bytes={self._kv_cache_config.max_gpu_total_bytes / (GB):.2f} GiB is provided. New max memory is {kv_cache_max_memory / (GB):.2f} GiB" ) logger.info( @@ -1423,8 +813,7 @@ def _create_kv_cache_manager( self, model_engine: PyTorchModelEngine, estimating_kv_cache: bool = False, - kv_cache_config_override: Optional[KvCacheConfig] = None, - cold_page_codec_provider: Optional[object] = None, + kv_cache_config_override: Optional[KvCacheConfig] = None ) -> KVCacheManager: mapping = self._mapping assert model_engine.model.model_config.is_generation, "Only construct KV cache for generation models." @@ -1457,16 +846,9 @@ def _create_kv_cache_manager( max_beam_width=self._max_beam_width, kv_connector_manager=self._kv_connector_manager, estimating_kv_cache=estimating_kv_cache, - enable_kv_cache_stats=self._enable_kv_cache_stats() - and not estimating_kv_cache, execution_stream=self._execution_stream, layer_mask=spec_dec_layer_mask, is_disagg=self._is_disagg, - kv_events_config=None - if estimating_kv_cache or model_engine.is_draft_model else - self._llm_args.kv_cache_config.kv_events_config, - cold_page_codec_provider=cold_page_codec_provider, - joint_kv_cache_reuse=self._joint_kv_cache_reuse, ) if not self._skip_est: @@ -1503,45 +885,13 @@ def _should_create_separate_draft_kv_cache(self) -> bool: in the target model and don't produce a separate ModelConfig. We fall back to the target model's config via _get_effective_draft_config(). """ - if self._speculative_config is None: - # No drafter at all, so there is nothing to give a manager to. - return False - # Narrower than is_external_drafter(): PARD and DRAFT_TARGET_ONE_MODEL - # never reach the arena this carve-out exists for. - spec_dec_mode = self._speculative_config.spec_dec_mode - is_standalone_drafter = (spec_dec_mode.is_dflash() - or spec_dec_mode.is_dspark()) - if self._mapping.enable_attention_dp and not is_standalone_drafter: - # This bail suits MTP, whose draft layers are target-shaped and - # appendable to the target pool. A standalone drafter has nothing to - # append, so it would be stranded on the private arena instead. + if self._mapping.enable_attention_dp: logger.info( "Attention DP is enabled, separate draft KV cache is not supported." ) return False - - sparse_cfg = self._sparse_attention_config - if (sparse_cfg is not None - and getattr(sparse_cfg, "algorithm", None) == "deepseek_v4" - and self._mapping.pp_size > 1): - logger.info( - "DeepSeek-V4 separate draft KV cache is only supported for PP=1; " - "folding draft layers into the unified manager for pp_size=%d.", - self._mapping.pp_size) - return False return should_use_separate_draft_kv_cache(self._speculative_config) - def _joint_reuse_supported(self) -> bool: - """Span half of the pairing decision; ``build_managers`` ANDs it with - ``has_separate_one_model_draft``. ``False`` = leave the run unpaired. - """ - if not self._is_kv_cache_manager_v2: - return False - if not getattr(self._kv_cache_manager_cls, - "_supports_reuse_match_backoff", False): - return False - return draft_prompt_lookahead(self._speculative_config) is not None - def _get_effective_draft_config(self) -> ModelConfig: """ Return the ModelConfig to use for draft KV cache creation. @@ -1569,88 +919,51 @@ def _get_num_draft_layers(self) -> int: return self._draft_config.pretrained_config.num_hidden_layers return get_num_spec_layers(self._speculative_config) - def _get_draft_max_attention_window( - self, - max_seq_len: int, - kv_cache_config: KvCacheConfig, - ) -> Optional[List[int]]: - """Derive the draft manager's per-layer attention windows.""" - effective_draft_config = self._get_effective_draft_config() - return _derive_draft_max_attention_window( - kv_cache_config, - effective_draft_config.pretrained_config, - max_seq_len, - self._get_num_draft_layers(), - ) - - def _get_one_model_draft_kv_cache_config( - self, - kv_cache_config: KvCacheConfig, - max_seq_len: int, - *, - estimating_kv_cache: bool = False, - ) -> KvCacheConfig: - """Return a clone with the draft manager's attention-window layout.""" - # Estimation uses a small max_tokens-sized temporary draft cache before - # the measured GPU budget is available to split. Applying VSWA there - # would size every window pool from the unsplit free-memory budget. - max_attention_window = (None if estimating_kv_cache else - self._get_draft_max_attention_window( - max_seq_len, kv_cache_config)) - return kv_cache_config.model_copy( - update={"max_attention_window": max_attention_window}) - def _create_one_model_draft_kv_cache_manager( self, - max_seq_len: int, estimating_kv_cache: bool = False, kv_cache_config_override: Optional[KvCacheConfig] = None, - cold_page_codec_provider: Optional[object] = None, ) -> Optional[KVCacheManager]: """ Create a KV cache manager for draft model layers in one-model mode when target and draft have different KV cache layouts. """ + # Get target model's num_hidden_layers to compute correct layer indices. + # Draft model layers in one-model mode start at target_num_layers. + target_pretrained_config = self._model_engine.model.model_config.pretrained_config + target_num_layers = target_pretrained_config.num_hidden_layers + + # PARD, External Drafter: draft is a separate model, layers start from 0. + # Other methods (EAGLE3, MTP): draft layers are appended after target layers. num_draft_layers = self._get_num_draft_layers() - spec_dec_layer_mask = self._get_one_model_draft_layer_mask() + if self._speculative_config.spec_dec_mode.is_external_drafter(): + spec_dec_layer_mask = [True] * num_draft_layers + else: + spec_dec_layer_mask = [False] * target_num_layers + [ + True + ] * num_draft_layers # Get the effective draft config (explicit draft_config if available, # otherwise fall back to target model config for MTP). effective_draft_config = self._get_effective_draft_config() - kv_cache_config = (kv_cache_config_override if kv_cache_config_override + draft_kv_config = (kv_cache_config_override if kv_cache_config_override is not None else self._kv_cache_config) - draft_kv_config = self._get_one_model_draft_kv_cache_config( - kv_cache_config, - max_seq_len, - estimating_kv_cache=estimating_kv_cache) - draft_kv_config.max_attention_window = ( - _expand_attention_window_pattern_to_global_layers( - draft_kv_config.max_attention_window, - spec_dec_layer_mask, - )) - if (not uses_vswa_kv_cache_layout(draft_kv_config.max_attention_window) - and draft_kv_config.pool_ratio is not None - and len(draft_kv_config.pool_ratio) != 1): - # pool_ratio describes one manager's layer-group layout. The - # target hybrid manager may have separate recurrent-state and - # attention layer groups, while a non-VSWA draft manager has one - # attention layer group. Reusing the target's ratios fails its arity - # check. - logger.info( - "Normalizing the separate one-model draft KV cache pool_ratio " - f"from {draft_kv_config.pool_ratio} to [1.0] for its single " - "layer group.") - draft_kv_config.pool_ratio = [1.0] - if uses_vswa_kv_cache_layout(draft_kv_config.max_attention_window): - logger.info( - f"Derived draft KV cache max_attention_window for separate " - f"draft manager: {draft_kv_config.max_attention_window}") # Get the appropriate KV cache manager class for the draft model draft_kv_cache_manager_cls = get_kv_cache_manager_cls( effective_draft_config, draft_kv_config, is_disagg=self._is_disagg) - draft_kv_cache_manager_cls = self._validate_or_fallback_kv_cache_manager_v2( - draft_kv_cache_manager_cls, effective_draft_config, draft_kv_config) + + # Use V2 if enabled and the base class is KVCacheManager + if draft_kv_cache_manager_cls == KVCacheManagerV2: + if self._kv_connector_manager is not None or ( + self._max_beam_width is not None and self._max_beam_width + > 1) or draft_kv_config.event_buffer_max_size > 0 or ( + self._cache_transceiver_config is not None + and self._cache_transceiver_config.backend is not None): + logger.warning( + "KVCacheManagerV2 is not supported with disaggregated serving or beam width > 1 or event buffer max size > 0 or disagg config. " + "Falling back to KVCacheManager for draft model.") + draft_kv_cache_manager_cls = KVCacheManager estimating_kv_cache = estimating_kv_cache and not self._skip_est # For MTP with models using sparse attention (e.g., DeepSeek V3 with DSA), @@ -1664,7 +977,7 @@ def _create_one_model_draft_kv_cache_manager( mapping=self._mapping, kv_cache_config=draft_kv_config, tokens_per_block=self._tokens_per_block, - max_seq_len=max_seq_len, + max_seq_len=self._max_seq_len, max_batch_size=self._max_batch_size, spec_config=self._speculative_config, sparse_attention_config=sparse_attn_config, @@ -1672,8 +985,6 @@ def _create_one_model_draft_kv_cache_manager( max_beam_width=self._max_beam_width, kv_connector_manager=self._kv_connector_manager, estimating_kv_cache=estimating_kv_cache, - enable_kv_cache_stats=self._enable_kv_cache_stats() - and not estimating_kv_cache, execution_stream=self._execution_stream, # One-model draft specific overrides model_config=effective_draft_config, @@ -1682,8 +993,6 @@ def _create_one_model_draft_kv_cache_manager( layer_mask=spec_dec_layer_mask, num_layers=num_draft_layers, is_disagg=self._is_disagg, - cold_page_codec_provider=cold_page_codec_provider, - joint_kv_cache_reuse=self._joint_kv_cache_reuse, ) def _get_target_and_draft_cache_costs( @@ -1694,18 +1003,14 @@ def _get_target_and_draft_cache_costs( target_kv_cache_config = (kv_cache_config if kv_cache_config is not None else self._kv_cache_config) total_kv = self._get_kv_size_per_token(target_kv_cache_config) - use_separate_draft_kv_cache = ( - self._should_create_separate_draft_kv_cache()) target_kv = self._per_manager_cache_cost( - self._kv_cache_manager_cls, - self._model_engine.model.model_config, - target_kv_cache_config, - use_separate_draft_kv_cache=use_separate_draft_kv_cache) + self._kv_cache_manager_cls, self._model_engine.model.model_config, + target_kv_cache_config) + # The draft contribution is whatever the aggregate has on top of the + # target. Both pieces are CacheCost; subtraction is component-wise. draft_kv = CacheCost(slope=total_kv.slope - target_kv.slope, intercept=total_kv.intercept - target_kv.intercept) - costs = (target_kv, draft_kv) - if any(cost.slope < 0 or cost.intercept < 0 or ( - cost.slope == 0 and cost.intercept == 0) for cost in costs): + if target_kv.slope <= 0 or draft_kv.slope <= 0: return None return target_kv, draft_kv @@ -1718,21 +1023,14 @@ def _compute_draft_budget_shares( """Split *total_budget* into (target_budget, draft_budget) byte shares.""" intercept_total = target_kv.intercept + draft_kv.intercept slope_budget = total_budget - intercept_total - slope_total = target_kv.slope + draft_kv.slope - if slope_budget < 0: + if slope_budget <= 0: logger.warning( f"KV cache budget {total_budget} is smaller than the fixed " - f"cache cost {intercept_total}; cannot split between " + f"mamba state cost {intercept_total}; cannot split between " f"target and draft.") return None - if slope_budget == 0 and slope_total > 0: - logger.warning( - f"KV cache budget {total_budget} leaves no capacity beyond " - f"the fixed cache cost {intercept_total}; cannot split " - f"between target and draft with a per-token cache cost.") - return None - draft_slope_share = (slope_budget * draft_kv.slope // - slope_total if slope_total > 0 else 0) + slope_total = target_kv.slope + draft_kv.slope + draft_slope_share = int(slope_budget * draft_kv.slope / slope_total) draft_budget = draft_kv.intercept + draft_slope_share target_budget = total_budget - draft_budget return target_budget, draft_budget @@ -1761,8 +1059,8 @@ def _split_kv_cache_budget_for_draft( GPU-resident state never occupies) the intercept is dropped so the split stays proportional to the per-token cost. - When the split is *infeasible* (the combined fixed cost exhausts the - budget while either manager has a per-token cost, or exceeds it) + When the split is *infeasible* (the combined fixed cost meets or exceeds + the budget — only possible for ``max_gpu_total_bytes`` after the above) the shortfall is fatal: both managers need their fixed state resident in GPU memory, so the run would OOM. It raises ``ValueError`` rather than silently producing an unusable config. A defensive degrade-to-zero path @@ -1794,7 +1092,7 @@ def _split_kv_cache_budget_for_draft( shares = self._compute_draft_budget_shares(total_budget, target_kv, draft_kv) if shares is None: - # The split cannot provide each manager with usable GPU capacity. + # The split is infeasible (combined fixed cost >= total budget). intercept_total = target_kv.intercept + draft_kv.intercept if budget_attr == "max_gpu_total_bytes": # A GPU budget that cannot even fit the combined fixed cost is @@ -1803,7 +1101,7 @@ def _split_kv_cache_budget_for_draft( # guidance rather than producing an unusable zero-budget draft. raise ValueError( f"KV cache GPU budget ({total_budget / GB:.2f} GiB) is " - f"insufficient after the combined fixed cost " + f"smaller than the combined fixed cost " f"({intercept_total / GB:.2f} GiB, e.g. mamba SSM state) " f"for target+draft. Increase free_gpu_memory_fraction or " f"max_gpu_total_bytes, or reduce max_batch_size (the fixed " @@ -2064,43 +1362,14 @@ def _create_cross_kv_cache_manager( def _needs_gpu_kv_cache_budget_split( self, - max_seq_len: int, kv_cache_config: Optional[KvCacheConfig] = None, ) -> bool: """Whether max_gpu_total_bytes must be split per manager.""" - if self._is_kv_cache_manager_v2: + if issubclass(self._kv_cache_manager_cls, KVCacheManagerV2): return self._should_create_separate_draft_kv_cache() kv_cache_config = (kv_cache_config if kv_cache_config is not None else self._kv_cache_config) - if uses_vswa_kv_cache_layout(kv_cache_config.max_attention_window): - return True - if not self._should_create_separate_draft_kv_cache(): - return False - draft_windows = self._get_draft_max_attention_window( - max_seq_len, kv_cache_config) - return uses_vswa_kv_cache_layout(draft_windows) - - @classmethod - def _drop_explicit_offload_tier_budgets( - cls, kv_cache_config: Optional[KvCacheConfig] - ) -> Optional[KvCacheConfig]: - """Return a copy of the config with explicit offload budgets unset. - - Host sizing then falls to the V2 auto host tier policy, which matches - the tier to the manager's own device quota; V1 builds no secondary pool. - The disk tier is V2 only and has no auto policy, so dropping its budget - leaves the tier out. - """ - if kv_cache_config is None: - return kv_cache_config - dropped = { - attr: None - for attr in cls._OFFLOAD_TIER_BUDGET_ATTRS - if getattr(kv_cache_config, attr) - } - if not dropped: - return kv_cache_config - return kv_cache_config.model_copy(update=dropped) + return is_vswa_enabled(kv_cache_config) def build_managers(self, resources: Dict, @@ -2120,72 +1389,36 @@ def build_managers(self, self_kv_cache_config, cross_kv_cache_config = self._split_kv_cache_budget_for_cross( ) - has_separate_one_model_draft = ( - self._draft_model_engine is None - and self._should_create_separate_draft_kv_cache()) - # One term per axis: a draft pool exists to pair with, block reuse is on, - # and the span is known. Anything outside keeps its unpaired path. - self._joint_kv_cache_reuse = (has_separate_one_model_draft and - self_kv_cache_config.enable_block_reuse - and self._joint_reuse_supported()) - - # Estimation managers are throwaway probes whose pools only hold dummy - # requests, so an explicit offload tier would reserve capacity the probe - # cannot fill. Encoder-decoder runs skip estimation, so dropping the - # cross budgets is defensive. - if estimating_kv_cache: - self_kv_cache_config = self._drop_explicit_offload_tier_budgets( - self_kv_cache_config) - cross_kv_cache_config = self._drop_explicit_offload_tier_budgets( - cross_kv_cache_config) - - # Split combined KV cache budgets before creating managers. + # Split combined KV cache budgets before creating managers. Skip during + # estimation — estimation uses max_tokens-based logic and must not + # mutate the config. has_draft = ( self._draft_model_engine is not None # two-model - or has_separate_one_model_draft) # one-model + or self._should_create_separate_draft_kv_cache()) # one-model draft_kv_cache_config = None - if has_draft: - # The GPU split applies when each manager sizes its pools from - # max_gpu_total_bytes (V2 and V1 VSWA). V1 non-VSWA and estimation - # size GPU pools from a shared max_tokens instead. - needs_gpu_split = (not estimating_kv_cache - and self._needs_gpu_kv_cache_budget_split( - original_max_seq_len, self_kv_cache_config)) - if needs_gpu_split: + if not estimating_kv_cache and has_draft: + # Used when each manager sizes pools from max_gpu_total_bytes (V2 + # and V1 VSWA). V1 non-VSWA GPU uses shared max_tokens instead. + if self._needs_gpu_kv_cache_budget_split(self_kv_cache_config): self_kv_cache_config, draft_kv_cache_config = ( self._split_kv_cache_budget_for_draft( "max_gpu_total_bytes", self_kv_cache_config, draft_kv_cache_config)) - for budget_attr in self._OFFLOAD_TIER_BUDGET_ATTRS: + # KVCacheManagerV2 does not support two-model draft budget splitting. + v2_two_model = (issubclass(self._kv_cache_manager_cls, + KVCacheManagerV2) + and self._draft_model_engine is not None) + if not v2_two_model: + # Each manager sizes its host pool from host_cache_size directly. self_kv_cache_config, draft_kv_cache_config = ( self._split_kv_cache_budget_for_draft( - budget_attr, self_kv_cache_config, + "host_cache_size", self_kv_cache_config, draft_kv_cache_config)) - compression_config = self._llm_args.kv_cache_compression_config - compression_manager = create_kv_cache_compression_manager( - compression_config, - model_engine=self._model_engine, - kv_cache_config=self_kv_cache_config, - estimating_kv_cache=estimating_kv_cache and not self._skip_est, - ) - cold_page_codec_provider = ( - compression_manager if compression_manager is not None - and compression_manager.provides_cold_page_codec else None) - kv_cache_manager = self._create_kv_cache_manager( self._model_engine, estimating_kv_cache, - kv_cache_config_override=self_kv_cache_config, - cold_page_codec_provider=cold_page_codec_provider, - ) - - # Carry the fp8 context-MLA workspace admission cap (computed in configure_kv_cache_capacity) onto - # the real KV manager so the scheduler reads it directly instead of re-deriving from pool layout. - # The estimation build reserves nothing and runs throwaway fresh-prefill dummies, so leave the - # attribute unset there (PyExecutor._get_ctx_mla_kv_len_cap does not cap during warmup). - if not estimating_kv_cache and kv_cache_manager is not None: - kv_cache_manager.fp8_ctx_mla_kv_len_cap = self._fp8_ctx_mla_kv_len_cap + kv_cache_config_override=self_kv_cache_config) if (not estimating_kv_cache and self._kv_connector_manager is not None and self._draft_model_engine is not None): @@ -2199,14 +1432,10 @@ def build_managers(self, # Two-model speculative decoding: draft model has separate engine if self._draft_model_engine is not None: - if (self._is_kv_cache_manager_v2 - and draft_kv_cache_config is not None): - # Offload budgets are divided per manager, GPU budgets are not. - assert (draft_kv_cache_config.max_gpu_total_bytes == - self_kv_cache_config.max_gpu_total_bytes), ( - "KVCacheManagerV2 does not support two-model " - "speculative decoding with separate draft GPU " - "budgets.") + if issubclass(self._kv_cache_manager_cls, KVCacheManagerV2): + assert draft_kv_cache_config is None, ( + "KVCacheManagerV2 does not support two-model speculative " + "decoding with separate draft KV cache budget splitting.") draft_kv_cache_manager = self._create_kv_cache_manager( self._draft_model_engine, estimating_kv_cache, @@ -2214,10 +1443,8 @@ def build_managers(self, # One-model speculative decoding with different KV layouts elif self._should_create_separate_draft_kv_cache(): draft_kv_cache_manager = self._create_one_model_draft_kv_cache_manager( - original_max_seq_len, estimating_kv_cache, - kv_cache_config_override=draft_build_kv_cache_config, - cold_page_codec_provider=cold_page_codec_provider) + kv_cache_config_override=draft_build_kv_cache_config) # Encoder-decoder cross-attention pool cross_kv_cache_manager = None @@ -2231,17 +1458,9 @@ def build_managers(self, ResourceManagerType.DRAFT_KV_CACHE_MANAGER] = draft_kv_cache_manager resources[ ResourceManagerType.CROSS_KV_CACHE_MANAGER] = cross_kv_cache_manager - if (compression_manager is not None - and compression_manager.uses_iteration_lifecycle): - resources[ResourceManagerType.KV_CACHE_COMPRESSION_MANAGER] = ( - compression_manager) def teardown_managers(self, resources: Dict) -> None: """Clean up KV caches for model, draft model, and cross pool.""" - compression_manager = resources.pop( - ResourceManagerType.KV_CACHE_COMPRESSION_MANAGER, None) - if compression_manager is not None: - compression_manager.shutdown() resources[ResourceManagerType.KV_CACHE_MANAGER].shutdown() del resources[ResourceManagerType.KV_CACHE_MANAGER] draft_kv_cache_manager = resources[ @@ -2290,67 +1509,6 @@ def _build_per_layer_num_kv_heads( ] * num_spec_layers -def _get_mamba_cache_layer_masks( - mamba_params: MambaKVCacheParams, - mapping: Mapping, - spec_config: Optional[SpeculativeConfig], - is_draft: bool, -) -> tuple[List[bool], List[bool]]: - use_separate_draft_kv_cache = ( - not mapping.enable_attention_dp - and should_use_separate_draft_kv_cache(spec_config)) - return mamba_params.get_layer_masks( - is_draft=is_draft, - use_separate_draft_kv_cache=use_separate_draft_kv_cache, - ) - - -# The V1 hybrid managers select the convolution-state layout by model_type; -# MambaHybridCacheManagerV2 takes the layout by name and rejects model_type. -_CONV_STATE_LAYOUT_BY_MODEL_TYPE = { - "nemotron_hybrid": "x_b_c", - "qwen3_next": "q_k_v", -} - - -def _mamba_conv_layout_kwargs(kv_cache_manager_cls: type, - model_type: str) -> dict: - """Constructor kwarg selecting the conv-state layout for a hybrid manager. - - Keeps the V1-vs-V2 dispatch in one place: a manager branch that forgets it - would previously get V2's silent "x_b_c" default (the Kimi K3 bug fixed in - this change). - """ - if issubclass(kv_cache_manager_cls, MambaHybridCacheManagerV2): - return { - "conv_state_layout": _CONV_STATE_LAYOUT_BY_MODEL_TYPE[model_type] - } - return {"model_type": model_type} - - -def _get_qwen4_exp_ple_cache_params(config, *, total_layers: int, - is_draft: bool): - """Align target-only PLE state with a target/draft cache layout.""" - if is_draft: - return None - - params = extract_qwen4_exp_ple_cache_params(config) - num_target_layers = len(params.ple_layer_mask) - if num_target_layers > total_layers: - raise ValueError( - "PLE layer mask cannot exceed the hybrid cache layout: " - f"got {num_target_layers}, expected at most {total_layers}") - if num_target_layers == total_layers: - return params - - # Unified one-model caches append attention-only MTP layers. - return dataclasses.replace( - params, - ple_layer_mask=params.ple_layer_mask + [False] * - (total_layers - num_target_layers), - ) - - def _create_kv_cache_manager( model_engine: Optional[PyTorchModelEngine], kv_cache_manager_cls, @@ -2365,7 +1523,6 @@ def _create_kv_cache_manager( max_beam_width: int, kv_connector_manager: Optional[KvCacheConnectorManager], estimating_kv_cache: bool = False, - enable_kv_cache_stats: bool = False, execution_stream: Optional[torch.cuda.Stream] = None, # Optional overrides for one-model draft case (when model_engine is None) model_config: Optional[ModelConfig] = None, @@ -2376,33 +1533,11 @@ def _create_kv_cache_manager( num_kv_heads: Optional[Union[int, List[int]]] = None, head_dim: Optional[int] = None, kv_cache_type=None, - is_disagg: bool = False, - cold_page_codec_provider: Optional[object] = None, - kv_events_config: Optional[KVEventsConfig] = None, - joint_kv_cache_reuse: bool = False) -> KVCacheManager: + is_disagg: bool = False) -> KVCacheManager: """ Returns: A KVCacheManager instance for the given model engine or model config """ - if cold_page_codec_provider is not None and not issubclass( - kv_cache_manager_cls, KVCacheManagerV2): - raise ValueError( - "Cold-page quantization requires the resolved KV cache manager " - f"to be KVCacheManagerV2; selected {kv_cache_manager_cls.__name__}") - - if (estimating_kv_cache - and issubclass(kv_cache_manager_cls, KVCacheManagerV2) - and kv_cache_config.pool_ratio is None - and kv_cache_config.avg_seq_len is not None - and kv_cache_config.avg_seq_len > max_seq_len): - # Estimation can build multiple managers from the same temporary - # config. The first manager may reduce max_seq_len to fit max_tokens, - # so later draft/cross managers need a per-manager workload length. - # Keep the shared config untouched because it is restored after - # estimation. - kv_cache_config = kv_cache_config.model_copy( - update={"avg_seq_len": max_seq_len}) - # Extract config from model_engine or use provided model_config if model_config is not None: config = model_config.pretrained_config @@ -2513,107 +1648,8 @@ def _create_kv_cache_manager( per_layer_num_kv_heads = _build_per_layer_num_kv_heads( num_key_value_heads, num_hidden_layers, spec_config, draft_config_for_kv) - manager_extra_kwargs = {} - if issubclass(kv_cache_manager_cls, KVCacheManagerV2): - manager_extra_kwargs["enable_stats"] = enable_kv_cache_stats - manager_extra_kwargs[ - "cold_page_codec_provider"] = cold_page_codec_provider - manager_extra_kwargs["kv_events_config"] = kv_events_config - manager_extra_kwargs["joint_kv_cache_reuse"] = joint_kv_cache_reuse - elif kv_events_config is not None and kv_events_config.enable_kv_cache_events: - logger.warning( - "kv_cache_config.kv_events_config is set but streaming KV event " - "publishing requires KV cache manager V2; events will not be " - f"published for {kv_cache_manager_cls.__name__}.") - if issubclass(kv_cache_manager_cls, MambaHybridCacheManagerV2): - manager_extra_kwargs["is_disagg"] = is_disagg - - if is_kimi_linear(config): - # Kimi K3 hybrid: KDA (Kimi Delta Attention) recurrent/conv states on - # the mamba side of the hybrid manager, absorbed-MQA MLA latent cache - # (num_kv_heads=1, head_dim = kv_lora_rank + qk_rope_head_dim, - # SELFKONLY) on the paged-KV side. Must come before the is_mla(...) - # route: the kimi_linear config carries MLA fields, but only 24 of - # its 93 layers are MLA. - if max_beam_width > 1: - raise ValueError( - "MambaHybridCacheManager + beam search is not supported yet.") - if not estimating_kv_cache and kv_connector_manager is not None: - raise NotImplementedError( - "Connector manager is not supported for MambaHybridCacheManager." - ) - mamba_params = extract_mamba_kv_cache_params( - config, - spec_config=spec_config, - quant_config=quant_config, - ) - mamba_layer_mask, full_attention_layer_mask = ( - _get_mamba_cache_layer_masks( - mamba_params, - mapping, - spec_config, - is_draft, - )) - num_mamba_layers = (0 if is_draft and mamba_params.num_draft_layers > 0 - else mamba_params.num_mamba_layers) - # Kimi K3 KDA state sharding follows the attention-family TP - # semantics (Qwen3-Next pattern): replicated under attention-DP, - # head-sharded across tp_size otherwise. That is exactly the cache - # manager's own internal gate (`tp_size = 1 if enable_attention_dp - # else tp_size`, then num_heads / n_groups / conv_dim divide by - # it), so the params pass through unscaled. - # KDA fused multi-token verify (trtllm::kda_mtp_decode): when the - # kernel can run here, allocate the per-slot replay caches instead - # of the legacy per-step intermediate verification buffers. The - # kernel replays accepted drafts from these caches and commits - # states in place, replacing the intermediate-buffer + promotion - # flow for KDA layers. - kimi_extra_kwargs = {} - kda_replay_manager_types = (MixedMambaHybridCacheManager, - MambaHybridCacheManagerV2) - if (spec_config is not None - and issubclass(kv_cache_manager_cls, kda_replay_manager_types)): - from ..modules.kimi_kda._kda_kernels import \ - is_kda_mtp_verify_available - if is_kda_mtp_verify_available(): - kimi_extra_kwargs["kda_replay_num_spec"] = ( - spec_config.tokens_per_gen_step - 1) - # KDA's conv state is a [Q | K | V] concatenation whose three sections - # have identical width, i.e. the qwen3_next section layout. - kimi_extra_kwargs.update( - _mamba_conv_layout_kwargs(kv_cache_manager_cls, "qwen3_next")) - kv_cache_manager = kv_cache_manager_cls( - # mamba (KDA) cache parameters - mamba_params.state_size, - mamba_params.conv_kernel, - mamba_params.num_heads, - mamba_params.n_groups, - mamba_params.head_dim, - num_mamba_layers, - mamba_layer_mask, - mamba_params.dtype, - mamba_params.mamba_ssm_cache_dtype, - # kv cache parameters (MLA latent cache) - kv_cache_config, - tensorrt_llm.bindings.internal.batch_manager.CacheType.SELFKONLY, - num_layers=sum(full_attention_layer_mask), - layer_mask=full_attention_layer_mask, - num_kv_heads=1, - head_dim=config.kv_lora_rank + config.qk_rope_head_dim, - tokens_per_block=tokens_per_block, - max_seq_len=max_seq_len, - max_num_tokens=max_num_tokens, - is_draft=is_draft, - max_batch_size=max_batch_size, - mapping=mapping, - dtype=kv_cache_dtype, - spec_config=spec_config, - is_estimating_kv_cache=estimating_kv_cache, - execution_stream=execution_stream, - **kimi_extra_kwargs, - **manager_extra_kwargs, - ) - elif is_mla(config): + + if is_mla(config): kv_cache_manager = kv_cache_manager_cls( kv_cache_config, tensorrt_llm.bindings.internal.batch_manager.CacheType.SELFKONLY, @@ -2627,7 +1663,6 @@ def _create_kv_cache_manager( dtype=kv_cache_dtype, spec_config=spec_config, vocab_size=config.vocab_size, - max_num_tokens=max_num_tokens, max_beam_width=max_beam_width, is_draft=is_draft, kv_connector_manager=kv_connector_manager @@ -2638,7 +1673,6 @@ def _create_kv_cache_manager( execution_stream=execution_stream, layer_mask=layer_mask, is_disagg=is_disagg, - **manager_extra_kwargs, ) elif is_nemotron_hybrid(config): if max_beam_width > 1: @@ -2652,18 +1686,10 @@ def _create_kv_cache_manager( mamba_params = extract_mamba_kv_cache_params( config, + layer_mask=layer_mask, spec_config=spec_config, quant_config=quant_config, ) - mamba_layer_mask, full_attention_layer_mask = ( - _get_mamba_cache_layer_masks( - mamba_params, - mapping, - spec_config, - is_draft, - )) - num_mamba_layers = (0 if is_draft and mamba_params.num_draft_layers > 0 - else mamba_params.num_mamba_layers) # Replay state update kernel for MTP: default on for sm >= 80; gates # below disable it for incompatible feature combinations. Cpp cache @@ -2722,9 +1748,6 @@ def _create_kv_cache_manager( mamba_ssm_stochastic_rounding = (stochastic_rounding and mamba_params.mamba_ssm_cache_dtype == torch.float16) - mamba_manager_extra_kwargs = dict(manager_extra_kwargs) - mamba_manager_extra_kwargs.update( - _mamba_conv_layout_kwargs(kv_cache_manager_cls, "nemotron_hybrid")) kv_cache_manager = kv_cache_manager_cls( # mamba cache parameters mamba_params.state_size, @@ -2732,20 +1755,19 @@ def _create_kv_cache_manager( mamba_params.num_heads, mamba_params.n_groups, mamba_params.head_dim, - num_mamba_layers, - mamba_layer_mask, + mamba_params.num_mamba_layers, + mamba_params.mamba_layer_mask, mamba_params.dtype, mamba_params.mamba_ssm_cache_dtype, # kv cache parameters kv_cache_config, tensorrt_llm.bindings.internal.batch_manager.CacheType.SELF, - num_layers=sum(full_attention_layer_mask), - layer_mask=full_attention_layer_mask, + num_layers=mamba_params.num_full_attention_layers, + layer_mask=mamba_params.full_attention_layer_mask, num_kv_heads=per_layer_num_kv_heads, head_dim=head_dim, tokens_per_block=tokens_per_block, max_seq_len=max_seq_len, - max_num_tokens=max_num_tokens, is_draft=is_draft, max_batch_size=max_batch_size, mapping=mapping, @@ -2753,11 +1775,11 @@ def _create_kv_cache_manager( spec_config=spec_config, is_estimating_kv_cache=estimating_kv_cache, execution_stream=execution_stream, + model_type="nemotron_hybrid", use_replay_state_update=use_replay, mamba_ssm_stochastic_rounding=mamba_ssm_stochastic_rounding, - **mamba_manager_extra_kwargs, ) - elif is_qwen3_hybrid(config) or is_qwen4_exp(config): + elif is_qwen3_hybrid(config): if max_beam_width > 1: raise ValueError( "MambaHybridCacheManager + beam search is not supported yet.") @@ -2768,88 +1790,10 @@ def _create_kv_cache_manager( ) mamba_params = extract_mamba_kv_cache_params( config, + layer_mask=layer_mask, spec_config=spec_config, quant_config=quant_config, ) - mamba_layer_mask, full_attention_layer_mask = ( - _get_mamba_cache_layer_masks( - mamba_params, - mapping, - spec_config, - is_draft, - )) - num_mamba_layers = (0 if is_draft and mamba_params.num_draft_layers > 0 - else mamba_params.num_mamba_layers) - # Replay state update for GDN MTP: mirrors the nemotron_hybrid gating - # above, minus the Mamba2-specific stochastic-rounding/Philox gate. - # The GDN replay kernel does a plain cast on checkpoint commit, so - # quantized SSM cache dtypes stay on the legacy path. - sm = get_sm_version() - use_replay = spec_config is not None and sm >= 80 - if spec_config is None: - logger.info( - "GDN replay kernel requires speculative decoding; using " - "non-replay path") - elif spec_config.tokens_per_gen_step > 8: - logger.info("GDN cached replay supports at most 8 tokens per " - "generation step; using non-replay path") - use_replay = False - - # Tree attention: replay assumes a linear token sequence. - if (spec_config is not None - and (getattr(spec_config, 'eagle_choices', None) is not None - or getattr(spec_config, 'use_dynamic_tree', False))): - logger.info("GDN replay kernel incompatible with tree attention; " - "using legacy MTP path") - use_replay = False - - if mamba_params.mamba_ssm_cache_dtype not in (torch.float32, - torch.bfloat16, - torch.float16): - logger.info( - "GDN replay kernel does not support quantized SSM cache " - f"dtype {mamba_params.mamba_ssm_cache_dtype}; using legacy " - "MTP path") - use_replay = False - - # Replay is enabled by default for eligible GDN MTP workloads. - if not is_gdn_replay_enabled(): - use_replay = False - - # GDN replay supports the contiguous C++ V1 state pool and the indirect - # per-layer state views exposed by V2. Mixed/Python does not expose an - # all-layer commit, so keep that manager but use non-replay MTP. - replay_manager_types = (CppMambaHybridCacheManager, - MambaHybridCacheManagerV2) - if use_replay and not issubclass(kv_cache_manager_cls, - replay_manager_types): - logger.info("GDN replay requires C++ V1 or V2 Mamba cache manager; " - f"{kv_cache_manager_cls.__name__} was selected, so the " - "non-replay MTP path will be used") - use_replay = False - logger.info("GDN replay state update: " + - ("ENABLED" if use_replay else "DISABLED")) - - mamba_manager_extra_kwargs = dict(manager_extra_kwargs) - mamba_manager_extra_kwargs.update( - _mamba_conv_layout_kwargs(kv_cache_manager_cls, "qwen3_next")) - if getattr(sparse_attention_config, "algorithm", None) == "qsa": - # Resolve the side-cache shape from the same checkpoint geometry - # used to construct the QSA index projection. - mamba_manager_extra_kwargs.update( - sparse_attention_config=sparse_attention_config, - pretrained_config=config, - ) - if is_qwen4_exp(config) and issubclass(kv_cache_manager_cls, - MambaHybridCacheManagerV2): - ple_cache_params = _get_qwen4_exp_ple_cache_params( - config, - total_layers=len(mamba_layer_mask), - is_draft=is_draft, - ) - if ple_cache_params is not None: - mamba_manager_extra_kwargs[ - "qwen4_exp_ple_cache_params"] = ple_cache_params kv_cache_manager = kv_cache_manager_cls( # mamba cache parameters mamba_params.state_size, @@ -2857,20 +1801,19 @@ def _create_kv_cache_manager( mamba_params.num_heads, mamba_params.n_groups, mamba_params.head_dim, - num_mamba_layers, - mamba_layer_mask, + mamba_params.num_mamba_layers, + mamba_params.mamba_layer_mask, mamba_params.dtype, mamba_params.mamba_ssm_cache_dtype, # kv cache parameters kv_cache_config, tensorrt_llm.bindings.internal.batch_manager.CacheType.SELF, - num_layers=sum(full_attention_layer_mask), - layer_mask=full_attention_layer_mask, + num_layers=mamba_params.num_full_attention_layers, + layer_mask=mamba_params.full_attention_layer_mask, num_kv_heads=per_layer_num_kv_heads, head_dim=head_dim, tokens_per_block=tokens_per_block, max_seq_len=max_seq_len, - max_num_tokens=max_num_tokens, is_draft=is_draft, max_batch_size=max_batch_size, mapping=mapping, @@ -2878,14 +1821,12 @@ def _create_kv_cache_manager( spec_config=spec_config, is_estimating_kv_cache=estimating_kv_cache, execution_stream=execution_stream, - use_replay_state_update=use_replay, - **mamba_manager_extra_kwargs, + model_type="qwen3_next", ) else: # NOTE: this is a workaround for VSWA to switch to calculate_max_num_blocks_for_vswa in KVCahceManager # Only needed for V1; V2 handles per-layer windows natively via life cycles. - is_vswa = uses_vswa_kv_cache_layout( - kv_cache_config.max_attention_window) + is_vswa = is_vswa_enabled(kv_cache_config) binding_model_config = None if is_vswa and kv_cache_manager_cls.__name__ == "KVCacheManager": binding_model_config = _model_config.get_bindings_model_config( @@ -2923,184 +1864,14 @@ def _create_kv_cache_manager( execution_stream=execution_stream, layer_mask=layer_mask, is_disagg=is_disagg, - **manager_extra_kwargs, ) # Note: Gemma4 KV sharing cache remapping is handled in Gemma4Attention # via cache_layer_idx — shared layers use target layer's index for # get_buffers(). No layer_offsets remapping needed here. - # Propagate the finalized chunked-prefill flag so KVCacheManager.fit_token_budget - # only shrinks context chunks when the attention backend can consume a - # partial context chunk. The flag is read from attn_runtime_features, which - # py_executor_creator finalizes (including the SM-version / - # attention-backend overrides) before build_managers runs. - if isinstance(kv_cache_manager, - KVCacheManager) and model_engine is not None: - kv_cache_manager.enable_chunked_prefill = bool( - model_engine.attn_runtime_features.chunked_prefill) - return kv_cache_manager -def validate_kv_cache_compression_compatibility( - config: KvCacheCompressionConfig, - kv_cache_config: KvCacheConfig, - spec_config: Optional[SpeculativeConfig], -) -> None: - """Reject unsupported KV-cache compression feature combinations.""" - if config.algorithm == "quantization_for_cold_page": - from tensorrt_llm.runtime.kv_cache_manager_v2 import _BACKEND - - if _BACKEND == "python": - raise ValueError( - "Cold-page quantization requires the C++ KVCacheManagerV2 backend" - ) - if config.quant == "nvfp4" and not is_sm_100f(): - raise RuntimeError( - "NVFP4 cold-page quantization requires an SM100-family device " - "(SM100 or SM103).") - elif config.algorithm == "triattention" and not is_sm_100f(): - raise RuntimeError( - "TriAttention requires an SM100-family device (SM100 or SM103).") - - if kv_cache_config.enable_block_reuse and not config.supports_block_reuse(): - raise ValueError( - f"KV-cache compression algorithm {config.algorithm!r} does not " - "support KV-cache block reuse. Set " - "KvCacheConfig.enable_block_reuse=False.") - if spec_config is None: - return - if not config.supports_speculative_decoding(): - guidance = ("; TriAttention requires eviction_mode='union'" - if config.algorithm == "triattention" else "") - raise ValueError( - f"KV-cache compression algorithm {config.algorithm!r} does not " - "support speculative decoding with its current configuration" - f"{guidance}") - mode = spec_config.spec_dec_mode - if config.algorithm == "quantization_for_cold_page": - supported = mode.is_mtp_eagle_one_model() or mode.is_eagle3_one_model() - guidance = "one-model MTP-EAGLE or EAGLE3" - else: - supported = mode.is_mtp_one_model() or mode.is_eagle3_one_model() - guidance = "one-model MTP or EAGLE3" - if not supported: - raise ValueError( - f"KV-cache compression does not support speculative decoding " - f"mode {mode.name}; use {guidance}") - - -def create_kv_cache_compression_manager( - config: Optional[KvCacheCompressionConfig], - *, - model_engine: PyTorchModelEngine, - kv_cache_config: KvCacheConfig, - estimating_kv_cache: bool = False, -) -> Optional[KVCacheCompressionManager]: - """Validate, select, and construct the configured manager before KVCM.""" - if config is None: - return None - if model_engine.mapping.has_cp_helix(): - # TODO: Revisit after KVCC validates HELIX-sharded Page ownership and migration. - raise ValueError( - "KV-cache compression does not support HELIX context parallelism.") - - if config.algorithm == "quantization_for_cold_page": - if config.quant != "nvfp4": - raise NotImplementedError( - f"Unsupported cold-page quantization format {config.quant!r}") - if estimating_kv_cache: - return None - quant_config = model_engine.model.model_config.quant_config - if (quant_config is not None and getattr( - quant_config, "kv_cache_quant_algo", None) == QuantAlgo.NVFP4): - logger.info( - "Skipping cold-page NVFP4 quantization because the active KV " - "cache already uses NVFP4; KVCM will migrate it losslessly.") - return None - - validate_kv_cache_compression_compatibility(config, kv_cache_config, - model_engine.spec_config) - from ..kv_cache_compression.quantization_for_cold_page.nvfp4_quantization import \ - Nvfp4ColdPageQuantizationCompression - - return Nvfp4ColdPageQuantizationCompression(config) - - if config.algorithm == "triattention": - validate_kv_cache_compression_compatibility(config, kv_cache_config, - model_engine.spec_config) - # TriAttention imports CuTe/CUTLASS; keep normal executor startup lazy. - from ..kv_cache_compression.triattention.triattention import \ - TriAttentionCompressionManager - - return TriAttentionCompressionManager( - config, - pretrained_config=model_engine.model.model_config.pretrained_config, - ) - - logger.warning( - "KV-cache compression algorithm '%s' is not registered; running without " - "a compression manager.", - config.algorithm, - ) - return None - - -def compute_max_num_sequences(mapping: Mapping, - max_batch_size: int, - disable_overlap_scheduler: bool, - enable_overlap_headroom: bool = False) -> int: - """Size the sequence-slot pool (and the sampler state it indexes). - - ``enable_overlap_headroom`` is intentionally opt-in. Disaggregated - attention-DP needs a second non-PP slot set because the V2 scheduler can - backfill seats before the overlap scheduler releases the previous - iteration's terminal slots. Pipeline parallelism already sizes the pool - by ``pp_size``. - """ - if mapping.has_pp(): - num_micro_batches = mapping.pp_size - else: - num_micro_batches = (2 if enable_overlap_headroom - and not disable_overlap_scheduler else 1) - return max_batch_size * num_micro_batches - - -def should_enable_adp_dummy_fixes(mapping: Mapping) -> bool: - """Enable transactional ADP dummy handling while PP remains follow-up.""" - return not mapping.has_pp() - - -_VALIDATED_OVERLAP_ADP_DUMMY_MODEL_TYPES = ("deepseek_v4", "qwen3_5_moe") - - -def should_enable_scheduler_aware_adp_dummy( - model_type: Optional[str], mapping: Mapping, - disable_overlap_scheduler: bool) -> bool: - """Enable scheduler-aware padding for validated lifecycle configurations.""" - return (should_enable_adp_dummy_fixes(mapping) - and (disable_overlap_scheduler - or model_type in _VALIDATED_OVERLAP_ADP_DUMMY_MODEL_TYPES)) - - -def should_enable_non_overlap_adp_forward_intent( - mapping: Mapping, disable_overlap_scheduler: bool) -> bool: - """Enable fresh cross-rank dummy intent for the generic non-overlap path.""" - return (should_enable_adp_dummy_fixes(mapping) - and disable_overlap_scheduler) - - -def should_enable_disagg_adp_overlap_headroom( - mapping: Mapping, - cache_transceiver_config: Optional[CacheTransceiverConfig], - disable_overlap_scheduler: bool) -> bool: - """Gate extra sequence slots to non-PP disaggregated attention-DP.""" - is_disagg = (cache_transceiver_config is not None - and cache_transceiver_config.backend is not None) - return (mapping.enable_attention_dp and is_disagg and not mapping.has_pp() - and not disable_overlap_scheduler) - - def create_py_executor_instance( *, dist, @@ -3127,18 +1898,12 @@ def create_py_executor_instance( virtual_memory_pools: Optional[dict] = None, execution_stream: Optional[torch.cuda.Stream] = None, dwdp_manager: Optional[DwdpManager] = None, - max_num_sequences: Optional[int] = None, ) -> PyExecutor: - set_low_latency_dispatch( - getattr(llm_args, 'enable_low_latency_host_dispatch', False)) - kv_cache_manager = resources.get(ResourceManagerType.KV_CACHE_MANAGER, None) spec_config = model_engine.spec_config - if max_num_sequences is None: - max_num_sequences = compute_max_num_sequences( - mapping, max_batch_size, llm_args.disable_overlap_scheduler) + max_num_sequences = max_batch_size * mapping.pp_size logger.info( f"max_seq_len={max_seq_len}, max_num_requests={max_num_sequences}, max_num_tokens={max_num_tokens}, max_batch_size={max_batch_size}" @@ -3157,11 +1922,9 @@ def create_py_executor_instance( # dataclass to avoid ad-hoc getattr + TP-division blocks per model type. from tensorrt_llm.bindings import LoraModule - initial_lora_data_type = None if len(lora_config.lora_dir) == 1: # Route to appropriate loader based on checkpoint source - initial_lora_data_type = _get_initial_lora_data_type( - load_torch_lora(lora_config)) + load_torch_lora(lora_config) else: assert len(lora_config.lora_target_modules ) >= 1, "Expecting at least one lora target module" @@ -3293,7 +2056,6 @@ def create_py_executor_instance( world_config=world_config, execution_stream=execution_stream, lora_target_modules=target_modules, - initial_data_type=initial_lora_data_type, ) resources[ResourceManagerType.PEFT_CACHE_MANAGER] = peft_cache_manager model_engine.set_lora_model_config( @@ -3305,37 +2067,23 @@ def create_py_executor_instance( resources[ResourceManagerType.SEQ_SLOT_MANAGER] = SeqSlotManager( max_num_sequences) - compression_manager = resources.get( - ResourceManagerType.KV_CACHE_COMPRESSION_MANAGER) - if compression_manager is not None: - compression_manager.bind_kv_cache_managers( - resources[ResourceManagerType.KV_CACHE_MANAGER], - resources.get(ResourceManagerType.DRAFT_KV_CACHE_MANAGER), - ) resource_manager = ResourceManager(resources) - # KV cache manager runs last (others may depend on it), except the - # compression manager which reconciles after it (below). + # Make sure the kv cache managers are always invoked last as they could + # depend on the results of other resource managers. if kv_cache_manager is not None: resource_manager.resource_managers.move_to_end( ResourceManagerType.KV_CACHE_MANAGER, last=True) + cross_kv_cache_manager = resources.get( ResourceManagerType.CROSS_KV_CACHE_MANAGER) if cross_kv_cache_manager is not None: resource_manager.resource_managers.move_to_end( ResourceManagerType.CROSS_KV_CACHE_MANAGER, last=True) - # Iteration-driven compression is the final reconciler after every native - # KV manager. Cold-page quantization runs only at native storage migration. - if (ResourceManagerType.KV_CACHE_COMPRESSION_MANAGER - in resource_manager.resource_managers): - resource_manager.resource_managers.move_to_end( - ResourceManagerType.KV_CACHE_COMPRESSION_MANAGER, last=True) # When scheduler_capacity == 1, attention dp dummy request will prevent the scheduling of DISAGG_GENERATION_INIT. # Enlarge scheduler capacity to avoid DISAGG_GENERATION_INIT stuck in the scheduler. - # V1 scheduler handles overlap via two_step_lookahead, so skip the - # slot-pool overlap factor here. - scheduler_capacity = max_batch_size * mapping.pp_size + scheduler_capacity = max_num_sequences if scheduler_capacity == 1 and mapping.enable_attention_dp and kv_cache_manager: scheduler_capacity += 1 @@ -3347,18 +2095,6 @@ def create_py_executor_instance( if cross_kv_cache_manager is not None else LlmRequestState.CONTEXT_INIT) - # V2 scheduler uses scheduler_capacity as the per-iteration request - # budget (BudgetTracker.max_num_requests). Unlike V1 which has a - # separate CapacityScheduler (needs pp_size * max_batch_size to hold - # requests across PP stages) and MicroBatchScheduler (uses - # max_batch_size for per-forward batch limit), V2 merges both into - # one loop. PP on-the-fly is handled by inflight_request_ids - # filtering, so its budget should be based on max_batch_size, not - # max_num_sequences (which includes the pp_size multiplier). - v2_scheduler_capacity = max_batch_size - if v2_scheduler_capacity == 1 and mapping.enable_attention_dp and kv_cache_manager: - v2_scheduler_capacity += 1 - if isinstance(kv_cache_manager, KVCacheManagerV2): # V2: interleaved scheduler handles both capacity and budget draft_kv_cache_manager = resources.get( @@ -3366,9 +2102,6 @@ def create_py_executor_instance( scheduler_policy = (scheduler_config.capacity_scheduler_policy if scheduler_config is not None else CapacitySchedulerPolicy.MAX_UTILIZATION) - enable_prefix_aware_scheduling = ( - scheduler_config.enable_prefix_aware_scheduling - if scheduler_config is not None else True) scheduler = KVCacheV2Scheduler( max_batch_size=max_batch_size, max_num_tokens=max_num_tokens, @@ -3377,17 +2110,13 @@ def create_py_executor_instance( ctx_chunk_config=ctx_chunk_config, peft_cache_manager=peft_cache_manager.impl if peft_cache_manager is not None else None, - scheduler_capacity=v2_scheduler_capacity, + scheduler_capacity=scheduler_capacity, draft_kv_cache_manager=draft_kv_cache_manager, cross_kv_cache_manager=cross_kv_cache_manager, no_schedule_until_state=no_schedule_until_state, - enable_prefix_aware_scheduling=enable_prefix_aware_scheduling, - # A disaggregated generation worker must not replay context locally. - enable_recompute_pause=not is_disagg, ) elif (scheduler_config is not None and scheduler_config.use_python_scheduler): - enable_prefix_aware_scheduling = scheduler_config.enable_prefix_aware_scheduling scheduler = SimpleUnifiedScheduler( max_batch_size=max_batch_size, max_num_tokens=max_num_tokens, @@ -3401,13 +2130,8 @@ def create_py_executor_instance( if cross_kv_cache_manager is not None else None, two_step_lookahead=mapping.has_pp(), scheduler_capacity=scheduler_capacity, - no_schedule_until_state=no_schedule_until_state, - enable_prefix_aware_scheduling=enable_prefix_aware_scheduling, - ) + no_schedule_until_state=no_schedule_until_state) else: - enable_prefix_aware_scheduling = ( - scheduler_config.enable_prefix_aware_scheduling - if scheduler_config is not None else True) capacity_scheduler = BindCapacityScheduler( scheduler_capacity, kv_cache_manager.impl if kv_cache_manager is not None else None, @@ -3416,16 +2140,10 @@ def create_py_executor_instance( cross_kv_cache_manager=cross_kv_cache_manager.impl if cross_kv_cache_manager is not None else None, two_step_lookahead=mapping.has_pp(), - no_schedule_until_state=no_schedule_until_state, - enable_prefix_aware_scheduling=enable_prefix_aware_scheduling, - ) + no_schedule_until_state=no_schedule_until_state) - mb_scheduler = BindMicroBatchScheduler( - max_batch_size, - max_num_tokens, - ctx_chunk_config, - no_schedule_until_state=no_schedule_until_state, - ) + mb_scheduler = BindMicroBatchScheduler(max_batch_size, max_num_tokens, + ctx_chunk_config) reorder_policy_config = llm_args.reorder_policy_config if reorder_policy_config is not None: @@ -3436,28 +2154,6 @@ def create_py_executor_instance( reorder_policy_config.policy_args.agent_inflight_seq_num) scheduler = SimpleScheduler(capacity_scheduler, mb_scheduler) - if getattr(model_engine, "mm_encoder_item_scheduling_enabled", False): - # Wrap the LLM scheduler with atomic MM item budgeting. ModelEngine - # already validated model-capability-dependent feature combinations. - multimodal_config = llm_args.multimodal_config - # `mm_encoder_item_scheduling_enabled` already excludes the DISABLED - # policy (a disabled model never reaches here and keeps the base LLM - # scheduler), so only the EAGER vs DEFAULT variant is selected here. - scheduler_cls = MultimodalScheduler - if (multimodal_config.encoder_scheduling_policy == - MultimodalEncoderSchedulingPolicy.EAGER): - logger.info("Eager multimodal encoder scheduling is enabled for " - "capacity-rejected active requests.") - scheduler_cls = MultimodalEagerEncoderScheduler - scheduler = scheduler_cls( - scheduler, - max_batch_size=model_engine.encoder_batch_size, - max_num_tokens=model_engine.encoder_max_num_tokens, - output_budget_bytes=model_engine.mm_encoder_output_budget_bytes, - bytes_per_encoder_embedding=( - model_engine.bytes_per_mm_encoder_embedding), - ) - config = model_engine.model.model_config.pretrained_config attention_type = AttentionTypeCpp.MLA if is_mla( config) else AttentionTypeCpp.DEFAULT @@ -3475,17 +2171,6 @@ def create_py_executor_instance( if scheduler_config is not None else WaitingQueuePolicy.FCFS) - # For enc-dec models max_seq_len covers the (longer) encoder sequence, so - # cap the executor's per-request max_tokens at the decoder position table - # (max_target_positions). - executor_max_seq_len = max_seq_len - if model_engine.model.model_config.is_encoder_decoder: - decoder_position_limit = getattr(config, "max_target_positions", None) - if (decoder_position_limit is not None - and executor_max_seq_len is not None): - executor_max_seq_len = min(executor_max_seq_len, - int(decoder_position_limit)) - return PyExecutor( resource_manager, scheduler, @@ -3509,7 +2194,7 @@ def create_py_executor_instance( garbage_collection_gen0_threshold=garbage_collection_gen0_threshold, kv_connector_manager=kv_connector_manager, resource_governor_queue=resource_governor_queue, - max_seq_len=executor_max_seq_len, + max_seq_len=max_seq_len, peft_cache_config=peft_cache_config, virtual_memory_pools=virtual_memory_pools, execution_stream=execution_stream, @@ -3528,15 +2213,11 @@ def create_torch_sampler_args( speculative_config: SpeculativeConfig, max_beam_width: int, disable_overlap_scheduler: bool, + disable_flashinfer_sampling: bool, enable_async_worker: bool, enable_speculative_beam_history_d2h: bool, - max_num_sequences: Optional[int] = None, ): - # The sampler's per-slot state is indexed by sequence slots, so it must - # be sized identically to the executor's slot pool. - if max_num_sequences is None: - max_num_sequences = compute_max_num_sequences( - mapping, max_batch_size, disable_overlap_scheduler) + max_num_sequences = max_batch_size * mapping.pp_size max_draft_len = (0 if speculative_config is None else speculative_config.max_draft_len) max_total_draft_tokens = (0 if speculative_config is None else @@ -3548,6 +2229,7 @@ def create_torch_sampler_args( max_total_draft_tokens=max_total_draft_tokens, max_num_sequences=max_num_sequences, max_beam_width=max_beam_width, + disable_flashinfer_sampling=disable_flashinfer_sampling, disable_overlap_scheduler=disable_overlap_scheduler, enable_async_worker=enable_async_worker, enable_speculative_beam_history_d2h=enable_speculative_beam_history_d2h, @@ -3561,9 +2243,12 @@ def instantiate_sampler( *, max_batch_size: int, max_beam_width: int, + max_seq_len: int, mm_encoder_only: bool, speculative_config: SpeculativeConfig, - max_num_sequences: Optional[int] = None, + decoding_config: trtllm.DecodingConfig, + kv_cache_config: KvCacheConfig, + disable_flashinfer_sampling: bool, ): enable_async_worker = (confidential_compute_enabled() or llm_args.sampler_force_async_worker) @@ -3575,11 +2260,16 @@ def instantiate_sampler( speculative_config=speculative_config, max_beam_width=max_beam_width, disable_overlap_scheduler=llm_args.disable_overlap_scheduler, + disable_flashinfer_sampling=disable_flashinfer_sampling, enable_async_worker=enable_async_worker, enable_speculative_beam_history_d2h=llm_args. enable_speculative_beam_history_d2h, - max_num_sequences=max_num_sequences, ) + decoding_mode = get_decoding_mode(decoding_config=decoding_config, + max_beam_width=max_beam_width) + if mapping.cp_config.get('cp_type') == CpType.STAR: + assert llm_args.attn_backend == "FLASHINFER_STAR_ATTENTION", "attention backend of star attention should be 'FLASHINFER_STAR_ATTENTION'" + return TorchSampler(sampler_args) if engine.spec_config is not None and engine.spec_config.spec_dec_mode.has_spec_decoder( ): return get_spec_decoder(sampler_args, engine.spec_config) @@ -3587,12 +2277,51 @@ def instantiate_sampler( if mm_encoder_only: # NOTE: handle model outputs specially for mm encoder executor/engine return EarlyStopWithMMResult() + if llm_args.sampler_type == SamplerType.TRTLLMSampler: + logger.warning( + "TRTLLMSampler is deprecated and will be removed in release 1.4. Please use TorchSampler instead." + ) + logger.debug(f"DecodingMode: {decoding_mode.name}") + return TRTLLMSampler(engine.model, + engine.dtype, + mapping, + decoding_mode, + llm_args.disable_overlap_scheduler, + max_seq_len=max_seq_len, + max_batch_size=max_batch_size, + max_beam_width=max_beam_width, + decoding_config=decoding_config, + kv_cache_config=kv_cache_config, + enable_async_worker=enable_async_worker) if not engine.model.model_config.is_generation: # NOTE: choose sampler based on model type return EarlyStopSampler() return TorchSampler(sampler_args) +def get_decoding_mode( + decoding_config: trtllm.DecodingConfig, + max_beam_width: int, +) -> DecodingMode: + '''This implementation is based off trtGptModelInflightBatching.cpp getDecodingMode().''' + if decoding_config and decoding_config.decoding_mode and not decoding_config.decoding_mode.isAuto( + ): + decoding_mode = decoding_config.decoding_mode + elif max_beam_width == 1: + decoding_mode = DecodingMode.TopKTopP() + else: + decoding_mode = DecodingMode.BeamSearch() + + # Override decoding mode when beam width is one + if max_beam_width == 1 and decoding_mode.isBeamSearch(): + logger.warning( + "Beam width is set to 1, but decoding mode is BeamSearch. Overwriting decoding mode to TopKTopP." + ) + decoding_mode = DecodingMode.TopKTopP() + + return decoding_mode + + _ATTN_MODULES = frozenset({ "attn_q", "attn_k", @@ -3667,8 +2396,7 @@ def _infer_shared_expert_size_from_adapter(adapter_dir: str) -> int: import json try: - from tensorrt_llm.models.convert_utils import (get_model_path, - load_state_dict) + from tensorrt_llm.lora_manager import get_model_path, load_state_dict model_path = get_model_path(adapter_dir, "adapter_model") if model_path is None: return 0 @@ -3763,7 +2491,7 @@ def _adjust_torch_mem_fraction(): torch.cuda.set_per_process_memory_fraction(mem_torch_fraction) -def validate_feature_combination(llm_args, model_engine): +def validate_feature_combination(llm_args, model_engine, sampler_type): # Validate the flags for features' combination def init_feature_status(llm_args) -> Dict[str, bool]: assert isinstance( @@ -3778,6 +2506,8 @@ def init_feature_status(llm_args) -> Dict[str, bool]: "mtp", "eagle3_one_model", "eagle3_two_model", + "torch_sampler", + "trtllm_sampler", "kv_cache_reuse", "slide_window_attention", "guided_decoding", @@ -3798,6 +2528,11 @@ def init_feature_status(llm_args) -> Dict[str, bool]: feature_status["eagle3_two_model"] = ( isinstance(llm_args.speculative_config, EagleDecodingConfig) and not llm_args.speculative_config.eagle3_one_model) + feature_status[ + "torch_sampler"] = sampler_type == SamplerType.TorchSampler + feature_status[ + "trtllm_sampler"] = sampler_type == SamplerType.TRTLLMSampler + feature_status[ "kv_cache_reuse"] = llm_args.kv_cache_config is not None and llm_args.kv_cache_config.enable_block_reuse feature_status["slide_window_attention"] = ( @@ -3815,9 +2550,29 @@ def init_feature_status(llm_args) -> Dict[str, bool]: feature_status: Dict[str, bool] = init_feature_status(llm_args) - # Kept as an extension point; there are currently no conflicting feature - # combinations to reject. - CONFLICT_RULES: list[dict[str, Any]] = [ + ERR_MSG_TMPL = "{feature1} and {feature2} enabled together is not supported yet." + + CONFLICT_RULES = [ + { + "features": ["trtllm_sampler", "mtp"], + "message": + ERR_MSG_TMPL.format(feature1="trtllm_sampler", feature2="mtp") + + " Please use sampler type auto instead." + }, + { + "features": ["trtllm_sampler", "eagle3_one_model"], + "message": + ERR_MSG_TMPL.format(feature1="trtllm_sampler", + feature2="eagle3_one_model") + + " Please use sampler type auto instead." + }, + { + "features": ["trtllm_sampler", "eagle3_two_model"], + "message": + ERR_MSG_TMPL.format(feature1="trtllm_sampler", + feature2="eagle3_two_model") + + " Please use sampler type auto instead." + }, # Add new conflict rules here in the future ] for rule in CONFLICT_RULES: diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py index 6bad069afaa4..149c1d677928 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py @@ -352,6 +352,32 @@ def test_regression_without_fix_would_overcount(): assert result != wrong +def test_integrated_gpu_estimation_ignores_affine_intercept(): + """Unified-memory estimation must not clamp to zero just because the + recurrent-state fixed cost exceeds the mem_get_info-derived budget.""" + tpb = 16 + c = _make_creator( + tpb, + [_make_mock_request(63)], + enable_attention_dp=False, + tp_size=1, + ) + + with ( + patch("torch.cuda.mem_get_info", return_value=(100, 100)), + patch.object( + KvCacheCreator, + "_get_kv_size_per_token", + return_value=CacheCost(slope=1, intercept=128), + ), + patch( + "tensorrt_llm._torch.pyexecutor._util.is_device_integrated", + return_value=True, + ), + ): + assert c._get_token_num_for_estimation() == 64 + + @pytest.mark.parametrize( ("model_cls", "encoder_cache_max_bytes", "expected_reserve"), [