From d713b1d9ea20436114195021ebdb3a6fc0a6cc9e Mon Sep 17 00:00:00 2001
From: Zheyu Fu <zheyuf@nvidia.com>
Date: Thu, 29 Oct 2026 19:46:32 -0700
Subject: [PATCH 01/14] [TRTLLM-17385][feat] MiniMax-M3: host explicit draft
 layer specs in the hybrid KV cache manager

MiniMaxM3KVCacheManagerV2 accepts a standalone DFlash/DSpark drafter's layers as
explicit draft layer specs: dense target layers stay FP8, index-K buffers follow the
spec-built layers, size-only coalescing keeps the target pools, and draft layers use
the base view and keep attention-op pools. A windowed draft lifecycle is funded with
its admissible sequences' windows plus one chunk instead of a bytes-per-token share.

Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com>
---
 .../sparse/minimax_m3/cache_manager.py        |  39 +++-
 .../kv_cache/kv_cache_manager_v2.py           |  76 ++++++-
 .../msa/test_minimax_m3_unified_draft_kv.py   | 199 ++++++++++++++++++
 3 files changed, 306 insertions(+), 8 deletions(-)
 create mode 100644 tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_unified_draft_kv.py

diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py
index 7722f76ea664..babb5aefee29 100644
--- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py
+++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py
@@ -28,12 +28,14 @@
 
 from __future__ import annotations
 
