[None][feat] Add the ported MiniMax-M3 decode kernels ahead of their wiring

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

@@ -13,6 +13,7 @@
 
 from __future__ import annotations
 
+import functools
 from typing import TYPE_CHECKING, Optional
 
 import torch
@@ -28,6 +29,76 @@
     from .common import MiniMaxM3SparseConfig
 
 
+@functools.lru_cache(maxsize=1)
+def cutedsl_score_runner():
+    """Return the CuTe DSL indexer scoring runner, or None if unavailable.
+
+    The CuTe DSL ops are registered only when the nvidia-cutlass-dsl package is
+    importable, so this stays a soft dependency.
+
+    Resolved once for the process: package availability cannot change under a
+    running model, and every sparse layer of every step scores through here.
+    """
+    try:
+        from tensorrt_llm._torch.custom_ops import cute_dsl_custom_ops
+    except ImportError:
+        return None
+    return getattr(cute_dsl_custom_ops, "CuteDSLMiniMaxM3IndexDecodeScoreRunner", None)
+
+
+def _cutedsl_score(
+    idx_q: torch.Tensor,
+    idx_k_paged: torch.Tensor,
+    max_score: torch.Tensor,
+    *,
+    block_table: torch.Tensor,
+    seq_lens_cuda: torch.Tensor,
+    decode_query_len: int,
+) -> bool:
+    """Try to fill `max_score` with the CuTe DSL scorer; report whether it ran.
+
+    `max_score` is the [num_index_heads, max_k_tiles, total_q] buffer the block
+    selector consumes. The kernel writes [head, token, block], so it is handed
+    the transposed view: same backing store, no copy, and the stores end up
+    coalesced across tokens rather than strided by max_k_tiles.
+
+    The buffer is deliberately not pre-filled with -inf. The kernel writes
+    blocks [0, ceil(seq_len / page_size)) for every token of a request, and the
+    selector reads only [0, n_valid_blocks[token])), which is bounded by that
+    same count for every token including the shorter ones in a multi-token
+    speculative step. So every entry the selector reads has just been written.
+    """
+    runner = cutedsl_score_runner()
+    if runner is None:
+        return False
+
+    total_q, num_index_heads, head_dim = idx_q.shape
+    page_size = int(idx_k_paged.shape[2])
+    if not runner.is_supported(
+        q_dtype=idx_q.dtype,
+        num_heads=num_index_heads,
+        head_dim=head_dim,
+        page_size=page_size,
+        max_decode_query_len=decode_query_len,
+    ):
+        return False
+    if idx_k_paged.dtype != idx_q.dtype or max_score.shape[2] != total_q:
+        return False
+
+    # The kernel wants MQA index-K as [num_pages, page_size, head_dim]; the
+    # squeeze is zero-copy and keeps the pool's real per-page stride, which the
+    # TMA descriptor reads at runtime.
+    torch.ops.trtllm.cute_dsl_minimax_m3_index_decode_score(
+        idx_q,
+        idx_k_paged.squeeze(1),
+        block_table,
+        seq_lens_cuda,
+        max_score.transpose(1, 2),
+        decode_query_len,
+    )
+    return True
+
+
 def _proxy_max_score(
     idx_q: torch.Tensor,
     idx_k_paged: torch.Tensor,