diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index f69fc8488766..88dfec13413c 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -39,7 +39,7 @@ from .guided_decoder import GuidedDecoder from .kv_cache_manager_v2 import KVCacheManagerV2 from .kv_cache_transceiver import AttentionTypeCpp, create_kv_cache_transceiver -from .llm_request import ExecutorResponse +from .llm_request import ExecutorResponse, LlmRequestState from .mamba_cache_manager import (BaseMambaCacheManager, CppMambaHybridCacheManager, MixedMambaHybridCacheManager, @@ -243,17 +243,23 @@ def __init__( self._draft_config = draft_config self._skip_est = skip_est - def _get_model_kv_cache_manager_cls(self, model_engine: PyTorchModelEngine): + def _get_model_kv_cache_manager_cls( + self, + model_engine: PyTorchModelEngine, + kv_cache_config_override: Optional[KvCacheConfig] = None, + ): + kv_cache_config = (kv_cache_config_override if kv_cache_config_override + is not None else self._kv_cache_config) config = model_engine.model.model_config.pretrained_config cls = get_kv_cache_manager_cls( model_engine.model.model_config, - self._kv_cache_config, + kv_cache_config, is_disagg=self._is_disagg, cache_transceiver_config=self._cache_transceiver_config) if 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 self._kv_cache_config.event_buffer_max_size > 0 or ( + > 1) or kv_cache_config.event_buffer_max_size > 0 or ( self._cache_transceiver_config is not None and self._cache_transceiver_config.backend is not None): # Per-layer head_dim models (e.g., Gemma4 hybrid) require V2's @@ -279,8 +285,9 @@ def _get_model_kv_cache_manager_cls(self, model_engine: PyTorchModelEngine): # the routing site so users see the warning where the decision is # actually made. if is_hybrid_linear(model_engine.model.model_config.pretrained_config) \ - and self._kv_cache_config.enable_block_reuse: - uses_v1_mamba_route = os.environ.get('TRTLLM_USE_CPP_MAMBA', '0') == '1' \ + 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: @@ -290,32 +297,44 @@ def _get_model_kv_cache_manager_cls(self, model_engine: PyTorchModelEngine): ) return cls - def _per_manager_cache_cost(self, manager_cls, model_config, + def _per_manager_cache_cost(self, + manager_cls, + model_config, + kv_cache_config: Optional[KvCacheConfig] = None, **extra_kwargs) -> CacheCost: + kv_cache_config = (kv_cache_config if kv_cache_config is not None else + self._kv_cache_config) return CacheCost.from_raw( manager_cls.get_cache_size_per_token( model_config, self._mapping, tokens_per_block=self._tokens_per_block, max_batch_size=self._max_batch_size, - kv_cache_config=self._kv_cache_config, + kv_cache_config=kv_cache_config, **extra_kwargs)) - def _get_kv_size_per_token(self) -> CacheCost: + def _get_kv_size_per_token(self, + kv_cache_config: Optional[KvCacheConfig] = None + ) -> CacheCost: """Aggregate KV cost across target + (optional) draft as a CacheCost. ``max_batch_size`` and ``kv_cache_config`` are passed unconditionally; managers that don't need them ignore via ``**kwargs``. """ + 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 total = self._per_manager_cache_cost(self._kv_cache_manager_cls, - model_config) + 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: draft_model_config = self._draft_model_engine.model.model_config draft_kv_cache_manager_cls = self._get_model_kv_cache_manager_cls( - self._draft_model_engine) + self._draft_model_engine, kv_cache_config) total += self._per_manager_cache_cost(draft_kv_cache_manager_cls, - draft_model_config) + draft_model_config, + kv_cache_config) 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 @@ -330,15 +349,17 @@ def _get_kv_size_per_token(self) -> CacheCost: # from target (e.g. hybrid target + plain transformer draft). draft_kv_cache_manager_cls = get_kv_cache_manager_cls( effective_draft_config, - self._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_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, + kv_cache_config, num_layers=self._get_num_draft_layers()) return total @@ -743,13 +764,17 @@ def configure_kv_cache_capacity(self, # ---------------------------handle max_gpu_total_bytes--------------------------------- def _create_kv_cache_manager( - self, - model_engine: PyTorchModelEngine, - estimating_kv_cache: bool = False) -> KVCacheManager: + self, + model_engine: PyTorchModelEngine, + estimating_kv_cache: bool = False, + 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." + kv_cache_config = (kv_cache_config_override if kv_cache_config_override + is not None else self._kv_cache_config) kv_cache_manager_cls = self._get_model_kv_cache_manager_cls( - model_engine) + model_engine, kv_cache_config) # When using separate draft KV cache in one-model speculative decoding, # use layer_mask to include only target layers. The draft layers should @@ -765,7 +790,7 @@ def _create_kv_cache_manager( model_engine=model_engine, kv_cache_manager_cls=kv_cache_manager_cls, mapping=mapping, - kv_cache_config=self._kv_cache_config, + kv_cache_config=kv_cache_config, tokens_per_block=self._tokens_per_block, max_seq_len=self._max_seq_len, max_batch_size=self._max_batch_size, @@ -876,17 +901,17 @@ def _create_one_model_draft_kv_cache_manager( # otherwise fall back to target model config for MTP). effective_draft_config = self._get_effective_draft_config() + draft_kv_config = (kv_cache_config_override if kv_cache_config_override + is not None else self._kv_cache_config) # Get the appropriate KV cache manager class for the draft model draft_kv_cache_manager_cls = get_kv_cache_manager_cls( - effective_draft_config, - self._kv_cache_config, - is_disagg=self._is_disagg) + effective_draft_config, draft_kv_config, is_disagg=self._is_disagg) # 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 self._kv_cache_config.event_buffer_max_size > 0 or ( + > 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( @@ -900,7 +925,6 @@ def _create_one_model_draft_kv_cache_manager( # the sparse_attention_config. Get it from effective_draft_config which # falls back to the target model's config for MTP mode. sparse_attn_config = effective_draft_config.sparse_attention_config - draft_kv_config = kv_cache_config_override if kv_cache_config_override is not None else self._kv_cache_config return _create_kv_cache_manager( model_engine=None, kv_cache_manager_cls=draft_kv_cache_manager_cls, @@ -926,11 +950,16 @@ def _create_one_model_draft_kv_cache_manager( ) def _get_target_and_draft_cache_costs( - self, ) -> Optional[tuple[CacheCost, CacheCost]]: + self, + kv_cache_config: Optional[KvCacheConfig] = None, + ) -> Optional[tuple[CacheCost, CacheCost]]: """Per-manager KV cache costs for target and draft layers.""" - total_kv = self._get_kv_size_per_token() + 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) target_kv = self._per_manager_cache_cost( - self._kv_cache_manager_cls, self._model_engine.model.model_config) + 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, @@ -963,21 +992,20 @@ def _compute_draft_budget_shares( def _split_kv_cache_budget_for_draft( self, budget_attr: str, + target_kv_cache_config: Optional[KvCacheConfig] = None, draft_kv_cache_config: Optional[KvCacheConfig] = None, - ) -> Optional[KvCacheConfig]: + ) -> tuple[KvCacheConfig, Optional[KvCacheConfig]]: """Split a byte budget (attribute on ``KvCacheConfig``) between target and draft KV caches. - Splits the value of ``self._kv_cache_config.`` using the - affine target/draft cache costs, updates the target config in-place, - and merges the draft share into ``draft_kv_cache_config`` (cloning the - target config if needed). + Splits the value of ``target_kv_cache_config.`` using the + affine target/draft cache costs, then returns cloned target and draft + configs containing their respective shares. - Returns the (possibly newly created) draft config. The input - ``draft_kv_cache_config`` is returned unchanged when the split is not - applicable (the budget is not set, or the per-manager cache costs are - unavailable) — in those cases sharing ``self._kv_cache_config`` is - correct. + The input target config and the creator's base config are not mutated. + When the split is not applicable (the budget is not set, or the + per-manager cache costs are unavailable), the input configs are returned + unchanged. The affine fixed (intercept) cost models GPU-resident state (e.g. mamba SSM state). It is only charged against ``max_gpu_total_bytes``; for any @@ -993,13 +1021,17 @@ def _split_kv_cache_budget_for_draft( for non-GPU budgets remains so the draft never silently inherits the full budget and double-allocates it. """ - total_budget = getattr(self._kv_cache_config, budget_attr) or 0 + target_kv_cache_config = (target_kv_cache_config + if target_kv_cache_config is not None else + self._kv_cache_config) + total_budget = getattr(target_kv_cache_config, budget_attr) or 0 if total_budget <= 0: - return draft_kv_cache_config + return target_kv_cache_config, draft_kv_cache_config - cache_costs = self._get_target_and_draft_cache_costs() + cache_costs = self._get_target_and_draft_cache_costs( + target_kv_cache_config) if cache_costs is None: - return draft_kv_cache_config + return target_kv_cache_config, draft_kv_cache_config target_kv, draft_kv = cache_costs # The fixed (intercept) cost models GPU-resident state such as mamba SSM @@ -1040,9 +1072,11 @@ def _split_kv_cache_budget_for_draft( f"assigning the draft a zero {budget_attr} budget to avoid " f"double-allocating the full budget.") if draft_kv_cache_config is None: - draft_kv_cache_config = self._kv_cache_config.model_copy() + draft_kv_cache_config = target_kv_cache_config.model_copy() + else: + draft_kv_cache_config = draft_kv_cache_config.model_copy() setattr(draft_kv_cache_config, budget_attr, 0) - return draft_kv_cache_config + return target_kv_cache_config, draft_kv_cache_config target_budget, draft_budget = shares logger.info( @@ -1050,17 +1084,248 @@ def _split_kv_cache_budget_for_draft( f"target={target_budget / GB:.2f} GiB ({target_kv}), " f"draft={draft_budget / GB:.2f} GiB ({draft_kv})") - setattr(self._kv_cache_config, budget_attr, target_budget) + split_target_kv_cache_config = target_kv_cache_config.model_copy() + setattr(split_target_kv_cache_config, budget_attr, target_budget) if draft_kv_cache_config is None: - draft_kv_cache_config = self._kv_cache_config.model_copy() - setattr(draft_kv_cache_config, budget_attr, draft_budget) - return draft_kv_cache_config + split_draft_kv_cache_config = target_kv_cache_config.model_copy() + else: + split_draft_kv_cache_config = draft_kv_cache_config.model_copy() + setattr(split_draft_kv_cache_config, budget_attr, draft_budget) + return split_target_kv_cache_config, split_draft_kv_cache_config + + def _is_encoder_decoder(self) -> bool: + return bool( + getattr(self._model_engine.model.model_config, "is_encoder_decoder", + False)) + + @staticmethod + def _get_config_int_attr(config, names: tuple[str, ...]) -> Optional[int]: + for name in names: + value = getattr(config, name, None) + if isinstance(value, int): + return value + return None + + def _get_cross_kv_cache_layout( + self, + fallback_max_seq_len: Optional[int] = None + ) -> tuple[int, int, int, int]: + """Return decoder-layer count and encoder KV geometry for cross cache.""" + config = self._model_engine.model.model_config.pretrained_config + + num_layers = self._get_config_int_attr( + config, + ("num_decoder_layers", "decoder_layers", "num_hidden_layers", + "num_layers"), + ) + if num_layers is None: + raise ValueError( + "Unable to determine decoder layer count for cross KV cache.") - def _needs_gpu_kv_cache_budget_split(self) -> bool: + encoder_num_heads = self._get_config_int_attr( + config, + ("encoder_num_heads", "encoder_attention_heads", "num_heads", + "num_attention_heads"), + ) + if encoder_num_heads is None: + raise ValueError( + "Unable to determine encoder attention head count for cross KV cache." + ) + + num_kv_heads = self._get_config_int_attr( + config, + ("encoder_num_kv_heads", "encoder_num_key_value_heads", + "encoder_attention_heads", "encoder_num_heads", + "num_key_value_heads", "num_heads", "num_attention_heads"), + ) + if num_kv_heads is None: + num_kv_heads = encoder_num_heads + + encoder_hidden_size = self._get_config_int_attr( + config, ("encoder_hidden_size", "d_model", "hidden_size")) + if encoder_hidden_size is None: + raise ValueError( + "Unable to determine encoder hidden size for cross KV cache.") + + head_dim = self._get_config_int_attr( + config, + ("encoder_head_size", "encoder_head_dim", "d_kv"), + ) + if head_dim is None: + head_dim = encoder_hidden_size // encoder_num_heads + + max_seq_len = fallback_max_seq_len or self._max_seq_len + max_input_len = getattr(self._llm_args, "max_input_len", None) + if isinstance(max_input_len, int) and max_input_len > 0: + max_seq_len = max_input_len + encoder_limit = self._get_config_int_attr( + config, + ("max_encoder_input_len", "encoder_max_input_length", + "max_encoder_position_embeddings", + "encoder_max_position_embeddings", "max_position_embeddings", + "n_positions"), + ) + if encoder_limit is not None: + max_seq_len = min(max_seq_len, encoder_limit) + + return num_layers, num_kv_heads, head_dim, max_seq_len + + def _get_cross_kv_size_per_token(self) -> int: + """Estimate bytes/token for the encoder-decoder cross-attention pool.""" + from types import SimpleNamespace + + model_config = self._model_engine.model.model_config + config = model_config.pretrained_config + (num_layers, num_kv_heads, head_dim, + _) = self._get_cross_kv_cache_layout() + num_attention_heads = self._get_config_int_attr( + config, + ("encoder_num_heads", "encoder_attention_heads", "num_heads", + "num_attention_heads"), + ) + hidden_size = self._get_config_int_attr( + config, ("encoder_hidden_size", "d_model", "hidden_size")) + proxy_model_config = SimpleNamespace( + pretrained_config=SimpleNamespace( + num_key_value_heads=num_kv_heads, + num_attention_heads=num_attention_heads, + hidden_size=hidden_size, + head_dim=head_dim, + ), + quant_config=model_config.quant_config, + ) + return self._kv_cache_manager_cls.get_cache_size_per_token( + proxy_model_config, + self._mapping, + tokens_per_block=self._tokens_per_block, + num_layers=num_layers, + ) + + def _split_kv_cache_budget_for_cross( + self, + kv_cache_config: Optional[KvCacheConfig] = None, + ) -> tuple[KvCacheConfig, KvCacheConfig]: + """Split enc-dec KV cache budgets between self and cross pools. + + The cross manager must exist for every encoder-decoder runtime. During + both estimation and final construction, split the same memory-derived + budget sources used by the legacy TRT path: the free-memory fraction, + any explicit ``max_gpu_total_bytes`` override, and any explicit host + cache budget. ``max_tokens`` is a logical cap, not a memory split knob, + so it is intentionally left unchanged. The creator's base config is not + mutated. + """ + base_kv_cache_config = (kv_cache_config if kv_cache_config is not None + else self._kv_cache_config) + fraction = base_kv_cache_config.cross_kv_cache_fraction + if fraction is None: + raise ValueError("Encoder-decoder models require " + "cross_kv_cache_fraction to size the cross " + "KV cache pool.") + + self_kv_cache_config = base_kv_cache_config.model_copy() + cross_kv_cache_config = base_kv_cache_config.model_copy() + split_any_budget = False + + free_fraction = base_kv_cache_config.free_gpu_memory_fraction + if free_fraction is not None: + cross_fraction = free_fraction * fraction + self_fraction = free_fraction - cross_fraction + logger.info( + "Splitting encoder-decoder free GPU memory fraction: " + f"total={free_fraction:.3f}, self={self_fraction:.3f}, cross={cross_fraction:.3f}" + ) + self_kv_cache_config.free_gpu_memory_fraction = self_fraction + cross_kv_cache_config.free_gpu_memory_fraction = cross_fraction + split_any_budget = True + + total_budget = base_kv_cache_config.max_gpu_total_bytes + if total_budget is not None and total_budget > 0: + cross_budget = int(total_budget * fraction) + self_budget = total_budget - cross_budget + logger.info( + f"Splitting KV cache budget for encoder-decoder: " + f"total={total_budget / GB:.2f} GiB, " + f"self={self_budget / GB:.2f} GiB ({1 - fraction:.0%}), " + f"cross={cross_budget / GB:.2f} GiB ({fraction:.0%})") + self_kv_cache_config.max_gpu_total_bytes = self_budget + cross_kv_cache_config.max_gpu_total_bytes = cross_budget + split_any_budget = True + + host_cache_size = base_kv_cache_config.host_cache_size + if host_cache_size is not None and host_cache_size > 0: + cross_host_cache_size = int(host_cache_size * fraction) + self_host_cache_size = host_cache_size - cross_host_cache_size + logger.info( + f"Splitting KV cache host budget for encoder-decoder: " + f"total={host_cache_size / GB:.2f} GiB, " + f"self={self_host_cache_size / GB:.2f} GiB ({1 - fraction:.0%}), " + f"cross={cross_host_cache_size / GB:.2f} GiB ({fraction:.0%})") + self_kv_cache_config.host_cache_size = self_host_cache_size + cross_kv_cache_config.host_cache_size = cross_host_cache_size + split_any_budget = True + + if not split_any_budget: + raise ValueError("Unable to size the encoder-decoder cross KV " + "cache pool: neither free_gpu_memory_fraction nor " + "max_gpu_total_bytes nor host_cache_size is " + "available.") + + return self_kv_cache_config, cross_kv_cache_config + + def _create_cross_kv_cache_manager( + self, + cross_kv_cache_config: KvCacheConfig, + estimating_kv_cache: bool = False, + fallback_max_seq_len: Optional[int] = None, + ) -> KVCacheManager: + """Create a KV cache manager for the cross-attention pool. + + The cross pool stores encoder K/V projections that are written once + during the first decoder context step and read on every subsequent + decoder generation step. It uses ``CacheType.CROSS`` with decoder + layer count but encoder-side KV geometry. + + The manager class mirrors the self pool (``KVCacheManager`` for V1, + ``KVCacheManagerV2`` for V2) so that both pools share the same + runtime ABI and scheduler integration. V1 is the default and the + production target for encoder-decoder models. + """ + (num_layers, num_kv_heads, head_dim, + max_seq_len) = self._get_cross_kv_cache_layout(fallback_max_seq_len) + estimating_kv_cache = estimating_kv_cache and not self._skip_est + return _create_kv_cache_manager( + model_engine=self._model_engine, + kv_cache_manager_cls=self._kv_cache_manager_cls, + mapping=self._mapping, + kv_cache_config=cross_kv_cache_config, + tokens_per_block=self._tokens_per_block, + max_seq_len=max_seq_len, + max_batch_size=self._max_batch_size, + spec_config=None, + sparse_attn_config=None, + max_num_tokens=self._max_num_tokens, + max_beam_width=1, + kv_connector_manager=None, + estimating_kv_cache=estimating_kv_cache, + execution_stream=self._execution_stream, + num_layers=num_layers, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + kv_cache_type=tensorrt_llm.bindings.internal.batch_manager. + CacheType.CROSS, + ) + + def _needs_gpu_kv_cache_budget_split( + self, + kv_cache_config: Optional[KvCacheConfig] = None, + ) -> bool: """Whether max_gpu_total_bytes must be split per manager.""" if issubclass(self._kv_cache_manager_cls, KVCacheManagerV2): return self._should_create_separate_draft_kv_cache() - return is_vswa_enabled(self._kv_cache_config) + kv_cache_config = (kv_cache_config if kv_cache_config is not None else + self._kv_cache_config) + return is_vswa_enabled(kv_cache_config) def build_managers(self, resources: Dict, @@ -1068,6 +1333,17 @@ def build_managers(self, """Construct KV caches for model and draft model (if applicable).""" if self._skip_est: self.configure_kv_cache_capacity() + original_max_seq_len = self._max_seq_len + + # For encoder-decoder models, split the self/cross budgets first so + # every enc-dec build creates a real cross pool. This must happen + # before any draft split so that the draft split operates on the + # already-reduced self-pool budget. + self_kv_cache_config = self._kv_cache_config + cross_kv_cache_config = None + if self._is_encoder_decoder(): + self_kv_cache_config, cross_kv_cache_config = self._split_kv_cache_budget_for_cross( + ) # Split combined KV cache budgets before creating managers. Skip during # estimation — estimation uses max_tokens-based logic and must not @@ -1079,26 +1355,36 @@ def build_managers(self, 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(): - draft_kv_cache_config = self._split_kv_cache_budget_for_draft( - "max_gpu_total_bytes", draft_kv_cache_config) + 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)) # 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. - draft_kv_cache_config = self._split_kv_cache_budget_for_draft( - "host_cache_size", draft_kv_cache_config) + self_kv_cache_config, draft_kv_cache_config = ( + self._split_kv_cache_budget_for_draft( + "host_cache_size", self_kv_cache_config, + draft_kv_cache_config)) kv_cache_manager = self._create_kv_cache_manager( - self._model_engine, estimating_kv_cache) + self._model_engine, + estimating_kv_cache, + 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: + if (not estimating_kv_cache and self._kv_connector_manager is not None + and self._draft_model_engine is not None): raise NotImplementedError( "Connector manager is not supported for draft model.") draft_kv_cache_manager = None + draft_build_kv_cache_config = (draft_kv_cache_config + if draft_kv_cache_config is not None else + self_kv_cache_config) # Two-model speculative decoding: draft model has separate engine if self._draft_model_engine is not None: @@ -1106,31 +1392,31 @@ def build_managers(self, assert draft_kv_cache_config is None, ( "KVCacheManagerV2 does not support two-model speculative " "decoding with separate draft KV cache budget splitting.") - # For V1, apply the draft's split budgets temporarily. - if draft_kv_cache_config is not None: - saved_budget = self._kv_cache_config.max_gpu_total_bytes - saved_host = self._kv_cache_config.host_cache_size - self._kv_cache_config.max_gpu_total_bytes = ( - draft_kv_cache_config.max_gpu_total_bytes) - self._kv_cache_config.host_cache_size = ( - draft_kv_cache_config.host_cache_size) draft_kv_cache_manager = self._create_kv_cache_manager( - self._draft_model_engine, estimating_kv_cache) - if draft_kv_cache_config is not None: - self._kv_cache_config.max_gpu_total_bytes = saved_budget - self._kv_cache_config.host_cache_size = saved_host + self._draft_model_engine, + estimating_kv_cache, + kv_cache_config_override=draft_build_kv_cache_config) # 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( estimating_kv_cache, - kv_cache_config_override=draft_kv_cache_config) + kv_cache_config_override=draft_build_kv_cache_config) + + # Encoder-decoder cross-attention pool + cross_kv_cache_manager = None + if cross_kv_cache_config is not None: + cross_kv_cache_manager = self._create_cross_kv_cache_manager( + cross_kv_cache_config, estimating_kv_cache, + original_max_seq_len) resources[ResourceManagerType.KV_CACHE_MANAGER] = kv_cache_manager resources[ ResourceManagerType.DRAFT_KV_CACHE_MANAGER] = draft_kv_cache_manager + resources[ + ResourceManagerType.CROSS_KV_CACHE_MANAGER] = cross_kv_cache_manager def teardown_managers(self, resources: Dict) -> None: - """Clean up KV caches for model and draft model (if applicable).""" + """Clean up KV caches for model, draft model, and cross pool.""" resources[ResourceManagerType.KV_CACHE_MANAGER].shutdown() del resources[ResourceManagerType.KV_CACHE_MANAGER] draft_kv_cache_manager = resources[ @@ -1138,6 +1424,12 @@ def teardown_managers(self, resources: Dict) -> None: if draft_kv_cache_manager: draft_kv_cache_manager.shutdown() del resources[ResourceManagerType.DRAFT_KV_CACHE_MANAGER] + cross_kv_cache_manager = resources.get( + ResourceManagerType.CROSS_KV_CACHE_MANAGER) + if cross_kv_cache_manager is not None: + cross_kv_cache_manager.shutdown() + if ResourceManagerType.CROSS_KV_CACHE_MANAGER in resources: + del resources[ResourceManagerType.CROSS_KV_CACHE_MANAGER] def _build_per_layer_num_kv_heads( @@ -1194,6 +1486,9 @@ def _create_kv_cache_manager( is_draft: Optional[bool] = None, layer_mask: Optional[List[bool]] = None, num_layers: Optional[int] = None, + num_kv_heads: Optional[Union[int, List[int]]] = None, + head_dim: Optional[int] = None, + kv_cache_type=None, is_disagg: bool = False) -> KVCacheManager: """ Returns: @@ -1215,11 +1510,15 @@ def _create_kv_cache_manager( if is_draft is None: is_draft = model_engine.is_draft_model + if kv_cache_type is None: + kv_cache_type = tensorrt_llm.bindings.internal.batch_manager.CacheType.SELF + hidden_size = config.hidden_size num_attention_heads = config.num_attention_heads - num_key_value_heads = getattr(config, 'num_key_value_heads', - num_attention_heads) - head_dim = getattr(config, "head_dim", None) + num_key_value_heads = num_kv_heads if num_kv_heads is not None else getattr( + config, 'num_key_value_heads', num_attention_heads) + if not isinstance(head_dim, int): + head_dim = getattr(config, "head_dim", None) if not isinstance(head_dim, int): head_dim = hidden_size // num_attention_heads @@ -1493,7 +1792,7 @@ def _create_kv_cache_manager( and kv_cache_manager_cls.__name__ == "KVCacheManager" else head_dim) kv_cache_manager = kv_cache_manager_cls( kv_cache_config, - tensorrt_llm.bindings.internal.batch_manager.CacheType.SELF, + kv_cache_type, num_layers=num_hidden_layers, num_kv_heads=per_layer_num_kv_heads, head_dim=effective_head_dim, @@ -1720,18 +2019,32 @@ def create_py_executor_instance( resource_manager = ResourceManager(resources) - # Make sure the kv cache manager is always invoked last as it could + # 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) + # 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. scheduler_capacity = max_num_sequences if scheduler_capacity == 1 and mapping.enable_attention_dp and kv_cache_manager: scheduler_capacity += 1 + # For encoder-decoder models, requests start in ENCODER_INIT and the + # capacity scheduler must admit them already at that state so the + # encoder loop can run. Decoder-only deployments keep the default + # CONTEXT_INIT gating. + no_schedule_until_state = (LlmRequestState.ENCODER_INIT + if cross_kv_cache_manager is not None else + LlmRequestState.CONTEXT_INIT) + if isinstance(kv_cache_manager, KVCacheManagerV2): # V2: interleaved scheduler handles both capacity and budget draft_kv_cache_manager = resources.get( @@ -1749,6 +2062,8 @@ def create_py_executor_instance( if peft_cache_manager is not None else None, 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, ) elif (scheduler_config is not None and scheduler_config.use_python_scheduler): @@ -1761,15 +2076,21 @@ def create_py_executor_instance( if peft_cache_manager is not None else None, scheduler_policy=scheduler_config.capacity_scheduler_policy, ctx_chunk_config=ctx_chunk_config, + 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(), - scheduler_capacity=scheduler_capacity) + scheduler_capacity=scheduler_capacity, + no_schedule_until_state=no_schedule_until_state) else: capacity_scheduler = BindCapacityScheduler( scheduler_capacity, kv_cache_manager.impl if kv_cache_manager is not None else None, peft_cache_manager.impl if peft_cache_manager is not None else None, scheduler_config.capacity_scheduler_policy, - two_step_lookahead=mapping.has_pp()) + 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) mb_scheduler = BindMicroBatchScheduler(max_batch_size, max_num_tokens, ctx_chunk_config) diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index 137007e1efbc..d3affad163a8 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -53,6 +53,7 @@ if TYPE_CHECKING: from tensorrt_llm._torch.attention_backend.interface import \ AttentionMetadata + from tensorrt_llm.llmapi.llm_args import DecodingBaseConfig BlocksPerWindow = Dict[int, Tuple[ int, @@ -78,6 +79,7 @@ class PoolConfiguration: class ResourceManagerType(enum.Enum): KV_CACHE_MANAGER = "KV_CACHE_MANAGER" DRAFT_KV_CACHE_MANAGER = "DRAFT_KV_CACHE_MANAGER" + CROSS_KV_CACHE_MANAGER = "CROSS_KV_CACHE_MANAGER" PEFT_CACHE_MANAGER = "PEFT_CACHE_MANAGER" SEQ_SLOT_MANAGER = "SEQ_SLOT_MANAGER" SPEC_RESOURCE_MANAGER = "SPEC_RESOURCE_MANAGER" @@ -668,49 +670,54 @@ def get_needed_resource_to_completion(self, request: LlmRequest) -> int: remaining_tokens / self.tokens_per_block) return need_blocks + def _context_seq_len(self, req: LlmRequest, is_cross: bool, + is_star_cp: bool) -> Optional[int]: + """Return the sequence length to pass to add_sequence_batch, or None to skip this request.""" + if is_cross: + if (getattr(req, "py_skip_cross_kv_projection", False) + or not req.is_first_context_chunk + or not self._kv_connector_should_add_sequence(req)): + return None + encoder_output_len = getattr(req, "encoder_output_len", None) + if encoder_output_len is None: + raise RuntimeError( + "Cross KV cache allocation requires " + f"encoder_output_len for request {req.py_request_id}.") + return int(encoder_output_len) + if is_star_cp: + if req.ctx_iters != 0: + return None + seq_len = sum(len(ctx_block) for ctx_block in req.ctx_blocks) + return seq_len + (len(req.query_id) if self.mapping.cp_rank + == self.mapping.cp_size - 1 else 0) + if not req.is_first_context_chunk or not self._kv_connector_should_add_sequence( + req): + return None + return req.prompt_len + def prepare_resources(self, scheduled_batch: ScheduledRequests): + # Cross/encoder K/V is allocated once and never grows; handle it on a + # dedicated path so the self-attention flow below stays unconditional. + if self.kv_cache_type == CacheTypeCpp.CROSS: + return self._prepare_cross_resources(scheduled_batch) + + is_star_cp = ('cp_type' in self.mapping.cp_config + and CpType.STAR == self.mapping.cp_config['cp_type']) with request_context(self.is_draft, scheduled_batch): # wait for all pending work to finish before launching offload/onboarding/partial copy self.impl.sync_transfer_manager_with_buffer_manager() - # Collect first-chunk requests eligible for add_sequence_batch. + # Collect first-chunk requests eligible for batch add_sequence_batch. # When block reuse is enabled, addSequenceBatch uses a two-phase # claim-then-onboard strategy that prevents host offloading from # evicting reusable blocks in the radix tree. - batch_request_infos = [] - batch_llm_requests = [] - batch_ctx_requests = [] - - # allocate KV Cache - is_star_cp = 'cp_type' in self.mapping.cp_config and CpType.STAR == self.mapping.cp_config[ - 'cp_type'] - - for req in scheduled_batch.context_requests: - req_beam_width = req.py_beam_width - if is_star_cp: - if req.ctx_iters == 0: - seq_len = sum( - len(ctx_block) for ctx_block in req.ctx_blocks) - prompt_len = seq_len + ( - len(req.query_id) if self.mapping.cp_rank - == self.mapping.cp_size - 1 else 0) - batch_request_infos.append( - (req.py_request_id, prompt_len, req_beam_width)) - batch_llm_requests.append(req) - batch_ctx_requests.append(req) - else: - if req.is_first_context_chunk and self._kv_connector_should_add_sequence( - req): - # Batch path: two-phase claim-then-onboard - batch_request_infos.append( - (req.py_request_id, req.prompt_len, req_beam_width)) - batch_llm_requests.append(req) - batch_ctx_requests.append(req) + batch_request_infos, batch_llm_requests = self._collect_context_sequences( + scheduled_batch, is_cross=False, is_star_cp=is_star_cp) if batch_request_infos: self.impl.add_sequence_batch(batch_request_infos, batch_llm_requests) - for req in batch_ctx_requests: + for req in batch_llm_requests: for _ in range(self.num_extra_kv_tokens): self.impl.add_token(req.py_request_id) for _ in range(get_draft_token_length(req)): @@ -750,6 +757,44 @@ def prepare_resources(self, scheduled_batch: ScheduledRequests): self.kv_connector_manager.build_scheduler_output( scheduled_batch, self) + def _collect_context_sequences(self, scheduled_batch: ScheduledRequests, + is_cross: bool, is_star_cp: bool): + """Build the (request_info, llm_request) lists for add_sequence_batch. + + Cross (encoder) sequences are sized from encoder_output_len with a beam + width of 1 (request-scoped); self-attention sequences use the request's + own beam width. + """ + batch_request_infos = [] + batch_llm_requests = [] + for req in scheduled_batch.context_requests: + seq_len = self._context_seq_len(req, is_cross, is_star_cp) + if seq_len is None: + continue + beam_width = 1 if is_cross else req.py_beam_width + batch_request_infos.append((req.py_request_id, seq_len, beam_width)) + batch_llm_requests.append(req) + return batch_request_infos, batch_llm_requests + + def _prepare_cross_resources(self, scheduled_batch: ScheduledRequests): + """Allocate cross (encoder) K/V blocks. + + Encoder K/V is written once at the first decoder context step and read + unchanged on every generation step, so it never grows: this skips the + decode-time token growth, draft-token reserve, and scheduler bookkeeping + that the self-attention path performs. + """ + with request_context(self.is_draft, scheduled_batch): + # wait for all pending work to finish before launching offload/onboarding/partial copy + self.impl.sync_transfer_manager_with_buffer_manager() + batch_request_infos, batch_llm_requests = self._collect_context_sequences( + scheduled_batch, is_cross=True, is_star_cp=False) + if batch_request_infos: + self.impl.add_sequence_batch(batch_request_infos, + batch_llm_requests) + # kernels wait for scheduled offload/onboard/partial copy work before launching + self.impl.refresh_blocks() + def extend_capacity_for_tokens(self, request: LlmRequest) -> None: """No-op for V1; interface kept consistent with V2.""" @@ -893,29 +938,41 @@ def update_resources(self, scheduled_batch: ScheduledRequests, attn_metadata: "AttentionMetadata" = None, kv_cache_dtype_byte_size: float = None): - # Rewind KV cache for requests with rejected draft tokens. - # Skip: - # - GENERATION_COMPLETE: finished requests - # - CONTEXT_INIT: requests whose state was reset after being paused with KV cache freed. - # With overlap scheduler, the scheduler pauses a request and frees KV cache at iteration N, - # while the previous batch (N-1) is still trying to update the KV cache after forward pass. - for request in scheduled_batch.generation_requests: - if request.state in (LlmRequestState.GENERATION_COMPLETE, - LlmRequestState.CONTEXT_INIT): - continue - if request.py_rewind_len > 0: - self.rewind_kv_cache(request, request.py_rewind_len) - # Symmetric companion to prepare_resources's reserve_slack - # add_token loop: when _kv_reserve_draft_tokens (e.g. dynamic - # tree's K*max_draft_len) exceeds the runtime draft length, - # those extra slots must also be rewound, otherwise the draft - # KV cache leaks reserve_slack tokens per generation iteration - # and eventually overflows mCacheBlockIndices. - runtime_draft_len = (request.py_rewind_len + - request.py_num_accepted_draft_tokens) - extra_rewind = self._kv_reserve_draft_tokens - runtime_draft_len - if extra_rewind > 0: - self.rewind_kv_cache(request, extra_rewind) + # Self-attention pools rewind rejected speculative tokens each step; + # cross/encoder K/V is immutable, so only the context-block commit below + # applies to it. + if self.kv_cache_type != CacheTypeCpp.CROSS: + if not self.is_draft: + from .kv_cache_manager_v2 import \ + _update_kv_cache_draft_token_location + + _update_kv_cache_draft_token_location(self, scheduled_batch, + attn_metadata, + kv_cache_dtype_byte_size) + + # Rewind KV cache for requests with rejected draft tokens. + # Skip: + # - GENERATION_COMPLETE: finished requests + # - CONTEXT_INIT: requests whose state was reset after being paused with KV cache freed. + # With overlap scheduler, the scheduler pauses a request and frees KV cache at iteration N, + # while the previous batch (N-1) is still trying to update the KV cache after forward pass. + for request in scheduled_batch.generation_requests: + if request.state in (LlmRequestState.GENERATION_COMPLETE, + LlmRequestState.CONTEXT_INIT): + continue + if request.py_rewind_len > 0: + self.rewind_kv_cache(request, request.py_rewind_len) + # Symmetric companion to prepare_resources's reserve_slack + # add_token loop: when _kv_reserve_draft_tokens (e.g. dynamic + # tree's K*max_draft_len) exceeds the runtime draft length, + # those extra slots must also be rewound, otherwise the draft + # KV cache leaks reserve_slack tokens per generation iteration + # and eventually overflows mCacheBlockIndices. + runtime_draft_len = (request.py_rewind_len + + request.py_num_accepted_draft_tokens) + extra_rewind = self._kv_reserve_draft_tokens - runtime_draft_len + if extra_rewind > 0: + self.rewind_kv_cache(request, extra_rewind) # For context requests, store completed context blocks for KV cache reuse. # We wait until context_remaining_length == 0 (all chunks processed) before diff --git a/tests/unittest/_torch/executor/test_dual_pool_kv_cache.py b/tests/unittest/_torch/executor/test_dual_pool_kv_cache.py new file mode 100644 index 000000000000..6bf94f2660d6 --- /dev/null +++ b/tests/unittest/_torch/executor/test_dual_pool_kv_cache.py @@ -0,0 +1,904 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Tests for dual-pool KV cache construction (enc-dec Steps 4 and 5). + +Validates budget splitting, ResourceManagerType.CROSS_KV_CACHE_MANAGER +registration, and the cross pool wiring for both the V1 ``KVCacheManager`` +(default and production target) and the V2 ``KVCacheManagerV2`` +(additive secondary path) scheduler integrations. +""" + +from types import SimpleNamespace +from unittest.mock import Mock, patch + +import pytest + +from tensorrt_llm._torch.pyexecutor._util import KvCacheCreator +from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2 +from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager, ResourceManagerType +from tensorrt_llm.llmapi.llm_args import CapacitySchedulerPolicy + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_mock_kv_cache_config( + cross_kv_cache_fraction=None, + max_gpu_total_bytes=None, + use_kv_cache_manager_v2=True, + max_tokens=None, + free_gpu_memory_fraction=0.9, + host_cache_size=None, +): + """Create a mock KvCacheConfig with the fields KvCacheCreator needs.""" + config = Mock() + config.cross_kv_cache_fraction = cross_kv_cache_fraction + config.max_gpu_total_bytes = max_gpu_total_bytes + config.use_kv_cache_manager_v2 = use_kv_cache_manager_v2 + config.max_tokens = max_tokens + config.free_gpu_memory_fraction = free_gpu_memory_fraction + config.host_cache_size = host_cache_size + config.max_attention_window = None + config.event_buffer_max_size = 0 + + def model_copy(): + c = Mock() + c.cross_kv_cache_fraction = config.cross_kv_cache_fraction + c.max_gpu_total_bytes = config.max_gpu_total_bytes + c.use_kv_cache_manager_v2 = config.use_kv_cache_manager_v2 + c.max_tokens = config.max_tokens + c.free_gpu_memory_fraction = config.free_gpu_memory_fraction + c.host_cache_size = config.host_cache_size + c.max_attention_window = config.max_attention_window + c.event_buffer_max_size = config.event_buffer_max_size + return c + + config.model_copy = model_copy + return config + + +def _make_mock_model_config( + is_encoder_decoder=False, + is_generation=True, + **pretrained_overrides, +): + """Minimal mock ModelConfig for KvCacheCreator.""" + model_config = Mock() + model_config.is_encoder_decoder = is_encoder_decoder + model_config.is_generation = is_generation + model_config.sparse_attention_config = None + + pretrained = Mock() + pretrained.num_hidden_layers = 6 + pretrained.num_attention_heads = 8 + pretrained.num_key_value_heads = 8 + pretrained.hidden_size = 512 + pretrained.head_dim = 64 + pretrained.vocab_size = 32000 + pretrained.quantization = Mock() + pretrained.quantization.quant_algo = None + pretrained.quantization.kv_cache_quant_algo = None + for key, value in pretrained_overrides.items(): + setattr(pretrained, key, value) + if "encoder_attention_heads" not in pretrained_overrides: + pretrained.encoder_attention_heads = pretrained.num_attention_heads + if "decoder_attention_heads" not in pretrained_overrides: + pretrained.decoder_attention_heads = pretrained.num_attention_heads + if "encoder_layers" not in pretrained_overrides: + pretrained.encoder_layers = pretrained.num_hidden_layers + if "decoder_layers" not in pretrained_overrides: + pretrained.decoder_layers = pretrained.num_hidden_layers + if "d_model" not in pretrained_overrides: + pretrained.d_model = pretrained.hidden_size + if "max_position_embeddings" not in pretrained_overrides: + pretrained.max_position_embeddings = 1024 + model_config.pretrained_config = pretrained + model_config.quant_config = None + return model_config + + +def _make_mock_model_engine(model_config): + """Minimal mock PyTorchModelEngine.""" + engine = Mock() + engine.model.model_config = model_config + engine.dtype = "bfloat16" + engine.is_draft_model = False + engine.kv_cache_manager_key = ResourceManagerType.KV_CACHE_MANAGER + return engine + + +def _make_creator(kv_cache_config, model_config=None, is_enc_dec=False, manager_cls=None): + """Create a KvCacheCreator with minimal mocking. + + ``manager_cls`` selects the KV cache manager class the creator binds to. + Defaults to ``KVCacheManagerV2`` when ``kv_cache_config.use_kv_cache_manager_v2`` + is True, otherwise the V1 ``KVCacheManager``. Tests can override + explicitly via ``manager_cls`` to exercise either path independently. + """ + if model_config is None: + model_config = _make_mock_model_config(is_encoder_decoder=is_enc_dec) + model_engine = _make_mock_model_engine(model_config) + + if manager_cls is None: + manager_cls = ( + KVCacheManagerV2 + if getattr(kv_cache_config, "use_kv_cache_manager_v2", True) + else KVCacheManager + ) + + with patch( + "tensorrt_llm._torch.pyexecutor._util.get_kv_cache_manager_cls", + return_value=manager_cls, + ): + creator = KvCacheCreator.__new__(KvCacheCreator) + creator._model_engine = model_engine + creator._draft_model_engine = None + creator._mapping = Mock() + creator._mapping.enable_attention_dp = False + creator._mapping.tp_size = 1 + creator._mapping.pp_size = 1 + creator._mapping.cp_config = {} + creator._mapping.is_last_pp_rank.return_value = True + creator._kv_cache_config = kv_cache_config + creator._max_kv_tokens_in = kv_cache_config.max_tokens + creator._max_num_tokens = 4096 + creator._max_beam_width = 1 + creator._kv_connector_manager = None + creator._llm_args = Mock() + creator._llm_args.extra_resource_managers = {} + creator._cache_transceiver_config = None + creator._speculative_config = None + creator._sparse_attention_config = None + creator._tokens_per_block = 64 + creator._max_seq_len = 2048 + creator._max_batch_size = 8 + creator._net_max_seq_len = 2048 + creator._dummy_reqs = None + creator._profiling_stage_data = None + creator._kv_cache_manager_cls = manager_cls + creator._execution_stream = None + creator._draft_config = None + creator._skip_est = True + return creator + + +# --------------------------------------------------------------------------- +# Tests: _split_kv_cache_budget_for_cross +# --------------------------------------------------------------------------- + + +class TestSplitKvCacheBudgetForCross: + """Test the budget splitting method directly.""" + + def test_split_50_50(self): + total = 10 * (1 << 30) # 10 GiB + config = _make_mock_kv_cache_config( + cross_kv_cache_fraction=0.5, + max_gpu_total_bytes=total, + free_gpu_memory_fraction=0.8, + ) + + creator = _make_creator(config, is_enc_dec=True) + self_config, cross_config = creator._split_kv_cache_budget_for_cross() + + assert self_config is not config + assert cross_config.max_gpu_total_bytes == total // 2 + assert self_config.max_gpu_total_bytes == total - total // 2 + assert cross_config.free_gpu_memory_fraction == pytest.approx(0.4) + assert self_config.free_gpu_memory_fraction == pytest.approx(0.4) + assert config.max_gpu_total_bytes == total + assert config.free_gpu_memory_fraction == pytest.approx(0.8) + + def test_split_30_70(self): + total = 10 * (1 << 30) + config = _make_mock_kv_cache_config( + cross_kv_cache_fraction=0.3, + max_gpu_total_bytes=total, + free_gpu_memory_fraction=0.8, + ) + + creator = _make_creator(config, is_enc_dec=True) + self_config, cross_config = creator._split_kv_cache_budget_for_cross() + + expected_cross = int(total * 0.3) + expected_self = total - expected_cross + assert cross_config.max_gpu_total_bytes == expected_cross + assert self_config.max_gpu_total_bytes == expected_self + assert cross_config.free_gpu_memory_fraction == pytest.approx(0.24) + assert self_config.free_gpu_memory_fraction == pytest.approx(0.56) + assert config.max_gpu_total_bytes == total + assert config.free_gpu_memory_fraction == pytest.approx(0.8) + + def test_no_split_when_fraction_is_none(self): + total = 10 * (1 << 30) + config = _make_mock_kv_cache_config(cross_kv_cache_fraction=None, max_gpu_total_bytes=total) + + creator = _make_creator(config, is_enc_dec=True) + with pytest.raises(ValueError, match="cross_kv_cache_fraction"): + creator._split_kv_cache_budget_for_cross() + + def test_split_free_fraction_when_budget_is_none(self): + config = _make_mock_kv_cache_config( + cross_kv_cache_fraction=0.5, + max_gpu_total_bytes=None, + max_tokens=1000, + free_gpu_memory_fraction=0.8, + ) + + creator = _make_creator(config, is_enc_dec=True) + self_config, cross_config = creator._split_kv_cache_budget_for_cross() + + assert cross_config.max_tokens == 1000 + assert self_config.max_tokens == 1000 + assert cross_config.free_gpu_memory_fraction == pytest.approx(0.4) + assert self_config.free_gpu_memory_fraction == pytest.approx(0.4) + assert config.free_gpu_memory_fraction == pytest.approx(0.8) + + def test_split_free_fraction_when_budget_is_zero(self): + config = _make_mock_kv_cache_config( + cross_kv_cache_fraction=0.5, + max_gpu_total_bytes=0, + max_tokens=1000, + free_gpu_memory_fraction=0.8, + ) + + creator = _make_creator(config, is_enc_dec=True) + self_config, cross_config = creator._split_kv_cache_budget_for_cross() + + assert cross_config.max_tokens == 1000 + assert self_config.max_tokens == 1000 + assert cross_config.free_gpu_memory_fraction == pytest.approx(0.4) + assert self_config.free_gpu_memory_fraction == pytest.approx(0.4) + assert config.free_gpu_memory_fraction == pytest.approx(0.8) + + def test_raises_when_no_budget_source_exists(self): + config = _make_mock_kv_cache_config( + cross_kv_cache_fraction=0.5, + max_gpu_total_bytes=0, + max_tokens=None, + free_gpu_memory_fraction=None, + ) + + creator = _make_creator(config, is_enc_dec=True) + with pytest.raises(ValueError, match="Unable to size"): + creator._split_kv_cache_budget_for_cross() + + def test_is_encoder_decoder_helper(self): + dec_config = _make_mock_model_config(is_encoder_decoder=False) + dec_creator = _make_creator(_make_mock_kv_cache_config(), model_config=dec_config) + assert not dec_creator._is_encoder_decoder() + + enc_dec_config = _make_mock_model_config(is_encoder_decoder=True) + enc_dec_creator = _make_creator(_make_mock_kv_cache_config(), model_config=enc_dec_config) + assert enc_dec_creator._is_encoder_decoder() + + def test_is_encoder_decoder_helper_handles_missing_attr(self): + pretrained = _make_mock_model_config().pretrained_config + model_config = SimpleNamespace( + is_generation=True, + sparse_attention_config=None, + pretrained_config=pretrained, + quant_config=None, + ) + creator = _make_creator(_make_mock_kv_cache_config(), model_config=model_config) + + assert not creator._is_encoder_decoder() + + def test_budgets_sum_to_total(self): + """Self + cross budgets always sum to the original total.""" + total = 7 * (1 << 30) + 123 # non-round number + config = _make_mock_kv_cache_config(cross_kv_cache_fraction=0.4, max_gpu_total_bytes=total) + + creator = _make_creator(config, is_enc_dec=True) + self_config, cross_config = creator._split_kv_cache_budget_for_cross() + + assert (self_config.max_gpu_total_bytes + cross_config.max_gpu_total_bytes) == total + assert config.max_gpu_total_bytes == total + + def test_host_cache_budget_is_split_without_mutating_base_config(self): + """Self + cross host cache budgets sum to the original host budget.""" + total_host = 7 * (1 << 30) + 123 + config = _make_mock_kv_cache_config( + cross_kv_cache_fraction=0.4, + max_gpu_total_bytes=8 * (1 << 30), + host_cache_size=total_host, + ) + + creator = _make_creator(config, is_enc_dec=True) + self_config, cross_config = creator._split_kv_cache_budget_for_cross() + + expected_cross_host = int(total_host * 0.4) + expected_self_host = total_host - expected_cross_host + assert cross_config.host_cache_size == expected_cross_host + assert self_config.host_cache_size == expected_self_host + assert (self_config.host_cache_size + cross_config.host_cache_size) == total_host + assert config.host_cache_size == total_host + + def test_host_cache_budget_counts_as_split_budget_source(self): + total_host = 4 * (1 << 30) + config = _make_mock_kv_cache_config( + cross_kv_cache_fraction=0.25, + max_gpu_total_bytes=None, + free_gpu_memory_fraction=None, + host_cache_size=total_host, + ) + + creator = _make_creator(config, is_enc_dec=True) + self_config, cross_config = creator._split_kv_cache_budget_for_cross() + + assert cross_config.host_cache_size == total_host // 4 + assert self_config.host_cache_size == total_host - total_host // 4 + assert config.host_cache_size == total_host + + +# --------------------------------------------------------------------------- +# Tests: ResourceManagerType enum +# --------------------------------------------------------------------------- + + +class TestResourceManagerType: + """Verify CROSS_KV_CACHE_MANAGER exists in the enum.""" + + def test_cross_kv_cache_manager_in_enum(self): + assert ResourceManagerType.CROSS_KV_CACHE_MANAGER.value == "CROSS_KV_CACHE_MANAGER" + + +# --------------------------------------------------------------------------- +# Tests: Cross-pool geometry and build_managers coverage +# --------------------------------------------------------------------------- + + +class TestCrossKvCacheConstruction: + """Exercise the Steps 4 and 5 construction path beyond helper math.""" + + @pytest.mark.parametrize("use_kv_cache_manager_v2", [False, True]) + def test_create_cross_kv_cache_manager_uses_encoder_geometry(self, use_kv_cache_manager_v2): + expected_cls = KVCacheManagerV2 if use_kv_cache_manager_v2 else KVCacheManager + config = _make_mock_kv_cache_config( + cross_kv_cache_fraction=0.5, + max_gpu_total_bytes=8 * (1 << 30), + use_kv_cache_manager_v2=use_kv_cache_manager_v2, + ) + model_config = _make_mock_model_config( + is_encoder_decoder=True, + num_hidden_layers=10, + num_attention_heads=16, + num_key_value_heads=16, + hidden_size=768, + head_dim=48, + encoder_layers=8, + decoder_layers=10, + encoder_attention_heads=12, + d_model=768, + max_position_embeddings=1024, + ) + creator = _make_creator(config, model_config=model_config, manager_cls=expected_cls) + cross_cfg = config.model_copy() + + with patch( + "tensorrt_llm._torch.pyexecutor._util._create_kv_cache_manager", + return_value=Mock(), + ) as create_mock: + creator._create_cross_kv_cache_manager(cross_cfg) + + kwargs = create_mock.call_args.kwargs + # Cross pool must use the same manager class as the self pool so + # both pools share the same runtime ABI. V1 is the default and + # production target; V2 is an additive secondary path. + assert kwargs["kv_cache_manager_cls"] is expected_cls + assert kwargs["num_layers"] == 10 + assert kwargs["num_kv_heads"] == 12 + assert kwargs["head_dim"] == 64 + assert kwargs["max_seq_len"] == 1024 + + import tensorrt_llm + + assert kwargs["kv_cache_type"] == ( + tensorrt_llm.bindings.internal.batch_manager.CacheType.CROSS + ) + + def test_cross_layout_uses_max_input_len_for_encoder_capacity(self): + config = _make_mock_kv_cache_config( + cross_kv_cache_fraction=0.5, + max_gpu_total_bytes=8 * (1 << 30), + ) + model_config = _make_mock_model_config( + is_encoder_decoder=True, + max_position_embeddings=4096, + ) + creator = _make_creator(config, model_config=model_config) + creator._llm_args.max_input_len = 1536 + creator._max_seq_len = 864 + + _, _, _, max_seq_len = creator._get_cross_kv_cache_layout(fallback_max_seq_len=2048) + + assert max_seq_len == 1536 + + def test_build_managers_cross_pool_ignores_mutated_self_max_seq_len(self): + config = _make_mock_kv_cache_config( + cross_kv_cache_fraction=0.5, + max_gpu_total_bytes=8 * (1 << 30), + ) + model_config = _make_mock_model_config( + is_encoder_decoder=True, + max_position_embeddings=4096, + ) + creator = _make_creator(config, model_config=model_config) + creator._llm_args.max_input_len = None + creator._max_seq_len = 2048 + creator.configure_kv_cache_capacity = Mock() + creator._should_create_separate_draft_kv_cache = Mock(return_value=False) + + def create_self_manager(*_args, **_kwargs): + creator._max_seq_len = 864 + manager = Mock() + manager.max_seq_len = 864 + return manager + + captured_cross_max_seq_lens = [] + + def create_cross_manager(*_args, **kwargs): + captured_cross_max_seq_lens.append(kwargs["max_seq_len"]) + manager = Mock() + manager.max_seq_len = kwargs["max_seq_len"] + return manager + + creator._create_kv_cache_manager = Mock(side_effect=create_self_manager) + with patch( + "tensorrt_llm._torch.pyexecutor._util._create_kv_cache_manager", + side_effect=create_cross_manager, + ): + creator.build_managers({}, estimating_kv_cache=False) + + assert creator._max_seq_len == 864 + assert captured_cross_max_seq_lens == [2048] + + def test_get_kv_size_per_token_includes_cross_pool_for_enc_dec(self): + config = _make_mock_kv_cache_config( + cross_kv_cache_fraction=0.5, max_gpu_total_bytes=8 * (1 << 30) + ) + model_config = _make_mock_model_config( + is_encoder_decoder=True, + num_hidden_layers=10, + num_attention_heads=16, + num_key_value_heads=16, + hidden_size=768, + head_dim=48, + encoder_attention_heads=12, + decoder_layers=10, + d_model=768, + ) + creator = _make_creator(config, model_config=model_config) + + with patch.object( + creator._kv_cache_manager_cls, + "get_cache_size_per_token", + side_effect=[100, 40], + ) as get_size_mock: + kv_size = creator._get_kv_size_per_token() + + assert kv_size.slope == 140 + assert kv_size.intercept == 0 + assert get_size_mock.call_count == 2 + + cross_call = get_size_mock.call_args_list[1] + proxy_model_config = cross_call.args[0] + assert proxy_model_config.pretrained_config.num_key_value_heads == 12 + assert proxy_model_config.pretrained_config.num_attention_heads == 12 + assert proxy_model_config.pretrained_config.head_dim == 64 + assert cross_call.kwargs["num_layers"] == 10 + + @pytest.mark.parametrize("use_kv_cache_manager_v2", [False, True]) + def test_build_managers_registers_cross_pool_for_enc_dec(self, use_kv_cache_manager_v2): + creator = _make_creator( + _make_mock_kv_cache_config( + cross_kv_cache_fraction=0.5, + max_gpu_total_bytes=8 * (1 << 30), + use_kv_cache_manager_v2=use_kv_cache_manager_v2, + ), + is_enc_dec=True, + ) + creator.configure_kv_cache_capacity = Mock() + creator._should_create_separate_draft_kv_cache = Mock(return_value=False) + creator._split_kv_cache_budget_for_cross = Mock(return_value=(Mock(), Mock())) + creator._create_kv_cache_manager = Mock(return_value=Mock()) + creator._create_cross_kv_cache_manager = Mock(return_value=Mock()) + + resources = {} + creator.build_managers(resources, estimating_kv_cache=False) + + creator._create_cross_kv_cache_manager.assert_called_once() + + @pytest.mark.parametrize("use_kv_cache_manager_v2", [False, True]) + def test_build_managers_registers_cross_pool_for_enc_dec_estimation( + self, use_kv_cache_manager_v2 + ): + creator = _make_creator( + _make_mock_kv_cache_config( + cross_kv_cache_fraction=0.5, + max_gpu_total_bytes=0, + max_tokens=1024, + free_gpu_memory_fraction=0.8, + use_kv_cache_manager_v2=use_kv_cache_manager_v2, + ), + is_enc_dec=True, + ) + creator.configure_kv_cache_capacity = Mock() + creator._should_create_separate_draft_kv_cache = Mock(return_value=False) + creator._split_kv_cache_budget_for_cross = Mock(return_value=(Mock(), Mock())) + creator._create_kv_cache_manager = Mock(return_value=Mock()) + creator._create_cross_kv_cache_manager = Mock(return_value=Mock()) + + resources = {} + creator.build_managers(resources, estimating_kv_cache=True) + + creator._create_cross_kv_cache_manager.assert_called_once() + + @pytest.mark.parametrize("use_kv_cache_manager_v2", [False, True]) + def test_build_managers_uses_split_cross_budget_without_mutating_base_config( + self, use_kv_cache_manager_v2 + ): + total_budget = 10 * (1 << 30) + creator = _make_creator( + _make_mock_kv_cache_config( + cross_kv_cache_fraction=0.5, + max_gpu_total_bytes=total_budget, + free_gpu_memory_fraction=0.9, + use_kv_cache_manager_v2=use_kv_cache_manager_v2, + ), + is_enc_dec=True, + ) + creator.configure_kv_cache_capacity = Mock() + creator._should_create_separate_draft_kv_cache = Mock(return_value=False) + + self_budgets = [] + cross_budgets = [] + + def create_self_manager(*_args, **kwargs): + self_cfg = kwargs["kv_cache_config_override"] + self_budgets.append( + ( + self_cfg.free_gpu_memory_fraction, + self_cfg.max_gpu_total_bytes, + ) + ) + return Mock() + + def create_cross_manager(cross_cfg, *_args, **_kwargs): + cross_budgets.append( + ( + cross_cfg.free_gpu_memory_fraction, + cross_cfg.max_gpu_total_bytes, + ) + ) + return Mock() + + creator._create_kv_cache_manager = Mock(side_effect=create_self_manager) + creator._create_cross_kv_cache_manager = Mock(side_effect=create_cross_manager) + + resources = {} + creator.build_managers(resources, estimating_kv_cache=True) + + assert creator._kv_cache_config.free_gpu_memory_fraction == pytest.approx(0.9) + assert creator._kv_cache_config.max_gpu_total_bytes == total_budget + + creator.build_managers(resources, estimating_kv_cache=False) + + expected_split = total_budget // 2 + assert self_budgets == [ + (pytest.approx(0.45), expected_split), + (pytest.approx(0.45), expected_split), + ] + assert cross_budgets == [ + (pytest.approx(0.45), expected_split), + (pytest.approx(0.45), expected_split), + ] + + def test_build_managers_skips_cross_pool_for_decoder_only(self): + creator = _make_creator( + _make_mock_kv_cache_config( + cross_kv_cache_fraction=None, + max_gpu_total_bytes=8 * (1 << 30), + ), + is_enc_dec=False, + ) + creator.configure_kv_cache_capacity = Mock() + creator._should_create_separate_draft_kv_cache = Mock(return_value=False) + creator._split_kv_cache_budget_for_cross = Mock() + creator._create_kv_cache_manager = Mock(return_value=Mock()) + creator._create_cross_kv_cache_manager = Mock() + + resources = {} + creator.build_managers(resources, estimating_kv_cache=False) + + creator._split_kv_cache_budget_for_cross.assert_not_called() + creator._create_cross_kv_cache_manager.assert_not_called() + assert resources[ResourceManagerType.CROSS_KV_CACHE_MANAGER] is None + + +# --------------------------------------------------------------------------- +# Tests: KVCacheV2Scheduler cross_kv_cache_manager parameter +# --------------------------------------------------------------------------- + + +class TestKVCacheV2SchedulerCrossParam: + """KVCacheV2Scheduler should accept and store cross_kv_cache_manager.""" + + def _make_mock_kv_mgr(self, tokens_per_block=64): + mgr = Mock(spec=KVCacheManagerV2) + mgr.tokens_per_block = tokens_per_block + return mgr + + def test_default_cross_is_none(self): + from tensorrt_llm._torch.pyexecutor.scheduler.scheduler_v2 import KVCacheV2Scheduler + + kv_mgr = self._make_mock_kv_mgr() + scheduler = KVCacheV2Scheduler( + max_batch_size=8, + max_num_tokens=4096, + kv_cache_manager=kv_mgr, + scheduler_policy=CapacitySchedulerPolicy.MAX_UTILIZATION, + ) + assert scheduler.cross_kv_cache_manager is None + + def test_cross_kv_cache_manager_is_stored(self): + from tensorrt_llm._torch.pyexecutor.scheduler.scheduler_v2 import KVCacheV2Scheduler + + kv_mgr = self._make_mock_kv_mgr() + cross_mgr = self._make_mock_kv_mgr() + scheduler = KVCacheV2Scheduler( + max_batch_size=8, + max_num_tokens=4096, + kv_cache_manager=kv_mgr, + scheduler_policy=CapacitySchedulerPolicy.MAX_UTILIZATION, + cross_kv_cache_manager=cross_mgr, + ) + assert scheduler.cross_kv_cache_manager is cross_mgr + + def test_factory_forwards_encoder_init_until_state_for_cross_pool(self): + """The executor factory must widen V2 scheduling to ENCODER_INIT. + + Without this, V2 enc-dec requests are filtered by the default + CONTEXT_INIT state gate before the encoder loop can see them. + """ + from tensorrt_llm._torch.pyexecutor._util import create_py_executor_instance + from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState + + kv_mgr = Mock() + kv_mgr.tokens_per_block = 64 + cross_mgr = Mock() + resources = { + ResourceManagerType.KV_CACHE_MANAGER: kv_mgr, + ResourceManagerType.CROSS_KV_CACHE_MANAGER: cross_mgr, + ResourceManagerType.DRAFT_KV_CACHE_MANAGER: None, + } + mapping = SimpleNamespace( + pp_size=1, + enable_attention_dp=False, + has_pp=lambda: False, + ) + model_engine = SimpleNamespace( + spec_config=None, + model=SimpleNamespace( + model_config=SimpleNamespace( + pretrained_config=SimpleNamespace( + kv_lora_rank=None, + qk_rope_head_dim=None, + ), + ), + ), + ) + llm_args = SimpleNamespace( + extra_resource_managers={}, + disable_overlap_scheduler=True, + enable_early_first_token_response=False, + kv_cache_config=SimpleNamespace(enable_kv_pool_rebalance=False), + ) + + with ( + patch( + "tensorrt_llm._torch.pyexecutor._util.KVCacheManagerV2", + new=Mock, + ), + patch( + "tensorrt_llm._torch.pyexecutor._util.KVCacheV2Scheduler", + ) as scheduler_cls, + patch( + "tensorrt_llm._torch.pyexecutor._util.create_kv_cache_transceiver", + return_value=None, + ), + patch( + "tensorrt_llm._torch.pyexecutor._util.PyExecutor", + ), + ): + scheduler_cls.return_value = Mock() + create_py_executor_instance( + dist=Mock(), + resources=resources, + mapping=mapping, + llm_args=llm_args, + ctx_chunk_config=None, + model_engine=model_engine, + start_worker=False, + sampler=Mock(), + drafter=None, + max_seq_len=128, + max_batch_size=8, + max_beam_width=1, + max_num_tokens=4096, + ) + + kwargs = scheduler_cls.call_args.kwargs + assert kwargs["cross_kv_cache_manager"] is cross_mgr + assert kwargs["no_schedule_until_state"] == LlmRequestState.ENCODER_INIT + + +# --------------------------------------------------------------------------- +# Tests: V1 scheduler cross_kv_cache_manager wiring. +# --------------------------------------------------------------------------- + + +class TestBindCapacitySchedulerCrossParam: + """C++-bound V1 ``BindCapacityScheduler`` exposes cross-KV wiring. + + The C++ ``CapacityScheduler`` already accepts a cross manager. The Python + wrapper forwards the cross pool and the ENCODER_INIT gating. + """ + + def test_default_cross_is_none_and_default_until_state(self): + from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState + from tensorrt_llm._torch.pyexecutor.scheduler.scheduler import BindCapacityScheduler + + with patch( + "tensorrt_llm._torch.pyexecutor.scheduler.scheduler.tb_internal.algorithms.CapacityScheduler" + ) as cap_cls: + cap_cls.return_value = Mock() + scheduler = BindCapacityScheduler( + max_num_requests=8, + kv_cache_manager=Mock(), + peft_cache_manager=None, + scheduler_policy=CapacitySchedulerPolicy.MAX_UTILIZATION, + ) + + assert scheduler.cross_kv_cache_manager is None + kwargs = cap_cls.call_args.kwargs + assert kwargs["no_schedule_until_state"] == LlmRequestState.CONTEXT_INIT + + def test_cross_kv_cache_manager_and_until_state_are_forwarded(self): + from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState + from tensorrt_llm._torch.pyexecutor.scheduler.scheduler import BindCapacityScheduler + + cross_mgr = Mock() + kv_mgr = Mock() + with patch( + "tensorrt_llm._torch.pyexecutor.scheduler.scheduler.tb_internal.algorithms.CapacityScheduler" + ) as cap_cls: + impl = Mock() + cap_cls.return_value = impl + scheduler = BindCapacityScheduler( + max_num_requests=8, + kv_cache_manager=kv_mgr, + peft_cache_manager=None, + scheduler_policy=CapacitySchedulerPolicy.MAX_UTILIZATION, + cross_kv_cache_manager=cross_mgr, + no_schedule_until_state=LlmRequestState.ENCODER_INIT, + ) + + # Construction forwarded the gating to the C++ binding. + ctor_kwargs = cap_cls.call_args.kwargs + assert ctor_kwargs["no_schedule_until_state"] == LlmRequestState.ENCODER_INIT + + # schedule_request must forward the cross manager to the C++ + # __call__ so the dual-pool scheduling logic activates. + impl.return_value = ([], [], []) + scheduler.schedule_request([]) + impl.assert_called_once_with([], kv_mgr, None, cross_mgr) + + +class TestSimpleUnifiedSchedulerCrossParam: + """V1 Python ``SimpleUnifiedScheduler`` exposes cross-KV wiring.""" + + def test_cross_kv_cache_manager_and_until_state_are_forwarded(self): + from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState + from tensorrt_llm._torch.pyexecutor.scheduler.scheduler import SimpleUnifiedScheduler + + kv_mgr = Mock() + kv_mgr.is_variable_window = False + kv_mgr.enable_block_reuse = False + cross_mgr = Mock() + cross_mgr.is_variable_window = False + cross_mgr.enable_block_reuse = False + + scheduler = SimpleUnifiedScheduler( + max_batch_size=8, + max_num_tokens=4096, + kv_cache_manager=kv_mgr, + peft_cache_manager=None, + scheduler_policy=CapacitySchedulerPolicy.MAX_UTILIZATION, + cross_kv_cache_manager=cross_mgr, + no_schedule_until_state=LlmRequestState.ENCODER_INIT, + ) + + assert scheduler.capacity_scheduler.cross_kv_cache_manager is cross_mgr + assert scheduler.capacity_scheduler.no_schedule_until_state == LlmRequestState.ENCODER_INIT + assert ( + scheduler.micro_batch_scheduler.no_schedule_until_state == LlmRequestState.ENCODER_INIT + ) + + +# --------------------------------------------------------------------------- +# Tests: V1 dual-pool smoke test. +# --------------------------------------------------------------------------- + + +class TestV1DualPoolSmoke: + """Smoke test exercising V1 dual-pool construction. + + Constructs both pools as V1 ``KVCacheManager`` instances with + ``CacheType.SELF`` / ``CacheType.CROSS`` (via mocked + ``_create_kv_cache_manager``) and verifies that ``build_managers`` + wires both pools into the resource map for the V1 production path. + + Running an actual encoder + decoder context iteration requires GPUs + and a full model engine; that lives in the integration suite. Here + we verify the V1 construction wiring with mocks consistent with the + rest of this file. + """ + + def test_build_managers_uses_v1_kv_cache_manager_for_both_pools(self): + kv_cache_config = _make_mock_kv_cache_config( + cross_kv_cache_fraction=0.5, + max_gpu_total_bytes=8 * (1 << 30), + use_kv_cache_manager_v2=False, + ) + creator = _make_creator(kv_cache_config, is_enc_dec=True, manager_cls=KVCacheManager) + creator.configure_kv_cache_capacity = Mock() + creator._should_create_separate_draft_kv_cache = Mock(return_value=False) + creator._split_kv_cache_budget_for_cross = Mock(return_value=(Mock(), Mock())) + + # Both _create_kv_cache_manager (self pool) and + # _create_cross_kv_cache_manager are exercised through the + # underlying free-function _create_kv_cache_manager so we can + # assert the manager_cls and CacheType for each call. + import tensorrt_llm + + cache_type_self = tensorrt_llm.bindings.internal.batch_manager.CacheType.SELF + cache_type_cross = tensorrt_llm.bindings.internal.batch_manager.CacheType.CROSS + + # Stub the self-pool path (_create_kv_cache_manager method) to + # avoid invoking the heavyweight free function. + self_mgr = Mock(spec=KVCacheManager) + self_mgr.kv_cache_type = cache_type_self + creator._create_kv_cache_manager = Mock(return_value=self_mgr) + + cross_mgr = Mock(spec=KVCacheManager) + cross_mgr.kv_cache_type = cache_type_cross + with patch( + "tensorrt_llm._torch.pyexecutor._util._create_kv_cache_manager", + return_value=cross_mgr, + ) as create_mock: + resources = {} + creator.build_managers(resources, estimating_kv_cache=False) + + # Self pool: registered as KV_CACHE_MANAGER. + assert resources[ResourceManagerType.KV_CACHE_MANAGER] is self_mgr + + # Cross pool: registered as CROSS_KV_CACHE_MANAGER and built + # with the V1 KVCacheManager class + CacheType.CROSS. + assert resources[ResourceManagerType.CROSS_KV_CACHE_MANAGER] is cross_mgr + cross_kwargs = create_mock.call_args.kwargs + assert cross_kwargs["kv_cache_manager_cls"] is KVCacheManager + assert cross_kwargs["kv_cache_type"] == cache_type_cross diff --git a/tests/unittest/_torch/executor/test_kv_cache_budget_split.py b/tests/unittest/_torch/executor/test_kv_cache_budget_split.py index 3c0cd8f3b491..8aa9ad22196d 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_budget_split.py +++ b/tests/unittest/_torch/executor/test_kv_cache_budget_split.py @@ -66,27 +66,39 @@ class TestSplitGpuBudgetForDraft: def test_gpu_budget_split_proportionally(self): total_gpu = 10 * GB c = _make_creator( - max_gpu_total_bytes=total_gpu, total_kv_per_token=100, target_kv_per_token=80 + max_gpu_total_bytes=total_gpu, + total_kv_per_token=100, + target_kv_per_token=80, ) - draft_config = c._split_kv_cache_budget_for_draft("max_gpu_total_bytes") + target_config, draft_config = c._split_kv_cache_budget_for_draft("max_gpu_total_bytes") assert draft_config is not None - assert c._kv_cache_config.max_gpu_total_bytes == 8 * GB + assert target_config.max_gpu_total_bytes == 8 * GB assert draft_config.max_gpu_total_bytes == 2 * GB + assert target_config.host_cache_size is None + assert c._kv_cache_config.max_gpu_total_bytes == total_gpu assert c._kv_cache_config.host_cache_size is None def test_returns_none_when_no_gpu_budget(self): c = _make_creator(max_gpu_total_bytes=0) - assert c._split_kv_cache_budget_for_draft("max_gpu_total_bytes") is None + target_config, draft_config = c._split_kv_cache_budget_for_draft("max_gpu_total_bytes") + + assert target_config is c._kv_cache_config + assert draft_config is None def test_returns_none_when_draft_kv_zero(self): c = _make_creator( - max_gpu_total_bytes=10 * GB, total_kv_per_token=100, target_kv_per_token=100 + max_gpu_total_bytes=10 * GB, + total_kv_per_token=100, + target_kv_per_token=100, ) - assert c._split_kv_cache_budget_for_draft("max_gpu_total_bytes") is None + target_config, draft_config = c._split_kv_cache_budget_for_draft("max_gpu_total_bytes") + + assert target_config is c._kv_cache_config + assert draft_config is None class TestSplitHostCacheBudgetForDraft: @@ -100,11 +112,13 @@ def test_host_budget_split_proportionally(self): target_kv_per_token=80, ) - draft_config = c._split_kv_cache_budget_for_draft("host_cache_size") + target_config, draft_config = c._split_kv_cache_budget_for_draft("host_cache_size") assert draft_config is not None - assert c._kv_cache_config.host_cache_size == 16 * GB + assert target_config.host_cache_size == 16 * GB assert draft_config.host_cache_size == 4 * GB + assert target_config.max_gpu_total_bytes == total_gpu + assert c._kv_cache_config.host_cache_size == total_host assert c._kv_cache_config.max_gpu_total_bytes == total_gpu def test_host_budget_not_doubled(self): @@ -117,10 +131,10 @@ def test_host_budget_not_doubled(self): target_kv_per_token=80, ) - draft_config = c._split_kv_cache_budget_for_draft("host_cache_size") + target_config, draft_config = c._split_kv_cache_budget_for_draft("host_cache_size") assert draft_config is not None - assert (c._kv_cache_config.host_cache_size + draft_config.host_cache_size) == total_host + assert (target_config.host_cache_size + draft_config.host_cache_size) == total_host def test_host_split_without_gpu_budget_uses_slope_ratio(self): """V1 non-VSWA: host split must not depend on max_gpu_total_bytes.""" @@ -132,10 +146,10 @@ def test_host_split_without_gpu_budget_uses_slope_ratio(self): target_kv_per_token=80, ) - draft_config = c._split_kv_cache_budget_for_draft("host_cache_size") + target_config, draft_config = c._split_kv_cache_budget_for_draft("host_cache_size") assert draft_config is not None - assert c._kv_cache_config.host_cache_size == 16 * GB + assert target_config.host_cache_size == 16 * GB assert draft_config.host_cache_size == 4 * GB def test_host_split_merges_into_existing_draft_config(self): @@ -148,13 +162,18 @@ def test_host_split_merges_into_existing_draft_config(self): target_kv_per_token=80, ) - draft_config = c._split_kv_cache_budget_for_draft("max_gpu_total_bytes") - draft_config = c._split_kv_cache_budget_for_draft("host_cache_size", draft_config) + target_config, draft_config = c._split_kv_cache_budget_for_draft("max_gpu_total_bytes") + target_config, draft_config = c._split_kv_cache_budget_for_draft( + "host_cache_size", target_config, draft_config + ) + assert draft_config is not None assert draft_config.max_gpu_total_bytes == 2 * GB - assert c._kv_cache_config.max_gpu_total_bytes == 8 * GB + assert target_config.max_gpu_total_bytes == 8 * GB assert draft_config.host_cache_size == 4 * GB - assert c._kv_cache_config.host_cache_size == 16 * GB + assert target_config.host_cache_size == 16 * GB + assert c._kv_cache_config.max_gpu_total_bytes == total_gpu + assert c._kv_cache_config.host_cache_size == total_host def test_host_split_after_gpu_split_is_unaffected_by_target_only_gpu_budget(self): """Regression: host split used to read max_gpu_total_bytes (already @@ -170,10 +189,13 @@ def test_host_split_after_gpu_split_is_unaffected_by_target_only_gpu_budget(self target_kv_per_token=80, ) - draft_config = c._split_kv_cache_budget_for_draft("max_gpu_total_bytes") - draft_config = c._split_kv_cache_budget_for_draft("host_cache_size", draft_config) + target_config, draft_config = c._split_kv_cache_budget_for_draft("max_gpu_total_bytes") + target_config, draft_config = c._split_kv_cache_budget_for_draft( + "host_cache_size", target_config, draft_config + ) - assert c._kv_cache_config.host_cache_size == 16 * GB + assert draft_config is not None + assert target_config.host_cache_size == 16 * GB assert draft_config.host_cache_size == 4 * GB def test_no_host_cache_leaves_none(self): @@ -184,7 +206,10 @@ def test_no_host_cache_leaves_none(self): target_kv_per_token=80, ) - assert c._split_kv_cache_budget_for_draft("host_cache_size") is None + target_config, draft_config = c._split_kv_cache_budget_for_draft("host_cache_size") + + assert target_config is c._kv_cache_config + assert draft_config is None def test_zero_host_cache_unchanged(self): c = _make_creator( @@ -194,25 +219,28 @@ def test_zero_host_cache_unchanged(self): target_kv_per_token=80, ) - assert c._split_kv_cache_budget_for_draft("host_cache_size") is None + target_config, draft_config = c._split_kv_cache_budget_for_draft("host_cache_size") + + assert target_config is c._kv_cache_config + assert draft_config is None @pytest.mark.parametrize("target_frac", [0.5, 0.75, 0.9, 0.95]) def test_various_ratios(self, target_frac): - total_gpu = 10 * GB total_host = 20 * GB total_kv = 1000 target_kv = int(total_kv * target_frac) c = _make_creator( - max_gpu_total_bytes=total_gpu, + max_gpu_total_bytes=10 * GB, host_cache_size=total_host, total_kv_per_token=total_kv, target_kv_per_token=target_kv, ) - draft_config = c._split_kv_cache_budget_for_draft("host_cache_size") + target_config, draft_config = c._split_kv_cache_budget_for_draft("host_cache_size") - assert (c._kv_cache_config.host_cache_size + draft_config.host_cache_size) == total_host + assert draft_config is not None + assert (target_config.host_cache_size + draft_config.host_cache_size) == total_host def test_budgets_sum_to_original_with_gpu_and_host(self): total_gpu = 15 * GB @@ -224,20 +252,20 @@ def test_budgets_sum_to_original_with_gpu_and_host(self): target_kv_per_token=700, ) - draft_config = c._split_kv_cache_budget_for_draft("max_gpu_total_bytes") - draft_config = c._split_kv_cache_budget_for_draft("host_cache_size", draft_config) + target_config, draft_config = c._split_kv_cache_budget_for_draft("max_gpu_total_bytes") + target_config, draft_config = c._split_kv_cache_budget_for_draft( + "host_cache_size", target_config, draft_config + ) - assert ( - c._kv_cache_config.max_gpu_total_bytes + draft_config.max_gpu_total_bytes - ) == total_gpu - assert (c._kv_cache_config.host_cache_size + draft_config.host_cache_size) == total_host + assert draft_config is not None + assert (target_config.max_gpu_total_bytes + draft_config.max_gpu_total_bytes) == total_gpu + assert (target_config.host_cache_size + draft_config.host_cache_size) == total_host + assert c._kv_cache_config.max_gpu_total_bytes == total_gpu + assert c._kv_cache_config.host_cache_size == total_host class TestHostSplitIgnoresGpuFixedCost: - """The fixed (intercept) cost models GPU-resident state (e.g. mamba SSM - state) and is not charged against host offload memory. The host split must - therefore stay proportional to the per-token (slope) cost even when the - GPU-resident fixed cost dwarfs the host budget.""" + """The fixed cost models GPU-resident state and is not host memory.""" def test_host_split_proportional_despite_large_intercept(self): total_host = 10 * GB @@ -246,14 +274,13 @@ def test_host_split_proportional_despite_large_intercept(self): host_cache_size=total_host, total_kv_per_token=100, target_kv_per_token=80, - total_kv_intercept=50 * GB, # huge GPU fixed cost, irrelevant to host + total_kv_intercept=50 * GB, ) - draft_config = c._split_kv_cache_budget_for_draft("host_cache_size") + target_config, draft_config = c._split_kv_cache_budget_for_draft("host_cache_size") - # Intercept ignored for host -> proportional on slope (draft 20/100). assert draft_config is not None - assert c._kv_cache_config.host_cache_size == 8 * GB + assert target_config.host_cache_size == 8 * GB assert draft_config.host_cache_size == 2 * GB def test_host_split_sums_to_original_despite_large_intercept(self): @@ -266,14 +293,14 @@ def test_host_split_sums_to_original_despite_large_intercept(self): total_kv_intercept=100 * GB, ) - draft_config = c._split_kv_cache_budget_for_draft("host_cache_size") + target_config, draft_config = c._split_kv_cache_budget_for_draft("host_cache_size") - assert (c._kv_cache_config.host_cache_size + draft_config.host_cache_size) == total_host + assert draft_config is not None + assert (target_config.host_cache_size + draft_config.host_cache_size) == total_host class TestGpuSplitChargesFixedCost: - """``max_gpu_total_bytes`` is where the GPU-resident fixed cost lives, so it - is charged the intercept and fails fast when the budget can't fit it.""" + """``max_gpu_total_bytes`` carries the GPU-resident fixed cost.""" def test_gpu_split_subtracts_intercept(self): total_gpu = 10 * GB @@ -281,44 +308,40 @@ def test_gpu_split_subtracts_intercept(self): max_gpu_total_bytes=total_gpu, total_kv_per_token=100, target_kv_per_token=80, - total_kv_intercept=5 * GB, # draft intercept = 5 - 0 = 5 GB + total_kv_intercept=5 * GB, target_kv_intercept=0, ) - draft_config = c._split_kv_cache_budget_for_draft("max_gpu_total_bytes") + target_config, draft_config = c._split_kv_cache_budget_for_draft("max_gpu_total_bytes") - # slope_budget = 10 - 5 = 5 GB; draft slope share = 5 * 20/100 = 1 GB; - # draft_budget = draft_intercept (5) + 1 = 6 GB; target = 4 GB. + assert draft_config is not None assert draft_config.max_gpu_total_bytes == 6 * GB - assert c._kv_cache_config.max_gpu_total_bytes == 4 * GB + assert target_config.max_gpu_total_bytes == 4 * GB def test_gpu_split_infeasible_raises(self): - """A GPU budget too small for the combined fixed cost is fatal (the run - would OOM), so the split must fail fast instead of degrading.""" - total_gpu = 1 * GB + """A GPU budget too small for fixed cost must fail fast.""" c = _make_creator( - max_gpu_total_bytes=total_gpu, + max_gpu_total_bytes=1 * GB, total_kv_per_token=100, target_kv_per_token=80, - total_kv_intercept=2 * GB, # fixed cost exceeds the gpu budget + total_kv_intercept=2 * GB, ) with pytest.raises(ValueError, match="GPU budget"): c._split_kv_cache_budget_for_draft("max_gpu_total_bytes") def test_gpu_raise_does_not_block_subsequent_host_split(self): - """GPU split raises on infeasible budget, but a host-only split with the - same large intercept still succeeds proportionally.""" total_host = 10 * GB c = _make_creator( - max_gpu_total_bytes=0, # no gpu split attempted + max_gpu_total_bytes=0, host_cache_size=total_host, total_kv_per_token=100, target_kv_per_token=80, total_kv_intercept=2 * GB, ) - draft_config = c._split_kv_cache_budget_for_draft("host_cache_size") + target_config, draft_config = c._split_kv_cache_budget_for_draft("host_cache_size") - assert c._kv_cache_config.host_cache_size == 8 * GB + assert draft_config is not None + assert target_config.host_cache_size == 8 * GB assert draft_config.host_cache_size == 2 * GB