# 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

import torch

import tensorrt_llm
import tensorrt_llm.bindings.executor as trtllm
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

# isort: off
from tensorrt_llm.llmapi.llm_args import (
    CacheTransceiverConfig, CapacitySchedulerPolicy, EagleDecodingConfig,
    KVEventsConfig, KvCacheCompressionConfig, KvCacheConfig, MTPDecodingConfig,
    MultimodalEncoderSchedulingPolicy, PeftCacheConfig, SchedulerConfig,
    SparseAttentionConfig, SpeculativeConfig, TorchLlmArgs, WaitingQueuePolicy,
    disagg_shares_draft_kv)
# 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.mapping import CpType, Mapping
from tensorrt_llm.quantization import QuantAlgo
from tensorrt_llm.runtime.kv_cache_manager_v2 import AttentionReusePolicy

from ..attention.backends import get_sparse_attn_kv_cache_manager
from ..attention.backends.sparse.params import \
    get_indexer_mqa_logits_workspace_bytes
from ..disaggregation.kv_cache_transceiver import (AttentionTypeCpp,
                                                   create_kv_cache_transceiver)
from ..hostfunc import set_low_latency_dispatch
from ..model_config import ModelConfig, _get_kv_cache_layer_specs
from ..models.modeling_multimodal_mixin import MultimodalModelMixin
from ..nccl_window_tensor_scope import discard_nccl_window_tensor_outputs
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 . import config_utils
from .config_utils import (MambaKVCacheParams, _is_sliding_attention_layer,
                           extract_mamba_kv_cache_params,
                           extract_qwen4_exp_ple_cache_params,
                           get_layer_attention_window, is_hybrid_linear,
                           is_kimi_linear, is_mla, is_nemotron_hybrid,
                           is_qwen3_hybrid, is_qwen4_exp, resolve_vocab_size,
                           uses_vswa_kv_cache_layout)
from .connectors.kv_cache_connector import KvCacheConnectorManager
from .dwdp import DwdpManager
from .engine.lora import make_lora_model_config
from .guided_decoder import GuidedDecoder
from .kv_cache.kv_cache_manager_v2 import (KVCacheLayerOwner, KVCacheLayerSpec,
                                           KVCacheManagerV2, Role,
                                           _resolve_v2_max_attention_window_vec)
from .kv_cache.mamba_cache_manager import (BaseMambaCacheManager,
                                           CppMambaHybridCacheManager,
                                           MambaHybridCacheManagerV2,
                                           MixedMambaHybridCacheManager,
                                           use_py_mamba_cache_manager)
from .llm_request import ExecutorResponse, LlmRequestState
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 .sampler.penalties import PenaltyStore
from .scheduler import (BindCapacityScheduler, BindMicroBatchScheduler,
                        KVCacheV2Scheduler, MultimodalEagerEncoderScheduler,
                        MultimodalScheduler, SimpleScheduler,
                        SimpleUnifiedScheduler)
from .seq_slot_manager import SeqSlotManager

GB = 1 << 30

# Page size of a standalone drafter's separate KV pool (DFlash, standalone
# DSpark). Deliberately not kv_cache_config.tokens_per_block: that describes
# the target pool, and target attention backends and deployment configs set it
# to values the drafter's block-decode kernels reject (FP4 MLA forces 128, which
# the trtllm-gen MLA decode does not take). The drafter adopts whatever page
# size its pool is carved with, so the pool picks one every drafter kernel
# accepts. Matches the KvCacheConfig.tokens_per_block default.
_STANDALONE_DRAFT_TOKENS_PER_BLOCK = 32


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 _has_variable_kv_cache_geometry(model_config: ModelConfig) -> bool:
    layer_specs = _get_kv_cache_layer_specs(model_config)
    return bool(layer_specs) and len({(spec.head_dim, spec.num_kv_heads)
                                      for spec in layer_specs}) > 1


def _non_hybrid_kv_cache_manager_cls(model_config: ModelConfig,
                                     kv_cache_config: KvCacheConfig):
    # Variable per-layer geometry requires V2's per-layer buffer sizes.
    needs_v2 = (kv_cache_config.use_kv_cache_manager_v2 is True
                or _has_variable_kv_cache_geometry(model_config))
    return KVCacheManagerV2 if needs_v2 else KVCacheManager


def kv_cache_manager_v2_incompatible_features(
        max_beam_width: Optional[int],
        *,
        enable_sleep: bool = False) -> List[str]:
    """Runtime features a V2 manager cannot serve.

    ``KvCacheCreator._validate_or_fallback_kv_cache_manager_v2`` demotes a plain
    V2 manager to ``KVCacheManager`` when this list is non-empty, and rejects the
    model families that require V2 outright. ``resolved_kv_cache_manager_is_v2``
    reads the same list so that callers sizing pools from the manager version
    cannot disagree with the selection itself.

    A KV connector is deliberately not a trigger: it is served through the pool
    layout registration path and no longer forces a fallback. It is not taken as
    a parameter either, so a future caller cannot reintroduce the demotion by
    passing it.
    """
    # The C++ backend supports beam search. Model-specific restrictions are
    # validated by KvCacheCreator before constructing the manager.
    # V2's native allocations are not tracked by release_with_tag.
    return ["SleepConfig"] if enable_sleep else []


def resolved_kv_cache_manager_is_v2(kv_cache_config: KvCacheConfig,
                                    max_beam_width: Optional[int],
                                    *,
                                    enable_sleep: bool = False) -> bool:
    """Whether the executor will actually hold a V2 manager.

    ``use_kv_cache_manager_v2`` is a request, not the outcome: model loading has
    already resolved ``"auto"``, but a plain model is still demoted to V1 at
    manager-selection time when its runtime features are V2-incompatible. Sizing
    a pool from the request would leave V2 geometry on a V1 executor.
    """
    return (kv_cache_config.use_kv_cache_manager_v2 is True
            and not kv_cache_manager_v2_incompatible_features(
                max_beam_width, enable_sleep=enable_sleep))


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 is_disagg_enabled(
        cache_transceiver_config: Optional[CacheTransceiverConfig]) -> bool:
    """Whether this executor participates in disaggregated serving."""
    return (cache_transceiver_config is not None
            and cache_transceiver_config.backend is not None)


def get_kv_cache_manager_cls(
        model_config: ModelConfig,
        kv_cache_config: KvCacheConfig,
        is_disagg: bool = False,
        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.

    Env-var overrides:
      * ``TRTLLM_USE_PY_MAMBA=1``  — Mixed manager in aggregated serving.
      * ``TLLM_MAMBA_MANAGER_PREFERENCE`` — explicit manager preference.
    """
    config = model_config.pretrained_config
    sparse_attn_config = model_config.sparse_attention_config
    sparse_attn_algorithm = getattr(sparse_attn_config, "algorithm", None)
    quant_config = getattr(model_config, "quant_config", None)
    if (sparse_attn_config is None and is_mla(config)
            and quant_config is not None
            and quant_config.quant_mode.has_fp4_kv_cache()):
        if kv_cache_config.use_kv_cache_manager_v2 is False:
            raise ValueError("FP4 MLA requires use_kv_cache_manager_v2=True.")
        if model_config.attn_backend != "TRTLLM":
            raise ValueError("FP4 MLA requires the TRTLLM attention backend.")
        if is_disagg:
            # Only the Python NIXL transceiver moves the V2 FP4 MLA roles.
            backend, runtime = _resolve_disagg_transceiver_route(
                cache_transceiver_config)
            if runtime != "PYTHON" or backend != "NIXL":
                raise ValueError(
                    "FP4 MLA disaggregated serving requires "
                    "cache_transceiver_config backend='NIXL' with "
                    "transceiver_runtime='PYTHON' (got "
                    f"backend={backend!r}, transceiver_runtime={runtime!r}).")
        if is_hybrid_linear(config):
            raise NotImplementedError(
                "FP4 MLA requires Fp4MlaKVCacheManagerV2, which does not "
                "support hybrid linear-attention models.")
        from ..attention.backends.fp4_mla.cache_manager import \
            Fp4MlaKVCacheManagerV2

        return Fp4MlaKVCacheManagerV2
    use_v2 = kv_cache_config.use_kv_cache_manager_v2 is True
    if 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.
        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)
            return _non_hybrid_kv_cache_manager_cls(model_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.")

        if config_utils.is_glm5_next(config):
            # glm5_next indexer state is a V2 extra buffer: reject any knob
            # that would select a non-V2 manager.
            if use_py_mamba_cache_manager() or os.environ.get(
                    'TLLM_MAMBA_MANAGER_PREFERENCE'):
                raise ValueError(
                    "glm5_next supports only its V2 cache manager; unset "
                    "TRTLLM_USE_PY_MAMBA / TLLM_MAMBA_MANAGER_PREFERENCE.")
            if is_disagg:
                # Only the Python NIXL transceiver moves KDA and indexer state.
                backend, runtime = _resolve_disagg_transceiver_route(
                    cache_transceiver_config)
                if runtime != "PYTHON" or backend != "NIXL":
                    raise ValueError(
                        "glm5_next disaggregated serving requires "
                        "cache_transceiver_config backend='NIXL' with "
                        "transceiver_runtime='PYTHON' (got "
                        f"backend={backend!r}, transceiver_runtime={runtime!r})."
                    )
            if not use_v2:
                raise ValueError(
                    "glm5_next requires KV cache manager V2 (the sparse-layer "
                    "indexer state is a V2 extra buffer). Leave "
                    "kv_cache_config.use_kv_cache_manager_v2='auto' or set it "
                    "to True.")
            from ..attention.backends.sparse.glm_kpool import \
                Glm5NextCacheManager
            return Glm5NextCacheManager

        # Helix x speculation bookkeeping (per-token verify groups on the
        # superblock ledger, py_helix_decode_group_index advancement) exists
        # only in KVCacheManagerV2. The V1-family hybrid managers account
        # helix decode one token per iteration and have no helix-x-spec
        # path, so selecting V1 would run silently wrong.
        if (model_config.mapping is not None
                and model_config.mapping.has_cp_helix()
                and model_config.spec_config is not None and not use_v2):
            raise ValueError(
                "Helix with speculative decoding requires "
                "kv_cache_config.use_kv_cache_manager_v2=True; the V1-family "
                "hybrid managers do not implement per-token verify-group "
                "bookkeeping.")

        # Kimi uses MambaHybridCacheManagerV2 with or without reuse when
        # use_kv_cache_manager_v2=True, the KvCacheConfig default. V1 can
        # be selected explicitly or resolved from "auto". Aggregated V1
        # uses Mixed without reuse and the C++ hybrid manager with reuse.
        # Disaggregated serving requires Python NIXL below. The Python
        # transceiver rejects a V1 manager at creation.
        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.
                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 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.")
            logger.info(
                "Using MixedMambaHybridCacheManager for hybrid mamba model")
            return MixedMambaHybridCacheManager
        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.")
                logger.warning(
                    "Environment variable TLLM_MAMBA_MANAGER_PREFERENCE=MIXED "
                    "overrides the default Mamba cache manager to "
                    "MixedMambaHybridCacheManager.")
                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.")
                logger.warning(
                    "Environment variable TLLM_MAMBA_MANAGER_PREFERENCE=CPP "
                    "overrides the default Mamba cache manager to "
                    "CppMambaHybridCacheManager.")
                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:
        return _non_hybrid_kv_cache_manager_cls(model_config, kv_cache_config)


# --- KV cache cost model ------------------------------------------------------
#
# 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.


@dataclasses.dataclass(frozen=True)
class CacheCost:
    """Affine KV cache budget: ``bytes = slope * tokens + intercept``.

    The legacy proportional case is just ``intercept = 0``.

    ``bytes_per_request`` is an optional estimate of the slope-funded
    per-request spend (``get_cache_bytes_per_request``): the worst-case bytes
    one request pins beyond what the intercept already reserves. It refines
    how the target/draft budget split weighs the two managers.
    ``quota_per_request`` is the optional allocator quota one request makes
    the manager insist on (``get_cache_quota_per_request``), which floors and
    caps the draft share. Neither is part of the affine cost model, so
    arithmetic (``__add__``) drops them.
    """
    slope: int
    intercept: int = 0
    bytes_per_request: Optional[int] = None
    quota_per_request: Optional[int] = None

    @classmethod
    def from_raw(cls, raw) -> "CacheCost":
        """Wrap an int / tuple / CacheCost result uniformly."""
        if isinstance(raw, CacheCost):
            return raw
        if isinstance(raw, tuple):
            slope, intercept = raw
            return cls(slope=int(slope), intercept=int(intercept))
        return cls(slope=int(raw))

    def __add__(self, other: "CacheCost") -> "CacheCost":
        if not isinstance(other, CacheCost):
            return NotImplemented
        return CacheCost(slope=self.slope + other.slope,
                         intercept=self.intercept + other.intercept)

    def __str__(self) -> str:
        text = f"{self.slope} bytes/token"
        if self.intercept != 0:
            text += f" + {self.intercept} bytes fixed cost"
        if self.bytes_per_request is not None:
            text += f", {self.bytes_per_request} bytes/request"
        if self.quota_per_request is not None:
            text += f", {self.quota_per_request} quota bytes/request"
        return text

    def tokens_for_budget(self, budget: int) -> int:
        """Memory budget -> max tokens. Clamps a negative result to 0."""
        if self.slope <= 0:
            return 0
        tokens = max((budget - self.intercept) // self.slope, 0)
        return tokens

    def bytes_for_tokens(self, tokens: int) -> int:
        """Token count -> memory bytes."""
        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_attention_workspace_fixed_bytes(model_config, mapping,
                                        **capacities) -> int:
    """Capacity-sized reserve declared by the model's selected backend."""
    from ..attention.backends.utils import get_attention_backend
    return get_attention_backend(
        model_config.attn_backend).runtime_workspace_fixed_bytes(
            model_config, mapping, **capacities)


def get_overlap_headroom_penalty_workspace_bytes(
        max_num_sequences: int, num_sequences_without_headroom: int,
        max_beam_width: int, vocab_size: Optional[int]) -> int:
    """Bytes the overlap headroom adds to the sampler's penalty workspace.

    ``PenaltyStore`` allocates its vocab-sized occurrence workspace -- int32
    ``counts_cuda`` with one row per beam of every seat, plus the bool
    ``presence_prefix_cuda`` with one row per seat -- only when the first
    penalized request arrives. That is after the KV cache has taken its share of
    free memory, and the profiling requests carry no penalties, so the
    estimation never measures it. The seats a configuration has without the
    overlap headroom come out of the ``free_gpu_memory_fraction`` slack, as they
    always have. The extra seats the headroom adds are reserved here instead.

    ``max_num_sequences`` and ``max_beam_width`` are the ``PenaltyStore``'s own
    dimensions; ``num_sequences_without_headroom`` is the seat pool the same
    configuration gets with the headroom off. Returns 0 when the store has no
    extra seats or when the vocabulary size is unknown.
    """
    extra_seats = max_num_sequences - num_sequences_without_headroom
    if extra_seats <= 0 or not vocab_size:
        return 0
    counts_bytes = extra_seats * max_beam_width * vocab_size * torch.int32.itemsize
    presence_prefix_bytes = extra_seats * vocab_size * torch.bool.itemsize
    return counts_bytes + presence_prefix_bytes


def get_penalty_store(sampler) -> Optional[PenaltyStore]:
    """Return the sampler's ``PenaltyStore``, or None if it does not own one."""
    penalty_handler = getattr(sampler, "_penalty_handler", None)
    store = getattr(penalty_handler, "store", None)
    return store if isinstance(store, PenaltyStore) else None


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,
        chunked_workspace_profiled=True,
        require_chunked_workspace_profile=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. A
    bounded chunk must also be exercised by profiling before dropping its
    reservation (``chunked_workspace_profiled``) when
    ``require_chunked_workspace_profile`` is enabled. Older hardware keeps
    the existing backend-declared reserve policy without this new requirement.

    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 or
             (require_chunked_workspace_profile
              and not chunked_workspace_profiled)))
    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_indexer_mqa_logits_workspace_reserve(kv_cache_memory: int,
                                             workspace_memory: int) -> int:
    """Bytes to carve out of the KV cache budget for the Indexer MQA-logits tile.

    The profiling forward does not reach the largest logits tile the runtime can
    allocate (fresh-prefill dummies attend only their own tokens), so reserve the
    full reachable tile. Returns 0 if either input is non-positive, and raises if
    the tile would leave no KV cache budget.
    """
    if kv_cache_memory <= 0 or workspace_memory <= 0:
        return 0
    if workspace_memory >= kv_cache_memory:
        raise ValueError(
            f"The DeepSeek-V4 Indexer MQA-logits workspace ({workspace_memory} bytes) "
            f"needs the whole KV cache budget ({kv_cache_memory} bytes). Lower "
            "TLLM_INDEXER_MQA_LOGITS_ELEM_BUDGET or max_num_tokens, or raise "
            "kv_cache_config.free_gpu_memory_fraction.")
    return workspace_memory


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 _derive_layer_type_attention_windows(
    config: object,
    max_seq_len: int,
) -> Optional[List[int]]:
    """Per-layer attention windows for a mixed sliding/full-attention config.

    Returns one window per decoder layer, in global model layer order
    (sliding layers get the model's ``sliding_window``, full-attention layers
    ``max_seq_len``), or ``None`` when the schedule is uniform and the
    single-window default is already right. Layer types that
    ``get_layer_attention_window`` does not recognize as sliding are treated
    as full attention.
    """
    num_layers = getattr(config, "num_hidden_layers", None)
    if not isinstance(num_layers, int) or num_layers <= 0:
        return None
    if not getattr(config, "layer_types", None):
        return None
    try:
        windows = [
            get_layer_attention_window(config, layer_idx)
            for layer_idx in range(num_layers)
        ]
    except (NotImplementedError, ValueError) as error:
        logger.warning(
            "Unable to derive per-layer attention windows from layer_types "
            f"({error}); falling back to the single-window default.")
        return None
    resolved = _normalize_attention_windows(
        [max_seq_len if window is None else window for window in windows],
        max_seq_len)
    if resolved is None or not uses_vswa_kv_cache_layout(resolved):
        # A schedule with <2 distinct positive windows can't select a VSWA
        # (multi-pool) layout, so a derived vector would only reshape the single
        # shared pool. Leave it to the single-window default.
        return None
    return resolved


