From a436008e8b2e5a5824ffc46b482bb2f3e34bc553 Mon Sep 17 00:00:00 2001 From: Zheyu Fu Date: Sat, 3 Oct 2026 13:48:57 -0700 Subject: [PATCH 1/3] [None][perf] MiniMax-M3: native P128 draft KV view for the NVFP4 hybrid cache With NVFP4 KV, the shared Eagle3 draft layer used a physical-P32 sub-page view (MiniMaxM3DraftSubpageView). Its copy_batch_block_offsets rebuilt the draft block table in Python/numpy on every decode prepare(), over each request's full padded page row (max_blocks_per_seq, 8196 for max_seq_len 1M), then expanded it x4 on the CPU and copied it H2D. That costs ~0.5 ms per request per decode step on the host and makes large-batch NVFP4 decode host-bound (B300 TP4, 128k context, batch 32: 38.4 ms/step vs 19.9 ms with FP8 KV). Run the Eagle layer at native FP8/P128 instead, as feat/m3_with_msa does since #18066: - _build_cache_config only strips the NVFP4 block-scale buffers from the dense/Eagle layers; no tokens_per_block_override for the Eagle layer. - New MiniMaxM3DraftKVCacheView: FP8, the trtllm-gen P128 opt-in, one pool rooted at the draft layer's K page, and a block table filled by a single copy_batch_block_offsets_to_device call with the view's own index_scales/kv_offset over the draft layer's lifecycle-pool slot rows. - Remove the P32 sub-page machinery (MiniMaxM3DraftSubpageView, DRAFT_PHYSICAL_TOKENS_PER_BLOCK, the P32 page-index override and helpers, the unused sub-paged scatter kernel) and inline the SWA-scratch expansion check back into KVCacheManagerV2._prepare_swa_scratch_copy_tensors, dropping the #19422-only hook. - Rename get_draft_subpage_view -> get_draft_kv_cache_view and is_fp8_subpaged_layer -> is_hybrid_draft_layer; update the unit tests. B300 (SM103), TP only, max_draft_len=3, measured on #19422 rev 4d7e78f3e3 plus this patch: batch-32 decode 38.4 -> 20.9 ms/step (FP8 KV 19.9). InferenceX AgentX P90 interactivity with forced acceptance: TP4 C40 79.5 -> 100.3, TP2 C45 44.6 -> 50.6. Rerun on main 00bdcb6ed (interleaved): TP4 C40 81.7 -> 100.4, TP2 C45 46.0 -> 50.4. Draft acceptance (24-prompt greedy probe) identical; GSM8K accuracy unchanged within run-to-run noise. Signed-off-by: Zheyu Fu --- .../sparse/minimax_m3/cache_manager.py | 412 ++++-------------- .../sparse/minimax_m3/kernels/msa_scatter.py | 112 ----- .../kernels/trtllm_gen_dense_decode.py | 6 +- .../_torch/models/modeling_minimaxm3.py | 2 +- .../kv_cache/kv_cache_manager_v2.py | 33 +- tensorrt_llm/_torch/speculative/interface.py | 2 +- tensorrt_llm/_torch/speculative/utils.py | 4 +- .../test_minimax_m3_shared_draft_layers.py | 200 ++++----- 8 files changed, 181 insertions(+), 590 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 a89bb47048a5..6cccf325173c 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 @@ -22,13 +22,14 @@ subclass that registers a per-sparse-layer ``Role.INDEX_KEY`` paged buffer alongside the standard K/V buffers; shared Eagle3 draft layers get their own virtual attention-op pools. + * :class:`MiniMaxM3DraftKVCacheView` — FP8/P128 view of the shared Eagle3 + draft layer when the target KV cache is NVFP4. """ from __future__ import annotations from typing import TYPE_CHECKING, List, Optional, Sequence, Tuple -import numpy as np import torch from tensorrt_llm._torch.disaggregation.resource.page import MapperKind @@ -268,12 +269,6 @@ 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() - # In the hybrid NVFP4/FP8 cache, shared FP8 Eagle layers require physical - # P32 pages for SM100/SM103 kernel compatibility. The KVCM config describes - # logical MSA blocks (P128), so allocation and the draft view share this - # physical page size instead. The subpage adapter is needed until Eagle - # kernels support the configured logical page size for this layout. - DRAFT_PHYSICAL_TOKENS_PER_BLOCK = 32 def __init__( self, @@ -343,10 +338,10 @@ def __init__( "[m3-kv] hybrid cache active: " f"{len(self.sparse_layer_ids)} sparse target layer(s)=NVFP4/P128, " f"{len(dense_target_layers)} dense target layer(s)=FP8/P128, " - f"{len(self._shared_draft_layer_ids)} shared Eagle layer(s)=FP8/P32" + f"{len(self._shared_draft_layer_ids)} shared Eagle layer(s)=FP8/P128" ) - self._draft_subpage_view_obj: Optional["MiniMaxM3DraftSubpageView"] = None + self._draft_kv_cache_view: Optional["MiniMaxM3DraftKVCacheView"] = None index_v_layer_ids = set(self.sparse_layer_ids) - self.disable_index_value_layer_ids if self.is_disagg and index_v_layer_ids: raise ValueError( @@ -406,35 +401,18 @@ def _build_cache_config(self, config): M3's 57 sparse target layers have a native MSA NVFP4 consumer. The three dense target layers and the appended one-model Eagle layer do not have a matching TRTLLM-Gen NVFP4 cubin, so those buffers retain - the proven FP8 representation. Target dense layers retain their - established P128 layout; only the shared Eagle layer uses physical - P32 pages because its SM100/SM103 kernels require that geometry. + the proven FP8 representation. """ if self.dtype != DataType.NVFP4: return super()._build_cache_config(config) - physical_page = self.DRAFT_PHYSICAL_TOKENS_PER_BLOCK - assert config.tokens_per_block % physical_page == 0, ( - f"M3 logical page P{config.tokens_per_block} must be divisible by " - f"the dense/Eagle physical page P{physical_page}." - ) scale_roles = {Role.KEY_BLOCK_SCALE, Role.VALUE_BLOCK_SCALE} for layer in config.layers: - local_layer_idx = int(layer.layer_id) - global_layer_idx = int(self.pp_layers[local_layer_idx]) - if global_layer_idx in self.sparse_layer_ids: + if int(self.pp_layers[int(layer.layer_id)]) in self.sparse_layer_ids: continue layer.buffers[:] = [ buffer for buffer in layer.buffers if buffer.role not in scale_roles ] - for buffer in layer.buffers: - if buffer.role not in (Role.KEY, Role.VALUE): - continue - if global_layer_idx in self._shared_draft_layer_ids: - buffer.size = ( - self.get_layer_bytes_per_token(local_layer_idx, buffer.role) * physical_page - ) - buffer.tokens_per_block_override = physical_page return super()._build_cache_config(config) def get_layer_bytes_per_token(self, local_layer_idx: int, data_role: Role): @@ -467,8 +445,8 @@ def is_fp8_dense_layer(self, layer_idx: int) -> bool: """Whether a dense target/Eagle layer is the FP8 half of hybrid KV.""" return self.dtype == DataType.NVFP4 and int(layer_idx) not in self.sparse_layer_ids - def is_fp8_subpaged_layer(self, layer_idx: int) -> bool: - """Whether the shared Eagle layer uses physical P32 FP8 pages.""" + def is_hybrid_draft_layer(self, layer_idx: int) -> bool: + """Whether a shared Eagle layer is read through :class:`MiniMaxM3DraftKVCacheView`.""" return self.is_fp8_dense_layer(layer_idx) and int(layer_idx) in self._shared_draft_layer_ids @property @@ -480,35 +458,16 @@ 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_attention_op_page_index_params( - self, layer_id: int, role: DataRole - ) -> Tuple[int, int, int]: - if ( - self.is_fp8_subpaged_layer(int(self.pp_layers[layer_id])) - and not self.enable_swa_scratch_reuse - ): - # The target table has one entry per P128 logical block. Record - # its first P32 page here; the drafter reads the separate expanded - # table supplied by MiniMaxM3DraftSubpageView. Heterogeneous target - # pools still need the base per-layer table conversion. - converter = self.impl.get_page_index_converter(layer_id, role) - return ( - int(converter.scale * converter.expansion), - int(converter.layer_offset * converter.expansion), - int(converter.scratch_pages_per_block * converter.expansion), - ) - return super()._get_attention_op_page_index_params(layer_id, role) + def get_draft_kv_cache_view(self) -> Optional["MiniMaxM3DraftKVCacheView"]: + """FP8/P128 view of the shared draft layer in a hybrid NVFP4 cache, or None. - def get_draft_subpage_view(self) -> Optional["MiniMaxM3DraftSubpageView"]: - """Sub-page view over the shared drafter pool, or None. + FP8/BF16 caches address the draft layer through this manager's own + attention-op pools instead. Only meaningful on a target manager carrying appended one-model draft layers; built lazily so the manager's page tables exist. A method rather than a property so ``getattr`` fetches it without - executing it (see ``resolve_draft_kv_cache_manager``). - - Retires with the P128 Eagle kernel fixes; see - ``DRAFT_PHYSICAL_TOKENS_PER_BLOCK``. + executing it (see ``get_draft_kv_cache_manager``). """ if self.dtype != DataType.NVFP4 or self.is_draft or not self._shared_draft_layer_ids: return None @@ -524,32 +483,23 @@ def get_draft_subpage_view(self) -> Optional["MiniMaxM3DraftSubpageView"]: ] if not draft_layers: return None - if self._draft_subpage_view_obj is None: - subpage_tokens = self.DRAFT_PHYSICAL_TOKENS_PER_BLOCK - self._draft_subpage_view_obj = MiniMaxM3DraftSubpageView( - self, - draft_layers, - subpage_tokens, - ) + if self._draft_kv_cache_view is None: + self._draft_kv_cache_view = MiniMaxM3DraftKVCacheView(self, draft_layers) logger.info( - f"[unified-kv] draft sub-page view active " - f"(tokens_per_block={self._draft_subpage_view_obj.tokens_per_block}, " - f"flat_page_bound={self._draft_subpage_view_obj.blocks_in_primary_pool})" + "[unified-kv] draft KV cache view active " + f"(flat_page_bound={self._draft_kv_cache_view.blocks_in_primary_pool})" ) - return self._draft_subpage_view_obj + return self._draft_kv_cache_view def add_dummy_requests(self, *args, **kwargs): - """Drop the draft sub-page view before delegating. + """Drop the draft KV cache view before delegating. 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. - - Retires with the P128 Eagle kernel fixes; see - ``DRAFT_PHYSICAL_TOKENS_PER_BLOCK``. """ - if isinstance(kwargs.get("draft_kv_cache_manager"), MiniMaxM3DraftSubpageView): + if isinstance(kwargs.get("draft_kv_cache_manager"), MiniMaxM3DraftKVCacheView): kwargs["draft_kv_cache_manager"] = None return super().add_dummy_requests(*args, **kwargs) @@ -572,7 +522,8 @@ def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None: "MiniMax-M3 shared Eagle3 draft layers do not support SWA scratch reuse." ) if self.dtype == DataType.NVFP4: - # Hybrid draft layers use MiniMaxM3DraftSubpageView. + # Hybrid draft layers use MiniMaxM3DraftKVCacheView, which also + # carries their P128 trtllm-gen opt-in. return # Draft layers run at the target's 128-token pages. trtllm-gen has P128 # kernels for their dense-GQA shapes but not for every shape, so opt in @@ -716,11 +667,6 @@ def _kv_slot_geometry( kv_layout = self._main_kv_layout_name() if kv_layout not in ("NHD", "HND"): raise ValueError(f"Unsupported kv_layout: {kv_layout}") - if self.is_fp8_subpaged_layer(layer_idx): - raise RuntimeError( - f"hybrid FP8 layer {layer_idx} uses four physical P32 pages; " - "use get_fp8_dense_buffers/get_dense_kv_subpage_pool instead" - ) if self.kv_cache_type == CacheTypeCpp.SELFKONLY: raise NotImplementedError( "MiniMaxM3KVCacheManagerV2 does not support the SELFKONLY cache type" @@ -795,23 +741,6 @@ def get_buffers( preserves the dim-0 stride (``scale * page_stride``), so ``view[s, 0/1, ...]`` lands on this layer's K/V at slot ``s``. """ - if self.is_fp8_subpaged_layer(layer_idx): - if kv_layout not in (None, "HND"): - raise ValueError( - "hybrid FP8 dense/Eagle buffers have a physical P32 HND layout; " - f"requested {kv_layout}" - ) - k, _v, slot_stride, pages_per_role = self._fp8_dense_data_buffers(layer_idx) - num_slots, _pages, num_heads, page_size, head_dim = k.shape - full = convert_to_torch_tensor( - TensorWrapper( - k.data_ptr(), - k.dtype, - [num_slots, slot_stride, num_heads, page_size, head_dim], - ) - ) - return full[:, : 2 * pages_per_role].unflatten(1, (2, pages_per_role)) - addr_key, torch_dtype, num_slots, scale, page_shape = self._kv_slot_geometry( layer_idx, kv_layout ) @@ -889,78 +818,6 @@ def get_block_scale_buffers( full_view = convert_to_torch_tensor(TensorWrapper(addr_key, torch_dtype, full_slot_shape)) return full_view[:, :2] - def _fp8_dense_data_buffers( - self, layer_idx: int - ) -> Tuple[torch.Tensor, torch.Tensor, int, int]: - """Return hybrid FP8 K/V views backed by physical P32 pages. - - Each logical P128 role page is laid out as four consecutive P32 HND - pages. ``slot_stride`` is measured in those physical pages and already - includes the V2 converter's expansion factor. - """ - if not self.is_fp8_subpaged_layer(layer_idx): - raise RuntimeError(f"layer {layer_idx} is not a physical-P32 hybrid FP8 layer") - local_layer_idx = self.layer_offsets[layer_idx] - physical_page = self.DRAFT_PHYSICAL_TOKENS_PER_BLOCK - pages_per_role = self.tokens_per_block // physical_page - addr_key = self.impl.get_mem_pool_base_address(local_layer_idx, Role.KEY) - addr_value = self.impl.get_mem_pool_base_address(local_layer_idx, Role.VALUE) - page_stride_key = self.impl.get_page_stride(local_layer_idx, Role.KEY) - page_stride_value = self.impl.get_page_stride(local_layer_idx, Role.VALUE) - assert page_stride_key == page_stride_value - assert addr_key + pages_per_role * page_stride_key == addr_value, ( - "M3 hybrid FP8 storage requires V immediately after K's physical " - f"P{physical_page} pages; layer={layer_idx} K={addr_key} " - f"stride={page_stride_key} V={addr_value}." - ) - - converter = self.impl.get_page_index_converter(local_layer_idx, Role.KEY) - assert int(converter.expansion) == pages_per_role, ( - f"layer {layer_idx} expected V2 expansion {pages_per_role}, got " - f"{int(converter.expansion)}" - ) - slot_stride = int(converter.scale) * pages_per_role - layer_offset_pages = int(converter.layer_offset) * pages_per_role - page_upper = self.impl.get_page_index_upper_bound(local_layer_idx, Role.KEY) - total_pages = int(page_upper) + layer_offset_pages - assert total_pages % slot_stride == 0 - num_slots = total_pages // slot_stride - - num_heads = self.num_kv_heads_per_layer[local_layer_idx] - head_dim = self.head_dim_per_layer[local_layer_idx] - full = convert_to_torch_tensor( - TensorWrapper( - addr_key, - torch.float8_e4m3fn, - [num_slots, slot_stride, num_heads, physical_page, head_dim], - ) - ) - return ( - full[:, :pages_per_role], - full[:, pages_per_role : 2 * pages_per_role], - slot_stride, - pages_per_role, - ) - - def get_fp8_dense_buffers(self, layer_idx: int) -> Tuple[torch.Tensor, torch.Tensor]: - """Return physical-P32 hybrid FP8 views as ``K, V``.""" - k, v, _slot_stride, _pages_per_role = self._fp8_dense_data_buffers(layer_idx) - return k, v - - def get_dense_kv_subpage_pool(self, layer_idx: int) -> Tuple[torch.Tensor, int, int]: - """Flat dense-attention pool, slot stride, and pages per K/V role.""" - if not self.is_fp8_subpaged_layer(layer_idx): - pool, slot_stride = self.get_kv_subpage_pool(layer_idx, "HND") - return pool, slot_stride, 1 - k, _v, slot_stride, pages_per_role = self._fp8_dense_data_buffers(layer_idx) - num_slots, _pages, num_heads, page_size, head_dim = k.shape - num_pages = (num_slots - 1) * slot_stride + 2 * pages_per_role - addr = k.data_ptr() - pool = convert_to_torch_tensor( - TensorWrapper(addr, k.dtype, [num_pages, num_heads, page_size, head_dim]) - ) - return pool, slot_stride, pages_per_role - def get_kv_subpage_pool( self, layer_idx: int, kv_layout: str = "HND" ) -> Tuple[torch.Tensor, int]: @@ -978,11 +835,6 @@ def get_kv_subpage_pool( spanning num_slots * scale, which would run off the pool by whatever this layer's K offset is inside a slot. """ - if self.is_fp8_subpaged_layer(layer_idx): - if kv_layout != "HND": - raise ValueError("hybrid FP8 dense/Eagle sub-pages are HND only") - pool, slot_stride, _pages_per_role = self.get_dense_kv_subpage_pool(layer_idx) - return pool, slot_stride addr_key, torch_dtype, num_slots, scale, page_shape = self._kv_slot_geometry( layer_idx, kv_layout ) @@ -1089,117 +941,74 @@ def get_block_ids_per_seq(self, request_ids): return padded_tensor -class MiniMaxM3DraftSubpageView: - """Present one shared draft layer's pool at a smaller kernel page size. +class MiniMaxM3DraftKVCacheView: + """The shared Eagle layer of a hybrid NVFP4 cache, presented as an FP8 cache. - With unified KV cache the drafter's KV lives inside the shared manager's - 128-token logical blocks, but the Eagle3 kernels are only healthy at - 32-token pages on this architecture. This view flows wherever a separate - draft manager would (``get_draft_kv_cache_manager`` and the - attention-metadata draft swap) and re-expresses the geometry only: + The draft layer runs the generic TRTLLM attention op, which reads the KV + dtype, pool pointers and block tables from the metadata's manager; the + shared manager's are NVFP4, with ``[data, scale]`` pool pointers. This view + flows wherever a separate draft manager would (``get_draft_kv_cache_manager`` + and the attention-metadata draft swap), with dtype FP8 and one pool rooted at + the draft layer's K page: slot ``s`` is K page ``s * scale`` and V page + ``s * scale + 1`` (see :meth:`MiniMaxM3KVCacheManagerV2.get_kv_subpage_pool`). + Everything else, block lifecycle included, stays with the shared manager. - * the single pool pointer is re-rooted at the drafter's K address, and - the draft layer's row in the pool mapping points at that pool; - * the block table expands each logical slot ``s`` into sub-pages — - K at ``s*scale*subdiv + j`` (``j < subdiv``), V at ``+subdiv`` - (``scale`` = the drafter's pages per mega-slot, from - ``_kv_slot_geometry``; same layout trick as the dense-layer - trtllm-gen adapter). - - The attention op reads ``tokens_per_block`` and the pool pointers from - the metadata's manager, so no attention-backend changes are needed. The - view owns no blocks: lifecycle stays entirely with the shared manager. - - Retires with the P128 Eagle kernel fixes; see the retirement plan on - ``MiniMaxM3KVCacheManagerV2.DRAFT_PHYSICAL_TOKENS_PER_BLOCK``. + Args: + manager: The shared target manager that owns the draft layer's storage. + draft_layer_ids: This rank's shared draft layers; exactly one is supported. """ - def __init__(self, manager, draft_layer_ids: Sequence[int], subpage_tokens: int): + dtype = DataType.FP8 + num_pools = num_attention_op_pools = 1 + # trtllm-gen has P128 kernels for the draft layer's dense-GQA shape but not + # for every shape, so only this view opts in. + trtllm_gen_extra_tokens_per_block = frozenset({128}) + + def __init__(self, manager: MiniMaxM3KVCacheManagerV2, draft_layer_ids: Sequence[int]) -> None: if len(draft_layer_ids) != 1: raise NotImplementedError( - "MiniMax-M3 draft subpage views support exactly one local shared draft layer; " - f"got {len(draft_layer_ids)}. Multiple layers require separate P32 pool roots." + "MiniMax-M3 draft KV cache views support exactly one local shared draft layer; " + f"got {len(draft_layer_ids)}. Multiple layers require separate pool roots." ) self._manager = manager - self.tokens_per_block = int(subpage_tokens) layer_id = draft_layer_ids[0] - is_hybrid_fp8 = bool( - getattr(manager, "is_fp8_subpaged_layer", lambda _layer_idx: False)(layer_id) - ) - if is_hybrid_fp8: - assert self.tokens_per_block == manager.DRAFT_PHYSICAL_TOKENS_PER_BLOCK, ( - "hybrid FP8 Eagle draft attention must use the physical draft-cache " - f"page size P{manager.DRAFT_PHYSICAL_TOKENS_PER_BLOCK}, got " - f"P{self.tokens_per_block}" - ) - assert manager.tokens_per_block % self.tokens_per_block == 0, ( - f"subpage size {subpage_tokens} must divide manager " - f"tokens_per_block {manager.tokens_per_block}" - ) - self._subdiv = manager.tokens_per_block // self.tokens_per_block - # The hybrid layout puts dense/Eagle FP8 pages in a separate physical - # pool from sparse NVFP4 data/scales. Root this single-pool view at the - # draft layer's K address, while sourcing raw logical slot IDs from the - # draft layer's actual V2 pool. - local = manager.layer_offsets[layer_id] - self._local_layer_idx = int(local) - self._source_pool_id = int(manager.impl.get_layer_group_id(int(local))) - if is_hybrid_fp8: - k, _v, slot_stride, pages_per_role = manager._fp8_dense_data_buffers(layer_id) - assert pages_per_role == self._subdiv - addr_key = k.data_ptr() - self._num_slots = int(k.shape[0]) - self._slot_units = slot_stride - self.dtype = DataType.FP8 - else: - addr_key, _dt, num_slots, scale, _shape = manager._kv_slot_geometry(layer_id, None) - self._num_slots = int(num_slots) - self._slot_units = scale * self._subdiv - self.num_pools = 1 - self.num_attention_op_pools = 1 - self.max_blocks_per_seq = manager.max_blocks_per_seq * self._subdiv - # Single-pool pointer rooted at the drafter's K; the op derives the - # page stride from tokens_per_block, so unit indices below address - # 32-token drafter pages directly. + local = int(manager.layer_offsets[layer_id]) + self._local_layer_idx = local + # get_kv_subpage_pool checks that V immediately follows K, hence kv_offset 1. + flat_pool, slot_stride = manager.get_kv_subpage_pool(layer_id) + # Set rather than delegated: the manager's bound counts from another pool base. + self.blocks_in_primary_pool = int(flat_pool.shape[0]) + pinned = prefer_pinned() self.kv_cache_pool_pointers = torch.tensor( - [[addr_key, 0]], dtype=torch.int64, pin_memory=prefer_pinned() + [[flat_pool.data_ptr(), 0]], dtype=torch.int64, pin_memory=pinned ) - mapping = manager.kv_cache_pool_mapping.clone() - mapping[int(local)] = torch.tensor([0, 0], dtype=mapping.dtype) - self.kv_cache_pool_mapping = mapping - # Placeholder host mirror: the dense TRTLLM path plans from device - # offsets; nothing reads the host table during the draft window. - self.host_kv_cache_block_offsets = torch.zeros( - (1, 1, 2, 1), dtype=torch.int32, pin_memory=prefer_pinned() - ) - self._slots_host: Optional[np.ndarray] = None - self._arange: Optional[torch.Tensor] = None - - @property - def blocks_in_primary_pool(self) -> int: - """Flattened sub-page index bound relative to the draft K pointer. - - ``FlashInferTrtllmGenFmha`` uses this value to size the flat paged-KV - tensor passed to FlashInfer. The wrapped V2 manager reports its bound - in 128-token page units relative to a different pool base, so - delegating that property through ``__getattr__`` under-describes this - 32-token, draft-K-rooted view. The final slot contributes only this - layer's K and V pages; inter-layer padding after V is not addressable - from the view and need not be included. - """ - return (self._num_slots - 1) * self._slot_units + 2 * self._subdiv + self.kv_cache_pool_mapping = manager.kv_cache_pool_mapping.clone() + self.kv_cache_pool_mapping[local] = 0 + # The source pool's own index_scales/kv_offset describe its + # representative layer, not this rerooted draft layer. + self.index_scales = torch.tensor([slot_stride], dtype=torch.int32, pin_memory=pinned) + self.kv_offset = torch.tensor([1], dtype=torch.int32, pin_memory=pinned) + # Slot ids come from the draft layer's lifecycle pool; with per-layer page + # tables its pool-mapping row holds the layer index instead. C++ keeps + # these rows current for every live request. + source_pool = int(manager.impl.get_layer_group_id(local)) + self.host_kv_cache_block_offsets = manager.host_kv_cache_block_offsets[ + source_pool : source_pool + 1 + ] - def get_attention_op_num_blocks(self, local_layer_idx: int) -> int: - """Bound relative to this view's draft-K root, in P32 page units.""" + def _check_layer(self, local_layer_idx: int) -> None: if local_layer_idx != self._local_layer_idx: raise ValueError("MiniMax-M3 draft cache view only addresses its shared draft layer") + + def get_attention_op_num_blocks(self, local_layer_idx: int) -> int: + """Page bound relative to this view's draft-K root.""" + self._check_layer(local_layer_idx) return self.blocks_in_primary_pool def get_attention_op_kv_page_offset(self, local_layer_idx: int) -> int: - """V follows this layer's K pages in each physical allocation slot.""" - if local_layer_idx != self._local_layer_idx: - raise ValueError("MiniMax-M3 draft cache view only addresses its shared draft layer") - return self._subdiv + """V immediately follows this layer's K page in each slot.""" + self._check_layer(local_layer_idx) + return 1 def __getattr__(self, name): manager = self.__dict__.get("_manager") @@ -1207,63 +1016,9 @@ def __getattr__(self, name): raise AttributeError(name) return getattr(manager, name) - @property - def host_kv_cache_pool_pointers(self): - return self.kv_cache_pool_pointers - def free_resources(self, request) -> None: """No-op: block lifecycle belongs to the shared manager.""" - def _host_block_table( - self, - slot_rows: Sequence[Sequence[int]], - num_seqs: int, - max_slots: int, - dtype: torch.dtype, - ) -> torch.Tensor: - """Expand this batch's slot ids into a freshly allocated pinned table. - - The buffer is allocated per call on purpose: it is the source of an - asynchronous H2D copy, whose source is read at copy execution time - rather than enqueue time. A persistent buffer refilled in place would - let the next iteration's refill clobber a still-pending copy — the - drafter would then index another batch's blocks (nvbug 6293536, whose - rationale the V1 manager spells out on - ``KVCacheManager._stage_block_offsets_for_copy``). The caching host - allocator keeps this block alive until the copy retires. The numpy - scratch below stays persistent: it is only ever read synchronously. - """ - sub = self._subdiv - if ( - self._slots_host is None - or self._slots_host.shape[0] < num_seqs - or self._slots_host.shape[1] != max_slots - ): - self._slots_host = np.zeros((num_seqs, max_slots), dtype=np.int32) - self._arange = torch.arange(sub, dtype=dtype) - slots_np = self._slots_host[:num_seqs] - slots_np.fill(0) - # Ragged fill is the only per-row work (numpy parses each row's list at - # C speed); the arithmetic below is one fused expansion across the - # batch. Pad/BAD_PAGE_INDEX entries clamp to slot 0 (safe pages: - # kernels never read past kv_lens). - for i, row in enumerate(slot_rows[:num_seqs]): - n = min(len(row), max_slots) - if n > 0: - slots_np[i, :n] = row[:n] - np.clip(slots_np, 0, None, out=slots_np) - slots = torch.from_numpy(slots_np).to(dtype) - host = torch.empty( - (num_seqs, 2, max_slots * sub), - dtype=dtype, - pin_memory=prefer_pinned(), - device="cpu", - ) - out = host.view(num_seqs, 2, max_slots, sub) - torch.add(slots.unsqueeze(-1) * self._slot_units, self._arange, out=out[:, 0]) - torch.add(out[:, 0], sub, out=out[:, 1]) - return host - def copy_batch_block_offsets( self, dst_tensor: torch.Tensor, @@ -1273,14 +1028,20 @@ def copy_batch_block_offsets( num_seqs: int, max_blocks: Optional[int] = None, ) -> None: - # Raw logical slot ids from the draft layer's physical pool. - slot_rows = self._manager._get_batch_cache_indices_by_pool_id( - request_ids, pool_id=self._source_pool_id - ) - host = self._host_block_table( - slot_rows, num_seqs, dst_tensor.shape[-1] // self._subdiv, dst_tensor.dtype + """Fill the draft block table with V2's single device copy. + + Not delegated to :meth:`KVCacheManagerV2.copy_batch_block_offsets`, which + would follow the target's per-layer page tables through ``__getattr__``. + """ + manager = self._manager + copy_batch_block_offsets_to_device( + self.host_kv_cache_block_offsets, + dst_tensor, + manager.index_mapper.get_copy_index(request_ids, num_contexts, beam_width), + self.index_scales, + self.kv_offset, + manager._stream.cuda_stream, ) - dst_tensor[0, :num_seqs].copy_(host, non_blocking=True) def get_minimax_m3_kv_cache_manager_cls(): @@ -1289,6 +1050,7 @@ def get_minimax_m3_kv_cache_manager_cls(): __all__ = [ + "MiniMaxM3DraftKVCacheView", "MiniMaxM3KVCacheManagerV2", "MiniMaxM3SparseIndexCache", "derive_shared_draft_layout", diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/msa_scatter.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/msa_scatter.py index 3df975f940ac..7f15cb3b2796 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/msa_scatter.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/msa_scatter.py @@ -71,61 +71,6 @@ def _fused_paged_scatter_kernel( tl.store(i_dst, i_vals.to(idx_cache.dtype.element_ty), mask=valid) -@triton.jit -def _fused_subpaged_scatter_kernel( - k_src, - v_src, - k_cache, - v_cache, - out_cache_loc, - k_src_row_stride, - v_src_row_stride, - kc_stride_page, - kc_stride_subpage, - kc_stride_head, - kc_stride_tok, - vc_stride_page, - vc_stride_subpage, - vc_stride_head, - vc_stride_tok, - logical_tokens_per_block, - physical_tokens_per_block, - H: tl.constexpr, - D: tl.constexpr, -): - """Scatter FP8 K/V into P32 pages inside one logical P128 slot.""" - t = tl.program_id(0).to(tl.int64) - slot = tl.load(out_cache_loc + t).to(tl.int64) - valid = slot >= 0 - page = slot // logical_tokens_per_block - logical_within = slot % logical_tokens_per_block - subpage = logical_within // physical_tokens_per_block - within = logical_within % physical_tokens_per_block - d = tl.arange(0, D) - for h in tl.static_range(H): - src = t * k_src_row_stride + h * D + d - k_vals = tl.load(k_src + src) - v_vals = tl.load(v_src + t * v_src_row_stride + h * D + d) - k_dst = ( - k_cache - + page * kc_stride_page - + subpage * kc_stride_subpage - + h * kc_stride_head - + within * kc_stride_tok - + d - ) - v_dst = ( - v_cache - + page * vc_stride_page - + subpage * vc_stride_subpage - + h * vc_stride_head - + within * vc_stride_tok - + d - ) - tl.store(k_dst, k_vals.to(k_cache.dtype.element_ty), mask=valid) - tl.store(v_dst, v_vals.to(v_cache.dtype.element_ty), mask=valid) - - @triton.jit def _fused_nvfp4_paged_scatter_kernel( k_data_src, @@ -330,62 +275,6 @@ def fused_write_layer_caches( return True -def fused_write_subpaged_layer_caches( - k_cache: torch.Tensor, - v_cache: torch.Tensor, - out_cache_loc: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, -) -> bool: - """Write ordinary K/V into physical sub-pages of a logical cache block. - - Hybrid M3 stores dense and shared-Eagle K/V as FP8 P32 pages while its - allocator and request lifecycle remain P128. ``k_cache``/``v_cache`` are - ``[logical_slots, pages_per_role, heads, P32, D]`` zero-copy views. - """ - if not (k.is_cuda and k_cache.is_cuda): - return False - if k_cache.dim() != 5 or v_cache.shape != k_cache.shape: - return False - if k_cache.stride(-1) != 1 or v_cache.stride(-1) != 1: - return False - _num_slots, pages_per_role, num_heads, physical_page, head_dim = k_cache.shape - if pages_per_role <= 0 or physical_page <= 0 or (head_dim & (head_dim - 1)) != 0: - return False - inner = num_heads * head_dim - k_stride = _row_stride_if_fusable(k, inner) - v_stride = _row_stride_if_fusable(v, inner) - if k_stride is None or v_stride is None: - return False - num_tokens = int(out_cache_loc.shape[0]) - if num_tokens == 0: - return True - - _fused_subpaged_scatter_kernel[(num_tokens,)]( - k, - v, - k_cache, - v_cache, - out_cache_loc, - k_stride, - v_stride, - k_cache.stride(0), - k_cache.stride(1), - k_cache.stride(2), - k_cache.stride(3), - v_cache.stride(0), - v_cache.stride(1), - v_cache.stride(2), - v_cache.stride(3), - pages_per_role * physical_page, - physical_page, - H=num_heads, - D=head_dim, - num_warps=2, - ) - return True - - def fused_write_layer_caches_nvfp4( k_data_cache: torch.Tensor, v_data_cache: torch.Tensor, @@ -537,5 +426,4 @@ def fused_write_layer_caches_nvfp4( __all__ = [ "fused_write_layer_caches", "fused_write_layer_caches_nvfp4", - "fused_write_subpaged_layer_caches", ] diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/trtllm_gen_dense_decode.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/trtllm_gen_dense_decode.py index 452863586356..075ebe04da12 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/trtllm_gen_dense_decode.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/trtllm_gen_dense_decode.py @@ -219,11 +219,13 @@ def uniform_subpages_per_slot(kv_cache_manager) -> int: if get_pool is None or not layer_offsets: return 0 is_nvfp4 = getattr(kv_cache_manager, "is_nvfp4_layer", lambda _: False) - is_subpaged = getattr(kv_cache_manager, "is_fp8_subpaged_layer", lambda _: False) + is_hybrid_draft = getattr(kv_cache_manager, "is_hybrid_draft_layer", lambda _: False) + # NVFP4 sparse layers and the hybrid draft layer (read through its own view) + # never read the staged table and may sit in pools with another factor. factors = { int(get_pool(layer_idx, "HND")[1]) for layer_idx in layer_offsets - if not is_nvfp4(layer_idx) and not is_subpaged(layer_idx) + if not is_nvfp4(layer_idx) and not is_hybrid_draft(layer_idx) } return factors.pop() if len(factors) == 1 else 0 diff --git a/tensorrt_llm/_torch/models/modeling_minimaxm3.py b/tensorrt_llm/_torch/models/modeling_minimaxm3.py index c422c9d6313c..f17b64016231 100644 --- a/tensorrt_llm/_torch/models/modeling_minimaxm3.py +++ b/tensorrt_llm/_torch/models/modeling_minimaxm3.py @@ -2884,7 +2884,7 @@ def __init__(self, model_config: "ModelConfig[PretrainedConfig]"): model_config = get_text_model_config(model_config) if model_config.quant_config.kv_cache_quant_algo == QuantAlgo.NVFP4: # M3's 57 sparse target layers have an MSA NVFP4 consumer, but the - # one-model Eagle layer has no shipped P32 NVFP4 decode cubin. + # one-model Eagle layer has no shipped NVFP4 trtllm-gen cubin. # Keep its modules and shared draft cache on their established FP8 # representation while the target remains NVFP4. model_config.extra_attrs["draft_kv_cache_quant_algo_override"] = QuantAlgo.FP8 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 371d7ed85327..2efa591f3f51 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 @@ -2479,23 +2479,6 @@ def _log_kv_cache_pool_lifecycle_mapping(self) -> None: for entry in entries: logger.info(entry) - def _get_attention_op_page_index_params( - self, layer_id: LayerId, role: DataRole - ) -> Tuple[int, int, int]: - """Scale, layer offset and scratch span for one entry per logical block.""" - converter = self.impl.get_page_index_converter(layer_id, role) - if converter.expansion != 1: - raise NotImplementedError( - "SWA scratch block-table conversion does not support " - f"expanded page indices yet: layer={layer_id}, role={role}, " - f"expansion={converter.expansion}" - ) - return ( - int(converter.scale), - int(converter.layer_offset), - int(converter.scratch_pages_per_block), - ) - def _prepare_swa_scratch_copy_tensors(self, index_mapper_capacity: int) -> None: pool_ids = torch.empty( self.num_attention_op_pools, @@ -2520,13 +2503,17 @@ def _prepare_swa_scratch_copy_tensors(self, index_mapper_capacity: int) -> None: role_a if role_b is None else role_b, ] for role_idx, role in enumerate(roles): - scale, layer_offset, scratch_span = self._get_attention_op_page_index_params( - layer_id, role - ) + converter = self.impl.get_page_index_converter(layer_id, role) + if converter.expansion != 1: + raise NotImplementedError( + "SWA scratch block-table conversion does not support " + f"expanded page indices yet: layer={layer_id}, role={role}, " + f"expansion={converter.expansion}" + ) pool_ids[local_layer_idx, role_idx] = pool_id - scales[local_layer_idx, role_idx] = scale - layer_offsets[local_layer_idx, role_idx] = layer_offset - scratch_pages[local_layer_idx, role_idx] = scratch_span + scales[local_layer_idx, role_idx] = int(converter.scale) + layer_offsets[local_layer_idx, role_idx] = int(converter.layer_offset) + scratch_pages[local_layer_idx, role_idx] = int(converter.scratch_pages_per_block) staging_capacity = index_mapper_capacity * self.max_beam_width device = torch.device("cuda", torch.cuda.current_device()) diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index 37992e2f8835..85d3fd1dd829 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -2990,7 +2990,7 @@ def get_draft_kv_cache_manager(self, resource_manager): if resource_manager is not None: target = resource_manager.get_resource_manager( ResourceManagerType.KV_CACHE_MANAGER) - get_view = getattr(target, "get_draft_subpage_view", None) + get_view = getattr(target, "get_draft_kv_cache_view", None) if get_view is not None: return get_view() return None diff --git a/tensorrt_llm/_torch/speculative/utils.py b/tensorrt_llm/_torch/speculative/utils.py index 053723280738..3145d5ba0fb9 100644 --- a/tensorrt_llm/_torch/speculative/utils.py +++ b/tensorrt_llm/_torch/speculative/utils.py @@ -787,7 +787,7 @@ def get_num_extra_kv_tokens(spec_config): def get_draft_kv_cache_manager(spec_config, resource_manager): """ - Return the one-model draft cache manager, including a shared subpage view + Return the one-model draft cache manager, including a shared draft KV view when the target manager owns the draft cache storage. """ from ..pyexecutor.resource_manager import ResourceManagerType @@ -802,7 +802,7 @@ def get_draft_kv_cache_manager(spec_config, resource_manager): return draft target = resource_manager.get_resource_manager( ResourceManagerType.KV_CACHE_MANAGER) - get_view = getattr(target, "get_draft_subpage_view", None) + get_view = getattr(target, "get_draft_kv_cache_view", None) return get_view() if get_view is not None else None diff --git a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_shared_draft_layers.py b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_shared_draft_layers.py index 125a1135c377..d6a9019571e6 100644 --- a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_shared_draft_layers.py +++ b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_shared_draft_layers.py @@ -31,12 +31,15 @@ cache_manager as m3_cache_manager, ) from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.cache_manager import ( - MiniMaxM3DraftSubpageView, + MiniMaxM3DraftKVCacheView, MiniMaxM3KVCacheManagerV2, derive_shared_draft_layout, extend_attention_op_pools_for_shared_draft_layers, shared_draft_layer_count, ) +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.trtllm_gen_dense_decode import ( + uniform_subpages_per_slot, +) from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention, TrtllmAttentionMetadata from tensorrt_llm._torch.pyexecutor.kv_cache import kv_cache_manager_v2 from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2, Role @@ -56,7 +59,9 @@ @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA cache pools") -def test_nvfp4_shared_draft_constructs_and_converts_heterogeneous_page_tables() -> None: +def test_nvfp4_shared_draft_constructs_and_converts_heterogeneous_page_tables( + monkeypatch: pytest.MonkeyPatch, +) -> None: """Exercise real pool allocation and both target/draft page-table consumers.""" torch.cuda.init() manager = MiniMaxM3KVCacheManagerV2( @@ -85,7 +90,8 @@ def test_nvfp4_shared_draft_constructs_and_converts_heterogeneous_page_tables() try: assert manager._use_per_layer_page_tables assert manager._shared_draft_layer_ids == [4] - assert manager.impl.get_page_index_converter(4, Role.KEY).expansion == 4 + # The shared Eagle layer keeps native P128 pages. + assert manager.impl.get_page_index_converter(4, Role.KEY).expansion == 1 # Fill two logical request rows; no context scratch is active. for row, slot in enumerate((2, 5)): manager.host_kv_cache_block_offsets[:, row].fill_(slot) @@ -104,19 +110,51 @@ def test_nvfp4_shared_draft_constructs_and_converts_heterogeneous_page_tables() for row, slot in enumerate((2, 5)): expected = converter( [slot] * manager.max_blocks_per_seq, PageIndexMode.PER_LAYER - )[:: converter.expansion] + ) assert actual[layer_id, row, role_idx].tolist() == expected - # The draft attention consumer gets every P32 page, not just the first. - draft = manager.get_draft_subpage_view() - assert draft is not None - table = draft._host_block_table([[2, 5]], 1, 2, torch.int32) + # The source pool's scale and K/V offset describe its representative + # layer; poison them so only the view's own values produce the pages below. + manager.index_scales.zero_() + manager.kv_offset.zero_() + + # The draft view is an FP8/P128 cache rooted at the draft layer's K page. + draft = manager.get_draft_kv_cache_view() + assert isinstance(draft, MiniMaxM3DraftKVCacheView) + assert draft.dtype == DataType.FP8 + assert draft.trtllm_gen_extra_tokens_per_block == frozenset({128}) + assert draft.kv_cache_pool_pointers.tolist() == [ + [manager.impl.get_mem_pool_base_address(4, Role.KEY, PageIndexMode.SHARED), 0] + ] + assert draft.kv_cache_pool_mapping[4].tolist() == [0, 0] + source_pool = manager.impl.get_layer_group_id(4) + assert ( + draft.host_kv_cache_block_offsets.data_ptr() + == manager.host_kv_cache_block_offsets[source_pool].data_ptr() + ) + + # One device copy converts the draft layer's slot ids to P128 pages. + copy_idx = torch.tensor([0, 1], dtype=torch.int32, pin_memory=True) + monkeypatch.setattr( + manager, + "index_mapper", + SimpleNamespace(get_copy_index=lambda ids, num_contexts, beam_width: copy_idx), + ) + device_copy = Mock(wraps=m3_cache_manager.copy_batch_block_offsets_to_device) + monkeypatch.setattr(m3_cache_manager, "copy_batch_block_offsets_to_device", device_copy) + draft_output = torch.full( + (1, 2, 2, draft.max_blocks_per_seq), -7, dtype=torch.int32, device="cuda" + ) + draft.copy_batch_block_offsets(draft_output, [101, 102], 1, 0, 2) + torch.cuda.synchronize() + device_copy.assert_called_once() converter = manager.impl.get_page_index_converter(4, Role.KEY) - expected_k = converter([2, 5], PageIndexMode.SHARED) - assert table[0, 0].tolist() == expected_k - assert table[0, 1].tolist() == [page + 4 for page in expected_k] + for row, slot in enumerate((2, 5)): + expected_k = converter([slot] * draft.max_blocks_per_seq, PageIndexMode.SHARED) + assert draft_output[0, row, 0].tolist() == expected_k + assert draft_output[0, row, 1].tolist() == [page + 1 for page in expected_k] - # FMHA must use the view's P32 extent and K/V displacement rather than + # FMHA must use the view's extent and K/V displacement rather than # rejecting the view or delegating to its differently rooted owner. attn = TrtllmAttention.__new__(TrtllmAttention) attn.local_layer_idx = 4 @@ -125,10 +163,14 @@ def test_nvfp4_shared_draft_constructs_and_converts_heterogeneous_page_tables() metadata = TrtllmAttentionMetadata( max_num_requests=2, max_num_tokens=256, kv_cache_manager=draft, mapping=manager.mapping ) - assert fmha._get_total_num_blocks(metadata) == draft.blocks_in_primary_pool - assert get_kv_page_offset(attn, metadata, 0) == 4 + assert fmha._get_total_num_blocks(metadata) == manager.get_kv_subpage_pool(4)[0].shape[0] + assert get_kv_page_offset(attn, metadata, 0) == 1 with pytest.raises(ValueError, match="only addresses its shared draft layer"): draft.get_attention_op_num_blocks(0) + + # Dense decode stages one sub-page factor for the dense target layers; the + # NVFP4 layer and the view-read draft layer may sit in other pools. + assert uniform_subpages_per_slot(manager) == manager.get_kv_subpage_pool(0)[1] finally: manager.shutdown() @@ -299,97 +341,9 @@ def test_shared_draft_layer_count_follows_the_base_manager(): assert shared_draft_layer_count(None, None) == 0 -NVFP4_DRAFT_LAYER = 60 -NVFP4_SCALE = 178 -NVFP4_ADDR = 0x7000_0000 - - -class _Nvfp4BaseManager: - tokens_per_block = 128 - max_blocks_per_seq = 16 - num_pools = 1 - # V2's flattened bound uses the target pool's 128-token page units and - # base pointer. The draft view must not delegate this value. - blocks_in_primary_pool = 1024 * NVFP4_SCALE - - def __init__(self): - self.layer_offsets = {NVFP4_DRAFT_LAYER: NVFP4_DRAFT_LAYER} - self.kv_cache_pool_mapping = torch.zeros((NVFP4_DRAFT_LAYER + 1, 2), dtype=torch.int32) - self.kv_cache_pool_mapping[NVFP4_DRAFT_LAYER] = torch.tensor([0, 7], dtype=torch.int32) - self.slot_rows = [[5, 7]] - self.impl = SimpleNamespace( - get_layer_group_id=lambda local: int(self.kv_cache_pool_mapping[local, 0]) - ) - - def _kv_slot_geometry(self, layer_idx, kv_layout): - assert layer_idx == NVFP4_DRAFT_LAYER - page_shape = [self.tokens_per_block, 16, 128] - return NVFP4_ADDR, torch.int8, 1024, NVFP4_SCALE, page_shape - - def _get_batch_cache_indices_by_pool_id(self, request_ids, *, pool_id): - assert pool_id == 0 - return self.slot_rows[: len(request_ids)] - - -class _Nvfp4HybridManager(_Nvfp4BaseManager): - dtype = DataType.NVFP4 - DRAFT_PHYSICAL_TOKENS_PER_BLOCK = 32 - num_pools = 2 - - def __init__(self): - super().__init__() - self.kv_cache_pool_mapping[NVFP4_DRAFT_LAYER] = torch.tensor([1, 7], dtype=torch.int32) - - def is_fp8_subpaged_layer(self, layer_idx): - assert layer_idx == NVFP4_DRAFT_LAYER - return True - - def _fp8_dense_data_buffers(self, layer_idx): - assert layer_idx == NVFP4_DRAFT_LAYER - - class _Pointer: - shape = (1024, NVFP4_SCALE * 4, 16, 32, 128) - - @staticmethod - def data_ptr(): - return NVFP4_ADDR - - return _Pointer(), None, NVFP4_SCALE * 4, 4 - - def _get_batch_cache_indices_by_pool_id(self, request_ids, *, pool_id): - assert pool_id == 1 - return self.slot_rows[: len(request_ids)] - - -def test_hybrid_view_publishes_an_fp8_pool_pointer_from_the_draft_pool(): - view = MiniMaxM3DraftSubpageView(_Nvfp4HybridManager(), [NVFP4_DRAFT_LAYER], 32) - expected = [[NVFP4_ADDR, 0]] - assert view.kv_cache_pool_pointers.tolist() == expected - assert view.host_kv_cache_pool_pointers.tolist() == expected - assert view.dtype == DataType.FP8 - assert view._source_pool_id == 1 - assert view.blocks_in_primary_pool == (1024 - 1) * NVFP4_SCALE * 4 + 8 - - dst = torch.full((1, 1, 2, view.max_blocks_per_seq), -7, dtype=torch.int32) - view.copy_batch_block_offsets(dst, request_ids=[123], beam_width=1, num_contexts=1, num_seqs=1) - unit = NVFP4_SCALE * 4 - expected_k = [5 * unit + j for j in range(4)] + [7 * unit + j for j in range(4)] - assert dst[0, 0, 0, :8].tolist() == expected_k - assert dst[0, 0, 1, :8].tolist() == [page + 4 for page in expected_k] - - -def test_hybrid_view_rejects_non_p32_draft_pages(): - try: - MiniMaxM3DraftSubpageView(_Nvfp4HybridManager(), [NVFP4_DRAFT_LAYER], 128) - except AssertionError as error: - assert "physical draft-cache page size P32" in str(error) - else: - raise AssertionError("expected NVFP4 Eagle draft view to require P32 pages") - - @pytest.mark.parametrize("local_draft_layers", [[], [60], [61]]) @pytest.mark.parametrize("swa_scratch_reuse", [False, True]) -def test_draft_subpage_accessor_respects_local_layers( +def test_draft_kv_cache_view_accessor_respects_local_layers( monkeypatch: pytest.MonkeyPatch, local_draft_layers: list[int], swa_scratch_reuse: bool ) -> None: """Only local draft layers may create a view; scratch reuse is unsupported on every rank.""" @@ -399,46 +353,44 @@ def test_draft_subpage_accessor_respects_local_layers( manager._shared_draft_layer_ids = [60, 61] manager.layer_offsets = {layer: local for local, layer in enumerate([30, *local_draft_layers])} manager.enable_swa_scratch_reuse = swa_scratch_reuse - manager._draft_subpage_view_obj = None - view = SimpleNamespace(tokens_per_block=32, blocks_in_primary_pool=1024) + manager._draft_kv_cache_view = None + view = SimpleNamespace(blocks_in_primary_pool=1024) create_view = Mock(return_value=view) - monkeypatch.setattr(m3_cache_manager, "MiniMaxM3DraftSubpageView", create_view) + monkeypatch.setattr(m3_cache_manager, "MiniMaxM3DraftKVCacheView", create_view) if swa_scratch_reuse: with pytest.raises(NotImplementedError, match="SWA scratch reuse"): - manager.get_draft_subpage_view() + manager.get_draft_kv_cache_view() create_view.assert_not_called() - assert manager._draft_subpage_view_obj is None + assert manager._draft_kv_cache_view is None elif not local_draft_layers: - assert manager.get_draft_subpage_view() is None + assert manager.get_draft_kv_cache_view() is None create_view.assert_not_called() - assert manager._draft_subpage_view_obj is None + assert manager._draft_kv_cache_view is None else: - assert manager.get_draft_subpage_view() is view - assert manager.get_draft_subpage_view() is view - create_view.assert_called_once_with(manager, local_draft_layers, 32) + assert manager.get_draft_kv_cache_view() is view + assert manager.get_draft_kv_cache_view() is view + create_view.assert_called_once_with(manager, local_draft_layers) assert manager._shared_draft_layer_ids == [60, 61] @pytest.mark.parametrize("via_accessor", [False, True]) -def test_draft_subpage_view_rejects_multiple_local_layers(via_accessor: bool) -> None: - """The real single-pool view must reject a second layer before publishing wrong pointers.""" - manager = _Nvfp4HybridManager() - manager.layer_offsets[61] = 61 - manager.kv_cache_pool_mapping = torch.cat( - [manager.kv_cache_pool_mapping, torch.tensor([[1, 9]], dtype=torch.int32)] - ) +def test_draft_kv_cache_view_rejects_multiple_local_layers(via_accessor: bool) -> None: + """The single-pool view must reject a second layer before publishing wrong pointers.""" + manager = MiniMaxM3KVCacheManagerV2.__new__(MiniMaxM3KVCacheManagerV2) + manager.dtype = DataType.NVFP4 manager.is_draft = False manager._shared_draft_layer_ids = [60, 61] + manager.layer_offsets = {60: 0, 61: 1} manager.enable_swa_scratch_reuse = False - manager._draft_subpage_view_obj = None + manager._draft_kv_cache_view = None with pytest.raises(NotImplementedError, match="exactly one local shared draft layer"): if via_accessor: - MiniMaxM3KVCacheManagerV2.get_draft_subpage_view(manager) + manager.get_draft_kv_cache_view() else: - MiniMaxM3DraftSubpageView(manager, [60, 61], 32) - assert manager._draft_subpage_view_obj is None + MiniMaxM3DraftKVCacheView(manager, [60, 61]) + assert manager._draft_kv_cache_view is None @pytest.mark.cpu_only From e54bd2610829d51b227ed3a552aa67d12c548844 Mon Sep 17 00:00:00 2001 From: Zheyu Fu Date: Mon, 5 Oct 2026 15:44:41 -0700 Subject: [PATCH 2/3] [None][fix] MiniMax-M3 draft KV view: check the copy index; CPU-only view tests - MiniMaxM3DraftKVCacheView.copy_batch_block_offsets now asserts beam_width == 1 and one copy_idx entry per sequence, as KVCacheManagerV2.copy_batch_block_offsets does. The view writes one draft block-table row per copy_idx entry, so a mismatch would otherwise leave stale rows without an error. - Mark test_draft_kv_cache_view_accessor_respects_local_layers and test_draft_kv_cache_view_rejects_multiple_local_layers cpu_only: both build the manager with __new__ and either mock the view or fail its single-layer check before any tensor is created. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Zheyu Fu --- .../attention/backends/sparse/minimax_m3/cache_manager.py | 5 ++++- .../sparse/msa/test_minimax_m3_shared_draft_layers.py | 2 ++ 2 files changed, 6 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 6cccf325173c..09aa58c3a723 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 @@ -1033,11 +1033,14 @@ def copy_batch_block_offsets( Not delegated to :meth:`KVCacheManagerV2.copy_batch_block_offsets`, which would follow the target's per-layer page tables through ``__getattr__``. """ + assert beam_width == 1, "beam_width must be 1 for KVCacheManagerV2" manager = self._manager + copy_idx = manager.index_mapper.get_copy_index(request_ids, num_contexts, beam_width) + assert copy_idx.shape[0] == num_seqs copy_batch_block_offsets_to_device( self.host_kv_cache_block_offsets, dst_tensor, - manager.index_mapper.get_copy_index(request_ids, num_contexts, beam_width), + copy_idx, self.index_scales, self.kv_offset, manager._stream.cuda_stream, diff --git a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_shared_draft_layers.py b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_shared_draft_layers.py index d6a9019571e6..3364da7089a5 100644 --- a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_shared_draft_layers.py +++ b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_shared_draft_layers.py @@ -341,6 +341,7 @@ def test_shared_draft_layer_count_follows_the_base_manager(): assert shared_draft_layer_count(None, None) == 0 +@pytest.mark.cpu_only @pytest.mark.parametrize("local_draft_layers", [[], [60], [61]]) @pytest.mark.parametrize("swa_scratch_reuse", [False, True]) def test_draft_kv_cache_view_accessor_respects_local_layers( @@ -374,6 +375,7 @@ def test_draft_kv_cache_view_accessor_respects_local_layers( assert manager._shared_draft_layer_ids == [60, 61] +@pytest.mark.cpu_only @pytest.mark.parametrize("via_accessor", [False, True]) def test_draft_kv_cache_view_rejects_multiple_local_layers(via_accessor: bool) -> None: """The single-pool view must reject a second layer before publishing wrong pointers.""" From ca7ef4ce0b4ab42c0bef92f70f31e2d5d963d5cb Mon Sep 17 00:00:00 2001 From: Zheyu Fu Date: Mon, 5 Oct 2026 16:04:14 -0700 Subject: [PATCH 3/3] [None][chore] Clarify why MiniMax-M3 opts in to trtllm-gen P128 per manager "Not for every shape" read as if the draft shape itself might lack P128 kernels. It means other shapes in the trtllm-gen artifact are P32-only, which is why P128 is opted in per manager rather than globally. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Zheyu Fu --- .../backends/sparse/minimax_m3/cache_manager.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 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 09aa58c3a723..841bb0366237 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 @@ -525,9 +525,9 @@ def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None: # Hybrid draft layers use MiniMaxM3DraftKVCacheView, which also # carries their P128 trtllm-gen opt-in. return - # Draft layers run at the target's 128-token pages. trtllm-gen has P128 - # kernels for their dense-GQA shapes but not for every shape, so opt in - # here rather than in the global allowlist. + # Draft layers run at the target's 128-token pages. Opt in to P128 here, not + # globally: trtllm-gen ships P128 kernels for these dense-GQA draft shapes, but + # many other shapes in its artifact are P32-only. if self.tokens_per_block == 128: self.trtllm_gen_extra_tokens_per_block = frozenset({128}) if self._use_per_layer_page_tables: @@ -960,8 +960,8 @@ class MiniMaxM3DraftKVCacheView: dtype = DataType.FP8 num_pools = num_attention_op_pools = 1 - # trtllm-gen has P128 kernels for the draft layer's dense-GQA shape but not - # for every shape, so only this view opts in. + # P128 is opted in here, not globally: trtllm-gen ships P128 kernels for this + # dense-GQA draft shape, but many other shapes in its artifact are P32-only. trtllm_gen_extra_tokens_per_block = frozenset({128}) def __init__(self, manager: MiniMaxM3KVCacheManagerV2, draft_layer_ids: Sequence[int]) -> None: