[None][perf] Fuse MiniMax-M3 QKV and index projection

cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.h

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 2025-2026, NVIDIA CORPORATION.  All rights reserved.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -17,6 +17,8 @@
 #pragma once
 
 #include "tensorrt_llm/common/config.h"
+
+#include <cstdint>
 #include <cuda_runtime.h>
 
 TRTLLM_NAMESPACE_BEGIN
@@ -61,6 +63,25 @@ void launchFusedQKNormRopeToFp8(void const* qkv_in, // BF16 input [num_tokens, t
     bool const interleave, int const* position_ids, float factor, float low, float high, float attention_factor,
     cudaStream_t stream, bool is_qk_norm, bool use_gemma, bool use_mrope, int mrope_section1, int mrope_section2);
 
+// MiniMax-M3-specific main-branch producer. It returns contiguous
+// FP8 Q and inserts normalized/RoPE'd FP8 K plus copy-cast FP8 V directly into
+// a paged HND pool [num_pages, 2, num_heads, page_size, head_dim].
+void launchMinimaxM3Fp8QKNormRopeKVInsert(void const* qkv_input, void* q_output, void* kv_cache,
+    int const* out_cache_loc, int64_t page_stride, int64_t plane_stride, int64_t head_stride, int64_t token_stride,
+    int64_t num_pages, int page_size, int num_tokens, int num_heads_q, int num_heads_k, int num_heads_v, int head_dim,
+    int rotary_dim, float eps, void const* q_weight, void const* k_weight, float base, int const* position_ids,
+    cudaStream_t stream);
+
+// MiniMax-M3 sparse producer for the packed [Q|K|V|index-Q|index-K]
+// projection. It uses a precomputed FP32 RoPE table, emits compact FP8 Q and
+// index-Q, and inserts main K/V plus index-K into their paged FP8 HND caches.
+void launchMinimaxM3Fp8QKVIndexerNormRopeKVInsert(void const* packed_input, void* q_output, void* index_q_output,
+    void* kv_cache, void* index_k_cache, int const* out_cache_loc, int64_t kv_page_stride, int64_t kv_plane_stride,
+    int64_t kv_head_stride, int64_t kv_token_stride, int64_t index_page_stride, int64_t index_token_stride,
+    int64_t num_pages, int page_size, int num_tokens, int num_heads_q, int num_heads_kv, int num_heads_index,
+    int head_dim, int rotary_dim, float eps, void const* q_weight, void const* k_weight, void const* index_q_weight,
+    void const* index_k_weight, float const* rotary_cos_sin, int const* position_ids, cudaStream_t stream);
+
 } // namespace kernels
 
 TRTLLM_NAMESPACE_END

docs/source/models/supported-models.md

@@ -84,7 +84,7 @@ Note: Support for other models may vary. Features marked "N/A" are not applicabl
 | `Gemma4ForConditionalGeneration` | Untested          | Yes        | Untested                   | No                    | Yes             | Yes | No               | No                | No     | Yes           | Untested         | No             | Yes                      | Untested              | Untested        |
 | `Gemma4UnifiedForConditionalGeneration` | Untested          | Untested   | Untested                   | No                    | Yes             | No  | No               | No                | No     | Yes           | Untested         | No             | Yes                      | Untested              | Untested        |
 | `Step3p7ForConditionalGeneration`| Yes               | Yes        | Yes                        | Untested              | Untested        | Yes | No               | No                | No     | Yes           | Untested         | Untested       | Yes                      | Untested              | Untested        |
-| `MiniMaxM3SparseForConditionalGeneration` [^12] | Yes               | Yes        | Yes                        | Untested              | Untested        | No  | No               | No                | No     | Yes           | Untested         | No             | N/A                      | Untested              | Untested        |
+| `MiniMaxM3SparseForConditionalGeneration` [^12] | Yes               | Yes        | Yes                        | Untested              | Untested        | No  | Yes              | No                | No     | Yes           | Untested         | No             | N/A                      | Untested              | Untested        |
 
 [^1]: Chunked Prefill for MLA can only be enabled on SM90/SM100/SM103/SM120.
 [^2]: KV cache reuse for MLA can only be enabled on SM90/SM100/SM103/SM120/SM121 and in BF16/FP8 KV cache dtype.