def _derive_v2_layer_type_attention_windows(
    kv_cache_config: KvCacheConfig,
    kv_cache_manager_cls: type,
    model_config: ModelConfig,
    max_seq_len: int,
) -> Optional[List[int]]:
    """Per-layer windows a KVCacheManagerV2 for `model_config` is built with.

    Shared by the static per-token cost model
    (`KvCacheCreator._per_manager_cache_cost`) and `_create_kv_cache_manager`,
    so the budget split sizes a manager from the same windows the manager
    receives. Returns `None`, leaving `kv_cache_config` as is, when the user
    supplied `max_attention_window`, for any class other than KVCacheManagerV2
    (KVCacheManager keeps the single-window default), when
    `_derive_layer_type_attention_windows` has nothing to derive, for hybrid
    linear-attention configs (their attention layers are interleaved with
    recurrent layers, and the static cost model indexes a window list by
    attention-layer position while the manager indexes it by global layer id,
    so the two would read a per-layer list differently), and when a
    user-supplied `pool_ratio` does not carry one entry per derived layer
    group (the manager would reject that arity at construction; a warning
    names the fix and the configuration keeps the single pool it was written
    for).
    """
    if kv_cache_config.max_attention_window is not None:
        return None
    if not (isinstance(kv_cache_manager_cls, type)
            and issubclass(kv_cache_manager_cls, KVCacheManagerV2)):
        return None
    config = model_config.pretrained_config
    layer_specs = _get_kv_cache_layer_specs(model_config)
    if layer_specs and any(spec.attention_window is not None
                           for spec in layer_specs):
        derived_windows = _normalize_attention_windows(
            [spec.attention_window for spec in layer_specs], max_seq_len)
        if derived_windows is not None and not uses_vswa_kv_cache_layout(
                derived_windows):
            derived_windows = None
    else:
        derived_windows = _derive_layer_type_attention_windows(
            config, max_seq_len)
    if derived_windows is None:
        return None
    if is_hybrid_linear(config):
        return None
    pool_ratio = kv_cache_config.pool_ratio
    num_layer_groups = len(set(derived_windows))
    if pool_ratio is not None and len(pool_ratio) != num_layer_groups:
        logger.warning_once(
            f"kv_cache_config.pool_ratio has {len(pool_ratio)} entries, but the "
            "per-layer attention windows derived from layer_types form "
            f"{num_layer_groups} layer groups; keeping the single-window "
            "default. Provide one pool_ratio entry per layer group to size the "
            "sliding-window and full-attention pools separately, or set "
            "max_attention_window explicitly.",
            key="derived_attention_windows_pool_ratio_arity")
        return None
    return derived_windows