+import dataclasses
 from typing import TYPE_CHECKING, List, Optional, Sequence, Tuple
 
 import torch
 
 from tensorrt_llm._torch.disaggregation.resource.page import MapperKind
 from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import (
+    KVCacheLayerOwner,
     KVCacheManagerV2,
     Role,
     _BasePageTableMaterializer,
@@ -336,6 +338,10 @@ class MiniMaxM3KVCacheManagerV2(KVCacheManagerV2):
     # Extra page sizes trtllm-gen may use with this manager (see
     # FlashInferTrtllmGenFmha); set with the virtual pools.
     trtllm_gen_extra_tokens_per_block: frozenset = frozenset()
+    # A standalone DFlash/DSpark drafter's layers may join as explicit draft layer specs: their own
+    # lifecycle keeps them out of the index-K pools, and size-only coalescing keeps the target pools.
+    hosts_unified_draft_layers = True
+    _separate_nvfp4_coalescing = False
 
     def __init__(
         self,
@@ -403,12 +409,23 @@ def __init__(
         # Filled at the end of __init__; the base __init__ may already call get_buffers.
         self._kv_buffers: dict[tuple[int, str], torch.Tensor] = {}
         self._block_scale_buffers: dict[tuple[int, str], torch.Tensor] = {}
+        layer_specs = kwargs.get("layer_specs")
+        if layer_specs is not None and kwargs.get("dtype") == DataType.NVFP4:
+            # The hybrid layout of _build_cache_config as specs: FP8 on the dense target layers.
+            kwargs["layer_specs"] = tuple(
+                dataclasses.replace(spec, dtype=DataType.FP8, buffer_roles=(Role.KEY, Role.VALUE))
+                if spec.owner is KVCacheLayerOwner.TARGET
+                and spec.dtype == DataType.NVFP4
+                and spec.global_layer_id not in self.sparse_layer_ids
+                else spec
+                for spec in layer_specs
+            )
         super().__init__(*args, **kwargs)
 
         if self.dtype == DataType.NVFP4:
             dense_target_layers = [
                 layer
-                for layer in self.layer_offsets
+                for layer in self.target_layer_ids
                 if layer not in self.sparse_layer_ids and layer not in self._shared_draft_layer_ids
             ]
             logger.info(
@@ -515,7 +532,7 @@ def _build_cache_config(self, config):
 
     def get_layer_bytes_per_token(self, local_layer_idx: int, data_role: Role):
         """Report the hybrid sparse-NVFP4 / dense-FP8 storage footprint."""
-        if self.dtype != DataType.NVFP4:
+        if self.layer_specs is not None or self.dtype != DataType.NVFP4:
             return super().get_layer_bytes_per_token(local_layer_idx, data_role)
         global_layer_idx = int(self.pp_layers[int(local_layer_idx)])
         if global_layer_idx in self.sparse_layer_ids:
@@ -567,6 +584,9 @@ def get_draft_kv_cache_view(self) -> Optional["MiniMaxM3DraftKVCacheView"]:
         method rather than a property so ``getattr`` fetches it without
         executing it (see ``get_draft_kv_cache_manager``).
         """
+        if self.layer_specs is not None:
+            # Explicit draft layers are addressed through this manager itself.
+            return super().get_draft_kv_cache_view()
         if self.dtype != DataType.NVFP4 or self.is_draft or not self._shared_draft_layer_ids:
             return None
         if self.enable_swa_scratch_reuse:
@@ -608,7 +628,8 @@ def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None:
         pool's slot ids; :meth:`copy_batch_block_offsets` scales them.
         """
         super()._prepare_page_table_tensor(index_mapper_capacity)
-        if self.dtype == DataType.NVFP4:
+        # Explicit draft layers keep the pools: the unified drafter reads their block offsets.
+        if self.dtype == DataType.NVFP4 and self.layer_specs is None:
             # Nothing reads a hybrid cache's own block offsets: MSA layers address pages
             # by slot id and the shared draft layer reads MiniMaxM3DraftKVCacheView's
             # table. Without pools, a new reader fails instead of reading unfilled offsets.
@@ -683,8 +704,8 @@ def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None:
         )
 
     def _prepare_swa_scratch_copy_tensors(self, index_mapper_capacity: int) -> None:
-        # Staging for the per-layer block-offset copy, which a hybrid NVFP4 cache never runs.
-        if self.dtype != DataType.NVFP4:
+        # Staging for the per-layer block-offset copy, which a pool-less hybrid NVFP4 cache never runs.
+        if self.dtype != DataType.NVFP4 or self.layer_specs is not None:
             super()._prepare_swa_scratch_copy_tensors(index_mapper_capacity)
 
     def copy_batch_block_offsets(
@@ -896,7 +917,15 @@ def get_buffers(
         at K's base, then slices ``[:, :2]`` to extract K+V. The slice
         preserves the dim-0 stride (``scale * page_stride``), so
         ``view[s, 0/1, ...]`` lands on this layer's K/V at slot ``s``.
+
+        Explicit draft layers keep the base view: their own lifecycle's pool holds only their K/V.
         """
+        if self.layer_specs is not None and self.get_layer_owner(layer_idx) is (
+            KVCacheLayerOwner.DRAFT
+        ):
+            return super().get_buffers(
+                layer_idx, kv_layout if kv_layout is not None else self._main_kv_layout_name()
+            )
         if kv_layout is None:
             kv_layout = self._main_kv_view_layout
         if kv_layout not in ("HND", "NHD"):
diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py
index 42b621bde4c1..deea92e10216 100644
--- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py
+++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py
@@ -1816,6 +1816,10 @@ class KVCacheManagerV2(BaseResourceManager):
     # need snapshotted (request_prefix_snapshot); kept apart from the request's
     # own expect_snapshot_points, which other code rewrites. Created on first use.
     _waiter_snapshot_points: Optional[Dict[int, set[int]]] = None
+    # A subclass whose extra buffers may join explicit layer specs (draft layers then get their own lifecycle).
+    hosts_unified_draft_layers: bool = False
+    # Layer specs coalesce NVFP4 data and scales apart; a subclass addressing its pools by size opts out.
+    _separate_nvfp4_coalescing: bool = True
 
     def __init__(
         self,
@@ -1861,6 +1865,7 @@ def __init__(
         if (
             self.layer_specs is not None
             and type(self)._extra_buffers_per_layer is not KVCacheManagerV2._extra_buffers_per_layer
+            and not self.hosts_unified_draft_layers
         ):
             raise ValueError(
                 f"{type(self).__name__} does not support layer_specs with "
@@ -3729,8 +3734,13 @@ def _build_base_config(
                         )
                     )
 
+        initial_pool_ratio = kv_cache_config.pool_ratio
         if self.layer_specs is not None:
             scale_roles = (Role.KEY_BLOCK_SCALE, Role.VALUE_BLOCK_SCALE)
+            # Subclass buffers beyond K/V/scales (e.g. MiniMax-M3's index-K) follow the spec-built ones.
+            extra_buffers_per_layer = (
+                self._extra_buffers_per_layer(tokens_per_block=tokens_per_block) or {}
+            )
             layer_configs = [
                 AttentionLayerConfig(
                     layer_id=LayerId(local_layer_id),
@@ -3741,18 +3751,25 @@ def _build_base_config(
                             # Keep NVFP4 data (1), its scales (2), and other data (0)
                             # separate while preserving matching K/V page order.
                             coalescing_group=(
-                                1 + int(role in scale_roles) if layer.dtype == DataType.NVFP4 else 0
+                                1 + int(role in scale_roles)
+                                if layer.dtype == DataType.NVFP4 and self._separate_nvfp4_coalescing
+                                else 0
                             ),
                         )
                         for role in (Role.KEY, Role.VALUE, *scale_roles)
                         if role in layer.buffer_roles
-                    ],
+                    ]
+                    + list(extra_buffers_per_layer.get(local_layer_id, ())),
                     sliding_window_size=self.max_attention_window_vec[local_layer_id],
                     num_sink_tokens=layer.num_sink_tokens,
                     reuse_policy=layer.reuse_policy,
                 )
                 for local_layer_id, layer in enumerate(self.layer_specs)
             ]
+            if initial_pool_ratio is None and typical_step is None:
+                initial_pool_ratio = self._windowed_draft_pool_ratio(
+                    layer_configs, cache_tiers[0].quota, tokens_per_block
+                )
         else:
             # Subclasses (e.g. MiniMax-M3 sparse cache) can register additional
             # per-layer BufferConfig entries — for example a sparse index-K
@@ -3827,7 +3844,7 @@ def _build_base_config(
                 kv_cache_config.enable_block_reuse
                 and self.block_reuse_policy != BlockReusePolicy.ALL_REUSABLE
             ),
-            initial_pool_ratio=kv_cache_config.pool_ratio,
+            initial_pool_ratio=initial_pool_ratio,
             pool_rebalance=PoolRebalanceConfig(
                 min_sampled_kv_caches=rebalance_cfg.min_sampled_kv_caches,
                 cooldown_secs=rebalance_cfg.cooldown_secs,
@@ -3891,6 +3908,59 @@ def _allocate_cache_pools(
             f"external cache storage ({external_bytes} bytes)."
         )
 
+    def _windowed_draft_pool_ratio(
+        self,
+        layer_configs: List[AttentionLayerConfig],
+        gpu_quota: int,
+        tokens_per_block: int,
+    ) -> Optional[List[float]]:
+        """Hot-tier ``[target, draft]`` byte ratios when the draft layers form one windowed lifecycle.
+
+        Without a typical step the backend splits by bytes per token, handing the draft layers
+        as many tokens as the target. They need only every admissible sequence's window plus
+        one context chunk, which is what the KV cache cost already charges for them.
+        """
+        owners = [spec.owner for spec in self.layer_specs]
+        keys = [(c.sliding_window_size, c.num_sink_tokens, c.reuse_policy) for c in layer_configs]
+        target_keys = {key for key, owner in zip(keys, owners) if owner is KVCacheLayerOwner.TARGET}
+        draft_keys = {key for key, owner in zip(keys, owners) if owner is KVCacheLayerOwner.DRAFT}
+        # Lifecycle ids follow first appearance, so the target's must come first.
+        if (
+            len(target_keys) != 1
+            or len(draft_keys) != 1
+            or target_keys == draft_keys
+            or not next(iter(draft_keys))[0]
+            or owners[0] is not KVCacheLayerOwner.TARGET
+        ):
+            return None
+        draft_specs = [spec for spec in self.layer_specs if spec.owner is KVCacheLayerOwner.DRAFT]
+        draft_sizes = [spec.bytes_per_token() for spec in draft_specs]
+        _, window_bytes_per_sequence = _estimate_swa_cache_size(
+            draft_sizes,
+            [next(iter(draft_keys))[0]] * len(draft_specs),
+            tokens_per_block,
+            context=False,
+            scratch=False,
+            generation_capacity_headroom=self._generation_kv_capacity_headroom,
+            layer_specs=draft_specs,
+        )
+        # Sequences that can hold KV at once (see max_admissible_sequences), counting the overlap
+        # scheduler's retiring request: the pool is sized before the scheduler mode is known.
+        num_sequences = (
+            self.max_batch_size
+            * self.mapping.pp_size
+            * (2 if self.is_disagg or not self.mapping.has_pp() else 1)
+        )
+        draft_bytes = window_bytes_per_sequence * num_sequences + sum(draft_sizes) * (
+            self.max_num_tokens or 0
+        )
+        bytes_per_block = {KVCacheLayerOwner.TARGET: 0, KVCacheLayerOwner.DRAFT: 0}
+        for config, owner in zip(layer_configs, owners):
+            bytes_per_block[owner] += sum(buffer.size for buffer in config.buffers)
+        token_share = bytes_per_block[KVCacheLayerOwner.DRAFT] / sum(bytes_per_block.values())
+        draft_ratio = min(draft_bytes / gpu_quota, token_share)
+        return [1.0 - draft_ratio, draft_ratio]
+
     def _build_cache_config(self, config: KVCacheManagerConfigPy) -> KVCacheManagerConfigPy:
         """Customize the general cache config for a specialized cache manager."""
         return config
diff --git a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_unified_draft_kv.py b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_unified_draft_kv.py
new file mode 100644
index 000000000000..b1a7679c2930
--- /dev/null
+++ b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_unified_draft_kv.py
@@ -0,0 +1,199 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""A standalone DSpark drafter's layers as explicit draft layer specs of the MiniMax-M3 manager."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.cache_manager import (
+    MiniMaxM3KVCacheManagerV2,
+)
+from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import (
+    KVCacheLayerOwner,
+    KVCacheLayerSpec,
+    KVCacheManagerV2,
+    Role,
+)
+from tensorrt_llm.bindings import DataType
+from tensorrt_llm.bindings.internal.batch_manager import CacheType
+from tensorrt_llm.llmapi.llm_args import KvCacheConfig, MiniMaxM3SparseAttentionConfig
+from tensorrt_llm.mapping import Mapping
+
+NUM_TARGET = 4  # M3 convention without a checkpoint: layers 0..2 dense, 3 sparse
+NUM_DRAFT = 2
+DRAFT_WINDOW = 256
+HEADS = 2
+HEAD_DIM = 128
+
+
+def _specs() -> tuple:
+    """Target layers in the NVFP4 cache dtype, then a windowed BF16 drafter."""
+    scale_roles = (Role.KEY_BLOCK_SCALE, Role.VALUE_BLOCK_SCALE)
+    target = [
+        KVCacheLayerSpec(
+            global_layer_id=layer,
+            owner=KVCacheLayerOwner.TARGET,
+            dtype=DataType.NVFP4,
+            num_kv_heads=HEADS,
+            head_dim=HEAD_DIM,
+            buffer_roles=(Role.KEY, Role.VALUE, *scale_roles),
+        )
+        for layer in range(NUM_TARGET)
+    ]
+    draft = [
+        KVCacheLayerSpec(
+            global_layer_id=NUM_TARGET + layer,
+            owner=KVCacheLayerOwner.DRAFT,
+            dtype=DataType.BF16,
+            num_kv_heads=HEADS,
+            head_dim=HEAD_DIM,
+            buffer_roles=(Role.KEY, Role.VALUE),
+            sliding_window_size=DRAFT_WINDOW,
+        )
+        for layer in range(NUM_DRAFT)
+    ]
+    return tuple(target + draft)
+
+
+def _manager(layer_specs) -> MiniMaxM3KVCacheManagerV2:
+    """A hybrid NVFP4 MSA manager, with or without explicit draft layer specs."""
+    return MiniMaxM3KVCacheManagerV2(
+        KvCacheConfig(
+            max_gpu_total_bytes=64 << 20,
+            host_cache_size=0,
+            enable_block_reuse=False,
+            enable_swa_scratch_reuse=False,
+        ),
+        CacheType.SELF,
+        num_layers=NUM_TARGET,
+        num_kv_heads=HEADS,
+        head_dim=HEAD_DIM,
+        tokens_per_block=128,
+        max_seq_len=1024,
+        max_batch_size=2,
+        max_num_tokens=256,
+        mapping=Mapping(world_size=1, rank=0, tp_size=1, pp_size=1),
+        dtype=DataType.NVFP4,
+        vocab_size=16,
+        sparse_attention_config=MiniMaxM3SparseAttentionConfig(
+            implementation="msa", indexer_kv_dtype="fp8"
+        ),
+        layer_specs=layer_specs,
+    )
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA cache pools")
+def test_draft_layer_specs_join_the_hybrid_nvfp4_manager() -> None:
+    """Draft layers join the hybrid manager without changing the target's pages."""
+    torch.cuda.init()
+    plain = _manager(None)
+    unified = _manager(_specs())
+    try:
+        assert unified.target_layer_ids == tuple(range(NUM_TARGET))
+        assert unified.draft_layer_ids == tuple(range(NUM_TARGET, NUM_TARGET + NUM_DRAFT))
+        # Dense target layers keep the FP8 half of the hybrid layout; draft layers their own dtype.
+        dtypes = [unified.get_layer_dtype(layer) for layer in unified.pp_layers]
+        assert dtypes == [DataType.FP8] * 3 + [DataType.NVFP4] + [DataType.BF16] * NUM_DRAFT
+        for layer in range(NUM_TARGET):
+            for role in (Role.KEY, Role.VALUE):
+                assert unified.impl.get_page_stride(
+                    unified.layer_offsets[layer], role
+                ) == plain.impl.get_page_stride(plain.layer_offsets[layer], role)
+        assert unified.get_index_k_buffer(3).shape[1:] == plain.get_index_k_buffer(3).shape[1:]
+
+        # Draft layers: their own lifecycle, the base view in their dtype.
+        groups = unified.layer_to_pool_mapping_dict
+        draft_groups = {groups[unified.layer_offsets[layer]] for layer in unified.draft_layer_ids}
+        target_groups = {groups[unified.layer_offsets[layer]] for layer in unified.target_layer_ids}
+        assert len(draft_groups) == 1 and not draft_groups & target_groups
+        view = unified.get_buffers(NUM_TARGET, kv_layout="HND")
+        assert view.dtype == torch.bfloat16
+        assert tuple(view.shape[1:]) == (2, HEADS, 128, HEAD_DIM)
+        assert unified.get_draft_kv_cache_view() is unified
+        assert plain.get_draft_kv_cache_view() is None
+
+        # The unified drafter reads its layers' block offsets, so they keep attention-op pools.
+        assert plain.num_attention_op_pools == 0
+        assert unified.num_attention_op_pools > 0
+        (request,) = unified.add_dummy_requests([7], token_nums=[200])
+        offsets = torch.zeros(
+            unified.num_attention_op_pools,
+            1,
+            2,
+            unified.max_blocks_per_seq,
+            dtype=torch.int32,
+            device="cuda",
+        )
+        unified.copy_batch_block_offsets(offsets, [7], 1, 1, 1)
+        torch.cuda.synchronize()
+        pool = int(unified.kv_cache_pool_mapping[unified.layer_offsets[NUM_TARGET], 0])
+        slots = [s for s in unified.get_batch_cache_indices([7], layer_idx=NUM_TARGET)[0] if s >= 0]
+        assert len(slots) == 2
+        # The row the drafter reads holds the draft layer's K page of each slot, V right after it.
+        scale = unified.get_layer_page_index_scale(NUM_TARGET)
+        pages = [[slot * scale + plane for slot in slots] for plane in (0, 1)]
+        assert offsets[pool, 0, :, : len(slots)].tolist() == pages
+        unified.free_resources(request)
+
+        # The draft lifecycle is sized by its window, not by its share of the bytes per token.
+        target_ratio, draft_ratio = unified.kv_cache_manager_py_config.initial_pool_ratio
+        assert target_ratio + draft_ratio == pytest.approx(1.0)
+        assert (
+            0
+            < draft_ratio
+            < 2 * HEADS * HEAD_DIM * 2 * NUM_DRAFT / (unified.get_cache_bytes_per_token())
+        )
+    finally:
+        unified.shutdown()
+        plain.shutdown()
+
+
+@pytest.mark.cpu_only
+@pytest.mark.parametrize("is_disagg", [False, True])
+def test_window_sized_draft_ratio(is_disagg: bool) -> None:
+    """A windowed draft lifecycle gets its windows plus one chunk, not a bytes-per-token share."""
+    manager = object.__new__(KVCacheManagerV2)
+    manager.layer_specs = _specs()[NUM_TARGET - 1 :]
+    manager.max_batch_size = 2
+    manager.max_num_tokens = 256
+    manager.mapping = Mapping()
+    manager.is_disagg = is_disagg
+    manager._generation_kv_capacity_headroom = 1
+    configs = [
+        SimpleNamespace(
+            sliding_window_size=spec.sliding_window_size,
+            num_sink_tokens=None,
+            reuse_policy=spec.reuse_policy,
+            buffers=[SimpleNamespace(size=spec.bytes_per_token() * 128)],
+        )
+        for spec in manager.layer_specs
+    ]
+    draft_bytes_per_token = 2 * HEADS * HEAD_DIM * 2 * NUM_DRAFT
+    # Four admissible sequences hold three pages of window each, plus one 256-token chunk.
+    draft_bytes = (4 * 3 * 128 + 256) * draft_bytes_per_token
+    assert manager._windowed_draft_pool_ratio(configs, 1 << 30, 128) == pytest.approx(
+        [1 - draft_bytes / (1 << 30), draft_bytes / (1 << 30)]
+    )
+    # A quota too small for the windows keeps the bytes-per-token split.
+    token_share = draft_bytes_per_token / sum(
+        spec.bytes_per_token() for spec in manager.layer_specs
+    )
+    assert manager._windowed_draft_pool_ratio(configs, 1 << 20, 128)[1] == pytest.approx(
+        token_share
+    )
+    # Without a windowed draft lifecycle the backend keeps its own split.
+    assert manager._windowed_draft_pool_ratio(configs[:1], 1 << 30, 128) is None

From 4979adfa59c7fee7887ddb64edc96ae478d0dec2 Mon Sep 17 00:00:00 2001
From: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com>
Date: Thu, 29 Oct 2026 20:19:40 -0700
Subject: [PATCH 02/14] [TRTLLM-17385][fix] Declare
 KVCacheManagerV2.layer_specs on the class

MiniMaxM3KVCacheManagerV2 now reads layer_specs in methods that tests and helpers call
on managers built without __init__, so the attribute defaults to None on the class.

Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com>
---
 tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py
index deea92e10216..81548aa75978 100644
--- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py
+++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py
@@ -1816,6 +1816,8 @@ class KVCacheManagerV2(BaseResourceManager):
     # need snapshotted (request_prefix_snapshot); kept apart from the request's
     # own expect_snapshot_points, which other code rewrites. Created on first use.
     _waiter_snapshot_points: Optional[Dict[int, set[int]]] = None
+    # Declared on the class so managers built without __init__ have no layer specs.
+    layer_specs: Optional[Tuple["KVCacheLayerSpec", ...]] = None
     # A subclass whose extra buffers may join explicit layer specs (draft layers then get their own lifecycle).
     hosts_unified_draft_layers: bool = False
     # Layer specs coalesce NVFP4 data and scales apart; a subclass addressing its pools by size opts out.

From 4efc6988009af616a916d751cb21be2d044e8db3 Mon Sep 17 00:00:00 2001
From: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com>
Date: Thu, 29 Oct 2026 21:29:44 -0700
Subject: [PATCH 03/14] [TRTLLM-17385][fix] MiniMax-M3: report index-K bytes
 for a manager with explicit layer specs

Layer specs describe K/V and their scales; index-K is the manager's own extra buffer.
With specs, get_layer_bytes_per_token(INDEX_KEY) asked the spec and raised, which fails
manager construction wherever the index-K buffer is sized through it.

Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com>
---
 .../backends/sparse/minimax_m3/cache_manager.py        | 10 +++++++++-
 .../sparse/msa/test_minimax_m3_unified_draft_kv.py     |  6 ++++++
 2 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py
index babb5aefee29..e4c68b5e8659 100644
--- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py
+++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py
@@ -532,7 +532,15 @@ def _build_cache_config(self, config):
 
     def get_layer_bytes_per_token(self, local_layer_idx: int, data_role: Role):
         """Report the hybrid sparse-NVFP4 / dense-FP8 storage footprint."""
-        if self.layer_specs is not None or self.dtype != DataType.NVFP4:
+        if self.layer_specs is not None:
+            # Explicit layer specs carry K/V and scales; index-K is this manager's extra buffer.
+            if data_role == Role.INDEX_KEY:
+                if int(self.pp_layers[int(local_layer_idx)]) not in self.sparse_layer_ids:
+                    return 0
+                index_dtype = self._torch_dtype_for_index_cache()
+                return self.sparse_index_dim * torch.tensor([], dtype=index_dtype).element_size()
+            return super().get_layer_bytes_per_token(local_layer_idx, data_role)
+        if self.dtype != DataType.NVFP4:
             return super().get_layer_bytes_per_token(local_layer_idx, data_role)
         global_layer_idx = int(self.pp_layers[int(local_layer_idx)])
         if global_layer_idx in self.sparse_layer_ids:
diff --git a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_unified_draft_kv.py b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_unified_draft_kv.py
index b1a7679c2930..c4de8e052149 100644
--- a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_unified_draft_kv.py
+++ b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_unified_draft_kv.py
@@ -114,6 +114,12 @@ def test_draft_layer_specs_join_the_hybrid_nvfp4_manager() -> None:
                     unified.layer_offsets[layer], role
                 ) == plain.impl.get_page_stride(plain.layer_offsets[layer], role)
         assert unified.get_index_k_buffer(3).shape[1:] == plain.get_index_k_buffer(3).shape[1:]
+        # Index-K is the manager's own buffer, not a layer-spec role.
+        index_k_bytes = [
+            unified.get_layer_bytes_per_token(unified.layer_offsets[layer], Role.INDEX_KEY)
+            for layer in unified.pp_layers
+        ]
+        assert index_k_bytes == [0, 0, 0, 128] + [0] * NUM_DRAFT
 
         # Draft layers: their own lifecycle, the base view in their dtype.
         groups = unified.layer_to_pool_mapping_dict

From 2d9090a56d2552b1734f10b462cd1a5488bfb703 Mon Sep 17 00:00:00 2001
From: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com>
Date: Fri, 30 Oct 2026 13:25:21 -0700
Subject: [PATCH 04/14] [TRTLLM-17385][fix] MiniMax-M3: keep draft-layer
 indices and pool ratio consistent

Draft-layer page indices take the base conversion that matches their base view.
The window-sized draft ratio keys layers by backend lifecycle and needs a chunk bound.
An NVFP4 draft layer is rejected: block scales belong to the sparse target layers.

Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com>
---
 .../sparse/minimax_m3/cache_manager.py        | 26 ++++++-
 .../kv_cache/kv_cache_manager_v2.py           | 16 +++-
 .../msa/test_minimax_m3_unified_draft_kv.py   | 76 +++++++++++++++----
 3 files changed, 99 insertions(+), 19 deletions(-)

diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py
index e4c68b5e8659..bc5d05b211f5 100644
--- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py
+++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py
@@ -410,6 +410,12 @@ def __init__(
         self._kv_buffers: dict[tuple[int, str], torch.Tensor] = {}
         self._block_scale_buffers: dict[tuple[int, str], torch.Tensor] = {}
         layer_specs = kwargs.get("layer_specs")
+        if layer_specs is not None and any(
+            spec.owner is KVCacheLayerOwner.DRAFT and spec.dtype == DataType.NVFP4
+            for spec in layer_specs
+        ):
+            # Block scales here belong to the MSA sparse target layers only.
+            raise ValueError(f"{type(self).__name__} does not host NVFP4 draft layers.")
         if layer_specs is not None and kwargs.get("dtype") == DataType.NVFP4:
             # The hybrid layout of _build_cache_config as specs: FP8 on the dense target layers.
             kwargs["layer_specs"] = tuple(
@@ -1077,6 +1083,13 @@ def _kv_pool_mapping_offset(self, layer_id, layer_group_id, key_base_addr) -> in
         )
         return layers_by_addr.index(int(layer_id))
 
+    def _draft_pool_ids(self) -> frozenset:
+        """Pools that hold the explicit draft layers."""
+        return frozenset(
+            self.layer_to_pool_mapping_dict[self.layer_offsets[layer]]
+            for layer in self.draft_layer_ids
+        )
+
     def _get_batch_cache_indices_by_pool_id(
         self,
         request_ids,
@@ -1102,7 +1115,8 @@ def _get_batch_cache_indices_by_pool_id(
         ``BAD_PAGE_INDEX`` slots remain ``-1`` here because disaggregation's
         :class:`KVRegionExtractorV1` filters ``region_ids >= 0``.
         :meth:`get_block_ids_per_seq` maps them to zero for the attention
-        metadata's padded tensor.
+        metadata's padded tensor. Pools of explicit draft layers keep the
+        base conversion, matching their base :meth:`get_buffers` view.
 
         Args:
             request_ids: Request IDs whose page-index rows are returned.
@@ -1117,6 +1131,16 @@ def _get_batch_cache_indices_by_pool_id(
             kv_factor: Kept for compatibility with the base virtual method;
                 M3 returns raw slot IDs without KV aggregation.
         """
+        if self.layer_specs is not None and pool_id in self._draft_pool_ids():
+            # Draft layers keep the base get_buffers() view, so their indices take the base conversion.
+            return super()._get_batch_cache_indices_by_pool_id(
+                request_ids,
+                pool_id=pool_id,
+                is_kv_aggregate=is_kv_aggregate,
+                num_blocks_per_seq=num_blocks_per_seq,
+                index_scale=index_scale,
+                kv_factor=kv_factor,
+            )
         res = []
         for req_idx, req_id in enumerate(request_ids):
             kv_cache = self.kv_cache_map[req_id]
diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py
index 81548aa75978..9d6dfa10f2ea 100644
--- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py
+++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py
@@ -3922,8 +3922,18 @@ def _windowed_draft_pool_ratio(
         as many tokens as the target. They need only every admissible sequence's window plus
         one context chunk, which is what the KV cache cost already charges for them.
         """
+        if self.max_num_tokens is None:
+            # Without a context-chunk bound the backend's split is the safe one.
+            return None
         owners = [spec.owner for spec in self.layer_specs]
-        keys = [(c.sliding_window_size, c.num_sink_tokens, c.reuse_policy) for c in layer_configs]
+        # Key layers as the backend groups them into lifecycles (sinks count in whole blocks).
+        life_cycles = [
+            AttnLifeCycle.make(
+                c.sliding_window_size, c.num_sink_tokens, tokens_per_block, c.reuse_policy
+            )
+            for c in layer_configs
+        ]
+        keys = [(lc.window_size, lc.num_sink_blocks, lc.reuse_policy) for lc in life_cycles]
         target_keys = {key for key, owner in zip(keys, owners) if owner is KVCacheLayerOwner.TARGET}
         draft_keys = {key for key, owner in zip(keys, owners) if owner is KVCacheLayerOwner.DRAFT}
         # Lifecycle ids follow first appearance, so the target's must come first.
@@ -3953,8 +3963,8 @@ def _windowed_draft_pool_ratio(
             * self.mapping.pp_size
             * (2 if self.is_disagg or not self.mapping.has_pp() else 1)
         )
-        draft_bytes = window_bytes_per_sequence * num_sequences + sum(draft_sizes) * (
-            self.max_num_tokens or 0
+        draft_bytes = (
+            window_bytes_per_sequence * num_sequences + sum(draft_sizes) * self.max_num_tokens
         )
         bytes_per_block = {KVCacheLayerOwner.TARGET: 0, KVCacheLayerOwner.DRAFT: 0}
         for config, owner in zip(layer_configs, owners):
diff --git a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_unified_draft_kv.py b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_unified_draft_kv.py
index c4de8e052149..a8b06fe4f2ca 100644
--- a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_unified_draft_kv.py
+++ b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_unified_draft_kv.py
@@ -14,7 +14,9 @@
 # limitations under the License.
 """A standalone DSpark drafter's layers as explicit draft layer specs of the MiniMax-M3 manager."""
 
+from dataclasses import replace
 from types import SimpleNamespace
+from typing import Optional
 
 import pytest
 import torch
@@ -40,8 +42,8 @@
 HEAD_DIM = 128
 
 
-def _specs() -> tuple:
-    """Target layers in the NVFP4 cache dtype, then a windowed BF16 drafter."""
+def _specs(draft_window: Optional[int] = DRAFT_WINDOW) -> tuple[KVCacheLayerSpec, ...]:
+    """Target layers in the NVFP4 cache dtype, then a BF16 drafter (windowed by default)."""
     scale_roles = (Role.KEY_BLOCK_SCALE, Role.VALUE_BLOCK_SCALE)
     target = [
         KVCacheLayerSpec(
@@ -62,14 +64,14 @@ def _specs() -> tuple:
             num_kv_heads=HEADS,
             head_dim=HEAD_DIM,
             buffer_roles=(Role.KEY, Role.VALUE),
-            sliding_window_size=DRAFT_WINDOW,
+            sliding_window_size=draft_window,
         )
         for layer in range(NUM_DRAFT)
     ]
     return tuple(target + draft)
 
 
-def _manager(layer_specs) -> MiniMaxM3KVCacheManagerV2:
+def _manager(layer_specs: Optional[tuple[KVCacheLayerSpec, ...]]) -> MiniMaxM3KVCacheManagerV2:
     """A hybrid NVFP4 MSA manager, with or without explicit draft layer specs."""
     return MiniMaxM3KVCacheManagerV2(
         KvCacheConfig(
@@ -147,12 +149,26 @@ def test_draft_layer_specs_join_the_hybrid_nvfp4_manager() -> None:
         unified.copy_batch_block_offsets(offsets, [7], 1, 1, 1)
         torch.cuda.synchronize()
         pool = int(unified.kv_cache_pool_mapping[unified.layer_offsets[NUM_TARGET], 0])
-        slots = [s for s in unified.get_batch_cache_indices([7], layer_idx=NUM_TARGET)[0] if s >= 0]
+        slots = [
+            s
+            for s in unified.get_batch_cache_indices([7], layer_idx=NUM_TARGET, raw_indices=True)[0]
+            if s >= 0
+        ]
         assert len(slots) == 2
         # The row the drafter reads holds the draft layer's K page of each slot, V right after it.
         scale = unified.get_layer_page_index_scale(NUM_TARGET)
         pages = [[slot * scale + plane for slot in slots] for plane in (0, 1)]
         assert offsets[pool, 0, :, : len(slots)].tolist() == pages
+        # Draft-layer indices address the base view's contiguous K/V rows, like the flat table.
+        rows = [page // 2 for page in pages[0]]
+        assert unified.get_batch_cache_indices([7], layer_idx=NUM_TARGET)[0][: len(slots)] == rows
+        flat = unified.get_batch_cache_indices_flat([7], [len(slots)], layer_idx=NUM_TARGET)
+        assert flat.tolist() == rows
+        # Target layers keep MiniMax's raw slot ids.
+        sparse_slots = unified.get_batch_cache_indices([7], layer_idx=3)[0]
+        assert (
+            sparse_slots == unified.get_batch_cache_indices([7], layer_idx=3, raw_indices=True)[0]
+        )
         unified.free_resources(request)
 
         # The draft lifecycle is sized by its window, not by its share of the bytes per token.
@@ -168,10 +184,35 @@ def test_draft_layer_specs_join_the_hybrid_nvfp4_manager() -> None:
         plain.shutdown()
 
 
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA cache pools")
+def test_nvfp4_draft_layers_are_rejected() -> None:
+    """Block scales belong to the sparse target layers, so an NVFP4 drafter cannot join."""
+    specs = _specs()
+    nvfp4_draft = replace(
+        specs[NUM_TARGET],
+        dtype=DataType.NVFP4,
+        buffer_roles=(Role.KEY, Role.VALUE, Role.KEY_BLOCK_SCALE, Role.VALUE_BLOCK_SCALE),
+    )
+    with pytest.raises(ValueError, match="does not host NVFP4 draft layers"):
+        _manager((*specs[:NUM_TARGET], nvfp4_draft, *specs[NUM_TARGET + 1 :]))
+
+
 @pytest.mark.cpu_only
 @pytest.mark.parametrize("is_disagg", [False, True])
 def test_window_sized_draft_ratio(is_disagg: bool) -> None:
     """A windowed draft lifecycle gets its windows plus one chunk, not a bytes-per-token share."""
+
+    def _configs(specs: tuple[KVCacheLayerSpec, ...], target_sinks: Optional[int] = None) -> list:
+        return [
+            SimpleNamespace(
+                sliding_window_size=spec.sliding_window_size,
+                num_sink_tokens=target_sinks if spec.owner is KVCacheLayerOwner.TARGET else None,
+                reuse_policy=spec.reuse_policy,
+                buffers=[SimpleNamespace(size=spec.bytes_per_token() * 128)],
+            )
+            for spec in specs
+        ]
+
     manager = object.__new__(KVCacheManagerV2)
     manager.layer_specs = _specs()[NUM_TARGET - 1 :]
     manager.max_batch_size = 2
@@ -179,15 +220,7 @@ def test_window_sized_draft_ratio(is_disagg: bool) -> None:
     manager.mapping = Mapping()
     manager.is_disagg = is_disagg
     manager._generation_kv_capacity_headroom = 1
-    configs = [
-        SimpleNamespace(
-            sliding_window_size=spec.sliding_window_size,
-            num_sink_tokens=None,
-            reuse_policy=spec.reuse_policy,
-            buffers=[SimpleNamespace(size=spec.bytes_per_token() * 128)],
-        )
-        for spec in manager.layer_specs
-    ]
+    configs = _configs(manager.layer_specs)
     draft_bytes_per_token = 2 * HEADS * HEAD_DIM * 2 * NUM_DRAFT
     # Four admissible sequences hold three pages of window each, plus one 256-token chunk.
     draft_bytes = (4 * 3 * 128 + 256) * draft_bytes_per_token
@@ -201,5 +234,18 @@ def test_window_sized_draft_ratio(is_disagg: bool) -> None:
     assert manager._windowed_draft_pool_ratio(configs, 1 << 20, 128)[1] == pytest.approx(
         token_share
     )
+    # Without a context-chunk bound the backend keeps its own split.
+    manager.max_num_tokens = None
+    assert manager._windowed_draft_pool_ratio(configs, 1 << 30, 128) is None
+    manager.max_num_tokens = 256
     # Without a windowed draft lifecycle the backend keeps its own split.
-    assert manager._windowed_draft_pool_ratio(configs[:1], 1 << 30, 128) is None
+    manager.layer_specs = _specs(draft_window=None)[NUM_TARGET - 1 :]
+    assert manager._windowed_draft_pool_ratio(_configs(manager.layer_specs), 1 << 30, 128) is None
+    # A target sharing the drafter's window shares its lifecycle too: sink tokens that round to
+    # no sink block do not make a second one.
+    manager.layer_specs = tuple(
+        replace(spec, sliding_window_size=DRAFT_WINDOW) for spec in _specs()[NUM_TARGET - 1 :]
+    )
+    assert (
+        manager._windowed_draft_pool_ratio(_configs(manager.layer_specs, 0), 1 << 30, 128) is None
+    )

From 1c1a77a61daa7a012aeafd61966c036a3a51e8d1 Mon Sep 17 00:00:00 2001
From: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com>
Date: Fri, 30 Oct 2026 13:51:23 -0700
Subject: [PATCH 05/14] [TRTLLM-17385][fix] MiniMax-M3: pick page-index
 conversion by layer owner

A draft layer can share its lifecycle, and so its pool, with target layers, so the
base conversion is chosen per requested draft layer instead of per pool. Managers
without layer specs log their dense layers from layer_offsets as before, and the
window-sized draft ratio also falls back without a quota.

Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com>
---
 .../sparse/minimax_m3/cache_manager.py        | 60 ++++++++++++-------
 .../kv_cache/kv_cache_manager_v2.py           |  4 +-
 .../msa/test_minimax_m3_unified_draft_kv.py   | 22 +++++++
 3 files changed, 63 insertions(+), 23 deletions(-)

diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py
index bc5d05b211f5..5cf39f53fb18 100644
--- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py
+++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py
@@ -29,7 +29,7 @@
 from __future__ import annotations
 
 import dataclasses
-from typing import TYPE_CHECKING, List, Optional, Sequence, Tuple
+from typing import TYPE_CHECKING, List, Optional, Sequence, Tuple, Union
 
 import torch
 
@@ -431,7 +431,9 @@ def __init__(
         if self.dtype == DataType.NVFP4:
             dense_target_layers = [
                 layer
-                for layer in self.target_layer_ids
+                for layer in (
+                    self.target_layer_ids if self.layer_specs is not None else self.layer_offsets
+                )
                 if layer not in self.sparse_layer_ids and layer not in self._shared_draft_layer_ids
             ]
             logger.info(
@@ -587,11 +589,14 @@ def _get_block_scale_role(self, role_a: DataRole, layer_id: int) -> Optional[Dat
             return None
         return super()._get_block_scale_role(role_a, layer_id)
 
-    def get_draft_kv_cache_view(self) -> Optional["MiniMaxM3DraftKVCacheView"]:
+    def get_draft_kv_cache_view(
+        self,
+    ) -> Optional[Union["MiniMaxM3DraftKVCacheView", "MiniMaxM3KVCacheManagerV2"]]:
         """FP8/P128 view of the shared draft layer in a hybrid NVFP4 cache, or None.
 
         FP8/BF16 caches address the draft layer through this manager's own
-        attention-op pools instead.
+        attention-op pools instead. With explicit draft layer specs, this
+        manager itself is the draft view.
 
         Only meaningful on a target manager carrying appended one-model
         draft layers; built lazily so the manager's page tables exist. A
@@ -1083,11 +1088,35 @@ def _kv_pool_mapping_offset(self, layer_id, layer_group_id, key_base_addr) -> in
         )
         return layers_by_addr.index(int(layer_id))
 
-    def _draft_pool_ids(self) -> frozenset:
-        """Pools that hold the explicit draft layers."""
-        return frozenset(
-            self.layer_to_pool_mapping_dict[self.layer_offsets[layer]]
-            for layer in self.draft_layer_ids
+    def get_batch_cache_indices(
+        self,
+        request_ids: List[int],
+        layer_idx: Optional[int] = None,
+        num_blocks_per_seq: Optional[Sequence[int]] = None,
+        *,
+        raw_indices: bool = False,
+    ) -> List[List[int]]:
+        """Return indices into a layer's :meth:`get_buffers` view.
+
+        Explicit draft layers keep the base view, so their indices take the base
+        conversion; every other layer gets raw slot ids.
+        """
+        if (
+            self.layer_specs is None
+            or layer_idx is None
+            or self.get_layer_owner(layer_idx) is not KVCacheLayerOwner.DRAFT
+        ):
+            return super().get_batch_cache_indices(
+                request_ids, layer_idx, num_blocks_per_seq, raw_indices=raw_indices
+            )
+        return KVCacheManagerV2._get_batch_cache_indices_by_pool_id(
+            self,
+            request_ids,
+            pool_id=self.layer_to_pool_mapping_dict[self.layer_offsets[layer_idx]],
+            is_kv_aggregate=not raw_indices,
+            num_blocks_per_seq=num_blocks_per_seq,
+            index_scale=1 if raw_indices else self.get_layer_page_index_scale(layer_idx),
+            kv_factor=self.get_layer_kv_factor(layer_idx),
         )
 
     def _get_batch_cache_indices_by_pool_id(
@@ -1115,8 +1144,7 @@ def _get_batch_cache_indices_by_pool_id(
         ``BAD_PAGE_INDEX`` slots remain ``-1`` here because disaggregation's
         :class:`KVRegionExtractorV1` filters ``region_ids >= 0``.
         :meth:`get_block_ids_per_seq` maps them to zero for the attention
-        metadata's padded tensor. Pools of explicit draft layers keep the
-        base conversion, matching their base :meth:`get_buffers` view.
+        metadata's padded tensor.
 
         Args:
             request_ids: Request IDs whose page-index rows are returned.
@@ -1131,16 +1159,6 @@ def _get_batch_cache_indices_by_pool_id(
             kv_factor: Kept for compatibility with the base virtual method;
                 M3 returns raw slot IDs without KV aggregation.
         """
-        if self.layer_specs is not None and pool_id in self._draft_pool_ids():
-            # Draft layers keep the base get_buffers() view, so their indices take the base conversion.
-            return super()._get_batch_cache_indices_by_pool_id(
-                request_ids,
-                pool_id=pool_id,
-                is_kv_aggregate=is_kv_aggregate,
-                num_blocks_per_seq=num_blocks_per_seq,
-                index_scale=index_scale,
-                kv_factor=kv_factor,
-            )
         res = []
         for req_idx, req_id in enumerate(request_ids):
             kv_cache = self.kv_cache_map[req_id]
diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py
index 9d6dfa10f2ea..a4e8c81cbf42 100644
--- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py
+++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py
@@ -3922,8 +3922,8 @@ def _windowed_draft_pool_ratio(
         as many tokens as the target. They need only every admissible sequence's window plus
         one context chunk, which is what the KV cache cost already charges for them.
         """
-        if self.max_num_tokens is None:
-            # Without a context-chunk bound the backend's split is the safe one.
+        if self.max_num_tokens is None or gpu_quota <= 0:
+            # Without a context-chunk bound or a quota the backend's split is the safe one.
             return None
         owners = [spec.owner for spec in self.layer_specs]
         # Key layers as the backend groups them into lifecycles (sinks count in whole blocks).
diff --git a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_unified_draft_kv.py b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_unified_draft_kv.py
index a8b06fe4f2ca..f2e334a7ed55 100644
--- a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_unified_draft_kv.py
+++ b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_unified_draft_kv.py
@@ -184,6 +184,28 @@ def test_draft_layer_specs_join_the_hybrid_nvfp4_manager() -> None:
         plain.shutdown()
 
 
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA cache pools")
+@pytest.mark.parametrize("draft_window", [DRAFT_WINDOW, None])
+def test_page_indices_follow_the_layer_owner(draft_window: Optional[int]) -> None:
+    """Target layers get raw slot ids and draft layers base view rows, even in one lifecycle."""
+    torch.cuda.init()
+    manager = _manager(_specs(draft_window))
+    try:
+        (request,) = manager.add_dummy_requests([7], token_nums=[200])
+        for layer in manager.pp_layers:
+            raw = manager.get_batch_cache_indices([7], layer_idx=layer, raw_indices=True)[0]
+            indices = manager.get_batch_cache_indices([7], layer_idx=layer)[0]
+            if layer in manager.draft_layer_ids:
+                scale = manager.get_layer_page_index_scale(layer)
+                expected = [s * scale // 2 if s >= 0 else s for s in raw]
+            else:
+                expected = raw
+            assert indices[: len(raw)] == expected, layer
+        manager.free_resources(request)
+    finally:
+        manager.shutdown()
+
+
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA cache pools")
 def test_nvfp4_draft_layers_are_rejected() -> None:
     """Block scales belong to the sparse target layers, so an NVFP4 drafter cannot join."""

From 645b0ca78e2f7f36cb0d6644eefb6dc62d1281a4 Mon Sep 17 00:00:00 2001
From: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com>
Date: Fri, 30 Oct 2026 14:00:47 -0700
Subject: [PATCH 06/14] [TRTLLM-17385][fix] MiniMax-M3: host draft layers only
 in their own lifecycle

Draft layers coalesced with target buffers would get target-relative pool offsets,
and a disaggregated Triton FP8 cache transfers token-major K/V while drafters write
head-major K/V, so both configurations are rejected at construction.

Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com>
---
 .../sparse/minimax_m3/cache_manager.py        | 24 +++++++++
 .../msa/test_minimax_m3_unified_draft_kv.py   | 54 ++++++++++++++-----
 2 files changed, 66 insertions(+), 12 deletions(-)

diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py
index 5cf39f53fb18..d7346be270d9 100644
--- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py
+++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py
@@ -416,6 +416,18 @@ def __init__(
         ):
             # Block scales here belong to the MSA sparse target layers only.
             raise ValueError(f"{type(self).__name__} does not host NVFP4 draft layers.")
+        hosts_draft_layers = layer_specs is not None and any(
+            spec.owner is KVCacheLayerOwner.DRAFT for spec in layer_specs
+        )
+        if (
+            hosts_draft_layers
+            and kwargs.get("is_disagg")
+            and self._main_kv_mapper_kind != MapperKind.HND
+        ):
+            # Drafters write head-major K/V, but transfer reads every K/V buffer with the main mapper.
+            raise ValueError(
+                f"{type(self).__name__} transfers draft layers only with the HND K/V layout."
+            )
         if layer_specs is not None and kwargs.get("dtype") == DataType.NVFP4:
             # The hybrid layout of _build_cache_config as specs: FP8 on the dense target layers.
             kwargs["layer_specs"] = tuple(
@@ -428,6 +440,18 @@ def __init__(
             )
         super().__init__(*args, **kwargs)
 
+        if hosts_draft_layers:
+            groups = self.layer_to_pool_mapping_dict
+            target_groups = {groups[self.layer_offsets[layer]] for layer in self.target_layer_ids}
+            if any(
+                groups[self.layer_offsets[layer]] in target_groups for layer in self.draft_layer_ids
+            ):
+                # Coalesced with target buffers, draft layers would get target-relative pool offsets.
+                self.shutdown()
+                raise ValueError(
+                    f"{type(self).__name__} hosts draft layers only in their own lifecycle."
+                )
+
         if self.dtype == DataType.NVFP4:
             dense_target_layers = [
                 layer
diff --git a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_unified_draft_kv.py b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_unified_draft_kv.py
index f2e334a7ed55..9117eebf4423 100644
--- a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_unified_draft_kv.py
+++ b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_unified_draft_kv.py
@@ -71,8 +71,13 @@ def _specs(draft_window: Optional[int] = DRAFT_WINDOW) -> tuple[KVCacheLayerSpec
     return tuple(target + draft)
 
 
-def _manager(layer_specs: Optional[tuple[KVCacheLayerSpec, ...]]) -> MiniMaxM3KVCacheManagerV2:
-    """A hybrid NVFP4 MSA manager, with or without explicit draft layer specs."""
+def _manager(
+    layer_specs: Optional[tuple[KVCacheLayerSpec, ...]],
+    dtype: DataType = DataType.NVFP4,
+    implementation: str = "msa",
+    is_disagg: bool = False,
+) -> MiniMaxM3KVCacheManagerV2:
+    """A MiniMax-M3 manager (hybrid NVFP4 MSA by default), with or without draft layer specs."""
     return MiniMaxM3KVCacheManagerV2(
         KvCacheConfig(
             max_gpu_total_bytes=64 << 20,
@@ -89,12 +94,14 @@ def _manager(layer_specs: Optional[tuple[KVCacheLayerSpec, ...]]) -> MiniMaxM3KV
         max_batch_size=2,
         max_num_tokens=256,
         mapping=Mapping(world_size=1, rank=0, tp_size=1, pp_size=1),
-        dtype=DataType.NVFP4,
+        dtype=dtype,
         vocab_size=16,
         sparse_attention_config=MiniMaxM3SparseAttentionConfig(
-            implementation="msa", indexer_kv_dtype="fp8"
+            implementation=implementation,
+            indexer_kv_dtype="fp8" if implementation == "msa" else "bf16",
         ),
         layer_specs=layer_specs,
+        is_disagg=is_disagg,
     )
 
 
@@ -185,11 +192,10 @@ def test_draft_layer_specs_join_the_hybrid_nvfp4_manager() -> None:
 
 
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA cache pools")
-@pytest.mark.parametrize("draft_window", [DRAFT_WINDOW, None])
-def test_page_indices_follow_the_layer_owner(draft_window: Optional[int]) -> None:
-    """Target layers get raw slot ids and draft layers base view rows, even in one lifecycle."""
+def test_page_indices_follow_the_layer_owner() -> None:
+    """Target layers get raw slot ids and draft layers base view rows."""
     torch.cuda.init()
-    manager = _manager(_specs(draft_window))
+    manager = _manager(_specs())
     try:
         (request,) = manager.add_dummy_requests([7], token_nums=[200])
         for layer in manager.pp_layers:
@@ -219,9 +225,31 @@ def test_nvfp4_draft_layers_are_rejected() -> None:
         _manager((*specs[:NUM_TARGET], nvfp4_draft, *specs[NUM_TARGET + 1 :]))
 
 
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA cache pools")
+def test_draft_layers_need_their_own_lifecycle() -> None:
+    """Unwindowed draft layers would share the target's lifecycle and pool offsets."""
+    torch.cuda.init()
+    with pytest.raises(ValueError, match="only in their own lifecycle"):
+        _manager(_specs(draft_window=None))
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA cache pools")
+def test_disagg_draft_layers_need_the_hnd_layout() -> None:
+    """A Triton FP8 cache transfers token-major K/V, but drafters write head-major K/V."""
+    specs = tuple(
+        replace(spec, dtype=DataType.FP8, buffer_roles=(Role.KEY, Role.VALUE))
+        if spec.owner is KVCacheLayerOwner.TARGET
+        else spec
+        for spec in _specs()
+    )
+    with pytest.raises(ValueError, match="only with the HND K/V layout"):
+        _manager(specs, dtype=DataType.FP8, implementation="triton", is_disagg=True)
+
+
 @pytest.mark.cpu_only
+@pytest.mark.parametrize("pp_size", [1, 2])
 @pytest.mark.parametrize("is_disagg", [False, True])
-def test_window_sized_draft_ratio(is_disagg: bool) -> None:
+def test_window_sized_draft_ratio(is_disagg: bool, pp_size: int) -> None:
     """A windowed draft lifecycle gets its windows plus one chunk, not a bytes-per-token share."""
 
     def _configs(specs: tuple[KVCacheLayerSpec, ...], target_sinks: Optional[int] = None) -> list:
@@ -239,13 +267,15 @@ def _configs(specs: tuple[KVCacheLayerSpec, ...], target_sinks: Optional[int] =
     manager.layer_specs = _specs()[NUM_TARGET - 1 :]
     manager.max_batch_size = 2
     manager.max_num_tokens = 256
-    manager.mapping = Mapping()
+    manager.mapping = Mapping(world_size=pp_size, rank=0, pp_size=pp_size)
     manager.is_disagg = is_disagg
     manager._generation_kv_capacity_headroom = 1
     configs = _configs(manager.layer_specs)
     draft_bytes_per_token = 2 * HEADS * HEAD_DIM * 2 * NUM_DRAFT
-    # Four admissible sequences hold three pages of window each, plus one 256-token chunk.
-    draft_bytes = (4 * 3 * 128 + 256) * draft_bytes_per_token
+    # Admissible sequences (batch x PP stages, doubled for the overlap scheduler's retiring
+    # request unless an aggregated PP pipeline) hold three pages of window each, plus one chunk.
+    num_sequences = 2 * pp_size * (2 if is_disagg or pp_size == 1 else 1)
+    draft_bytes = (num_sequences * 3 * 128 + 256) * draft_bytes_per_token
     assert manager._windowed_draft_pool_ratio(configs, 1 << 30, 128) == pytest.approx(
         [1 - draft_bytes / (1 << 30), draft_bytes / (1 << 30)]
     )

From 3a11d12be87facaf312b867b0fa6822f4a6cd0c6 Mon Sep 17 00:00:00 2001
From: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com>
Date: Fri, 30 Oct 2026 14:19:36 -0700
Subject: [PATCH 07/14] [TRTLLM-17385][fix] MiniMax-M3: drop the self draft
 view when adding dummy requests

With explicit draft layers the draft view is the manager itself, and CUDA-graph
padding passes it back as the draft manager, which created each dummy request twice.

Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com>
---
 .../attention/backends/sparse/minimax_m3/cache_manager.py | 8 ++++++--
 .../_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py     | 7 ++++---
 .../sparse/msa/test_minimax_m3_unified_draft_kv.py        | 5 ++++-
 3 files changed, 14 insertions(+), 6 deletions(-)

diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py
index d7346be270d9..2838898afe16 100644
--- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py
+++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py
@@ -658,9 +658,13 @@ def add_dummy_requests(self, *args, **kwargs):
         The base method mirrors dummy KV caches into a *separate* draft
         manager. With shared draft layers a dummy request's blocks already
         span the drafter's pool (pools allocate in lockstep per logical
-        block), and the view owns no block lifecycle.
+        block), and the view owns no block lifecycle. Explicit draft layers
+        are this manager's own, so their view is this manager itself.
         """
-        if isinstance(kwargs.get("draft_kv_cache_manager"), MiniMaxM3DraftKVCacheView):
+        draft_kv_cache_manager = kwargs.get("draft_kv_cache_manager")
+        if draft_kv_cache_manager is self or isinstance(
+            draft_kv_cache_manager, MiniMaxM3DraftKVCacheView
+        ):
             kwargs["draft_kv_cache_manager"] = None
         return super().add_dummy_requests(*args, **kwargs)
 
diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py
index a4e8c81cbf42..9642bad8f7c6 100644
--- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py
+++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py
@@ -3918,9 +3918,10 @@ def _windowed_draft_pool_ratio(
     ) -> Optional[List[float]]:
         """Hot-tier ``[target, draft]`` byte ratios when the draft layers form one windowed lifecycle.
 
-        Without a typical step the backend splits by bytes per token, handing the draft layers
-        as many tokens as the target. They need only every admissible sequence's window plus
-        one context chunk, which is what the KV cache cost already charges for them.
+        Without a pool ratio or typical step the backend splits the hot tier for one 2048-token
+        request, which says nothing about how many sequences hold a draft window at once. The
+        draft layers need every admissible sequence's window plus one context chunk, which is
+        what the KV cache cost already charges for them, capped at an equal-token share.
         """
         if self.max_num_tokens is None or gpu_quota <= 0:
             # Without a context-chunk bound or a quota the backend's split is the safe one.
diff --git a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_unified_draft_kv.py b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_unified_draft_kv.py
index 9117eebf4423..be01af6adeaf 100644
--- a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_unified_draft_kv.py
+++ b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_unified_draft_kv.py
@@ -144,7 +144,10 @@ def test_draft_layer_specs_join_the_hybrid_nvfp4_manager() -> None:
         # The unified drafter reads its layers' block offsets, so they keep attention-op pools.
         assert plain.num_attention_op_pools == 0
         assert unified.num_attention_op_pools > 0
-        (request,) = unified.add_dummy_requests([7], token_nums=[200])
+        # CUDA-graph padding passes the draft view (this manager) back in.
+        (request,) = unified.add_dummy_requests(
+            [7], token_nums=[200], draft_kv_cache_manager=unified.get_draft_kv_cache_view()
+        )
         offsets = torch.zeros(
             unified.num_attention_op_pools,
             1,

From 6b8c95e26248da1ad0027678a9dd2cfdb12012c7 Mon Sep 17 00:00:00 2001
From: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com>
Date: Fri, 30 Oct 2026 14:26:57 -0700
Subject: [PATCH 08/14] [TRTLLM-17385][fix] MiniMax-M3: take the target layers
 from explicit layer specs

The sparse-layer default counts the target-owned specs, so a per-layer head list
that spans the drafter does not make draft layers sparse. Target specs must use the
cache dtype, which the target views and block scales follow.

Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com>
---
 .../sparse/minimax_m3/cache_manager.py        | 14 ++++++-
 .../msa/test_minimax_m3_unified_draft_kv.py   | 42 +++++++++++++++++++
 2 files changed, 55 insertions(+), 1 deletion(-)

diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py
index 2838898afe16..0358a2283637 100644
--- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py
+++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py
@@ -383,6 +383,12 @@ def __init__(
             kwargs.get("num_kv_heads"),
             shared_draft_layer_count(kwargs.get("spec_config"), kwargs.get("layer_mask")),
         )
+        layer_specs = kwargs.get("layer_specs")
+        if layer_specs is not None:
+            # Explicit specs declare which layers are the target's.
+            num_target_layers = sum(
+                1 for spec in layer_specs if spec.owner is KVCacheLayerOwner.TARGET
+            )
         if sparse_layer_ids is None:
             sparse_layer_ids = default_sparse_layer_ids(
                 kwargs.get("pretrained_config"), num_target_layers
@@ -409,7 +415,13 @@ def __init__(
         # Filled at the end of __init__; the base __init__ may already call get_buffers.
         self._kv_buffers: dict[tuple[int, str], torch.Tensor] = {}
         self._block_scale_buffers: dict[tuple[int, str], torch.Tensor] = {}
-        layer_specs = kwargs.get("layer_specs")
+        if layer_specs is not None and any(
+            spec.owner is KVCacheLayerOwner.TARGET
+            and spec.dtype != kwargs.get("dtype", DataType.HALF)
+            for spec in layer_specs
+        ):
+            # Target views and block scales follow the manager dtype.
+            raise ValueError(f"{type(self).__name__} needs target layer specs in the cache dtype.")
         if layer_specs is not None and any(
             spec.owner is KVCacheLayerOwner.DRAFT and spec.dtype == DataType.NVFP4
             for spec in layer_specs
diff --git a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_unified_draft_kv.py b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_unified_draft_kv.py
index be01af6adeaf..2d2e88d2e0ac 100644
--- a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_unified_draft_kv.py
+++ b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_unified_draft_kv.py
@@ -228,6 +228,48 @@ def test_nvfp4_draft_layers_are_rejected() -> None:
         _manager((*specs[:NUM_TARGET], nvfp4_draft, *specs[NUM_TARGET + 1 :]))
 
 
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA cache pools")
+def test_specs_decide_the_sparse_target_layers() -> None:
+    """A per-layer head list spanning the drafter does not make draft layers sparse targets."""
+    torch.cuda.init()
+    manager = MiniMaxM3KVCacheManagerV2(
+        KvCacheConfig(
+            max_gpu_total_bytes=64 << 20,
+            host_cache_size=0,
+            enable_block_reuse=False,
+            enable_swa_scratch_reuse=False,
+        ),
+        CacheType.SELF,
+        num_layers=NUM_TARGET,
+        num_kv_heads=[HEADS] * (NUM_TARGET + NUM_DRAFT),
+        head_dim=HEAD_DIM,
+        tokens_per_block=128,
+        max_seq_len=1024,
+        max_batch_size=2,
+        max_num_tokens=256,
+        mapping=Mapping(world_size=1, rank=0, tp_size=1, pp_size=1),
+        dtype=DataType.NVFP4,
+        vocab_size=16,
+        sparse_attention_config=MiniMaxM3SparseAttentionConfig(
+            implementation="msa", indexer_kv_dtype="fp8"
+        ),
+        layer_specs=_specs(),
+    )
+    try:
+        assert manager.sparse_layer_ids == [3]
+    finally:
+        manager.shutdown()
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA cache pools")
+def test_target_specs_need_the_cache_dtype() -> None:
+    """Target views and block scales follow the manager dtype, so target specs must use it."""
+    specs = _specs()
+    fp8_target = replace(specs[3], dtype=DataType.FP8, buffer_roles=(Role.KEY, Role.VALUE))
+    with pytest.raises(ValueError, match="target layer specs in the cache dtype"):
+        _manager((*specs[:3], fp8_target, *specs[NUM_TARGET:]))
+
+
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA cache pools")
 def test_draft_layers_need_their_own_lifecycle() -> None:
     """Unwindowed draft layers would share the target's lifecycle and pool offsets."""

From c54e0f0783896ebe09e5a70c3af5717517bbf298 Mon Sep 17 00:00:00 2001
From: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com>
Date: Fri, 30 Oct 2026 14:37:13 -0700
Subject: [PATCH 09/14] [TRTLLM-17385][fix] MiniMax-M3: keep the global sparse
 split and drop draft ids from it

Layer specs are rank-local, so counting target specs broke the sparse default under
PP. The default split stays global and explicit draft ids are removed from it and
from the native dense layers. A test admits a full context chunk on the draft pool.

Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com>
---
 .../sparse/minimax_m3/cache_manager.py        | 16 +++---
 .../msa/test_minimax_m3_unified_draft_kv.py   | 49 +++++++++++++++++++
 2 files changed, 59 insertions(+), 6 deletions(-)

diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py
index 0358a2283637..1b738f62d096 100644
--- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py
+++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py
@@ -384,15 +384,17 @@ def __init__(
             shared_draft_layer_count(kwargs.get("spec_config"), kwargs.get("layer_mask")),
         )
         layer_specs = kwargs.get("layer_specs")
-        if layer_specs is not None:
-            # Explicit specs declare which layers are the target's.
-            num_target_layers = sum(
-                1 for spec in layer_specs if spec.owner is KVCacheLayerOwner.TARGET
-            )
+        # Explicit specs declare the draft layers, which are never sparse or native target layers.
+        draft_ids = {
+            spec.global_layer_id
+            for spec in (layer_specs or ())
+            if spec.owner is KVCacheLayerOwner.DRAFT
+        }
         if sparse_layer_ids is None:
             sparse_layer_ids = default_sparse_layer_ids(
                 kwargs.get("pretrained_config"), num_target_layers
             )
+        sparse_layer_ids = [i for i in sparse_layer_ids if int(i) not in draft_ids]
         if disable_index_value_layer_ids is None:
             disable_index_value_layer_ids = list(sparse_layer_ids)
 
@@ -402,7 +404,9 @@ def __init__(
         self.sparse_layer_ids = sorted(int(i) for i in sparse_layer_ids)
         if implementation == "triton" and native_dense_cache and num_target_layers is not None:
             self._dense_native_layer_ids = tuple(
-                i for i in range(num_target_layers) if i not in self.sparse_layer_ids
+                i
+                for i in range(num_target_layers)
+                if i not in self.sparse_layer_ids and i not in draft_ids
             )
         self.disable_index_value_layer_ids = set(int(i) for i in disable_index_value_layer_ids)
         self.sparse_index_dim = sparse_index_dim
diff --git a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_unified_draft_kv.py b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_unified_draft_kv.py
index 2d2e88d2e0ac..966b3f9e8aa2 100644
--- a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_unified_draft_kv.py
+++ b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_unified_draft_kv.py
@@ -291,6 +291,55 @@ def test_disagg_draft_layers_need_the_hnd_layout() -> None:
         _manager(specs, dtype=DataType.FP8, implementation="triton", is_disagg=True)
 
 
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA cache pools")
+def test_window_sized_draft_pool_admits_a_full_context_chunk() -> None:
+    """Allocation rounds the window-sized draft share, but one full chunk still fits."""
+    torch.cuda.init()
+    max_num_tokens = 8320
+    specs = tuple(
+        KVCacheLayerSpec(
+            global_layer_id=layer,
+            owner=owner,
+            dtype=DataType.BF16,
+            num_kv_heads=HEADS,
+            head_dim=HEAD_DIM,
+            buffer_roles=(Role.KEY, Role.VALUE),
+            sliding_window_size=window,
+        )
+        for layer, owner, window in (
+            (0, KVCacheLayerOwner.TARGET, None),
+            (1, KVCacheLayerOwner.DRAFT, DRAFT_WINDOW),
+        )
+    )
+    manager = KVCacheManagerV2(
+        KvCacheConfig(
+            max_gpu_total_bytes=1 << 30,
+            host_cache_size=0,
+            enable_block_reuse=False,
+            enable_swa_scratch_reuse=False,
+        ),
+        CacheType.SELF,
+        num_layers=1,
+        num_kv_heads=HEADS,
+        head_dim=HEAD_DIM,
+        tokens_per_block=128,
+        max_seq_len=16384,
+        max_batch_size=1,
+        max_num_tokens=max_num_tokens,
+        mapping=Mapping(world_size=1, rank=0, tp_size=1, pp_size=1),
+        dtype=DataType.BF16,
+        vocab_size=16,
+        layer_specs=specs,
+    )
+    try:
+        assert manager.kv_cache_manager_py_config.initial_pool_ratio is not None
+        requests = manager.add_dummy_requests([7], token_nums=[max_num_tokens])
+        assert requests is not None
+        manager.free_resources(requests[0])
+    finally:
+        manager.shutdown()
+
+
 @pytest.mark.cpu_only
 @pytest.mark.parametrize("pp_size", [1, 2])
 @pytest.mark.parametrize("is_disagg", [False, True])

From e227da002c033e782543b19b1ab8545b22281c94 Mon Sep 17 00:00:00 2001
From: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com>
Date: Fri, 30 Oct 2026 14:46:09 -0700
Subject: [PATCH 10/14] [TRTLLM-17385][fix] KVCacheManagerV2: floor the
 window-sized draft pool at one context chunk

Allocation rounds the window-sized ratio's byte target, which could leave a separately
sized draft pool a block short of a full context chunk. The ratio now comes with the
context-chunk constraint the typical-step path already declares.

Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com>
---
 .../backends/sparse/minimax_m3/cache_manager.py     |  7 ++++---
 .../pyexecutor/kv_cache/kv_cache_manager_v2.py      | 13 +++++++++++++
 .../sparse/msa/test_minimax_m3_unified_draft_kv.py  | 12 +++++++-----
 3 files changed, 24 insertions(+), 8 deletions(-)

diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py
index 1b738f62d096..5e1366dd04e0 100644
--- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py
+++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py
@@ -29,7 +29,7 @@
 from __future__ import annotations
 
 import dataclasses
-from typing import TYPE_CHECKING, List, Optional, Sequence, Tuple, Union
+from typing import TYPE_CHECKING, Any, List, Optional, Sequence, Tuple, Union
 
 import torch
 
@@ -65,6 +65,7 @@
 if TYPE_CHECKING:
     from transformers import PretrainedConfig
 
+    from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest
     from tensorrt_llm.llmapi.llm_args import DecodingBaseConfig
 
 
@@ -578,7 +579,7 @@ def _build_cache_config(self, config):
             ]
         return super()._build_cache_config(config)
 
-    def get_layer_bytes_per_token(self, local_layer_idx: int, data_role: Role):
+    def get_layer_bytes_per_token(self, local_layer_idx: int, data_role: Role) -> int:
         """Report the hybrid sparse-NVFP4 / dense-FP8 storage footprint."""
         if self.layer_specs is not None:
             # Explicit layer specs carry K/V and scales; index-K is this manager's extra buffer.
@@ -668,7 +669,7 @@ def get_draft_kv_cache_view(
             )
         return self._draft_kv_cache_view
 
-    def add_dummy_requests(self, *args, **kwargs):
+    def add_dummy_requests(self, *args: Any, **kwargs: Any) -> Optional[List[LlmRequest]]:
         """Drop the draft KV cache view before delegating.
 
         The base method mirrors dummy KV caches into a *separate* draft
diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py
index 9642bad8f7c6..c36a9038bc20 100644
--- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py
+++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py
@@ -3772,6 +3772,19 @@ def _build_base_config(
                 initial_pool_ratio = self._windowed_draft_pool_ratio(
                     layer_configs, cache_tiers[0].quota, tokens_per_block
                 )
+                if initial_pool_ratio is not None:
+                    # Allocation rounds the ratio's byte target; this floor keeps one fresh
+                    # context chunk admissible in every lifecycle.
+                    constraints.append(
+                        BatchDesc(
+                            [
+                                KVCacheDesc(
+                                    capacity=self.max_num_tokens + self.num_extra_kv_tokens,
+                                    history_length=0,
+                                )
+                            ]
+                        )
+                    )
         else:
             # Subclasses (e.g. MiniMax-M3 sparse cache) can register additional
             # per-layer BufferConfig entries — for example a sparse index-K
diff --git a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_unified_draft_kv.py b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_unified_draft_kv.py
index 966b3f9e8aa2..174df019c417 100644
--- a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_unified_draft_kv.py
+++ b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_unified_draft_kv.py
@@ -301,14 +301,14 @@ def test_window_sized_draft_pool_admits_a_full_context_chunk() -> None:
             global_layer_id=layer,
             owner=owner,
             dtype=DataType.BF16,
-            num_kv_heads=HEADS,
+            num_kv_heads=heads,
             head_dim=HEAD_DIM,
             buffer_roles=(Role.KEY, Role.VALUE),
             sliding_window_size=window,
         )
-        for layer, owner, window in (
-            (0, KVCacheLayerOwner.TARGET, None),
-            (1, KVCacheLayerOwner.DRAFT, DRAFT_WINDOW),
+        for layer, owner, heads, window in (
+            (0, KVCacheLayerOwner.TARGET, 2 * HEADS, None),
+            (1, KVCacheLayerOwner.DRAFT, HEADS, DRAFT_WINDOW),
         )
     )
     manager = KVCacheManagerV2(
@@ -320,7 +320,7 @@ def test_window_sized_draft_pool_admits_a_full_context_chunk() -> None:
         ),
         CacheType.SELF,
         num_layers=1,
-        num_kv_heads=HEADS,
+        num_kv_heads=[2 * HEADS, HEADS],
         head_dim=HEAD_DIM,
         tokens_per_block=128,
         max_seq_len=16384,
@@ -333,6 +333,8 @@ def test_window_sized_draft_pool_admits_a_full_context_chunk() -> None:
     )
     try:
         assert manager.kv_cache_manager_py_config.initial_pool_ratio is not None
+        # Different page sizes keep the draft lifecycle in its own physical pool.
+        assert int(manager.kv_cache_pool_mapping[0, 0]) != int(manager.kv_cache_pool_mapping[1, 0])
         requests = manager.add_dummy_requests([7], token_nums=[max_num_tokens])
         assert requests is not None
         manager.free_resources(requests[0])

From d8273c07160595e15bf3e4c2b8b6ae19c967a8c8 Mon Sep 17 00:00:00 2001
From: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com>
Date: Fri, 30 Oct 2026 14:55:14 -0700
Subject: [PATCH 11/14] [TRTLLM-17385][fix] MiniMax-M3: count target layers
 from the first draft layer

The first explicit draft layer id is the global target count, so a head list that spans
the drafter keeps the checkpoint's sparse split. The window-sized draft share is sized
over the resume watermark.

Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com>
---
 .../sparse/minimax_m3/cache_manager.py        |  3 +++
 .../kv_cache/kv_cache_manager_v2.py           | 11 ++++++---
 .../msa/test_minimax_m3_unified_draft_kv.py   | 24 ++++++++++++++++---
 3 files changed, 32 insertions(+), 6 deletions(-)

diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py
index 5e1366dd04e0..f0690f833ecc 100644
--- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py
+++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py
@@ -391,6 +391,9 @@ def __init__(
             for spec in (layer_specs or ())
             if spec.owner is KVCacheLayerOwner.DRAFT
         }
+        if draft_ids:
+            # Draft layers follow every target layer, so the first one is the global target count.
+            num_target_layers = min(draft_ids)
         if sparse_layer_ids is None:
             sparse_layer_ids = default_sparse_layer_ids(
                 kwargs.get("pretrained_config"), num_target_layers
diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py
index c36a9038bc20..4880347f279f 100644
--- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py
+++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py
@@ -3770,7 +3770,10 @@ def _build_base_config(
             ]
             if initial_pool_ratio is None and typical_step is None:
                 initial_pool_ratio = self._windowed_draft_pool_ratio(
-                    layer_configs, cache_tiers[0].quota, tokens_per_block
+                    layer_configs,
+                    cache_tiers[0].quota,
+                    tokens_per_block,
+                    kv_cache_config.max_util_for_resume,
                 )
                 if initial_pool_ratio is not None:
                     # Allocation rounds the ratio's byte target; this floor keeps one fresh
@@ -3928,13 +3931,15 @@ def _windowed_draft_pool_ratio(
         layer_configs: List[AttentionLayerConfig],
         gpu_quota: int,
         tokens_per_block: int,
+        max_util_for_resume: float = 1.0,
     ) -> Optional[List[float]]:
         """Hot-tier ``[target, draft]`` byte ratios when the draft layers form one windowed lifecycle.
 
         Without a pool ratio or typical step the backend splits the hot tier for one 2048-token
         request, which says nothing about how many sequences hold a draft window at once. The
         draft layers need every admissible sequence's window plus one context chunk, which is
-        what the KV cache cost already charges for them, capped at an equal-token share.
+        what the KV cache cost already charges for them, over the resume watermark and capped
+        at an equal-token share.
         """
         if self.max_num_tokens is None or gpu_quota <= 0:
             # Without a context-chunk bound or a quota the backend's split is the safe one.
@@ -3979,7 +3984,7 @@ def _windowed_draft_pool_ratio(
         )
         draft_bytes = (
             window_bytes_per_sequence * num_sequences + sum(draft_sizes) * self.max_num_tokens
-        )
+        ) / max_util_for_resume
         bytes_per_block = {KVCacheLayerOwner.TARGET: 0, KVCacheLayerOwner.DRAFT: 0}
         for config, owner in zip(layer_configs, owners):
             bytes_per_block[owner] += sum(buffer.size for buffer in config.buffers)
diff --git a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_unified_draft_kv.py b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_unified_draft_kv.py
index 174df019c417..1c10d7b8126b 100644
--- a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_unified_draft_kv.py
+++ b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_unified_draft_kv.py
@@ -229,9 +229,22 @@ def test_nvfp4_draft_layers_are_rejected() -> None:
 
 
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA cache pools")
-def test_specs_decide_the_sparse_target_layers() -> None:
-    """A per-layer head list spanning the drafter does not make draft layers sparse targets."""
+@pytest.mark.parametrize(
+    "sparse_attention_freq,expected",
+    [(None, [3]), ([0, 1, 0, 1], [1, 3])],
+    ids=["default", "config"],
+)
+def test_specs_decide_the_sparse_target_layers(
+    sparse_attention_freq: Optional[list[int]], expected: list[int]
+) -> None:
+    """A per-layer head list spanning the drafter keeps the target's sparse split."""
     torch.cuda.init()
+    pretrained_config = None
+    if sparse_attention_freq is not None:
+        pretrained_config = SimpleNamespace(
+            num_hidden_layers=NUM_TARGET,
+            sparse_attention_config={"sparse_attention_freq": sparse_attention_freq},
+        )
     manager = MiniMaxM3KVCacheManagerV2(
         KvCacheConfig(
             max_gpu_total_bytes=64 << 20,
@@ -254,9 +267,10 @@ def test_specs_decide_the_sparse_target_layers() -> None:
             implementation="msa", indexer_kv_dtype="fp8"
         ),
         layer_specs=_specs(),
+        pretrained_config=pretrained_config,
     )
     try:
-        assert manager.sparse_layer_ids == [3]
+        assert manager.sparse_layer_ids == expected
     finally:
         manager.shutdown()
 
@@ -375,6 +389,10 @@ def _configs(specs: tuple[KVCacheLayerSpec, ...], target_sinks: Optional[int] =
     assert manager._windowed_draft_pool_ratio(configs, 1 << 30, 128) == pytest.approx(
         [1 - draft_bytes / (1 << 30), draft_bytes / (1 << 30)]
     )
+    # Resume refuses a pool above its watermark, so the share covers the windows below it.
+    assert manager._windowed_draft_pool_ratio(configs, 1 << 30, 128, 0.5)[1] == pytest.approx(
+        2 * draft_bytes / (1 << 30)
+    )
     # A quota too small for the windows keeps the bytes-per-token split.
     token_share = draft_bytes_per_token / sum(
         spec.bytes_per_token() for spec in manager.layer_specs

From 7ba0bb17bbb13dbd8f496430c8418244b1867121 Mon Sep 17 00:00:00 2001
From: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com>
Date: Fri, 30 Oct 2026 15:04:35 -0700
Subject: [PATCH 12/14] [TRTLLM-17385][fix] MiniMax-M3: take the global target
 count from the checkpoint

Layer specs may be one PP rank's subset, so the checkpoint's layer count sets the target
range when there is one. The window-sized draft ratio, and the context-chunk floor that
comes with it, apply only when the quota can hold that floor.

Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com>
---
 .../attention/backends/sparse/minimax_m3/cache_manager.py | 8 ++++++--
 .../_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py     | 4 ++++
 .../sparse/msa/test_minimax_m3_unified_draft_kv.py        | 2 ++
 3 files changed, 12 insertions(+), 2 deletions(-)

diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py
index f0690f833ecc..164c2ce4beb9 100644
--- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py
+++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py
@@ -392,8 +392,12 @@ def __init__(
             if spec.owner is KVCacheLayerOwner.DRAFT
         }
         if draft_ids:
-            # Draft layers follow every target layer, so the first one is the global target count.
-            num_target_layers = min(draft_ids)
+            # Layer specs may be a PP rank's subset: the checkpoint's layer count is the global
+            # target count, and without it draft layers still follow every target layer.
+            num_target_layers = int(
+                getattr(kwargs.get("pretrained_config"), "num_hidden_layers", None)
+                or min(draft_ids)
+            )
         if sparse_layer_ids is None:
             sparse_layer_ids = default_sparse_layer_ids(
                 kwargs.get("pretrained_config"), num_target_layers
diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py
index 4880347f279f..d7d4ca731bf1 100644
--- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py
+++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py
@@ -3988,6 +3988,10 @@ def _windowed_draft_pool_ratio(
         bytes_per_block = {KVCacheLayerOwner.TARGET: 0, KVCacheLayerOwner.DRAFT: 0}
         for config, owner in zip(layer_configs, owners):
             bytes_per_block[owner] += sum(buffer.size for buffer in config.buffers)
+        chunk_blocks = -(-self.max_num_tokens // tokens_per_block)
+        if chunk_blocks * sum(bytes_per_block.values()) > gpu_quota * max_util_for_resume:
+            # The context-chunk floor declared with this ratio would not fit the quota.
+            return None
         token_share = bytes_per_block[KVCacheLayerOwner.DRAFT] / sum(bytes_per_block.values())
         draft_ratio = min(draft_bytes / gpu_quota, token_share)
         return [1.0 - draft_ratio, draft_ratio]
diff --git a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_unified_draft_kv.py b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_unified_draft_kv.py
index 1c10d7b8126b..d1b337448531 100644
--- a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_unified_draft_kv.py
+++ b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_unified_draft_kv.py
@@ -400,6 +400,8 @@ def _configs(specs: tuple[KVCacheLayerSpec, ...], target_sinks: Optional[int] =
     assert manager._windowed_draft_pool_ratio(configs, 1 << 20, 128)[1] == pytest.approx(
         token_share
     )
+    # A quota that cannot hold the context-chunk floor keeps the backend's split.
+    assert manager._windowed_draft_pool_ratio(configs, 1 << 19, 128) is None
     # Without a context-chunk bound the backend keeps its own split.
     manager.max_num_tokens = None
     assert manager._windowed_draft_pool_ratio(configs, 1 << 30, 128) is None

From 15d4b903509e50d1398b6fbe177da6611d3f6ff7 Mon Sep 17 00:00:00 2001
From: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com>
Date: Fri, 30 Oct 2026 15:15:03 -0700
Subject: [PATCH 13/14] [TRTLLM-17385][fix] MiniMax-M3: size the draft-ratio
 guard like the context-chunk floor

The guard counts the chunk with extra KV tokens and its slots over the resume watermark,
as the floor does. Every rank with layer specs, target-only ones included, takes the
global target count from the checkpoint.

Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com>
---
 .../sparse/minimax_m3/cache_manager.py        | 11 +++++-----
 .../kv_cache/kv_cache_manager_v2.py           |  8 ++++---
 .../msa/test_minimax_m3_unified_draft_kv.py   | 22 +++++++++++++------
 3 files changed, 26 insertions(+), 15 deletions(-)

diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py
index 164c2ce4beb9..db0f24b29a6a 100644
--- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py
+++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py
@@ -391,13 +391,14 @@ def __init__(
             for spec in (layer_specs or ())
             if spec.owner is KVCacheLayerOwner.DRAFT
         }
-        if draft_ids:
+        if layer_specs is not None:
             # Layer specs may be a PP rank's subset: the checkpoint's layer count is the global
             # target count, and without it draft layers still follow every target layer.
-            num_target_layers = int(
-                getattr(kwargs.get("pretrained_config"), "num_hidden_layers", None)
-                or min(draft_ids)
-            )
+            num_hidden_layers = getattr(kwargs.get("pretrained_config"), "num_hidden_layers", None)
+            if num_hidden_layers is not None:
+                num_target_layers = int(num_hidden_layers)
+            elif draft_ids:
+                num_target_layers = min(draft_ids)
         if sparse_layer_ids is None:
             sparse_layer_ids = default_sparse_layer_ids(
                 kwargs.get("pretrained_config"), num_target_layers
diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py
index d7d4ca731bf1..4d0e1589a7ea 100644
--- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py
+++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py
@@ -3988,9 +3988,11 @@ def _windowed_draft_pool_ratio(
         bytes_per_block = {KVCacheLayerOwner.TARGET: 0, KVCacheLayerOwner.DRAFT: 0}
         for config, owner in zip(layer_configs, owners):
             bytes_per_block[owner] += sum(buffer.size for buffer in config.buffers)
-        chunk_blocks = -(-self.max_num_tokens // tokens_per_block)
-        if chunk_blocks * sum(bytes_per_block.values()) > gpu_quota * max_util_for_resume:
-            # The context-chunk floor declared with this ratio would not fit the quota.
+        # The context-chunk floor declared with this ratio: every lifecycle holds one fresh chunk
+        # below the resume watermark. Without room for it the backend keeps its own split.
+        chunk_blocks = -(-(self.max_num_tokens + self.num_extra_kv_tokens) // tokens_per_block)
+        floor_slots = math.ceil(chunk_blocks / max_util_for_resume)
+        if floor_slots * sum(bytes_per_block.values()) > gpu_quota:
             return None
         token_share = bytes_per_block[KVCacheLayerOwner.DRAFT] / sum(bytes_per_block.values())
         draft_ratio = min(draft_bytes / gpu_quota, token_share)
diff --git a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_unified_draft_kv.py b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_unified_draft_kv.py
index d1b337448531..3af4c8a45102 100644
--- a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_unified_draft_kv.py
+++ b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_unified_draft_kv.py
@@ -230,12 +230,16 @@ def test_nvfp4_draft_layers_are_rejected() -> None:
 
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA cache pools")
 @pytest.mark.parametrize(
-    "sparse_attention_freq,expected",
-    [(None, [3]), ([0, 1, 0, 1], [1, 3])],
-    ids=["default", "config"],
+    "sparse_attention_freq,num_specs,expected",
+    [
+        (None, NUM_TARGET + NUM_DRAFT, [3]),
+        ([0, 1, 0, 1], NUM_TARGET + NUM_DRAFT, [1, 3]),
+        ([0, 1, 0, 1], NUM_TARGET, [1, 3]),
+    ],
+    ids=["default", "config", "config-target-only"],
 )
 def test_specs_decide_the_sparse_target_layers(
-    sparse_attention_freq: Optional[list[int]], expected: list[int]
+    sparse_attention_freq: Optional[list[int]], num_specs: int, expected: list[int]
 ) -> None:
     """A per-layer head list spanning the drafter keeps the target's sparse split."""
     torch.cuda.init()
@@ -266,7 +270,7 @@ def test_specs_decide_the_sparse_target_layers(
         sparse_attention_config=MiniMaxM3SparseAttentionConfig(
             implementation="msa", indexer_kv_dtype="fp8"
         ),
-        layer_specs=_specs(),
+        layer_specs=_specs()[:num_specs],
         pretrained_config=pretrained_config,
     )
     try:
@@ -380,6 +384,7 @@ def _configs(specs: tuple[KVCacheLayerSpec, ...], target_sinks: Optional[int] =
     manager.mapping = Mapping(world_size=pp_size, rank=0, pp_size=pp_size)
     manager.is_disagg = is_disagg
     manager._generation_kv_capacity_headroom = 1
+    manager.num_extra_kv_tokens = 0
     configs = _configs(manager.layer_specs)
     draft_bytes_per_token = 2 * HEADS * HEAD_DIM * 2 * NUM_DRAFT
     # Admissible sequences (batch x PP stages, doubled for the overlap scheduler's retiring
@@ -400,8 +405,11 @@ def _configs(specs: tuple[KVCacheLayerSpec, ...], target_sinks: Optional[int] =
     assert manager._windowed_draft_pool_ratio(configs, 1 << 20, 128)[1] == pytest.approx(
         token_share
     )
-    # A quota that cannot hold the context-chunk floor keeps the backend's split.
-    assert manager._windowed_draft_pool_ratio(configs, 1 << 19, 128) is None
+    # The quota must hold the context-chunk floor: two blocks per lifecycle, over the resume
+    # watermark; otherwise the backend keeps its own split.
+    floor_bytes = 3 * sum(spec.bytes_per_token() * 128 for spec in manager.layer_specs)
+    assert manager._windowed_draft_pool_ratio(configs, floor_bytes - 1, 128, 0.95) is None
+    assert manager._windowed_draft_pool_ratio(configs, floor_bytes, 128, 0.95) is not None
     # Without a context-chunk bound the backend keeps its own split.
     manager.max_num_tokens = None
     assert manager._windowed_draft_pool_ratio(configs, 1 << 30, 128) is None

From 567997406ca1b1a75009328ef681032e8214eefe Mon Sep 17 00:00:00 2001
From: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com>
Date: Fri, 30 Oct 2026 15:24:18 -0700
Subject: [PATCH 14/14] [TRTLLM-17385][fix] KVCacheManagerV2: round the
 draft-ratio guard to the allocation grain

The allocator rounds each pool of the context-chunk floor to its allocation granularity
and reads the resume watermark as float32. The guard does the same, bounding pools by
buffers, so the floor never raises allocation past the GPU quota.

Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com>
---
 .../pyexecutor/kv_cache/kv_cache_manager_v2.py      | 13 ++++++++++---
 .../sparse/msa/test_minimax_m3_unified_draft_kv.py  | 13 ++++++++-----
 2 files changed, 18 insertions(+), 8 deletions(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py
index 4d0e1589a7ea..bc1107bcda8d 100644
--- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py
+++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py
@@ -3989,10 +3989,17 @@ def _windowed_draft_pool_ratio(
         for config, owner in zip(layer_configs, owners):
             bytes_per_block[owner] += sum(buffer.size for buffer in config.buffers)
         # The context-chunk floor declared with this ratio: every lifecycle holds one fresh chunk
-        # below the resume watermark. Without room for it the backend keeps its own split.
+        # below the resume watermark, each pool rounded to the allocation granularity (bounded
+        # here per buffer). Without room for it the backend keeps its own split.
         chunk_blocks = -(-(self.max_num_tokens + self.num_extra_kv_tokens) // tokens_per_block)
-        floor_slots = math.ceil(chunk_blocks / max_util_for_resume)
-        if floor_slots * sum(bytes_per_block.values()) > gpu_quota:
+        floor_slots = math.ceil(chunk_blocks / float(np.float32(max_util_for_resume)))
+        granularity = _gpu_alloc_granularity(gpu_quota)
+        floor_bytes = sum(
+            -(-floor_slots * buffer.size // granularity) * granularity
+            for config in layer_configs
+            for buffer in config.buffers
+        )
+        if floor_bytes > gpu_quota:
             return None
         token_share = bytes_per_block[KVCacheLayerOwner.DRAFT] / sum(bytes_per_block.values())
         draft_ratio = min(draft_bytes / gpu_quota, token_share)
diff --git a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_unified_draft_kv.py b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_unified_draft_kv.py
index 3af4c8a45102..fd26e42e5447 100644
--- a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_unified_draft_kv.py
+++ b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_unified_draft_kv.py
@@ -398,16 +398,19 @@ def _configs(specs: tuple[KVCacheLayerSpec, ...], target_sinks: Optional[int] =
     assert manager._windowed_draft_pool_ratio(configs, 1 << 30, 128, 0.5)[1] == pytest.approx(
         2 * draft_bytes / (1 << 30)
     )
-    # A quota too small for the windows keeps the bytes-per-token split.
+    # A quota too small for every window keeps the bytes-per-token split.
     token_share = draft_bytes_per_token / sum(
         spec.bytes_per_token() for spec in manager.layer_specs
     )
-    assert manager._windowed_draft_pool_ratio(configs, 1 << 20, 128)[1] == pytest.approx(
+    manager.max_batch_size = 64
+    assert manager._windowed_draft_pool_ratio(configs, 8 << 20, 128)[1] == pytest.approx(
         token_share
     )
-    # The quota must hold the context-chunk floor: two blocks per lifecycle, over the resume
-    # watermark; otherwise the backend keeps its own split.
-    floor_bytes = 3 * sum(spec.bytes_per_token() * 128 for spec in manager.layer_specs)
+    manager.max_batch_size = 2
+    # The quota must hold the context-chunk floor: two blocks per lifecycle over the resume
+    # watermark, each buffer rounded to the 2 MiB allocation granularity; otherwise the backend
+    # keeps its own split.
+    floor_bytes = len(configs) * (2 << 20)
     assert manager._windowed_draft_pool_ratio(configs, floor_bytes - 1, 128, 0.95) is None
     assert manager._windowed_draft_pool_ratio(configs, floor_bytes, 128, 0.95) is not None
     # Without a context-chunk bound the backend keeps its own split.