[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

tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/common.py

@@ -32,12 +32,39 @@
 _LOCAL_SCORE = 1e29
 
 
+def index_head_range(
+    num_index_heads: int, num_kv_heads: int, mapping: Optional["Mapping"] = None
+) -> Tuple[int, int]:
+    """Return this rank's global index-head interval, preserving KV groups.
+
+    Whole KV groups (including all their index heads) replicate when TP
+    exceeds the KV-head count. Attention-DP retains every head on each rank.
+    """
+    tp_size = 1 if mapping is None or mapping.enable_attention_dp else mapping.tp_size
+    if num_index_heads <= 0:
+        raise ValueError("MiniMax-M3 requires positive index heads")
+    # Metadata without model geometry is usable only for an unsharded view.
+    if num_kv_heads == 0 and tp_size == 1:
+        return 0, num_index_heads
+    if num_kv_heads <= 0 or num_index_heads % num_kv_heads != 0:
+        raise ValueError("MiniMax-M3 index heads must be divisible by global KV heads")
+    shard_count = min(tp_size, num_kv_heads)
+    if tp_size % shard_count != 0 or num_kv_heads % shard_count != 0:
+        raise ValueError("MiniMax-M3 TP and KV heads must divide one another")
+    rank = 0 if tp_size == 1 else mapping.tp_rank // (tp_size // shard_count)
+    count = num_index_heads // shard_count
+    return rank * count, (rank + 1) * count
+
+
 @dataclass(frozen=True)
 class MiniMaxM3SparseParams(SparseParams):
     """Lowered runtime parameters for the MiniMax-M3 sparse backend."""
 
     algorithm: Literal["minimax_m3"] = field(init=False, default="minimax_m3")
     num_index_heads: int = 4
+    # None keeps explicit rank-local backend construction supported. Model
+    # lowering supplies the global count so both backends preserve KV groups.
+    global_num_kv_heads: Optional[int] = None
     sparse_index_dim: int = 128
     block_size: int = 128
     topk: int = 16
@@ -47,6 +74,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:
@@ -83,6 +111,10 @@ 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:
+        start, end = index_head_range(self.num_index_heads, self.global_num_kv_heads, mapping)
+        return end - start
+
 
 @dataclass(frozen=True)
 class MiniMaxM3SparseConfig:
@@ -149,11 +181,17 @@ def from_sparse_params(
         """Build a kernel param bundle from lowered ``MiniMaxM3SparseParams``
         and the per-rank model geometry.
         """
+        num_index_heads = int(sparse_params.num_index_heads)
+        if sparse_params.global_num_kv_heads is not None:
+            global_kv_heads = int(sparse_params.global_num_kv_heads)
+            if global_kv_heads <= 0 or num_index_heads % global_kv_heads != 0:
+                raise ValueError("MiniMax-M3 index heads must be divisible by global KV heads")
+            num_index_heads = num_index_heads // global_kv_heads * int(num_kv_heads)
         return cls(
             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=num_index_heads,
             sparse_index_dim=int(sparse_params.sparse_index_dim),
             block_size=int(sparse_params.block_size),
             topk=int(sparse_params.topk),

tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_backend.py

@@ -368,14 +368,15 @@ def _validate_decode_kernel_support(self) -> None:
         # num_index_heads * query tokens into one Q block, which bounds the
         # draft length it can verify.
         decode_query_len = self._msa_max_decode_query_len()
+        num_index_heads = params.sharded_index_head_count(self.mapping)
         if not self._cutedsl_indexer_supported(
-            num_index_heads=params.num_index_heads,
+            num_index_heads=num_index_heads,
             page_size=page_size,
             decode_query_len=decode_query_len,
         ):
             raise RuntimeError(
                 "The MiniMax-M3 CuTe DSL indexer scorer does not support this "
-                f"configuration: {params.num_index_heads} index heads, page size "
+                f"configuration: {num_index_heads} index heads, page size "
                 f"{page_size}, index dtype {self._msa_index_kv_dtype()}, up to "
                 f"{decode_query_len} query tokens per generation request."
             )
@@ -475,13 +476,13 @@ def _create_msa_buffers(self) -> None:
             fmha_sm100 = require_msa_module()
             max_k_tiles = _worst_case_proxy_max_k_tiles(
                 fmha_sm100,
-                num_index_heads=params.num_index_heads,
+                num_index_heads=params.sharded_index_head_count(self.mapping),
                 kv_cache_manager=kv_cache_manager,
                 max_batch=max_num_sequences,
             )
             self._msa_worst_case_max_k_tiles = int(max_k_tiles)
             self._alloc_msa_proxy_scratch(
-                num_index_heads=params.num_index_heads,
+                num_index_heads=params.sharded_index_head_count(self.mapping),
                 max_tokens=self._msa_max_decode_tokens(),
                 max_k_tiles=max_k_tiles,
                 capture_graph=capture_graph,
@@ -863,7 +864,7 @@ def _build_step_plans(self) -> None:
         params = self._msa_params
         if params is None:
             return
-        num_index_heads = params.num_index_heads
+        num_index_heads = params.sharded_index_head_count(self.mapping)
         qo_lens_cpu = self.msa_qo_lens_cpu
         kv_lens_cpu = self.msa_kv_lens_cpu
         qo_offset_cpu = self.msa_qo_offset_cpu