def _get_num_pool_groups_for_estimation(
    model_config: ModelConfig,
    max_seq_len: int,
    kv_cache_config: KvCacheConfig,
) -> int:
    """Estimate pool groups for the temporary V2 profiling manager.

    Hybrid recurrent/attention models also retain their distinct layer types
    when automatic window derivation does not apply.
    """
    config = model_config.pretrained_config
    attention_windows = kv_cache_config.max_attention_window
    if attention_windows is None and _has_variable_kv_cache_geometry(
            model_config):
        attention_windows = _derive_v2_layer_type_attention_windows(
            kv_cache_config, KVCacheManagerV2, model_config, max_seq_len)
    if attention_windows is not None and not is_hybrid_linear(config):
        normalized_windows = _normalize_attention_windows(
            attention_windows, max_seq_len)
        return 1 if normalized_windows is None else len(set(normalized_windows))

    layer_types = getattr(config, "layer_types", None)
    if isinstance(layer_types, (list, tuple)):
        # Without distinct storage windows, sliding/full computation and
        # different page geometries do not create separate window groups.
        pool_types = set()
        for layer_type in layer_types:
            if attention_windows is None and (
                    _is_sliding_attention_layer(layer_type)
                    or getattr(layer_type, "name",
                               str(layer_type)).lower() == "full_attention"):
                layer_type = "full_attention"
            pool_types.add(layer_type)
        if len(pool_types) > 1:
            return len(pool_types)

    if attention_windows is not None:
        normalized_windows = _normalize_attention_windows(
            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


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,
        *,
        model_engine: PyTorchModelEngine,
        mapping: Mapping,
        net_max_seq_len: int,
        kv_connector_manager: Optional[KvCacheConnectorManager],
        max_num_tokens: int,
        max_beam_width: int,
        tokens_per_block: int,
        max_seq_len: int,
        max_batch_size: int,
        kv_cache_config: KvCacheConfig,
        llm_args: TorchLlmArgs,
        speculative_config: SpeculativeConfig,
        sparse_attention_config: SparseAttentionConfig,
        profiling_stage_data: Optional[dict],
        is_disagg: bool,
        execution_stream: Optional[torch.cuda.Stream] = None,
        draft_config: Optional[ModelConfig] = None,
        skip_est: bool = False,
    ):
        self._model_engine = model_engine
        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
        self._speculative_config = speculative_config
        self._sparse_attention_config = sparse_attention_config
        self._tokens_per_block = tokens_per_block
        self._max_seq_len = max_seq_len
        # Estimation can shrink _max_seq_len; cache ownership must stay fixed.
        self._unification_max_seq_len = max_seq_len
        self._max_batch_size = max_batch_size
        self._net_max_seq_len = net_max_seq_len
        self._dummy_reqs = None
        self._mla_chunked_profile_length: int | None = 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._disable_overlap_scheduler = llm_args.disable_overlap_scheduler
        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

    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)
        model_config = model_engine.model.model_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
            quant_config = getattr(model_config, "quant_config", None)
            fp4_mla = (sparse_attn_config is None and is_mla(config)
                       and quant_config is not None
                       and quant_config.quant_mode.has_fp4_kv_cache())
            incompat = kv_cache_manager_v2_incompatible_features(
                self._max_beam_width,
                enable_sleep=self._llm_args.sleep_config is not None)
            # Sparse attention: ModelEngine only forwards cache_indirection when
            # the metadata type is exactly TrtllmAttentionMetadata, and every
            # sparse backend uses a subclass, so beams would read beam 0's
            # unmapped prompt rows. The sparse managers' own block tables
            # (indexer K-cache, pool block indices) are beam-0 only as well.
            # Disaggregated serving transfers the shared prompt in beam 0.
            # The C++ V2 cache expands beams after receive completion, copying
            # the prompt's partial tail before the first generation step.
            if (self._max_beam_width is not None and self._max_beam_width > 1
                    and (is_hybrid_linear(config)
                         or sparse_attn_config is not None or fp4_mla)
                    and "max_beam_width > 1" not in incompat):
                incompat.append("max_beam_width > 1")
            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.
                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.")
                if _has_variable_kv_cache_geometry(model_config):
                    raise NotImplementedError(
                        f"Variable per-layer KV geometry requires KVCacheManagerV2, "
                        f"which is not yet supported with {incompat_str}. "
                        f"Disable these features to run this model.")
                if fp4_mla:
                    raise NotImplementedError(
                        "FP4 MLA requires Fp4MlaKVCacheManagerV2, which is "
                        f"not yet supported with {incompat_str}. Disable these "
                        "features to run FP4 MLA.")
                if is_hybrid_linear(config):
                    raise NotImplementedError(
                        "Hybrid Mamba KVCacheManagerV2 does not support "
                        f"{incompat_str}; automatic fallback is unavailable "
                        "for hybrid linear models. Disable the "
                        "incompatible features to run hybrid linear models.")
                if kv_cache_config is None:
                    kv_cache_config = self._kv_cache_config
                if kv_cache_config.block_reuse_config.policy != "all_reusable":
                    raise NotImplementedError(
                        "Non-default block_reuse_config.policy requires "
                        "KVCacheManagerV2, which is not supported with "
                        f"{incompat_str}.")
                if kv_cache_config.block_reuse_config.swa_endpoint_rewind_tokens > 0:
                    raise NotImplementedError(
                        "block_reuse_config.swa_endpoint_rewind_tokens requires "
                        "KVCacheManagerV2, which is not supported with "
                        f"{incompat_str}.")
                # 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.
                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))

    def _per_manager_cache_cost(self,
                                manager_cls,
                                model_config,
                                kv_cache_config: Optional[KvCacheConfig] = None,
                                *,
                                is_draft: bool = False,
                                mapping=None,
                                tokens_per_block: Optional[int] = None,
                                include_request_cost: bool = False,
                                **extra_kwargs) -> CacheCost:
        kv_cache_config = (kv_cache_config if kv_cache_config is not None else
                           self._kv_cache_config)
        if tokens_per_block is None:
            tokens_per_block = self._tokens_per_block
        if not is_draft and self._should_unify_dflash_kv_cache():
            extra_kwargs = dict(extra_kwargs)
            extra_kwargs[
                "layer_specs"] = self._build_unified_dflash_layer_specs(
                    kv_cache_config)
        if not is_draft:
            derived_windows = _derive_v2_layer_type_attention_windows(
                kv_cache_config, manager_cls, model_config, self._max_seq_len)
            if derived_windows is not None:
                # Cost the manager with the per-layer windows
                # `_create_kv_cache_manager` builds it with, so the budget
                # split between the target and draft managers matches their
                # pools. The draft manager never derives (see there).
                kv_cache_config = kv_cache_config.model_copy(
                    update={"max_attention_window": derived_windows})
        cost = CacheCost.from_raw(
            manager_cls.get_cache_size_per_token(
                model_config,
                mapping if mapping is not None else self._mapping,
                tokens_per_block=tokens_per_block,
                max_seq_len=self._max_seq_len,
                max_batch_size=self._max_batch_size,
                max_num_tokens=self._max_num_tokens,
                kv_cache_config=kv_cache_config,
                spec_config=self._speculative_config,
                is_draft=is_draft,
                is_disagg=self._is_disagg,
                **extra_kwargs))
        if include_request_cost:
            if cost.bytes_per_request is None:
                bytes_per_request = self._per_manager_request_estimate(
                    manager_cls,
                    "get_cache_bytes_per_request",
                    model_config,
                    kv_cache_config,
                    mapping=mapping,
                    tokens_per_block=tokens_per_block,
                    is_draft=is_draft,
                    # Zero is a valid weight: a manager whose entire retention
                    # is intercept-funded spends nothing per request from the
                    # slope budget.
                    allow_zero=True,
                    **extra_kwargs)
                if bytes_per_request is not None:
                    cost = dataclasses.replace(
                        cost, bytes_per_request=bytes_per_request)
            if cost.quota_per_request is None:
                quota_per_request = self._per_manager_request_estimate(
                    manager_cls,
                    "get_cache_quota_per_request",
                    model_config,
                    kv_cache_config,
                    mapping=mapping,
                    tokens_per_block=tokens_per_block,
                    is_draft=is_draft,
                    **extra_kwargs)
                if quota_per_request is not None:
                    cost = dataclasses.replace(
                        cost, quota_per_request=quota_per_request)
        return cost

    def _per_manager_request_estimate(self,
                                      manager_cls,
                                      method_name: str,
                                      model_config,
                                      kv_cache_config: KvCacheConfig,
                                      *,
                                      mapping=None,
                                      tokens_per_block: Optional[int] = None,
                                      is_draft: bool = False,
                                      allow_zero: bool = False,
                                      **extra_kwargs) -> Optional[int]:
        """Per-request byte estimate from a manager classmethod, or None.

        None when the manager class does not model per-request spend (only
        ``KVCacheManagerV2`` and subclasses do) or declines (see the manager
        method's docstring for when).
        """
        estimate_method = getattr(manager_cls, method_name, None)
        if estimate_method is None:
            return None
        estimate = estimate_method(
            model_config,
            mapping if mapping is not None else self._mapping,
            tokens_per_block=(tokens_per_block if tokens_per_block is not None
                              else 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,
            kv_cache_config=kv_cache_config,
            spec_config=self._speculative_config,
            is_draft=is_draft,
            **extra_kwargs)
        # Mocked manager classes in tests return non-int sentinels; treat
        # anything but a usable int as "no estimate".
        if isinstance(estimate, int) and (estimate > 0 or
                                          (allow_zero and estimate == 0)):
            return estimate
        return None

    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:
        """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
        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)
        if self._is_encoder_decoder():
            total += CacheCost.from_raw(self._get_cross_kv_size_per_token())
        draft_cost = self._get_draft_cache_cost(
            kv_cache_config,
            use_separate_draft_kv_cache=use_separate_draft_kv_cache,
        )
        if draft_cost is not None:
            total += draft_cost
        return total

    def _get_draft_cache_cost(
        self,
        kv_cache_config: KvCacheConfig,
        *,
        use_separate_draft_kv_cache: bool,
        include_request_cost: bool = False,
    ) -> Optional[CacheCost]:
        """Return the draft manager's standalone cache cost, if it has one.

        Under helix CP the drafter is dense rather than helix-sharded, so it is
        costed with the same repurposed mapping runtime construction uses, then
        the slope is multiplied by cp_size to express it per rank-LOCAL target
        token (the target stores only every cp_size-th page per rank).
        Intercepts are per-request rank-local bytes and stay unscaled.
        """
        draft_mapping = self._mapping
        helix_cp_scale = 1
        if self._mapping.has_cp_helix():
            draft_mapping = self._mapping.repurpose_helix_cp_to_tp()
            helix_cp_scale = self._mapping.cp_size

        def scaled(cost: CacheCost) -> CacheCost:
            # Rescale only the slope: every other cost component (e.g. the
            # fixed intercept) is per-request rank-local bytes and stays
            # unscaled. dataclasses.replace keeps any component the rescale
            # does not touch, where rebuilding the CacheCost from named
            # fields would silently drop it.
            return dataclasses.replace(cost, slope=cost.slope * helix_cp_scale)

        if use_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
            # (e.g. EAGLE3: config says 1, runtime uses 4).
            # For PP, draft layers are only on the last rank (see
            # get_pp_layers), so only that rank should include draft cost.
            # _get_draft_kv_model_config(), not _get_effective_draft_config():
            # the cost charged here must be the cost of the pool that
            # _create_one_model_draft_kv_cache_manager actually allocates.
            effective_draft_config = self._get_draft_kv_model_config()
            draft_kv_cache_config = self._get_one_model_draft_kv_cache_config(
                kv_cache_config, self._max_seq_len)
            # Charge the page size the pool is allocated with, not the
            # target's: per-request quotas round up to whole pages.
            draft_tokens_per_block = self._get_draft_tokens_per_block()
            # 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,
                is_disagg=self._is_disagg,
                cache_transceiver_config=self._cache_transceiver_config)
            draft_kv_cache_manager_cls = self._validate_or_fallback_kv_cache_manager_v2(
                draft_kv_cache_manager_cls, effective_draft_config,
                draft_kv_cache_config)
            if self._speculative_config.spec_dec_mode.is_external_drafter():
                # External drafter: layers start from 0, normal PP distribution
                return scaled(
                    self._per_manager_cache_cost(
                        draft_kv_cache_manager_cls,
                        effective_draft_config,
                        draft_kv_cache_config,
                        mapping=draft_mapping,
                        tokens_per_block=draft_tokens_per_block,
                        is_draft=True,
                        include_request_cost=include_request_cost))
            elif self._mapping.is_last_pp_rank():
                # EAGLE3/MTP: draft layers only on last PP rank
                return scaled(
                    self._per_manager_cache_cost(
                        draft_kv_cache_manager_cls,
                        effective_draft_config,
                        draft_kv_cache_config,
                        mapping=draft_mapping,
                        tokens_per_block=draft_tokens_per_block,
                        num_layers=self._get_num_draft_layers(),
                        is_draft=True,
                        include_request_cost=include_request_cost))
        return None

    def _cal_max_memory(self, peak_memory, total_gpu_memory, fraction,
                        allocated_bytes: int) -> int:
        """
        Calculate the max KV cache capacity.

        NOTE: `allocated_bytes` is the total KV-cache memory that must be pre-allocated during the estimation phase (for both the main and draft models) so the estimation run can complete successfully. When computing `available_kv_mem`, add this amount back in.
        """
        kv_size_per_token = self._get_kv_size_per_token()

        available_kv_mem = (total_gpu_memory - peak_memory +
                            allocated_bytes) * fraction
        logger.info(
            f"Peak memory during memory usage profiling (torch + non-torch): {peak_memory / (GB):.2f} GiB, "
            f"available KV cache memory when calculating max tokens: {available_kv_mem / (GB):.2f} GiB, "
            f"fraction is set {fraction}, kv size per token is {kv_size_per_token}. device total memory {total_gpu_memory / (GB):.2f} GiB, "
            f"temporary kv cache memory during profiling {allocated_bytes / (GB):.2f} GiB"
        )
        return int(available_kv_mem)

    def _get_mla_chunked_profile_length(self, input_seq_len: int) -> int | None:
        """Length needed to profile two full cached-KV chunks with a full query."""
        model_config = self._model_engine.model.model_config
        # Skip-softmax preserves dense MLA. DSA-style hooks can switch to
        # absorption before this request reaches a full cached-KV chunk.
        if (not is_mla(model_config.pretrained_config)
                or model_config.attn_backend != "TRTLLM"
                or getattr(model_config.sparse_attention_config, "algorithm",
                           None) not in (None, "skip_softmax")):
            return None
        # Beam search and speculative decoding use this same dense context
        # path. Their extra KV capacity is accounted for by
        # _get_token_num_for_estimation; neither needs a separate exclusion.
        features = self._model_engine.attn_runtime_features
        if not features.chunked_prefill or features.chunk_size <= 0:
            return None
        # MLA.forward_context uses full-gather on Hopper even when scheduler
        # chunking is enabled. Only the SM100+ path bounds cached-KV staging.
        if get_sm_version() < 100:
            return None

        kv_chunk_tokens = (features.chunk_size *
                           features.chunked_prefill_buffer_batch_size)
        # The previous loop's K/V tensors can remain live while the next
        # chunk is expanded. Exercise two full cached-KV chunks in one forward
        # to include this overlap, alongside a full query budget. Round the
        # prefix up to a scheduler-step boundary so the final query is full.
        cached_tokens = (ceil_div(2 * kv_chunk_tokens, self._max_num_tokens) *
                         self._max_num_tokens)
        profile_length = cached_tokens + self._max_num_tokens
        if profile_length > input_seq_len:
            # A short-context workload may fill the KV chunk through fan-out;
            # a single long request cannot cover that case. Keep its reserve.
            return None
        return profile_length

    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 = []
        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

        self._mla_chunked_profile_length = self._get_mla_chunked_profile_length(
            input_seq_len)
        if self._mla_chunked_profile_length is not None:
            input_seq_len = self._mla_chunked_profile_length
            remaining_tokens = input_seq_len
            logger.info(
                "Profiling chunked MLA with a cached prefix: "
                f"prompt length {input_seq_len}, query budget {max_num_tokens}."
            )
        else:
            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)
            input_tokens = torch.randint(low=0,
                                         high=vocab_size,
                                         size=(input_seq_len, )).tolist()
            request = trtllm.Request(input_tokens,
                                     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)
                    }
                }
            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 equivalent store output."""
        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,
                        ),
                    )
                # The clone is retained only for memory accounting; no tensor
                # data is consumed after this profiling forward returns.
                with discard_nccl_window_tensor_outputs(encoder_inputs):
                    output = self._model_engine.model.encode_multimodal_inputs(
                        encoder_inputs)
                    # The runtime store owns detached copies rather than views
                    # of an encoder batch. Reproduce that allocation boundary.
                    return output.detach().clone()
        finally:
            self._dummy_encoder_inputs = []

    def _get_multimodal_encoder_memory_reserve(self,
                                               profiled_output_bytes: int = 0
                                               ) -> int:
        """Reserve the unprofiled encoder store and chunk-assembly copies."""
        encoder_cache = self._model_engine.mm_encoder_cache
        if encoder_cache is None:
            return 0
        reserve = encoder_cache.max_bytes
        if self._model_engine.mm_encoder_item_scheduling_enabled:
            # Item outputs live only in the store, so the retained profiling
            # output already occupies part of it. Inline encoding keeps its
            # batch output beside the store's copies; the retained output
            # stands in for that batch, so the whole store stays unprofiled.
            reserve = max(0, reserve - profiled_output_bytes)
            # The text-only dummy does not assemble cached MM rows for LLM
            # prefill. Per-request joins and their batch-wide join can coexist.
            # Both contain only current-chunk rows, bounded by the LLM token
            # budget even when many requests reuse the same cached item.
            chunk_output_bytes = (
                self._max_num_tokens *
                self._model_engine.bytes_per_mm_encoder_embedding)
            reserve += 2 * chunk_output_bytes
        return reserve

    def _get_penalty_workspace_reserve(self, py_executor) -> int:
        """Reserve the overlap-headroom seats' penalty workspace.

        Sized from the profiling executor's own ``PenaltyStore``, so it follows
        the sampler that will serve. See
        ``get_overlap_headroom_penalty_workspace_bytes``.
        """
        store = get_penalty_store(getattr(py_executor, "sampler", None))
        if store is None:
            return 0
        return get_overlap_headroom_penalty_workspace_bytes(
            max_num_sequences=store.max_num_sequences,
            num_sequences_without_headroom=compute_max_num_sequences(
                self._mapping,
                self._max_batch_size,
                self._llm_args.disable_overlap_scheduler,
                enable_overlap_headroom=False),
            max_beam_width=store.max_beam_width,
            vocab_size=resolve_vocab_size(
                self._model_engine.model.model_config.pretrained_config))

    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:
            raise ValueError(
                "KV cache size estimation not supported with context parallelism."
            )
        # estimate_max_kv_cache_tokens submits self._dummy_reqs
        num_cache_blocks = 0
        num_extra_tokens_per_seq = 1  # account for generated tokens
        spec_cfg = self._speculative_config
        if not self._llm_args.disable_overlap_scheduler and spec_cfg is not None:
            num_extra_tokens_per_seq += spec_cfg.tokens_per_gen_step - 1

        if spec_cfg is not None:
            num_extra_tokens_per_seq += spec_cfg.tokens_per_gen_step - 1
            num_extra_tokens_per_seq += get_num_extra_kv_tokens(spec_cfg)

        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.
            num_cache_blocks += ceil_div(num_req_tokens, self._tokens_per_block)

        # With ADP enabled, _create_dummy_context_requests produces tp_size
        # copies so each rank gets work during the estimation warmup. But the
        # scheduler distributes them evenly (1 per rank), so each rank's KV
        # cache only needs capacity for its own share, not all of them.
        if self._mapping.enable_attention_dp and self._mapping.tp_size > 1:
            num_cache_blocks = (num_cache_blocks + self._mapping.tp_size -
                                1) // self._mapping.tp_size

        # Max cuda graph warmup required tokens
        max_cuda_graph_bs = min(self._model_engine.batch_size,
                                self._model_engine._max_cuda_graph_batch_size)
        # Round up the max seq len to the block size
        max_seq_len_blocks = ceil_div(self._model_engine.max_seq_len + 1,
                                      self._tokens_per_block)
        cuda_graph_warmup_block = max_seq_len_blocks + max_cuda_graph_bs - 1
        num_cache_blocks = max(cuda_graph_warmup_block, num_cache_blocks)

        # This is the minimal blocks required to run with max bs
        # 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).
        num_pool_groups = 1
        if self._is_kv_cache_manager_v2:
            model_cfg = self._model_engine.model.model_config
            num_pool_groups = _get_num_pool_groups_for_estimation(
                model_cfg,
                self._model_engine.max_seq_len,
                self._kv_cache_config,
            )
        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()
        max_memory = self._kv_cache_config.free_gpu_memory_fraction * free_mem
        kv_size_per_token = self._get_kv_size_per_token()
        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)

    def try_prepare_estimation(self) -> bool:
        """Prepare for possible KV cache capacity estimation.

        This updates `kv_cache_config` and returns a boolean indicating whether KV cache
        estimation is to be performend.
        """
        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
        model_config = self._model_engine.model.model_config
        if model_config.attn_backend == "VANILLA":
            # Make build_managers call configure_kv_cache_capacity() for
            # the V2 memory quota. Non-HELIX CP has no sizing path there;
            # HELIX is handled above.
            if (self._is_kv_cache_manager_v2
                    and 'cp_type' not in self._mapping.cp_config):
                self._skip_est = True
            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:
            # User-provided pool sizing can underprovision the temporary
            # estimation cache and cause warmup to hang or fail. Override it
            # before budgeting, then restore it in configure_kv_cache_capacity().
            self._kv_cache_config.pool_ratio = None
            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
            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.
        NOTE: for VSWA case, we calculate and set kv cache memory instead of using max_tokens in kv_cache_config.

        This updates `kv_cache_config`.
        """
        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

        torch.cuda.empty_cache()
        torch.cuda.reset_peak_memory_stats()
        end, total_gpu_memory = torch.cuda.mem_get_info()
        total_used_bytes = total_gpu_memory - end
        model_bytes = torch.cuda.memory_stats()["allocated_bytes.all.current"]
        logger.info(
            f"Memory used after loading model weights (inside torch) in memory usage profiling: {model_bytes / (GB):.2f} GiB"
        )
        logger.info(
            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
            # equivalent runtime-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
            req_ids = []
            if py_executor.dist.mapping.rank == 0:
                req_ids = py_executor.enqueue_requests(self._dummy_reqs)
            req_ids = py_executor.dist.broadcast(req_ids, root=0)
            py_executor.is_warmup = True
            py_executor.start_worker()
            try:
                responses = py_executor.await_responses(req_ids)
                for response_or_list in responses:
                    response_list = [response_or_list] if isinstance(
                        response_or_list,
                        ExecutorResponse) else response_or_list
                    for response in response_list:
                        if response.has_error():
                            raise RuntimeError(response.error_msg)

                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()
                # Get draft KV cache stats if present (one-model mode with a separate draft KV cache)
                draft_kv_cache_manager = py_executor.resource_manager.resource_managers.get(
                    ResourceManagerType.DRAFT_KV_CACHE_MANAGER)
                kv_stats_draft = draft_kv_cache_manager.get_kv_cache_stats(
                ) if draft_kv_cache_manager is not None else None

                # get total allocated bytes
                allocated_bytes = kv_stats.allocated_bytes + (
                    kv_stats_draft.allocated_bytes
                    if kv_stats_draft is not None else 0)
                py_executor.is_warmup = False
                py_executor.shutdown_for_kv_cache_estimation()
                py_executor.enable_iter_perf_stats = origin_iter_stats
                py_executor.set_gather_responses(False)

            total_used_bytes = total_gpu_memory - end
            activation_bytes = torch_peak_memory - model_bytes
            extra_cost = max(total_used_bytes - torch_used_bytes, 0)
            peak_memory = torch_peak_memory + extra_cost
            logger.info(
                f"Memory dynamically allocated during inference (inside torch) in memory usage profiling: {activation_bytes / (GB):.2f} GiB"
            )
            logger.info(
                f"Memory used outside torch (e.g., NCCL and CUDA graphs) in memory usage profiling: {extra_cost / (GB):.2f} GiB"
            )

        else:
            peak_memory = total_used_bytes
            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.")

        fixed_workspace_reserve = get_attention_workspace_fixed_bytes(
            self._model_engine.model.model_config,
            self._mapping,
            max_batch_size=self._max_batch_size,
            max_num_tokens=self._max_num_tokens,
            max_seq_len=self._max_seq_len,
            enable_cuda_graph=self._llm_args.cuda_graph_config is not None,
        )
        # A decode arena can first appear after the peak prefill step. Reserve
        # its full capacity for coexistence with the profiled transient peak.
        peak_memory += fixed_workspace_reserve
        if fixed_workspace_reserve:
            logger.info(
                f"Reserving {fixed_workspace_reserve / GB:.2f} GiB for capacity-sized "
                "attention workspace across serving steps.")

        penalty_workspace_reserve = self._get_penalty_workspace_reserve(
            py_executor)
        peak_memory += penalty_workspace_reserve
        if penalty_workspace_reserve:
            logger.info(
                f"Reserving {penalty_workspace_reserve / GB:.2f} GiB for the penalty "
                "workspace of the overlap-headroom seats, which is allocated on "
                "the first penalized request.")

        # 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

        if getattr(self._sparse_attention_config, "algorithm",
                   None) == "deepseek_v4":
            workspace_memory = get_indexer_mqa_logits_workspace_bytes(
                self._max_num_tokens, self._max_seq_len)
            budget_before = kv_cache_max_memory
            workspace_reserve = get_indexer_mqa_logits_workspace_reserve(
                budget_before, workspace_memory)
            if workspace_reserve > 0:
                kv_cache_max_memory = int(budget_before - workspace_reserve)
                logger.info(
                    f"Reserving {workspace_reserve / (GB):.2f} GiB for the DeepSeek-V4 Indexer MQA "
                    f"logits workspace (workspace cap {workspace_memory / (GB):.2f} GiB): KV cache "
                    f"budget {budget_before / (GB):.2f} -> {kv_cache_max_memory / (GB):.2f} GiB."
                )

        # 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.
        # A completed chunk-aware profiling run already prices the cached-KV
        # staging buffers in the measured peak. Do not subtract an additional
        # workspace reserve or install an attended-KV admission cap for it.
        profiled_mla_chunks = (py_executor is not None and not self._skip_est
                               and self._mla_chunked_profile_length is not None)
        w_bytes_per_token = (0 if profiled_mla_chunks else
                             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,
            chunked_workspace_profiled=profiled_mla_chunks,
            # Chunk-aware profiling is SM100+ only. Preserve Hopper's existing
            # backend-declared policy rather than adding a new reserve/cap to
            # its dense full-gather path just because it cannot profile chunks.
            require_chunked_workspace_profile=(
                w_bytes_per_token > 0 and self._llm_args.enable_chunked_prefill
                and get_sm_version() >= 100))
        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
        #   • max_gpu_total_bytes
        # For KVCacheManagerV2, KvCacheCreator controls capacity using max_gpu_total_bytes only.
        # This leaves max_tokens as a user-defined constraint.

        # ---------------------------handle max_tokens---------------------------------
        if self._is_kv_cache_manager_v2:
            # 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)

                # raise error if it is VSWA case
                if is_vswa:
                    logger.warning(
                        "max_tokens should not be set for VSWA case as it is ambiguous concept for VSWA."
                    )
                # calculate max memory from max_tokens
                kv_size_per_token = self._get_kv_size_per_token()
                kv_cache_max_memory_from_max_tokens = (
                    kv_size_per_token.bytes_for_tokens(self._max_kv_tokens_in))
                kv_cache_max_memory = min(kv_cache_max_memory,
                                          kv_cache_max_memory_from_max_tokens)
                logger.info(
                    f"max_tokens={self._max_kv_tokens_in} is provided. It limits max memory to {kv_cache_max_memory_from_max_tokens / (GB):.2f} GiB. "
                    f"New max_memory is set to {kv_cache_max_memory / (GB):.2f} GiB"
                )
            # For KvCacheManager, its logic still relies on max_tokens to control capacity
            self._kv_cache_config.max_tokens = (self._get_kv_size_per_token(
            ).tokens_for_budget(kv_cache_max_memory))
        # ---------------------------handle max_tokens---------------------------------

        # ---------------------------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):
            kv_cache_max_memory = min(kv_cache_max_memory,
                                      self._max_gpu_total_bytes_in)
            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"
            )

        logger.info(
            f"Estimated max memory in KV cache : {kv_cache_max_memory / (GB):.2f} GiB"
        )
        # set max_gpu_total_bytes
        self._kv_cache_config.max_gpu_total_bytes = kv_cache_max_memory
        if isinstance(self._profiling_stage_data, dict):
            self._profiling_stage_data["activation_bytes"] = activation_bytes
        # ---------------------------handle max_gpu_total_bytes---------------------------------

    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,
    ) -> 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, 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
        # only be in the separate draft KV cache manager.
        # We still pass spec_config so that num_extra_kv_tokens is calculated.
        spec_dec_layer_mask = None
        if self._should_create_separate_draft_kv_cache():
            num_target_layers = model_engine.model.model_config.pretrained_config.num_hidden_layers
            spec_dec_layer_mask = [True] * num_target_layers

        estimating_kv_cache = estimating_kv_cache and not self._skip_est
        unified_layer_specs = None
        if self._should_unify_dflash_kv_cache():
            unified_layer_specs = self._build_unified_dflash_layer_specs(
                kv_cache_config)
            logger.info(
                "DFlash: target and draft KV layers share one KV cache manager "
                f"({len(unified_layer_specs)} layers, disagg={self._is_disagg})."
            )
        kv_cache_manager = _create_kv_cache_manager(
            model_engine=model_engine,
            kv_cache_manager_cls=kv_cache_manager_cls,
            mapping=mapping,
            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,
            spec_config=self._speculative_config,
            sparse_attention_config=self._sparse_attention_config,
            max_num_tokens=self._max_num_tokens,
            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,
            disable_overlap_scheduler=self._disable_overlap_scheduler,
            kv_events_config=None if estimating_kv_cache 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,
            layer_specs=unified_layer_specs,
        )

        if not self._skip_est:
            # KVCacheManager (Non-draft) modifies the max_seq_len field, update it to self
            if model_engine.kv_cache_manager_key == ResourceManagerType.KV_CACHE_MANAGER:
                # When SWA is enabled, max_seq_len is updated inside kv_cache_manager.
                if kv_cache_manager is not None:
                    if kv_cache_manager.max_seq_len < self._max_seq_len:
                        self._dummy_reqs = self._create_dummy_context_requests(
                            max(
                                1, self._net_max_seq_len - 1 -
                                (self._max_seq_len -
                                 kv_cache_manager.max_seq_len)))
                    self._max_seq_len = kv_cache_manager.max_seq_len

                # When SWA is enabled, max_seq_len is updated inside kv_cache_manager.
                if kv_cache_manager is not None:
                    if kv_cache_manager.max_seq_len < self._max_seq_len:
                        self._dummy_reqs = self._create_dummy_context_requests(
                            max(1, kv_cache_manager.max_seq_len - 1))
                    self._max_seq_len = kv_cache_manager.max_seq_len
        else:
            if kv_cache_manager is not None:
                self._max_seq_len = kv_cache_manager.max_seq_len

        return kv_cache_manager

    def _should_create_separate_draft_kv_cache(self) -> bool:
        """
        Check if we need a separate draft KV cache manager for one-model mode.
        Returns True if the speculative config has use_separate_draft_kv_cache=True.

        Note: For MTP, _draft_config may be None since MTP layers are embedded
        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
        if self._should_unify_dflash_kv_cache():
            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.
            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 _should_unify_dflash_kv_cache(self) -> bool:
        """Whether this run can place target and DFlash layers in one V2 manager."""
        spec_config = self._speculative_config
        if (spec_config is None or not spec_config.spec_dec_mode.is_dflash()
                or spec_config.attention_backend != "TRTLLM"
                or self._tokens_per_block not in (16, 32, 64)
                or not spec_config._allow_separate_draft_kv_cache
                or spec_config.draft_len_schedule is not None
                or self._draft_config is None or self._mapping.pp_size != 1
                or self._mapping.cp_size != 1
                or self._llm_args.attn_backend != "TRTLLM"
                or _has_variable_kv_cache_geometry(
                    self._model_engine.model.model_config)
                or getattr(self._kv_cache_config, "mixed_precision_kv",
                           None) is not None
                or self._llm_args.kv_cache_compression_config is not None):
            return False
        # Specialized managers add cache roles and lifecycle protocols that an
        # external drafter cannot yet describe. Keep their existing fallback.
        if self._kv_cache_manager_cls is not KVCacheManagerV2:
            return False
        target_layout = self._dflash_kv_layout(
            self._model_engine.model.model_config,
            self._model_engine.model.config.torch_dtype)
        draft_model_config = self._get_draft_kv_model_config()
        draft_layout = self._dflash_kv_layout(
            draft_model_config,
            draft_model_config.pretrained_config.torch_dtype)
        if target_layout is None or draft_layout is None:
            logger.info_once(
                "DFlash: a per-layer KV-head count is not supported in a shared "
                "target and draft KV cache; keeping a separate draft KV cache.",
                key="dflash_per_layer_kv_heads_fallback")
            return False
        draft_windows = _resolve_v2_max_attention_window_vec(
            self._get_draft_max_attention_window(self._unification_max_seq_len,
                                                 self._kv_cache_config),
            self._unification_max_seq_len,
            range(draft_model_config.pretrained_config.num_hidden_layers))
        if len(set(draft_windows)) > 1:
            logger.info_once(
                "DFlash: draft layers have different effective attention "
                "windows; keeping a separate draft KV cache.",
                key="dflash_mixed_draft_windows_fallback")
            return False
        if self._is_disagg and not disagg_shares_draft_kv(spec_config):
            return False
        if self._is_disagg and not self._dflash_disagg_layouts_compatible(
                target_layout, draft_layout):
            logger.info_once(
                "DFlash: draft KV layout (kv_heads, head_dim, dtype, k_only)="
                f"{draft_layout} cannot be transferred with the target's "
                f"{target_layout}; disaggregated serving keeps a separate "
                "draft KV cache.",
                key="dflash_disagg_draft_layout_fallback")
            return False
        return True

    @staticmethod
    def _dflash_disagg_layouts_compatible(
            target_layout: tuple[int, int, object, bool],
            draft_layout: tuple[int, int, object, bool]) -> bool:
        """Whether the transceiver can move draft layers with the target's.

        Heads are resharded in bytes: a view's bytes per head are its own
        per-layer bytes over the rank's KV-head count, so dtype and head_dim
        may differ. The KV-head count and K-only flag come from the rank's
        single AttentionInfo and must match. NVFP4 adds block-scale views, so
        an NVFP4 side is only paired with an identical layout.
        """
        target_heads, _, target_dtype, target_k_only = target_layout
        draft_heads, _, draft_dtype, draft_k_only = draft_layout
        if (target_heads, target_k_only) != (draft_heads, draft_k_only):
            return False
        nvfp4 = tensorrt_llm.bindings.DataType.NVFP4
        return target_dtype == draft_dtype or nvfp4 not in (target_dtype,
                                                            draft_dtype)

    def _dflash_kv_layout(
            self, model_config: ModelConfig,
            fallback_dtype) -> Optional[tuple[int, int, object, bool]]:
        """Rank-local (kv_heads, head_dim, dtype, k_only) of a model's KV layers.

        Returns None when the KV-head count varies per layer.
        """
        config = model_config.pretrained_config
        k_only = is_mla(config)
        heads = 1 if k_only else getattr(config, "num_key_value_heads",
                                         config.num_attention_heads)
        if not isinstance(heads, int):
            return None
        head_dim = (config.kv_lora_rank + config.qk_rope_head_dim
                    if k_only else getattr(config, "head_dim", None))
        if not isinstance(head_dim, int):
            head_dim = config.hidden_size // config.num_attention_heads
        dtype = self._layer_spec_dtype(model_config, fallback_dtype)
        # Attention DP replicates attention weights, so every rank holds all heads.
        tp_size = 1 if self._mapping.enable_attention_dp else self._mapping.tp_size
        return ceil_div(heads, tp_size), head_dim, dtype, k_only

    @staticmethod
    def _layer_spec_dtype(model_config: ModelConfig, fallback_dtype) -> object:
        quant_config = model_config.quant_config
        if quant_config is not None and quant_config.quant_mode.has_fp8_kv_cache(
        ):
            return tensorrt_llm.bindings.DataType.FP8
        if quant_config is not None and quant_config.quant_mode.has_fp4_kv_cache(
        ):
            return tensorrt_llm.bindings.DataType.NVFP4
        return str_dtype_to_binding(torch_dtype_to_str(fallback_dtype))

    @staticmethod
    def _layer_spec_roles(dtype, *, k_only: bool) -> tuple:
        roles = [Role.KEY]
        if not k_only:
            roles.append(Role.VALUE)
        if dtype == tensorrt_llm.bindings.DataType.NVFP4:
            roles.append(Role.KEY_BLOCK_SCALE)
            if not k_only:
                roles.append(Role.VALUE_BLOCK_SCALE)
        return tuple(roles)

    def _build_unified_dflash_layer_specs(
            self,
            kv_cache_config: KvCacheConfig) -> tuple[KVCacheLayerSpec, ...]:
        """Build the rank-local target+DFlash storage layout used by cost and allocation."""
        if not self._should_unify_dflash_kv_cache():
            return ()

        target_model_config = self._model_engine.model.model_config
        target_config = target_model_config.pretrained_config
        target_layers = int(target_config.num_hidden_layers)
        target_local_ids = list(self._mapping.pp_layers(target_layers))
        target_heads, target_head_dim, target_dtype, target_k_only = (
            self._dflash_kv_layout(target_model_config,
                                   self._model_engine.model.config.torch_dtype))
        target_window_config = kv_cache_config.max_attention_window
        if target_window_config is None:
            target_window_config = _derive_v2_layer_type_attention_windows(
                kv_cache_config,
                self._kv_cache_manager_cls,
                target_model_config,
                self._max_seq_len,
            )
        target_windows = _resolve_v2_max_attention_window_vec(
            target_window_config,
            self._max_seq_len,
            target_local_ids,
        )

        draft_model_config = self._get_draft_kv_model_config()
        draft_config = draft_model_config.pretrained_config
        draft_layers = int(draft_config.num_hidden_layers)
        draft_local_ids = list(self._mapping.pp_layers(draft_layers))
        draft_heads, draft_head_dim, draft_dtype, draft_k_only = (
            self._dflash_kv_layout(draft_model_config,
                                   draft_config.torch_dtype))
        draft_windows = _resolve_v2_max_attention_window_vec(
            self._get_draft_max_attention_window(self._max_seq_len,
                                                 kv_cache_config),
            self._max_seq_len,
            draft_local_ids,
        )

        specs = [
            KVCacheLayerSpec(
                global_layer_id=layer_id,
                owner=KVCacheLayerOwner.TARGET,
                dtype=target_dtype,
                num_kv_heads=target_heads,
                head_dim=target_head_dim,
                buffer_roles=self._layer_spec_roles(target_dtype,
                                                    k_only=target_k_only),
                sliding_window_size=window,
                reuse_policy=AttentionReusePolicy.REQUIRED,
            ) for layer_id, window in zip(target_local_ids, target_windows)
        ]
        specs.extend(
            KVCacheLayerSpec(
                global_layer_id=target_layers + layer_id,
                owner=KVCacheLayerOwner.DRAFT,
                dtype=draft_dtype,
                num_kv_heads=draft_heads,
                head_dim=draft_head_dim,
                buffer_roles=self._layer_spec_roles(draft_dtype,
                                                    k_only=draft_k_only),
                sliding_window_size=window,
                reuse_policy=AttentionReusePolicy.REQUIRED,
            ) for layer_id, window in zip(draft_local_ids, draft_windows))
        return tuple(specs)

    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
        lookahead = draft_prompt_lookahead(self._speculative_config)
        if lookahead is None:
            return False
        if lookahead > 0 and not getattr(self._kv_cache_manager_cls,
                                         "_supports_reuse_match_backoff",
                                         False):
            # The opt-out is about backing the match off by `lookahead` tokens,
            # which a specialized commit/history protocol (recurrent snapshots,
            # DSA) cannot express. A zero span asks for no backoff at all, so
            # every backoff-sized path stays a no-op and the pairing is safe.
            return False
        return True

    def _get_effective_draft_config(self) -> ModelConfig:
        """
        Return the ModelConfig to use for draft KV cache creation.

        For Eagle3 and draft-target one-model modes, a dedicated draft config
        is provided at construction time.  For MTP one-model mode, no separate
        draft config exists because the MTP layers share the same architecture
        as the target model.  In that case we fall back to the target model's
        config so that the draft KV cache is created with the correct layout.
        """
        if self._draft_config is not None:
            return self._draft_config
        # MTP: MTP layers reuse the target model architecture, so the target
        # model's config describes the correct KV cache layout for the draft
        # layers as well.
        return self._model_engine.model.model_config

    def _get_draft_kv_model_config(self) -> ModelConfig:
        """The draft ModelConfig describing the KV pool as it is ALLOCATED.

        The args-level ``kv_cache_config.dtype`` sync stamps the TARGET's fp8 KV
        algo onto every loaded model, including a standalone drafter. The drafter
        stores and reads its pool in its weights dtype (DFlash validates a bf16 pool
        and otherwise falls back to the max_seq_len-dense private arena, which OOMs at
        long context), so the pool dtype must follow the drafter.

        Every consumer of draft KV bytes must go through here. If the budget split
        and the allocation read different dtypes, the split charges fp8 bytes for a
        bf16 pool and the draft manager gets HALF the target's tokens. The capacity
        scheduler admits on the target pool alone, so past ~50% target utilization it
        raises "Draft KV cache context resize failed", fatal to every rank.
        """
        effective_draft_config = self._get_effective_draft_config()
        # Narrower than is_external_drafter(), matching
        # _should_create_separate_draft_kv_cache. PARD and DRAFT_TARGET_ONE_MODEL
        # reach here too and can carry a genuine fp8 KV algo of their own, which
        # dtype="auto" keeps; dropping it would allocate bf16 under attention
        # modules that still read and write fp8.
        spec_dec_mode = self._speculative_config.spec_dec_mode
        if not (spec_dec_mode.is_dflash() or spec_dec_mode.is_dspark()):
            return effective_draft_config
        quant_config = getattr(effective_draft_config, "quant_config", None)
        if quant_config is None or not quant_config.quant_mode.has_kv_cache_quant(
        ):
            return effective_draft_config
        # A drafter that writes and reads FP8 codes with its own per-layer scales
        # (DFlashForCausalLM.set_kv_cache_scales) keeps the fp8 algo: its pool is FP8.
        draft_model = getattr(getattr(self._model_engine, "model", None),
                              "draft_model", None)
        if getattr(draft_model, "kv_cache_storage_dtype",
                   None) == torch.float8_e4m3fn:
            logger.info("External drafter KV pool is FP8: the drafter stores "
                        "and reads it with its own per-layer scales.")
            return effective_draft_config
        logger.info(
            "External drafter KV pool keeps the drafter dtype; dropping "
            "the KV quantization inherited from the target.")
        neutral_quant = copy.copy(quant_config)
        neutral_quant.kv_cache_quant_algo = None
        # QuantConfig.quant_mode and .layer_quant_mode are both cached_property
        # and the copy carries the already-computed caches, so BOTH must be
        # dropped for the mutation to take: _create_kv_cache_manager reads
        # quant_mode off this copy, and layer_quant_mode is the pair's other
        # half, stale in the same way.
        neutral_quant.__dict__.pop("quant_mode", None)
        neutral_quant.__dict__.pop("layer_quant_mode", None)
        # No _frozen dance: ModelConfig.__setattr__ exempts quant_config by
        # name, and restoring _frozen to True would freeze a copy whose source
        # may not have been frozen.
        effective_draft_config = copy.copy(effective_draft_config)
        effective_draft_config.quant_config = neutral_quant
        return effective_draft_config

    def _get_draft_tokens_per_block(self) -> int:
        """Page size of the separate one-model draft KV pool.

        A standalone drafter (DFlash, standalone DSpark) gets its own page
        size, ``_STANDALONE_DRAFT_TOKENS_PER_BLOCK``, NOT self._tokens_per_block:
        the target's value can be one the drafter's kernels reject (FP4 MLA
        forces 128; trtllm-gen MLA decode takes only 32 or 64). Those drafters
        read their pool geometry from the manager, so nothing else has to agree.

        Every other separate draft pool (MTP, Eagle3, ...) keeps the target's
        page size: its draft layers run through the target's attention
        metadata, which is not re-planned for a different page size.

        Both the budget split (_get_draft_cache_cost) and the allocation
        (_create_one_model_draft_kv_cache_manager) must go through here, or
        the split charges one page granularity and the pool uses another.
        """
        spec_dec_mode = self._speculative_config.spec_dec_mode
        if spec_dec_mode.is_dflash() or spec_dec_mode.is_dspark():
            return _STANDALONE_DRAFT_TOKENS_PER_BLOCK
        return self._tokens_per_block

    def _get_num_draft_layers(self) -> int:
        """Return the actual number of draft KV cache layers.

        This must stay in sync with the num_layers passed to the draft KV
        cache manager constructor in _create_one_model_draft_kv_cache_manager.
        """
        if self._speculative_config.spec_dec_mode.is_external_drafter():
            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.
        """
        num_draft_layers = self._get_num_draft_layers()
        spec_dec_layer_mask = self._get_one_model_draft_layer_mask()

        # Get the effective draft config (explicit draft_config if available,
        # otherwise fall back to target model config for MTP), with the
        # target's inherited fp8 KV algo dropped for a standalone drafter. The
        # budget split in _get_kv_size_per_token resolves it through the SAME
        # helper, so the bytes/token it charges match the pool allocated here.
        effective_draft_config = self._get_draft_kv_model_config()

        kv_cache_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,
            cache_transceiver_config=self._cache_transceiver_config)
        draft_kv_cache_manager_cls = self._validate_or_fallback_kv_cache_manager_v2(
            draft_kv_cache_manager_cls, effective_draft_config, draft_kv_config)

        estimating_kv_cache = estimating_kv_cache and not self._skip_est
        # For MTP with models using sparse attention (e.g., DeepSeek V3 with DSA),
        # the draft layers share the same architecture as the target model and need
        # 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
        # Under helix the standalone drafter is built against the repurposed
        # mapping (CP ranks become TP ranks) and every rank keeps its full
        # drafter KV, so its paged manager needs the CP-free mapping: the
        # round-robin ledger applies to the TARGET KV alone, and
        # KVCacheManagerV2 rejects is_draft x helix outright.
        draft_mapping = self._mapping
        if draft_mapping.has_cp_helix():
            draft_mapping = draft_mapping.repurpose_helix_cp_to_tp()
        draft_tokens_per_block = self._get_draft_tokens_per_block()
        if draft_tokens_per_block != self._tokens_per_block:
            # Logged because nothing fails when the two differ: only the draft
            # pool's block count, block-table width and fragmentation change.
            logger.info(
                f"Separate draft KV cache page size {draft_tokens_per_block} "
                f"(target pool uses {self._tokens_per_block}).")
        return _create_kv_cache_manager(
            model_engine=None,
            max_cuda_graph_batch_size=self._model_engine.
            _max_cuda_graph_batch_size,
            kv_cache_manager_cls=draft_kv_cache_manager_cls,
            mapping=draft_mapping,
            kv_cache_config=draft_kv_config,
            tokens_per_block=draft_tokens_per_block,
            max_seq_len=max_seq_len,
            max_batch_size=self._max_batch_size,
            spec_config=self._speculative_config,
            sparse_attention_config=sparse_attn_config,
            max_num_tokens=self._max_num_tokens,
            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,
            dtype=effective_draft_config.pretrained_config.torch_dtype,
            is_draft=True,
            layer_mask=spec_dec_layer_mask,
            num_layers=num_draft_layers,
            is_disagg=self._is_disagg,
            disable_overlap_scheduler=self._disable_overlap_scheduler,
            cold_page_codec_provider=cold_page_codec_provider,
            joint_kv_cache_reuse=self._joint_kv_cache_reuse,
        )

    def _get_target_and_draft_cache_costs(
        self,
        kv_cache_config: Optional[KvCacheConfig] = None,
        *,
        include_request_cost: bool = False,
    ) -> Optional[tuple[CacheCost, CacheCost]]:
        """Per-manager KV cache costs for target and draft layers.

        ``include_request_cost`` additionally attaches the optional per-request
        estimates (``CacheCost.bytes_per_request`` and
        ``CacheCost.quota_per_request``) to each cost. Only the GPU budget
        split uses them; default callers keep the plain affine costs.
        """
        target_kv_cache_config = (kv_cache_config if kv_cache_config is not None
                                  else self._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,
            include_request_cost=include_request_cost)
        # Estimate the draft component directly so its independently modelled
        # affine intercept is preserved exactly.
        draft_kv = self._get_draft_cache_cost(
            target_kv_cache_config,
            use_separate_draft_kv_cache=use_separate_draft_kv_cache,
            include_request_cost=include_request_cost,
        )
        if draft_kv is None:
            return None
        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):
            return None
        return target_kv, draft_kv

    def _compute_draft_budget_shares(
        self,
        total_budget: int,
        target_kv: CacheCost,
        draft_kv: CacheCost,
    ) -> Optional[tuple[int, int]]:
        """Split *total_budget* into (target_budget, draft_budget) byte shares.

        Both managers' affine intercepts are reserved first; the remaining
        (slope) budget is weighed by slope-funded per-request spend
        (``CacheCost.bytes_per_request``, which excludes the retention the
        intercept already pays) when both managers report it, and by per-token
        cost otherwise. The per-token weighting hands both pools the same token
        capacity, which under-budgets the draft pool whenever the target bounds
        retention with windows the intercept does not carry while every draft
        mirror is reserved at the request's full context length: the draft pool
        then caps concurrency while the target pool sits mostly idle. Weighting
        by the uncovered per-request spend sizes both pools for the same
        worst-case request count instead.

        Under per-request weighting the draft share is additionally capped at
        ``max_batch_size`` full-length mirrors — the pool can never hold more
        than the scheduler admits, so any excess share stays with the target —
        using the per-mirror allocator quota (``CacheCost.quota_per_request``:
        resume-utilization headroom and allocation granularity included) when
        the manager reports it. That quota also floors the draft share: the
        allocator raises any smaller quota at construction regardless of the
        split, so assigning less would silently overcommit the target's share.
        A floor that does not fit in the budget remaining after both
        intercepts makes the split infeasible (``None``), like an unfundable
        fixed cost; flooring anyway would push the target below its intercept.
        """
        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:
            logger.warning(
                f"KV cache budget {total_budget} is smaller than the fixed "
                f"cache 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
        per_request_weighted = (target_kv.bytes_per_request is not None
                                and draft_kv.bytes_per_request is not None
                                and target_kv.bytes_per_request +
                                draft_kv.bytes_per_request > 0)
        if per_request_weighted:
            target_weight = target_kv.bytes_per_request
            draft_weight = draft_kv.bytes_per_request
        else:
            target_weight = target_kv.slope
            draft_weight = draft_kv.slope
        weight_total = target_weight + draft_weight
        draft_slope_share = (slope_budget * draft_weight //
                             weight_total if weight_total > 0 else 0)
        if per_request_weighted and self._max_batch_size > 0:
            # Cap in allocator-quota units when known: the usable-bytes
            # estimate under-reserves by the resume-utilization headroom and
            # the allocation granularity the runtime rounds each pool to.
            draft_bytes_per_mirror = (draft_kv.quota_per_request
                                      if draft_kv.quota_per_request is not None
                                      else draft_kv.bytes_per_request)
            max_useful_draft_bytes = (self._max_batch_size *
                                      draft_bytes_per_mirror)
            draft_slope_share = min(draft_slope_share, max_useful_draft_bytes)
        if draft_kv.quota_per_request is not None:
            if draft_kv.quota_per_request > slope_budget:
                # The allocator will raise the draft quota to this one-request
                # floor at construction regardless of the split, so a budget
                # that cannot fund it alongside both intercepts is infeasible,
                # exactly like an unfundable fixed cost. Flooring anyway would
                # silently push the target below its intercept.
                logger.warning(
                    f"KV cache budget {total_budget} cannot fund the draft "
                    f"manager's minimum allocator quota "
                    f"({draft_kv.quota_per_request} bytes for one max-length "
                    f"request) beyond the fixed cache cost {intercept_total}; "
                    f"cannot split between target and draft.")
                return None
            # The allocator raises any smaller quota to this one-request floor
            # at construction regardless of the split, so assigning less would
            # silently overcommit the target's share.
            draft_slope_share = max(draft_slope_share,
                                    draft_kv.quota_per_request)
        draft_budget = draft_kv.intercept + draft_slope_share
        target_budget = total_budget - draft_budget
        return target_budget, draft_budget

    def _bound_unsplit_v2_draft_by_tokens(
        self,
        total_budget: int,
        target_kv_cache_config: KvCacheConfig,
        draft_kv_cache_config: Optional[KvCacheConfig],
    ) -> Optional[KvCacheConfig]:
        """Backstop for a one-model V2 draft whose GPU budget cannot be split.

        When the per-manager costs are unavailable the split returns the
        configs unchanged, so the draft would read the target's whole
        max_gpu_total_bytes and claim the budget a second time. Bound the draft
        alone by the token capacity of that budget instead: its V2 quota
        becomes min(max_gpu_total_bytes, max_tokens * draft bytes/token), as V1
        does with a shared max_tokens. The target config is never touched, so a
        successful split (which already sizes both managers) is unaffected.
        """
        if (not self._is_kv_cache_manager_v2
                or not self._should_create_separate_draft_kv_cache()
                or target_kv_cache_config.max_tokens is not None):
            return draft_kv_cache_config
        max_tokens = self._get_kv_size_per_token(
            target_kv_cache_config).tokens_for_budget(total_budget)
        # When every attention layer is windowed the per-token slope is 0 and
        # tokens_for_budget returns 0. V2 treats an unset max_tokens as
        # unbounded, so leave it unset rather than impose a cap of zero.
        if max_tokens <= 0:
            return draft_kv_cache_config
        bounded = (draft_kv_cache_config if draft_kv_cache_config is not None
                   else target_kv_cache_config).model_copy()
        bounded.max_tokens = max_tokens
        logger.info(
            f"Cannot split KV cache max_gpu_total_bytes between target and "
            f"draft; bounding the draft to max_tokens={max_tokens}.")
        return bounded

    # Warn only when admissible concurrency exceeds the graph-covered range
    # by this factor, so a small overshoot off-graph does not trip it.
    _CUDA_GRAPH_CAPTURE_WARN_FACTOR = 1.5

    def _warn_if_admitted_concurrency_outruns_graph_capture(self) -> None:
        """Warn when admitted decode concurrency far outruns CUDA-graph coverage.

        Decode above the covered range runs eager, which can cost more
        throughput than the extra concurrency buys; still valid at low load,
        hence a warning that names the knobs rather than a refusal.
        """
        capture_sizes = getattr(self._model_engine, "_cuda_graph_batch_sizes",
                                None)
        padding = getattr(self._model_engine, "_cuda_graph_padding_enabled",
                          None)
        if (not isinstance(capture_sizes, (list, tuple)) or not capture_sizes
                or not all(isinstance(s, int) for s in capture_sizes)):
            # CUDA graphs disabled, or the engine does not expose its capture
            # set.
            return
        if padding is True:
            # Padding rounds a decode batch up to the nearest captured size,
            # so everything up to the largest capture replays a graph.
            covered = max(capture_sizes)
        else:
            # Without padding a decode batch replays a graph only at an exactly
            # captured size, so the reliably covered range is the dense {1..n}
            # prefix of the capture set; above it ragged batches rarely match.
            covered = 0
            for size in sorted(set(capture_sizes)):
                if size != covered + 1:
                    break
                covered = size
        if covered <= 0:
            return
        if not (isinstance(self._max_batch_size, int)
                and self._max_batch_size > 0):
            return
        admissible = self._max_batch_size
        if admissible <= self._CUDA_GRAPH_CAPTURE_WARN_FACTOR * covered:
            return
        logger.warning(
            f"The engine admits up to ~{admissible} concurrent decode "
            f"requests (scheduler max_batch_size={self._max_batch_size}), but "
            f"captured CUDA graphs only cover decode "
            f"batches up to {covered} (largest capture "
            f"{max(capture_sizes)}, enable_padding="
            f"{bool(padding)}). Decode batches above {covered} run in eager "
            f"mode with a large per-iteration latency penalty, which can "
            f"cost more throughput than the added concurrency buys. Raise "
            f"cuda_graph_config.max_batch_size (or add larger "
            f"cuda_graph_config.batch_sizes) and set "
            f"cuda_graph_config.enable_padding=True so ragged batches round "
            f"up to a captured size, or lower max_batch_size if decode above "
            f"{covered} concurrent requests is not intended.")

    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,
    ) -> tuple[KvCacheConfig, Optional[KvCacheConfig]]:
        """Split a byte budget (attribute on ``KvCacheConfig``) between target
        and draft KV caches.

        Splits the value of ``target_kv_cache_config.<budget_attr>`` using the
        affine target/draft cache costs, then returns cloned target and draft
        configs containing their respective shares.

        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
        other budget (e.g. ``host_cache_size``, which is host offload memory the
        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)
        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
        for non-GPU budgets remains so the draft never silently inherits the full
        budget and double-allocates it.
        """
        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 target_kv_cache_config, draft_kv_cache_config

        cache_costs = self._get_target_and_draft_cache_costs(
            target_kv_cache_config, include_request_cost=True)
        if cache_costs is None:
            if budget_attr == "max_gpu_total_bytes":
                draft_kv_cache_config = self._bound_unsplit_v2_draft_by_tokens(
                    total_budget, target_kv_cache_config, 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
        # state; it does not consume host offload memory. When splitting a
        # non-GPU budget (e.g. host_cache_size), drop the intercept so the split
        # stays proportional to the per-token (slope) cost instead of being
        # spuriously starved by a GPU-only fixed cost. The per-request
        # estimates are dropped with it: only the GPU pool caps live
        # concurrency and only its allocator enforces quota floors, so offload
        # tiers keep the per-token proportional split.
        if budget_attr != "max_gpu_total_bytes":
            target_kv = CacheCost(slope=target_kv.slope)
            draft_kv = CacheCost(slope=draft_kv.slope)

        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.
            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
                # fatal: both managers need their fixed state resident in GPU
                # memory, so the run would OOM. Fail fast with actionable
                # 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"({intercept_total / GB:.2f} GiB, e.g. SWA or mamba state) "
                    f"for target+draft. Increase free_gpu_memory_fraction or "
                    f"max_gpu_total_bytes, or reduce max_batch_size (the fixed "
                    f"cost scales with batch size).")
            # Defensive: non-GPU budgets zero out the intercept above, so with a
            # positive budget this branch is currently unreachable for them. It
            # remains as a safety net guaranteeing that, should a non-GPU budget
            # ever carry a fixed cost it cannot fit, we degrade gracefully rather
            # than letting both managers inherit the full budget and
            # double-allocate it: keep the full budget on the target and zero the
            # draft's share for this attribute.
            logger.warning(
                f"Cannot split KV cache {budget_attr} between target and 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 = 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 target_kv_cache_config, draft_kv_cache_config
        target_budget, draft_budget = shares

        # The proportional split assumes the draft pool holds as many tokens as
        # the target pool. A disaggregated generation worker never runs the
        # prompt forward, so an external drafter there only stores the tokens
        # generated on this worker; TLLM_DRAFT_KV_GPU_BUDGET_GIB caps the
        # draft's GPU share and hands the rest to the target pool.
        draft_cap_gib = os.environ.get("TLLM_DRAFT_KV_GPU_BUDGET_GIB")
        if draft_cap_gib and budget_attr == "max_gpu_total_bytes":
            try:
                cap_gib = float(draft_cap_gib)
            except ValueError as e:
                raise ValueError(
                    "TLLM_DRAFT_KV_GPU_BUDGET_GIB must be a number of GiB, got "
                    f"{draft_cap_gib!r}") from e
            # Reject NaN too: it fails every comparison below.
            if not cap_gib > 0:
                raise ValueError(
                    "TLLM_DRAFT_KV_GPU_BUDGET_GIB must be positive, got "
                    f"{draft_cap_gib!r}")
            draft_cap = draft_kv.intercept + int(cap_gib * GB)
            if draft_kv.quota_per_request is not None:
                # The allocator raises the draft quota to one request's floor
                # regardless of the split (see _compute_draft_budget_shares),
                # so a lower cap would silently overcommit the target's share.
                draft_cap = max(draft_cap,
                                draft_kv.intercept + draft_kv.quota_per_request)
            if draft_cap < draft_budget:
                logger.info(
                    f"Capping the draft KV cache GPU budget at "
                    f"{draft_cap / GB:.2f} GiB (TLLM_DRAFT_KV_GPU_BUDGET_GIB="
                    f"{draft_cap_gib}); the proportional share was "
                    f"{draft_budget / GB:.2f} GiB.")
                draft_budget = draft_cap
                target_budget = total_budget - draft_budget

        logger.info(
            f"Splitting KV cache {budget_attr}: total={total_budget / GB:.2f} GiB, "
            f"target={target_budget / GB:.2f} GiB ({target_kv}), "
            f"draft={draft_budget / GB:.2f} GiB ({draft_kv})")

        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:
            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 self._model_engine.model.model_config.is_encoder_decoder

    @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.")

        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_attention_config=None,
            max_num_tokens=self._max_num_tokens,
            max_beam_width=1,
            max_copy_beam_width=self._max_beam_width,
            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,
            disable_overlap_scheduler=self._disable_overlap_scheduler,
            kv_cache_type=tensorrt_llm.bindings.internal.batch_manager.
            CacheType.CROSS,
        )

    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:
            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)

    def build_managers(self,
                       resources: Dict,
                       estimating_kv_cache: bool = False) -> None:
        """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(
            )

        has_separate_one_model_draft = 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.
        draft_kv_cache_config = None
        if has_separate_one_model_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:
                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:
                self_kv_cache_config, draft_kv_cache_config = (
                    self._split_kv_cache_budget_for_draft(
                        budget_attr, self_kv_cache_config,
                        draft_kv_cache_config))

        # Estimation managers are throwaway probes; the real build that
        # follows runs the check against the same engine and scheduler
        # settings.
        if not estimating_kv_cache:
            self._warn_if_admitted_concurrency_outruns_graph_capture()

        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

        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)

        # One-model speculative decoding with different KV layouts
        if has_separate_one_model_draft:
            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)

        # 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
        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[
            ResourceManagerType.DRAFT_KV_CACHE_MANAGER]
        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(
    num_key_value_heads: int,
    num_hidden_layers: int,
    spec_config: Optional[SpeculativeConfig] = None,
    draft_config: Optional[ModelConfig] = None,
) -> Union[int, List[int]]:
    """
    Returns:
        An int when all layers share the same num_kv_heads (common case),
        or a list of num_kv_heads (one entry per layer) when target and
        draft models differ.
    """
    if spec_config is None or draft_config is None:
        return num_key_value_heads

    from ..speculative.utils import get_num_spec_layers
    draft_pretrained = draft_config.pretrained_config
    draft_num_kv_heads = getattr(
        draft_pretrained, 'num_key_value_heads',
        getattr(draft_pretrained, 'num_attention_heads', None))

    if draft_num_kv_heads is None or draft_num_kv_heads == num_key_value_heads:
        return num_key_value_heads

    num_spec_layers = get_num_spec_layers(spec_config)
    logger.info(f"Per-layer KV heads for speculative decoding: "
                f"target={num_key_value_heads} x {num_hidden_layers} layers, "
                f"draft={draft_num_kv_heads} x {num_spec_layers} layers, "
                f"total={num_hidden_layers + num_spec_layers} layers")
    return [num_key_value_heads] * num_hidden_layers + [draft_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,
        mapping: Mapping,
        kv_cache_config: KvCacheConfig,
        tokens_per_block: int,
        max_seq_len: int,
        max_batch_size: int,
        spec_config: Optional[SpeculativeConfig],
        sparse_attention_config: Optional[SparseAttentionConfig],
        max_num_tokens: int,
        max_beam_width: int,
        kv_connector_manager: Optional[KvCacheConnectorManager],
        max_copy_beam_width: Optional[int] = None,
        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,
        dtype: Optional[torch.dtype] = None,
        is_draft: bool = False,
        layer_mask: Optional[List[bool]] = None,
        num_layers: Optional[int] = None,
        num_kv_heads: Optional[Union[int, List[int]]] = None,
        head_dim: Optional[Union[int, List[int]]] = None,
        kv_cache_type=None,
        is_disagg: bool = False,
        disable_overlap_scheduler: bool = False,
        cold_page_codec_provider: Optional[object] = None,
        kv_events_config: Optional[KVEventsConfig] = None,
        joint_kv_cache_reuse: bool = False,
        max_cuda_graph_batch_size: Optional[int] = None,
        layer_specs: Optional[Sequence[KVCacheLayerSpec]] = None
) -> 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
        quant_config = model_config.quant_config
        _model_config = model_config
    else:
        config = model_engine.model.model_config.pretrained_config
        quant_config = model_engine.model.model_config.quant_config
        _model_config = model_engine.model.model_config

    if dtype is None:
        dtype = model_engine.dtype

    if kv_cache_type is None:
        kv_cache_type = tensorrt_llm.bindings.internal.batch_manager.CacheType.SELF

    hidden_size = config.hidden_size

    model_layer_specs = _get_kv_cache_layer_specs(_model_config)
    if model_layer_specs:
        head_dim = [spec.head_dim for spec in model_layer_specs]
        num_key_value_heads = [spec.num_kv_heads for spec in model_layer_specs]
    else:
        num_attention_heads = config.num_attention_heads
        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

    # Derive per-layer max_attention_window for any model that publishes a
    # mixed sliding/full `layer_types` schedule (Gemma4 hybrid included) so V2
    # creates separate pool groups for sliding vs full-attention layers;
    # otherwise every layer lands in one full-context pool and the bounded
    # layers keep blocks they can never read.  Sliding layers get the model's
    # `sliding_window`, full layers `max_seq_len`; V2 then evicts old blocks
    # once kv_len exceeds the window (only ~ceil(sliding_window/page_size)
    # blocks per sequence for sliding layers), and FlashInfer's prepare()
    # picks up the smaller per-pool block count automatically.  Gemma4 hybrid
    # always resolves to KVCacheManagerV2 (see _non_hybrid_kv_cache_manager_cls),
    # so the V2 guard never excludes it.  A user-supplied max_attention_window
    # always wins.  `_derive_v2_layer_type_attention_windows` applies the same
    # rules in the creator's static cost model, so the budget split sizes the
    # manager from the windows it is built with; it also keeps the default
    # when a configured `pool_ratio` does not match the derived layer groups.
    # Skip derivation for the one-model draft manager: (1) in one-model
    # spec-decode the KV memory budget is already split between target and
    # draft, so VSWA sizing here would size each window pool from the unsplit
    # free-memory budget; (2) draft layers live at global indices past the
    # target's num_hidden_layers, so _project_max_attention_window_vec would
    # wrap them back onto pattern[0]. The draft config sets
    # max_attention_window=None to opt out, not to request derivation.
    # The cross-attention pool holds encoder-side KV that the decoder's
    # `layer_types` do not describe, so it keeps the default too.
    # Keep estimation storage full-context unless distinct page layouts require
    # windowed storage (Gemma4). Enable automatic windowing for the final manager.
    derived_windows = None
    if (not is_draft and kv_cache_type
            == tensorrt_llm.bindings.internal.batch_manager.CacheType.SELF
            and (not estimating_kv_cache
                 or _has_variable_kv_cache_geometry(_model_config))):
        derived_windows = _derive_v2_layer_type_attention_windows(
            kv_cache_config, kv_cache_manager_cls, _model_config, max_seq_len)
    if derived_windows is not None:
        if layer_mask is None and spec_config is not None:
            # `get_pp_layers` appends the one-model speculative layers after
            # the decoder layers when they share this manager, and V2 resolves
            # each layer's window by global layer id: give them the full
            # context the single-window default gave them instead of wrapping
            # them onto the decoder pattern.
            derived_windows = derived_windows + [int(
                max_seq_len)] * get_num_spec_layers(spec_config)
        assert uses_vswa_kv_cache_layout(derived_windows), (
            "derived per-layer windows must select a VSWA layout; a non-VSWA "
            "vector would reshape the single pool instead of splitting it")
        logger.info(
            "Derived per-layer max_attention_window from layer_types for "
            f"{kv_cache_manager_cls.__name__}: {derived_windows} "
            f"({len(set(derived_windows))} distinct windows)")
        kv_cache_config = copy.copy(kv_cache_config)
        kv_cache_config.max_attention_window = derived_windows

    if quant_config is not None and quant_config.quant_mode.has_fp8_kv_cache():
        kv_cache_dtype = tensorrt_llm.bindings.DataType.FP8
    elif quant_config is not None and quant_config.quant_mode.has_fp4_kv_cache(
    ):
        kv_cache_dtype = tensorrt_llm.bindings.DataType.NVFP4
    else:
        kv_cache_dtype = str_dtype_to_binding(torch_dtype_to_str(dtype))

    # Use provided num_layers if available, otherwise use config.
    # When layer_mask is set (e.g., KV sharing), num_layers for the cache
    # manager must equal the number of enabled (True) layers in the mask.
    if num_layers is not None:
        num_hidden_layers = num_layers
    elif layer_mask is not None:
        num_hidden_layers = sum(layer_mask)
    else:
        num_hidden_layers = config.num_hidden_layers
    # Only include draft KV heads in the per-layer list when draft layers
    # are NOT handled by a separate draft KV cache manager.  When layer_mask
    # is provided from the caller, it means the main KV cache covers only
    # the masked (target) layers and draft layers live in their own manager.
    draft_config_for_kv = None
    if layer_mask is None:
        draft_config_for_kv = (getattr(model_engine.model, 'draft_config', None)
                               if model_engine is not None else None)
    # Use model-provided per-layer KV heads directly; otherwise expand the
    # scalar value for speculative decoding when needed.
    if isinstance(num_key_value_heads, list):
        per_layer_num_kv_heads = num_key_value_heads
    else:
        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["max_cuda_graph_batch_size"] = (
            model_engine._max_cuda_graph_batch_size
            if model_engine is not None else max_cuda_graph_batch_size)
        manager_extra_kwargs["enable_stats"] = enable_kv_cache_stats
        manager_extra_kwargs["max_copy_beam_width"] = max_copy_beam_width
        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
        manager_extra_kwargs[
            "disable_overlap_scheduler"] = disable_overlap_scheduler
        manager_extra_kwargs["layer_specs"] = layer_specs
        # Vocab size also enables multimodal event decoding and its per-block
        # digest scan. Leave it unset for text-only engines. Without an engine
        # (e.g. separate one-model draft caches), preserve config resolution:
        # a text sub-config alone cannot identify a multimodal deployment.
        needs_multimodal_keys = model_engine is None or model_engine.is_multimodal
        manager_extra_kwargs["vocab_size"] = (resolve_vocab_size(config) if
                                              needs_multimodal_keys else None)
        if (needs_multimodal_keys and manager_extra_kwargs["vocab_size"] is None
                and kv_cache_config.enable_block_reuse):
            logger.warning(
                "Could not resolve vocab_size from the model config; "
                "multimodal requests will fail when block reuse is on. "
                "Disable block reuse to serve them.")
    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 config_utils.is_glm5_next(config) or 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. glm5_next has the same layout (11 of its 45
        # layers are sparse MLA) plus paged indexer state; both families use
        # sequential verify when replay is unavailable.
        is_glm5_next = config_utils.is_glm5_next(config)
        text_config = config_utils.unwrap_glm5_next_text_config(
            config) if is_glm5_next else config
        if max_beam_width > 1:
            raise ValueError(
                "glm5_next + beam search is not supported yet."
                if is_glm5_next else
                "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 glm5_next."
                if is_glm5_next else
                "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)
        # 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.
        kda_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 kda_mtp_replay_num_spec
            kda_replay_num_spec = kda_mtp_replay_num_spec(
                spec_config.tokens_per_gen_step - 1)
            if kda_replay_num_spec is not None:
                kda_extra_kwargs["kda_replay_num_spec"] = kda_replay_num_spec
                # A model whose KDA verify replays the accepted drafts from
                # its own per-draft records (instead of the replay caches)
                # asks for a buffer of them with `kda_draft_records`.
                backbone = getattr(getattr(model_engine, "model", None),
                                   "model", None)
                if (issubclass(kv_cache_manager_cls, MambaHybridCacheManagerV2)
                        and getattr(backbone, "kda_draft_records", False)):
                    kda_extra_kwargs["kda_draft_records"] = True
        if is_glm5_next:
            # The manager places an indexer buffer on every attention layer.
            from ..attention.backends.sparse.glm_kpool import \
                glm_kpool_cache_row_dim
            kda_extra_kwargs["index_state_dim"] = glm_kpool_cache_row_dim(
                int(text_config.index_head_dim))
        # KDA's conv state is a [Q | K | V] concatenation whose three sections
        # have identical width, i.e. the qwen3_next section layout.
        kda_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=(int(text_config.kv_lora_rank) +
                      int(getattr(text_config, "qk_rope_head_dim", 0) or 0)
                      if is_glm5_next else 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,
            **kda_extra_kwargs,
            **manager_extra_kwargs,
        )
    elif (getattr(sparse_attention_config, "algorithm", None) == "csa2"
          or is_mla(config)):
        if getattr(sparse_attention_config, "algorithm", None) == "csa2":
            manager_extra_kwargs["context_swa_layer_limit"] = (
                _model_config.extra_attrs.get("csa2_context_swa_layer_limit"))
            manager_extra_kwargs["bounded_replay_on_generation"] = (
                _model_config.extra_attrs.get("bounded_replay_on_generation",
                                              False))
        kv_cache_manager = kv_cache_manager_cls(
            kv_cache_config,
            tensorrt_llm.bindings.internal.batch_manager.CacheType.SELFKONLY,
            num_layers=num_hidden_layers,
            num_kv_heads=1,
            head_dim=(getattr(config, "head_dim", 512)
                      if sparse_attention_config is not None
                      and sparse_attention_config.algorithm == "csa2" else
                      config.kv_lora_rank + config.qk_rope_head_dim),
            tokens_per_block=tokens_per_block,
            max_seq_len=max_seq_len,
            max_batch_size=max_batch_size,
            mapping=mapping,
            dtype=kv_cache_dtype,
            spec_config=spec_config,
            max_num_tokens=max_num_tokens,
            max_beam_width=max_beam_width,
            is_draft=is_draft,
            kv_connector_manager=kv_connector_manager
            if not estimating_kv_cache else None,
            sparse_attention_config=sparse_attention_config,
            pretrained_config=config,
            is_estimating_kv_cache=estimating_kv_cache,
            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:
            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)

        # Replay state update kernel for MTP: default on for sm >= 80; gates
        # below disable it for incompatible feature combinations.  Cpp cache
        # manager doesn't expose use_replay_state_update, so the wrapper
        # property's getattr default keeps replay off there automatically.
        sm = get_sm_version()
        stochastic_rounding = getattr(
            quant_config, 'mamba_ssm_stochastic_rounding',
            False) if quant_config is not None else False

        use_replay = spec_config is not None and sm >= 80
        if spec_config is None:
            logger.info(
                "Replay kernel requires speculative decoding; using non-replay path"
            )

        # Block reuse (prefix caching): replay leaves SSM state at a
        # checkpoint after speculation. The next decode step replays forward
        # to correct it. If block reuse feeds that stale state into a new
        # prefill, the correction never happens.
        # Currently we only save and reuse context tokens so this does not affect.

        # Tree attention: replay assumes linear token sequence.
        if (spec_config is not None
                and getattr(spec_config, 'use_dynamic_tree', False)):
            logger.info("Replay kernel incompatible with tree attention; "
                        "using legacy MTP path")
            use_replay = False

        # Replay Philox uses PTX cvt.rs.f16x2.f32 which needs 100 <= sm < 120.
        # Flashinfer has a SW fallback at any SM.
        if (stochastic_rounding
                and mamba_params.mamba_ssm_cache_dtype == torch.float16
                and (sm < 100 or sm in (120, 121))):
            logger.info("Replay kernel Philox requires 100 <= sm < 120; "
                        "using legacy MTP path for stochastic rounding support")
            use_replay = False

        # Use replay algorithm for mamba (default is on).
        enforce_disable_replay = os.environ.get('TRTLLM_USE_MAMBA_REPLAY',
                                                '1') == '0'
        if enforce_disable_replay:
            logger.info(
                "Replay kernel is disabled by TRTLLM_USE_MAMBA_REPLAY=0")
            use_replay = False
        else:
            logger.info(
                "Replay kernel is not changed since TRTLLM_USE_MAMBA_REPLAY=1")

        # Stochastic-rounding seeds must live on the cache manager (not be
        # re-created with torch.randint per forward) whenever SR can fire
        # on the fp16 SSM cache.  This mirrors the predicate the mixer uses
        # internally (`_stochastic_rounding_for_flashinfer` /
        # `_stochastic_rounding_for_replay`) so allocation matches consumption.
        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,
            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
            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_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,
            dtype=kv_cache_dtype,
            spec_config=spec_config,
            is_estimating_kv_cache=estimating_kv_cache,
            execution_stream=execution_stream,
            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):
        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)
        # 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, '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,
            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
            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_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,
            dtype=kv_cache_dtype,
            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,
        )
    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)
        binding_model_config = None
        if is_vswa and kv_cache_manager_cls.__name__ == "KVCacheManager":
            binding_model_config = _model_config.get_bindings_model_config(
                tokens_per_block=tokens_per_block,
                kv_cache_config=kv_cache_config,
                spec_config=spec_config)

        # KVCacheManager (V1) doesn't support per-layer head_dim lists;
        # use max for estimation. KVCacheManagerV2 handles lists natively.
        effective_head_dim = (
            max(head_dim) if isinstance(head_dim, list)
            and kv_cache_manager_cls.__name__ == "KVCacheManager" else head_dim)
        kv_cache_manager = kv_cache_manager_cls(
            kv_cache_config,
            kv_cache_type,
            num_layers=num_hidden_layers,
            num_kv_heads=per_layer_num_kv_heads,
            head_dim=effective_head_dim,
            tokens_per_block=tokens_per_block,
            max_seq_len=max_seq_len,
            max_batch_size=max_batch_size,
            mapping=mapping,
            dtype=kv_cache_dtype,
            spec_config=spec_config,
            max_num_tokens=max_num_tokens,
            model_config=binding_model_config,
            max_beam_width=max_beam_width,
            is_draft=is_draft,
            kv_connector_manager=kv_connector_manager
            if not estimating_kv_cache else None,
            sparse_attention_config=sparse_attention_config,
            pretrained_config=config,
            is_estimating_kv_cache=estimating_kv_cache,
            execution_stream=execution_stream,
            layer_mask=layer_mask,
            is_disagg=is_disagg,
            **manager_extra_kwargs,
        )
    # 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":
        if config.quant == "nvfp4" and not is_sm_100f():
            raise RuntimeError(
                "NVFP4 cold-page quantization requires an SM100-family device "
                "(SM100, SM103 or SM107).")
    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,
            pretrained_config=model_engine.model.model_config.pretrained_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; see
    ``should_enable_overlap_headroom``. 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 resolve_max_num_sequences(model_engine,
                              mapping: Mapping,
                              max_batch_size: int,
                              llm_args,
                              max_num_sequences: Optional[int] = None) -> int:
    """Resolve the seat-pool size, preferring an explicit value, then the
    engine's published pool, then a fresh ``compute_max_num_sequences``."""
    if max_num_sequences is not None:
        return max_num_sequences
    engine_seats = getattr(model_engine, "max_num_seq_slots", None)
    if engine_seats is not None:
        return engine_seats
    return compute_max_num_sequences(mapping,
                                     max_batch_size,
                                     llm_args.disable_overlap_scheduler,
                                     enable_overlap_headroom=getattr(
                                         model_engine,
                                         "_enable_overlap_headroom", False))


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_overlap_headroom(mapping: Mapping,
                                   disable_overlap_scheduler: bool,
                                   kv_cache_manager_is_v2: bool,
                                   is_hybrid: bool = False,
                                   has_mrope_delta_cache: bool = False) -> bool:
    """Gate the extra micro-batch of sequence slots.

    True only where a retiring request and the replacement that took its place
    can own a seat at the same time: non-PP, overlap-on, V2 and non-hybrid.

    Under the overlap scheduler a request that emitted its last token keeps its
    seat until the next iteration's ``_process_previous_batch``, which runs
    after that iteration has already admitted new work. Admission excludes such
    retiring requests whenever this gate is on (``ADPRouter`` on the
    attention-DP path, ``_fetch_new_requests`` otherwise), so the replacement
    cohort needs seats of its own. Without the extra micro-batch, a workload
    whose requests retire after one iteration -- a context server, where every
    request finishes in its prefill step -- can only refill the seats the
    previous iteration did not use, and runs at roughly half its batch.

    Widening the pool is only safe when every ``py_seq_slot``-indexed pool is
    sized from ``compute_max_num_sequences``. Two model families size one from
    something else instead, so they keep the single-micro-batch pool:

    * ``is_hybrid``: ``MambaCacheManager`` re-derives its own capacity as
      ``max_batch_size * pp_size``, which a doubled non-PP pool would exhaust.
    * ``has_mrope_delta_cache``: Qwen2/2.5-VL and Qwen3-VL hold
      ``max_num_tokens * pp_size + 1`` MRoPE deltas while indexing them by
      ``py_seq_slot``, relying on ``max_batch_size <= max_num_tokens`` to stay in
      bounds. The top entry is the reserved dummy slot, so a doubled pool first
      aliases the dummy -- silently giving padded requests a real request's
      delta -- and then indexes past the end.
    """
    if is_hybrid or has_mrope_delta_cache or not kv_cache_manager_is_v2:
        return False
    return not mapping.has_pp() and not disable_overlap_scheduler