@@ -96,7 +96,7 @@ Note: Support for other models may vary. Features marked "N/A" are not applicabl
 [^9]: Audio modality only supported on E2B/E4B variants.
 [^10]: Audio requires a checkpoint with a `sound_config` and is supported only on the full (non-disaggregated) model path, not the EPD disaggregated path.
 [^11]: DeepSeek-V4 is only supported on Blackwell GPUs (`SM100+`). See the [DeepSeek-V4 example README](../../../examples/models/core/deepseek_v4/README.md) for setup and parallelism.
-[^12]: Supports text, image, and video inputs over the block-sparse attention path. The published MXFP8 checkpoint is dequantized on load so the runtime sees an effectively BF16 model. The text decoder is also usable standalone (text-only) via the `MiniMaxM3SparseForCausalLM` architecture. KV cache reuse and MTP are not supported on the sparse-attention path in this release.
+[^12]: Supports text, image, and video inputs over the block-sparse attention path. The published MXFP8 checkpoint is dequantized on load so the runtime sees an effectively BF16 model. The text decoder is also usable standalone (text-only) via the `MiniMaxM3SparseForCausalLM` architecture. KV cache reuse and MTP are not supported on the sparse-attention path in this release. One-model linear EAGLE-3 is supported; combining it with CUDA graphs requires the MSA implementation on SM100.
 [^13]: The Cosmos 3 family also supports visual generation through the VisualGen API. See [Visual Generation Models](#visual-generation-models).
 [^14]: Requires `transformers>=5.7.0`: MiniCPM-V 4.6 was upstreamed into transformers as a native model type (`minicpmv4_6`) and the checkpoint ships no remote code (`auto_map`) to fall back on. The Qwen3.5-hybrid text tower runs in BF16. Image, video, and text inputs are supported in this release (video reuses the same NaViT-packed vision path as image via `MiniCPMV4_6InputProcessor`).
 

tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/cache_manager.py

@@ -155,6 +155,15 @@ class MiniMaxM3KVCacheManagerV2(KVCacheManagerV2):
       * ``sparse_index_dim`` — width of the index-K/V vectors.
     """
 
+    # INDEX_KEY is coalesced into the target's V2 pool. Dense Eagle3 draft
+    # layers cannot consume the synthetic AttentionOp view of that layout, so
+    # they require their own ordinary KV manager even under attention DP.
+    supports_shared_draft_layers = False
+
+    # MSA requires 128-token target pages. The dense Eagle3 generation kernel
+    # uses its validated 32-token page geometry in the separate draft manager.
+    draft_manager_tokens_per_block = 32
+
     def __init__(
         self,
         *args,

tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/common.py

@@ -43,6 +43,7 @@ class MiniMaxM3SparseParams(SparseParams):
     disable_index_value: bool = True
     implementation: Literal["triton", "msa"] = "triton"
     indexer_kv_dtype: Literal["bf16", "fp8"] = "bf16"
+    fuse_qkv_index_projection: bool = False
 
     @property
     def indices_block_size(self) -> int:
@@ -62,6 +63,7 @@ class MiniMaxM3SparseMetadataParams(SparseMetadataParams):
     global_num_kv_heads: int = 0
     num_index_heads: int = 4
     topk: int = 16
+    fuse_qkv_index_projection: bool = False
 
     def sharded_head_counts(self, mapping: Optional["Mapping"] = None) -> Tuple[int, int]:
         """Return per-rank (num_q_heads, num_kv_heads) for mapping.
@@ -79,6 +81,13 @@ def _shard(num_heads: int) -> int:
 
         return _shard(self.global_num_q_heads), _shard(self.global_num_kv_heads)
 
+    def sharded_index_head_count(self, mapping: Optional["Mapping"] = None) -> int:
+        """Return the index-head count used by this rank's proxy attention."""
+        if not self.fuse_qkv_index_projection:
+            return int(self.num_index_heads)
+        _, num_kv_heads = self.sharded_head_counts(mapping)
+        return num_kv_heads
+
 
 @dataclass(frozen=True)
 class MiniMaxM3SparseConfig:
@@ -149,7 +158,11 @@ def from_sparse_params(
             num_q_heads=int(num_q_heads),
             num_kv_heads=int(num_kv_heads),
             head_dim=int(head_dim),
-            num_index_heads=int(sparse_params.num_index_heads),
+            num_index_heads=(
+                int(num_kv_heads)
+                if sparse_params.fuse_qkv_index_projection
+                else int(sparse_params.num_index_heads)
+            ),
             sparse_index_dim=int(sparse_params.sparse_index_dim),
             block_size=int(sparse_params.block_size),
             topk=int(sparse_params.topk),