[None][perf] MiniMax-M3: native P128 draft KV view for the NVFP4 hybrid cache

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",
 ]

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
 

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

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

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

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