def validate_seq_slot_pool_covers_admission(max_num_sequences: int,
                                            kv_cache_manager) -> None:
    """Fail at startup if the KV index pool cannot cover the seat pool.

    The check is one-sided on purpose: an index pool narrower than the seat pool
    silently defers admitted requests, while a wider one is legitimate. Managers
    that do not publish an integer ``max_admissible_sequences`` are skipped.
    """
    admissible = getattr(kv_cache_manager, "max_admissible_sequences", None)
    if not isinstance(admissible, int):
        return
    if admissible >= max_num_sequences:
        return
    raise ValueError(
        f"{type(kv_cache_manager).__name__} can lease KV cache indices for "
        f"{admissible} concurrent sequences but the executor's sequence-slot "
        f"pool holds {max_num_sequences}: the index pool is smaller than the "
        "seat pool, so admitted requests would be silently deferred one at a "
        "time (nvbug 6627795). The seat pool must come from "
        "_util.compute_max_num_sequences and the index pool must cover it; a "
        "shortfall means one of them was re-derived from max_batch_size.")


def create_py_executor_instance(
    *,
    dist,
    resources,
    mapping,
    llm_args,
    ctx_chunk_config,
    model_engine,
    start_worker,
    sampler,
    drafter,
    guided_decoder: Optional[GuidedDecoder] = None,
    lora_config: Optional[LoraConfig] = None,
    garbage_collection_gen0_threshold: Optional[int] = None,
    kv_connector_manager: Optional[KvCacheConnectorManager] = None,
    resource_governor_queue=None,
    max_seq_len: Optional[int] = None,
    max_batch_size: Optional[int] = None,
    max_beam_width: Optional[int] = None,
    max_num_tokens: Optional[int] = None,
    peft_cache_config: Optional[PeftCacheConfig] = None,
    scheduler_config: Optional[SchedulerConfig] = None,
    cache_transceiver_config: Optional[CacheTransceiverConfig] = None,
    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

    is_disagg = is_disagg_enabled(cache_transceiver_config)

    max_num_sequences = resolve_max_num_sequences(
        model_engine,
        mapping,
        max_batch_size,
        llm_args,
        max_num_sequences=max_num_sequences)

    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}"
    )
    for key, value in llm_args.extra_resource_managers.items():
        if key in resources:
            raise ValueError(
                f"Cannot overwrite existing resource manager {key}.")
        resources[key] = value

    peft_cache_manager = None
    if lora_config is not None:
        lora_owner = model_engine.initialize_lora_owner()
        # TODO: Refactor dimension resolution into a LoraModuleDimensions
        # 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, owner=lora_owner))
        else:
            assert len(lora_config.lora_target_modules
                       ) >= 1, "Expecting at least one lora target module"
            if not bool(lora_config.trtllm_modules_to_hf_modules):
                lora_config.trtllm_modules_to_hf_modules = get_default_trtllm_modules_to_hf_modules(
                )

        model_binding_config = model_engine.model.model_config.get_bindings_model_config(
            is_disagg=is_disagg)

        num_experts = _try_infer_num_experts(model_engine.model.model_config)

        num_kv_attention_heads_per_layer = model_binding_config.num_kv_heads_per_layer
        if max(num_kv_attention_heads_per_layer) != min(
                num_kv_attention_heads_per_layer):
            logger.warning(
                "Defining LORA with per-layer KV heads is not supported for LORA, using the max number of KV heads per layer"
            )
            num_kv_attention_heads = max(num_kv_attention_heads_per_layer)
        else:
            # all layers have the same number of KV heads
            num_kv_attention_heads = num_kv_attention_heads_per_layer[0]

        pretrained_config = model_engine.model.model_config.pretrained_config

        # Derive shared expert intermediate size from the LoRA adapter
        # weights, which are the source of truth for dimension validation.
        # The model config's shared_expert_intermediate_size may not match
        # the adapter (e.g., upcycled models).
        shared_expert_hidden_size = 0
        if lora_config.lora_dir:
            shared_expert_global = _infer_shared_expert_size_from_adapter(
                lora_config.lora_dir[0])
            if shared_expert_global > 0:
                shared_expert_hidden_size = shared_expert_global // mapping.tp_size

        moe_hidden_size = 0
        moe_intermediate = getattr(pretrained_config, 'moe_intermediate_size',
                                   None)
        if moe_intermediate is not None and moe_intermediate > 0:
            moe_hidden_size = moe_intermediate // mapping.tp_size

        # Mamba dimensions for hybrid models (e.g., Nemotron-H)
        # d_inner = mamba_head_dim * mamba_num_heads
        # d_in_proj = 2 * d_inner + 2 * n_groups * d_state + mamba_num_heads
        mamba_in_proj_size = 0
        mamba_inner_size = 0
        mamba_head_dim = getattr(pretrained_config, 'mamba_head_dim', 0)
        mamba_num_heads = getattr(pretrained_config, 'mamba_num_heads', 0)
        if mamba_head_dim > 0 and mamba_num_heads > 0:
            d_inner = mamba_head_dim * mamba_num_heads
            mamba_inner_size = d_inner // mapping.tp_size
            n_groups = getattr(pretrained_config, 'n_groups', 1)
            d_state = getattr(pretrained_config, 'ssm_state_size', 128)
            d_in_proj = 2 * d_inner + 2 * n_groups * d_state + mamba_num_heads
            mamba_in_proj_size = d_in_proj // mapping.tp_size

        # MoE latent size for latent MoE models (e.g., Nemotron-H SuperV3).
        # Latent projections are replicated (not TP-sharded), so pass the
        # raw config value without dividing by tp_size.
        moe_latent_size = getattr(pretrained_config, 'moe_latent_size', 0) or 0

        # For MoE models with shared experts: replace mlp_* target modules with
        # shared_expert_* equivalents. The shared expert uses different LoRA
        # module types with their own intermediate size. For pure MoE models
        # (no dense MLP layers), mlp_* modules don't exist in the model.
        target_modules = list(lora_config.lora_target_modules)
        if shared_expert_hidden_size > 0:
            has_dense_mlp = bool(
                getattr(pretrained_config, 'mlp_only_layers', None))
            mlp_to_shared_expert = {
                'mlp_h_to_4h': 'shared_expert_h_to_4h',
                'mlp_4h_to_h': 'shared_expert_4h_to_h',
                'mlp_gate': 'shared_expert_gate',
            }
            for mlp_mod, se_mod in mlp_to_shared_expert.items():
                if mlp_mod in target_modules:
                    if se_mod not in target_modules:
                        target_modules.append(se_mod)
                    if not has_dense_mlp:
                        target_modules.remove(mlp_mod)

        lora_modules = LoraModule.create_lora_modules(
            lora_module_names=target_modules,
            hidden_size=model_binding_config.hidden_size,
            mlp_hidden_size=model_binding_config.mlp_hidden_size,
            num_attention_heads=model_binding_config.num_heads,
            num_kv_attention_heads=num_kv_attention_heads,
            attention_head_size=model_binding_config.head_size,
            tp_size=mapping.tp_size,
            num_experts=num_experts,
            shared_expert_hidden_size=shared_expert_hidden_size,
            moe_hidden_size=moe_hidden_size,
            mamba_in_proj_size=mamba_in_proj_size,
            mamba_inner_size=mamba_inner_size,
            moe_latent_size=moe_latent_size)

        model_binding_config.use_lora_plugin = True
        model_binding_config.lora_modules = lora_modules
        model_binding_config.max_lora_rank = lora_config.max_lora_rank

        max_lora_rank = lora_config.max_lora_rank
        num_lora_modules = _compute_num_lora_modules(
            pretrained_config,
            target_modules + lora_config.missing_qkv_modules,
        )

        peft_cache_config_model = PeftCacheConfig(
        ) if peft_cache_config is None else peft_cache_config
        if lora_config.max_loras is not None:
            peft_cache_config_model.num_device_module_layer = \
                max_lora_rank * num_lora_modules * lora_config.max_loras
        if lora_config.max_cpu_loras is not None:
            peft_cache_config_model.num_host_module_layer = \
                max_lora_rank * num_lora_modules * lora_config.max_cpu_loras

        from tensorrt_llm.bindings import WorldConfig
        world_config = WorldConfig(
            tensor_parallelism=mapping.tp_size,
            pipeline_parallelism=mapping.pp_size,
            context_parallelism=mapping.cp_size,
            rank=dist.mapping.rank,
            gpus_per_node=dist.mapping.gpus_per_node,
        )
        lora_model_config = make_lora_model_config(
            model_engine.model,
            target_modules,
            lora_config.trtllm_modules_to_hf_modules,
            lora_config.swap_gate_up_proj_lora_b_weight,
            owner=lora_owner,
        )
        model_engine.set_lora_model_config(lora_model_config)
        peft_cache_manager = PeftCacheManager(
            peft_cache_config=peft_cache_config_model,
            lora_config=lora_config,
            model_config=model_binding_config,
            lora_model_config=lora_model_config,
            world_config=world_config,
            execution_stream=execution_stream,
            initial_data_type=initial_lora_data_type,
        )
        resources[ResourceManagerType.PEFT_CACHE_MANAGER] = peft_cache_manager
        if isinstance(model_engine, PyTorchModelEngine):
            model_engine._init_cuda_graph_lora_manager(lora_config)

    validate_seq_slot_pool_covers_admission(max_num_sequences, kv_cache_manager)
    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).
    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
    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)

    # 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(
            ResourceManagerType.DRAFT_KV_CACHE_MANAGER)
        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,
            kv_cache_manager=kv_cache_manager,
            scheduler_policy=scheduler_policy,
            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,
            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,
            kv_cache_manager=kv_cache_manager.impl
            if kv_cache_manager is not None else None,
            peft_cache_manager=peft_cache_manager.impl
            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,
            no_schedule_until_state=no_schedule_until_state,
            enable_prefix_aware_scheduling=enable_prefix_aware_scheduling,
        )
    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,
            peft_cache_manager.impl if peft_cache_manager is not None else None,
            scheduler_config.capacity_scheduler_policy,
            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,
        )

        mb_scheduler = BindMicroBatchScheduler(
            max_batch_size,
            max_num_tokens,
            ctx_chunk_config,
            no_schedule_until_state=no_schedule_until_state,
        )

        reorder_policy_config = llm_args.reorder_policy_config
        if reorder_policy_config is not None:
            assert reorder_policy_config.policy_name == "AgentTree", "Reorder policy only supports AgentTree for now"
            capacity_scheduler.impl.set_agent_tree_reorder_policy(
                reorder_policy_config.policy_args.agent_percentage,
                reorder_policy_config.policy_args.agent_types,
                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)
            and model_engine.mapping.is_first_pp_rank()):
        # 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
        encoder_cache = model_engine.mm_encoder_cache
        if encoder_cache is None:
            raise RuntimeError(
                "MM encoder item scheduling requires its unified output cache")
        scheduler = scheduler_cls(
            scheduler,
            max_batch_size=model_engine.encoder_batch_size,
            max_num_tokens=model_engine.encoder_max_num_tokens,
            encoder_cache=encoder_cache,
            get_item_cache_keys=model_engine.get_mm_encoder_item_cache_keys,
            bytes_per_encoder_embedding=(
                model_engine.bytes_per_mm_encoder_embedding),
            retain_cache_entries=model_engine.model.encoder_cache_active,
            # V2 resolves actual KV reuse before encoder item selection.
            # V1 keeps whole-request encoder scheduling.
            context_chunk_unit_size=(
                ctx_chunk_config[1]
                if isinstance(kv_cache_manager, KVCacheManagerV2)
                and ctx_chunk_config is not None else None),
        )

    config = model_engine.model.model_config.pretrained_config
    attention_type = AttentionTypeCpp.MLA if is_mla(
        config) else AttentionTypeCpp.DEFAULT

    # For hybrid models, this has both impl and mamba_impl
    mamba_cache_manager = None
    if isinstance(kv_cache_manager, BaseMambaCacheManager):
        mamba_cache_manager = kv_cache_manager

    # A model that declares a disaggregated handoff contributes its semantic cache
    # identity, the record it freezes for each finished context request and, when its
    # drafter's state lives in the separate draft manager, that manager's pages to every
    # transfer. Models without one keep the target-only transfer.
    transfer_kwargs = {}
    cache_identity_fn = getattr(model_engine.model, "disagg_cache_identity",
                                None)
    if (cache_identity_fn is not None and cache_transceiver_config is not None
            and cache_transceiver_config.backend is not None):
        transfer_kwargs["cache_identity"] = cache_identity_fn()
        transfer_kwargs["context_state_exporter"] = getattr(
            model_engine.model, "export_disagg_context_request", None)
        transfer_kwargs["context_state_validator"] = getattr(
            model_engine.model, "validate_disagg_generation_request", None)
        # The drafter's context travels only in draft layers the target's KV cache
        # manager owns; they move with the target's pages.
        get_draft_view = getattr(kv_cache_manager, "get_draft_kv_cache_view",
                                 None)
        if (getattr(model_engine.model, "transfers_draft_kv_cache", False)
                and (get_draft_view is None
                     or get_draft_view() is not kv_cache_manager)):
            raise ValueError(
                "the model's disaggregated handoff carries the drafter's context, "
                "which needs the drafter's KV layers in the target's KV cache "
                "manager; this engine keeps a separate draft KV cache. Run DFlash "
                "with attention_backend TRTLLM and without "
                "TRTLLM_DISAGG_SHARED_DRAFT_KV=0 on both workers.")

    kv_cache_transceiver = create_kv_cache_transceiver(
        mapping, dist, kv_cache_manager, attention_type,
        cache_transceiver_config, mamba_cache_manager, **transfer_kwargs)

    waiting_queue_policy = (scheduler_config.waiting_queue_policy
                            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,
        model_engine=model_engine,
        sampler=sampler,
        drafter=drafter,
        dist=dist,
        max_num_sequences=max_num_sequences,
        disable_overlap_scheduler=llm_args.disable_overlap_scheduler,
        enable_early_first_token_response=llm_args.
        enable_early_first_token_response,
        max_batch_size=max_batch_size,
        max_beam_width=max_beam_width,
        max_draft_len=spec_config.max_draft_len
        if spec_config is not None else 0,
        max_total_draft_tokens=(spec_config.tokens_per_gen_step -
                                1) if spec_config is not None else 0,
        kv_cache_transceiver=kv_cache_transceiver,
        guided_decoder=guided_decoder,
        start_worker=start_worker,
        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,
        peft_cache_config=peft_cache_config,
        virtual_memory_pools=virtual_memory_pools,
        execution_stream=execution_stream,
        waiting_queue_policy=waiting_queue_policy,
        dwdp_manager=dwdp_manager,
        enable_kv_pool_rebalance=llm_args.kv_cache_config.
        enable_kv_pool_rebalance,
        kv_pool_rebalance_check_interval=llm_args.kv_cache_config.
        kv_pool_rebalance_config.check_interval,
    )


