[None][perf] Optimize MiniMax-M3 MSA block selection

cpp/tensorrt_llm/kernels/minimaxM3SelectBlocks.h

@@ -0,0 +1,36 @@
+/*
+ * 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.
+ */
+
+#pragma once
+
+#include "tensorrt_llm/common/config.h"
+
+#include <cstdint>
+#include <cuda_runtime.h>
+
+TRTLLM_NAMESPACE_BEGIN
+
+namespace kernels
+{
+
+void invokeMinimaxM3SelectBlocks(float const* scores, int64_t headStride, int64_t blockStride, int64_t queryStride,
+    int32_t const* nValidBlocks, int32_t* output, int32_t numKvHeads, int32_t numBlocks, int32_t totalQueries,
+    int32_t initBlocks, int32_t localBlocks, bool headMajorOutput, cudaStream_t stream);
+
+} // namespace kernels
+
+TRTLLM_NAMESPACE_END

cpp/tensorrt_llm/thop/CMakeLists.txt

@@ -116,6 +116,7 @@ add_library(
   IndexerKCacheScatterOp.cpp
   IndexerTopKOp.cpp
   sparseKvCacheCompactOp.cpp
+  minimaxM3SelectBlocksOp.cpp
   mlaRopeInplaceOp.cpp
   ncclCommunicatorOp.cpp
   allocateOutput.cpp

cpp/tensorrt_llm/thop/minimaxM3SelectBlocksOp.cpp

@@ -0,0 +1,103 @@
+/*
+ * 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.
+ */
+
+#include "tensorrt_llm/kernels/minimaxM3SelectBlocks.h"
+
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <limits>
+#include <torch/extension.h>
+
+TRTLLM_NAMESPACE_BEGIN
+
+namespace torch_ext
+{
+
+//! Select the highest-scoring MiniMax-M3 MSA blocks for each query and KV head.
+//!
+//! @param scores Float32 CUDA tensor shaped [num_kv_heads, num_blocks, total_queries]. Strided tensors are supported.
+//! @param nValidBlocks Contiguous int32 CUDA tensor shaped [total_queries], on the same device as scores.
+//! @param topK Number of blocks to select. The operator currently requires topK == 16.
+//! @param initBlocks Number of initial blocks forced into the selection.
+//! @param localBlocks Number of blocks immediately preceding each query's valid-block boundary forced into selection.
+//! @param headMajorOutput Whether to use head-major backing storage while preserving the logical return shape.
+//! @return Int32 CUDA tensor shaped [total_queries, num_kv_heads, 16]. Valid block IDs are ascending, followed by -1
+//! padding when fewer than 16 valid blocks exist.
+torch::Tensor minimaxM3SelectBlocks(torch::Tensor const& scores, torch::Tensor const& nValidBlocks, int64_t topK,
+    int64_t initBlocks, int64_t localBlocks, bool headMajorOutput)
+{
+    constexpr int64_t kRequiredTopK = 16;
+    constexpr int64_t kMaxBlockIndex = 65'535;
+
+    TORCH_CHECK(scores.is_cuda(), "minimax_m3_select_blocks expects CUDA scores");
+    TORCH_CHECK(scores.scalar_type() == torch::kFloat32, "minimax_m3_select_blocks expects float32 scores, got ",
+        scores.scalar_type());
+    TORCH_CHECK(scores.dim() == 3, "scores must be [num_kv_heads, num_blocks, total_queries]");
+    TORCH_CHECK(scores.stride(0) >= 0 && scores.stride(1) >= 0 && scores.stride(2) >= 0,
+        "scores must have non-negative strides");
+
+    TORCH_CHECK(nValidBlocks.is_cuda(), "minimax_m3_select_blocks expects CUDA n_valid_blocks");
+    TORCH_CHECK(nValidBlocks.device() == scores.device(), "scores and n_valid_blocks must be on the same device");
+    TORCH_CHECK(nValidBlocks.scalar_type() == torch::kInt32,
+        "minimax_m3_select_blocks expects int32 n_valid_blocks, got ", nValidBlocks.scalar_type());
+    TORCH_CHECK(
+        nValidBlocks.dim() == 1 && nValidBlocks.size(0) == scores.size(2), "n_valid_blocks must be [total_queries]");
+    TORCH_CHECK(nValidBlocks.is_contiguous(), "n_valid_blocks must be contiguous");
+
+    TORCH_CHECK(topK == kRequiredTopK, "minimax_m3_select_blocks supports topk=16, got ", topK);
+    TORCH_CHECK(initBlocks >= 0, "init_blocks must be non-negative");
+    TORCH_CHECK(localBlocks >= 0, "local_blocks must be non-negative");
+    TORCH_CHECK(scores.size(1) <= kMaxBlockIndex, "minimax_m3_select_blocks supports at most ", kMaxBlockIndex,
+        " blocks, got ", scores.size(1));
+    TORCH_CHECK(scores.size(0) <= std::numeric_limits<int32_t>::max()
+            && scores.size(1) <= std::numeric_limits<int32_t>::max()
+            && scores.size(2) <= std::numeric_limits<int32_t>::max(),
+        "minimax_m3_select_blocks dimensions exceed int32 range");
+    TORCH_CHECK(scores.size(0) * scores.size(2) <= std::numeric_limits<int32_t>::max(),
+        "minimax_m3_select_blocks output rows exceed int32 range");
+    TORCH_CHECK(initBlocks <= std::numeric_limits<int32_t>::max() && localBlocks <= std::numeric_limits<int32_t>::max(),
+        "minimax_m3_select_blocks forcing ranges exceed int32 range");
+
+    c10::cuda::CUDAGuard const deviceGuard{scores.device()};
+    auto const outputOptions = scores.options().dtype(torch::kInt32);
+    auto const output = headMajorOutput
+        ? torch::empty({scores.size(0), scores.size(2), topK}, outputOptions).permute({1, 0, 2})
+        : torch::empty({scores.size(2), scores.size(0), topK}, outputOptions);
+    auto const stream = at::cuda::getCurrentCUDAStream(scores.get_device());
+    tensorrt_llm::kernels::invokeMinimaxM3SelectBlocks(scores.data_ptr<float>(), scores.stride(0), scores.stride(1),
+        scores.stride(2), nValidBlocks.data_ptr<int32_t>(), output.data_ptr<int32_t>(),
+        static_cast<int32_t>(scores.size(0)), static_cast<int32_t>(scores.size(1)),
+        static_cast<int32_t>(scores.size(2)), static_cast<int32_t>(initBlocks), static_cast<int32_t>(localBlocks),
+        headMajorOutput, stream);
+    return output;
+}
+
+} // namespace torch_ext
+
+TRTLLM_NAMESPACE_END
+
+TORCH_LIBRARY_FRAGMENT(trtllm, m)
+{
+    m.def(
+        "minimax_m3_select_blocks(Tensor scores, Tensor n_valid_blocks, int topk, int init_blocks, int "
+        "local_blocks, bool head_major_output=False) -> Tensor");
+}
+
+TORCH_LIBRARY_IMPL(trtllm, CUDA, m)
+{
+    m.impl("minimax_m3_select_blocks", &tensorrt_llm::torch_ext::minimaxM3SelectBlocks);
+}

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

@@ -824,6 +824,9 @@ def run_indexer(
         config = self.m3_config
         idx_sm_scale = idx_sm_scale if idx_sm_scale is not None else config.sparse_index_dim**-0.5
         num_tokens = int(idx_q.shape[0])
+        head_major_output = (
+            int(metadata.num_contexts or 0) > 0 and int(metadata.num_generations or 0) == 0
+        )
         idx_q_view = idx_q.view(num_tokens, config.num_index_heads, config.sparse_index_dim)
         idx_k_view = idx_k.view(num_tokens, 1, config.sparse_index_dim)
 
@@ -861,6 +864,7 @@ def run_indexer(
             proxy_plan=proxy_plan,
             max_score=max_score,
             n_valid_blocks=n_valid_blocks,
+            head_major_output=head_major_output,
         )
 
     def sparse_attn_predict(

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

@@ -131,6 +131,7 @@ def select_blocks(
         proxy_plan: Optional[tuple] = None,
         max_score: Optional[torch.Tensor] = None,
         n_valid_blocks: Optional[torch.Tensor] = None,
+        head_major_output: bool = False,
     ) -> torch.Tensor:
         """Return [total_q, num_kv_heads, topk] selected block indices.
 
@@ -183,6 +184,7 @@ def select_blocks(
             n_valid_blocks=n_valid_blocks,
             init_blocks=config.init_blocks,
             local_blocks=config.local_blocks,
+            head_major_output=head_major_output,
         )
 
 

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

@@ -11,7 +11,7 @@
 
 import torch
 
-from .common import _INIT_SCORE, _LOCAL_SCORE, write_kv_slots
+from .common import write_kv_slots
 
 # fmha_sm100 ships only head_dim 128 variants and the MiniMax-M3 checkpoint
 # selects topk 16. Callers enforce these early so a misconfiguration fails
@@ -187,42 +187,26 @@ def select_blocks_from_maxscore(
     n_valid_blocks: torch.Tensor,
     init_blocks: int,
     local_blocks: int,
+    head_major_output: bool = False,
 ) -> torch.Tensor:
     """Select per-query top-k blocks from per-KV-head block scores.
 
     Applies init and local forced blocks and per-query valid-block masking
     on the amax-reduced scores [num_kv_heads, n_blocks, total_q]. Returns
     [total_q, num_kv_heads, topk] int32 ascending block ids with -1 tail
-    padding.
+    padding. When ``head_major_output`` is set, the logical result uses a
+    head-major backing so ``result.permute(1, 0, 2)`` is contiguous without a
+    copy.
     """
-    num_kv_heads, n_blocks, total_q = max_score_kv.shape
-    device = max_score_kv.device
-    scores = max_score_kv.permute(2, 0, 1).to(torch.float32).clone()
-    block_ids = torch.arange(n_blocks, device=device, dtype=torch.long)
-    nvb = n_valid_blocks.to(device=device, dtype=torch.long)
-
-    if init_blocks > 0:
-        init_mask = block_ids.view(1, 1, -1) < init_blocks
-        scores = torch.where(init_mask, torch.full_like(scores, _INIT_SCORE), scores)
-    if local_blocks > 0:
-        local_start = (nvb - local_blocks).clamp_min(0)
-        local_mask = (block_ids.view(1, -1) >= local_start.view(-1, 1)) & (
-            block_ids.view(1, -1) < nvb.view(-1, 1)
-        )
-        scores = torch.where(local_mask.unsqueeze(1), torch.full_like(scores, _LOCAL_SCORE), scores)
-    block_valid = block_ids.view(1, -1) < nvb.view(-1, 1)
-    scores = scores.masked_fill(~block_valid.unsqueeze(1), float("-inf"))
-
-    k = min(topk, n_blocks)
-    vals, idx = scores.topk(k=k, dim=-1)
-    idx = torch.where(vals != float("-inf"), idx, torch.full_like(idx, -1))
-    sort_key = torch.where(idx < 0, torch.full_like(idx, n_blocks), idx)
-    sort_key, _ = torch.sort(sort_key, dim=-1)
-    idx = torch.where(sort_key >= n_blocks, torch.full_like(sort_key, -1), sort_key)
-    if k < topk:
-        pad = torch.full((total_q, num_kv_heads, topk - k), -1, dtype=idx.dtype, device=device)
-        idx = torch.cat([idx, pad], dim=-1)
-    return idx.to(torch.int32)
+    nvb = n_valid_blocks.to(device=max_score_kv.device, dtype=torch.int32).contiguous()
+    return torch.ops.trtllm.minimax_m3_select_blocks(
+        max_score_kv,
+        nvb,
+        topk,
+        init_blocks,
+        local_blocks,
+        head_major_output,
+    )
 
 
 __all__ = [

tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py

@@ -276,6 +276,22 @@ def _(logits,
         # In-place operation, no return value (void function)
         pass
 
+    @torch.library.register_fake("trtllm::minimax_m3_select_blocks")
+    def _(
+        scores,
+        n_valid_blocks,
+        topk,
+        init_blocks,
+        local_blocks,
+        head_major_output=False,
+    ):
+        del n_valid_blocks, init_blocks, local_blocks
+        if head_major_output:
+            return scores.new_empty((scores.shape[0], scores.shape[2], topk),
+                                    dtype=torch.int32).permute(1, 0, 2)
+        return scores.new_empty((scores.shape[2], scores.shape[0], topk),
+                                dtype=torch.int32)
+
     @torch.library.register_fake("trtllm::kda_decode")
     def _(x_q, x_k, x_v, w_q_t, w_k_t, w_v_t, bias_q, bias_k, bias_v,
           conv_state_q, conv_state_k, conv_state_v, a_log, g, dt_bias, beta,

tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py

@@ -7,6 +7,8 @@
 by the SM100 integration accuracy test.
 """
 
+from types import SimpleNamespace
+
 import pytest
 import torch
 
@@ -165,6 +167,66 @@ def fake_select_blocks_from_maxscore(*args, **kwargs):
     assert result is expected
 
 
+@pytest.mark.parametrize(
+    ("num_contexts", "num_generations", "expected_head_major"),
+    [(2, 0, True), (1, 1, False), (0, 2, False)],
+)
+def test_run_indexer_routes_head_major_output_by_batch_mode(
+    num_contexts, num_generations, expected_head_major
+):
+    num_tokens, num_index_heads, sparse_index_dim = 3, 4, 128
+    captured = {}
+
+    class FakeIndexer:
+        def select_blocks(self, *args, **kwargs):
+            del args
+            captured["head_major_output"] = kwargs["head_major_output"]
+            return torch.zeros(num_tokens, 1, 16, dtype=torch.int32)
+
+    class FakeMetadata:
+        msa_decode_proxy_plan = None
+        msa_eager_proxy_plan = ("eager",)
+        msa_eager_n_valid_blocks = torch.ones(num_tokens, dtype=torch.int32)
+        msa_kv_indices = torch.arange(num_tokens, dtype=torch.int32)
+        msa_qo_lens_cpu = torch.tensor([num_tokens], dtype=torch.int32)
+        msa_kv_lens_cpu = torch.tensor([num_tokens], dtype=torch.int32)
+        msa_qo_offset_cpu = torch.tensor([0], dtype=torch.int32)
+
+        def __init__(self):
+            self.num_contexts = num_contexts
+            self.num_generations = num_generations
+            self.idx_k_cache = None
+
+        def msa_write_idx_k(self, layer_idx, idx_k):
+            del layer_idx
+            self.idx_k_cache = idx_k
+
+        def msa_idx_k_cache(self, layer_idx):
+            del layer_idx
+            return self.idx_k_cache
+
+    attention = SimpleNamespace(
+        layer_idx=0,
+        m3_config=SimpleNamespace(
+            sparse_index_dim=sparse_index_dim,
+            num_index_heads=num_index_heads,
+            num_kv_heads=1,
+        ),
+        indexer=FakeIndexer(),
+    )
+    metadata = FakeMetadata()
+
+    result = MiniMaxM3MsaSparseAttention.run_indexer(
+        attention,
+        torch.zeros(num_tokens, num_index_heads * sparse_index_dim),
+        torch.zeros(num_tokens, sparse_index_dim),
+        metadata,
+    )
+
+    assert result.shape == (num_tokens, 1, 16)
+    assert captured["head_major_output"] is expected_head_major
+
+
 def test_msa_proxy_max_score_strided_index_k_matches_packed():
     if not torch.cuda.is_available():
         pytest.skip("CUDA required")