def create_torch_sampler_args(
    *,
    max_seq_len: int,
    speculative_config: SpeculativeConfig,
    max_beam_width: int,
    disable_overlap_scheduler: bool,
    enable_async_worker: bool,
    enable_speculative_beam_history_d2h: bool,
    max_num_sequences: int,
):
    # The sampler's per-slot state is indexed by sequence slots, so it must
    # be sized identically to the executor's slot pool.
    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
                              speculative_config.tokens_per_gen_step - 1)

    return TorchSampler.Args(
        max_seq_len=max_seq_len,
        max_draft_len=max_draft_len,
        max_total_draft_tokens=max_total_draft_tokens,
        max_num_sequences=max_num_sequences,
        max_beam_width=max_beam_width,
        disable_overlap_scheduler=disable_overlap_scheduler,
        enable_async_worker=enable_async_worker,
        enable_speculative_beam_history_d2h=enable_speculative_beam_history_d2h,
    )


def instantiate_sampler(
    engine: PyTorchModelEngine,
    llm_args: TorchLlmArgs,
    mapping: Mapping,
    *,
    max_batch_size: int,
    max_beam_width: int,
    mm_encoder_only: bool,
    speculative_config: SpeculativeConfig,
    max_num_sequences: Optional[int] = None,
):
    enable_async_worker = (confidential_compute_enabled()
                           or llm_args.sampler_force_async_worker)

    max_num_sequences = resolve_max_num_sequences(
        engine,
        mapping,
        max_batch_size,
        llm_args,
        max_num_sequences=max_num_sequences)

    sampler_args = create_torch_sampler_args(
        max_seq_len=engine.max_seq_len,
        speculative_config=speculative_config,
        max_beam_width=max_beam_width,
        disable_overlap_scheduler=llm_args.disable_overlap_scheduler,
        enable_async_worker=enable_async_worker,
        enable_speculative_beam_history_d2h=llm_args.
        enable_speculative_beam_history_d2h,
        max_num_sequences=max_num_sequences,
    )
    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)

    if mm_encoder_only:
        # NOTE: handle model outputs specially for mm encoder executor/engine.
        # Only export results on ranks whose responses reach the frontend
        # (see PyExecutor._enqueue_responses): rank 0, plus the attention-DP
        # ranks gathered into it. A CUDA IPC handle exported on any other rank
        # is never opened, so PyTorch would never release its allocation.
        returns_responses = mapping.rank == 0 or (mapping.enable_attention_dp
                                                  and 0 in mapping.tp_group)
        return EarlyStopWithMMResult(return_mm_results=returns_responses)
    if not engine.model.model_config.is_generation:
        # NOTE: choose sampler based on model type
        return EarlyStopSampler()
    return TorchSampler(sampler_args)


_ATTN_MODULES = frozenset({
    "attn_q",
    "attn_k",
    "attn_v",
    "attn_qkv",
    "attn_dense",
    "cross_attn_q",
    "cross_attn_k",
    "cross_attn_v",
})
_MLP_MODULES = frozenset({
    "mlp_h_to_4h",
    "mlp_4h_to_h",
    "mlp_gate",
    "mlp_gate_up",
})


def _compute_num_lora_modules(pretrained_config,
                              all_target_modules: list[str]) -> int:
    """Compute the total number of LoRA module-layer slots for cache sizing.

    For models with per-layer block_configs (e.g. Nemotron-NAS / DeciLM),
    layers with no_op or replace_with_linear attention/FFN cannot host LoRA
    adapters, so they are excluded from the count.  For all other models,
    falls back to the uniform num_hidden_layers x len(target_modules).
    """
    num_layers = pretrained_config.num_hidden_layers
    block_configs = getattr(pretrained_config, "block_configs", None)

    if block_configs is None:
        return num_layers * len(all_target_modules)

    attn_modules = [m for m in all_target_modules if m in _ATTN_MODULES]
    mlp_modules = [m for m in all_target_modules if m in _MLP_MODULES]
    other_modules = [
        m for m in all_target_modules
        if m not in _ATTN_MODULES and m not in _MLP_MODULES
    ]

    def _has_lora_capable_attn(bc):
        return not bc.attention.no_op and not bc.attention.replace_with_linear

    def _has_lora_capable_ffn(bc):
        return not bc.ffn.no_op and not bc.ffn.replace_with_linear

    layers_with_attn = sum(1 for bc in block_configs
                           if _has_lora_capable_attn(bc))
    layers_with_mlp = sum(1 for bc in block_configs
                          if _has_lora_capable_ffn(bc))

    total = (layers_with_attn * len(attn_modules) +
             layers_with_mlp * len(mlp_modules) +
             num_layers * len(other_modules))

    logger.info(f"LoRA module-layer count: {total} "
                f"(attn: {layers_with_attn}x{len(attn_modules)}, "
                f"mlp: {layers_with_mlp}x{len(mlp_modules)}, "
                f"other: {num_layers}x{len(other_modules)}, "
                f"uniform would be {num_layers * len(all_target_modules)})")
    return total


def _infer_shared_expert_size_from_adapter(adapter_dir: str) -> int:
    """Infer shared expert intermediate size from LoRA adapter weights.

    Scans the adapter for shared_expert.down_proj lora_A weights and
    returns the global (unsharded) intermediate size. This is more reliable
    than the model config, which may not match the adapter (e.g., upcycled
    models).
    """
    import json

    try:
        from tensorrt_llm.models.convert_utils import (get_model_path,
                                                       load_state_dict)
        model_path = get_model_path(adapter_dir, "adapter_model")
        if model_path is None:
            return 0
        adapter_weights = load_state_dict(model_path)
        if adapter_weights is None:
            return 0
        for key, tensor in adapter_weights.items():
            if 'shared_expert' in key and 'down_proj' in key and 'lora_A' in key:
                adapter_config_path = os.path.join(adapter_dir,
                                                   "adapter_config.json")
                with open(adapter_config_path) as f:
                    rank = json.load(f).get("r", 0)
                if rank > 0:
                    return tensor.shape[1] if tensor.shape[
                        0] == rank else tensor.shape[0]
    except Exception as e:
        logger.debug(f"Failed to infer shared expert size from adapter: {e}")
    return 0


def _try_infer_num_experts(model_config: ModelConfig) -> int:
    """
    Attempt to infer the number of experts from the model configuration.

    Different MoE models use different attribute names for storing the number of experts,
    so this function checks for various possible names and returns a default of 1 if none are found.
    However, this function is not exhaustive and may miss some cases, so it should be revised.
    """
    config = getattr(model_config, 'pretrained_config', model_config)

    expert_attr_names = [
        'num_experts', 'num_local_experts', 'moe_num_experts',
        'experts_per_router'
    ]
    num_experts = None
    for attr_name in expert_attr_names:
        if hasattr(config, attr_name):
            num_experts = getattr(config, attr_name)
            break

    # Default to 1 for non-MoE models or if no experts attribute is found
    if num_experts is None:
        return 1

    return num_experts


def _adjust_torch_mem_fraction():
    # If true, adjust PyTorch CUDA memory fraction to correspond to the
    # total GPU memory minus the statically allocated engine memory.
    # If false, set the PyTorch CUDA memory fraction to 1.0.
    _limit_torch_cuda_mem_fraction: bool = True

    # FIXME: PyTorch only uses the garbage_collection_threshold setting
    #        if a memory fraction is set, cf.
    #   https://github.com/pytorch/pytorch/blob/cd995bfb2aac8891465809be3ce29543bd524287/c10/cuda/CUDACachingAllocator.cpp#L1357
    logger.debug("Setting PyTorch memory fraction to 1.0")
    torch.cuda.set_per_process_memory_fraction(1.0)

    # FIXME: As soon as
    #     torch.cuda._set_allocator_settings (added in PyTorch 2.8.0-rc1)
    #   or a similar API is available, the warning below should be removed
    #   and the allocator GC threshold be set via the new API instead.
    torch_allocator_config = os.environ.get("PYTORCH_ALLOC_CONF", "")
    torch_mem_threshold_advised = (
        torch.cuda.get_allocator_backend() == "native"
        and "expandable_segments:True" not in torch_allocator_config)
    torch_mem_threshold_set = "garbage_collection_threshold:" in torch_allocator_config
    if torch_mem_threshold_advised and not torch_mem_threshold_set:
        logger.warning(
            "It is recommended to incl. 'garbage_collection_threshold:0.???' or 'backend:cudaMallocAsync'"
            " or 'expandable_segments:True' in PYTORCH_ALLOC_CONF.")

    # NOTE: Even if a memory threshold was not set (cf. warning above), setting a memory
    #       fraction < 1.0 is beneficial, because
    #         https://github.com/pytorch/pytorch/blob/5228986c395dc79f90d2a2b991deea1eef188260/c10/cuda/CUDACachingAllocator.cpp#L2719
    #       and
    #         https://github.com/pytorch/pytorch/blob/5228986c395dc79f90d2a2b991deea1eef188260/c10/cuda/CUDACachingAllocator.cpp#L1240
    #       lead PyTorch to release all unused memory before hitting the set fraction. This
    #       still mitigates OOM, although at a higher performance impact, because it
    #       effectively resets the allocator cache.
    if not _limit_torch_cuda_mem_fraction:
        return
    mem_reserved = torch.cuda.memory_reserved()
    mem_free, mem_total = torch.cuda.mem_get_info()
    safety_margin = 32 * 1024**2
    mem_torch_max = mem_free + mem_reserved - safety_margin
    mem_torch_fraction = mem_torch_max / mem_total
    logger.info(
        f"Setting PyTorch memory fraction to {mem_torch_fraction} ({mem_torch_max / 1024**3} GiB)"
    )
    torch.cuda.set_per_process_memory_fraction(mem_torch_fraction)


def validate_feature_combination(llm_args, model_engine):
    # Validate the flags for features' combination
    def init_feature_status(llm_args) -> Dict[str, bool]:
        assert isinstance(
            llm_args, TorchLlmArgs
        ), "Expect TorchLlmArgs used for feature status validation."
        feature_list = [
            "overlap_scheduler",
            "cuda_graph",
            "attention_dp",
            "disaggregated_serving",
            "chunked_prefill",
            "mtp",
            "eagle3_one_model",
            "kv_cache_reuse",
            "slide_window_attention",
            "guided_decoding",
        ]
        feature_status: Dict[str, bool] = dict.fromkeys(feature_list)
        feature_status[
            "overlap_scheduler"] = not llm_args.disable_overlap_scheduler
        feature_status["cuda_graph"] = llm_args.cuda_graph_config is not None
        feature_status["attention_dp"] = llm_args.enable_attention_dp
        feature_status[
            "disaggregated_serving"] = llm_args.cache_transceiver_config is not None
        feature_status["chunked_prefill"] = llm_args.enable_chunked_prefill
        feature_status["mtp"] = isinstance(llm_args.speculative_config,
                                           MTPDecodingConfig)
        feature_status["eagle3_one_model"] = isinstance(
            llm_args.speculative_config, EagleDecodingConfig)
        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"] = (
            hasattr(model_engine.model.model_config.pretrained_config,
                    "layer_types") and "sliding_attention"
            in model_engine.model.model_config.pretrained_config.layer_types)
        feature_status[
            "guided_decoding"] = llm_args.guided_decoding_backend is not None
        assert all(v is not None for v in feature_status.values()
                   ), "feature status has not been fully initialized."
        assert all(
            k in feature_list for k in
            feature_status.keys()), "unexpected feature type in feature_status."
        return feature_status

    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]] = [
        # Add new conflict rules here in the future
    ]
    for rule in CONFLICT_RULES:
        if all(feature_status[feature] for feature in rule["features"]):
            raise ValueError(rule["message"])