[None][perf] Use FP8 MiniMax-M3 MSA indexer QK

cpp/tensorrt_llm/kernels/minimaxM3Fp8IndexerKernel.cu

@@ -0,0 +1,181 @@
+/*
+ * Copyright (c) 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.
+ * 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 "minimaxM3Fp8IndexerKernel.h"
+
+#include "tensorrt_llm/common/cudaUtils.h"
+#include "tensorrt_llm/common/mathUtils.h"
+#include "tensorrt_llm/common/reduceKernelUtils.cuh"
+
+#include <cuda_bf16.h>
+#include <cuda_fp8.h>
+#include <cuda_runtime.h>
+
+#include <cstdint>
+
+TRTLLM_NAMESPACE_BEGIN
+
+namespace kernels
+{
+
+namespace
+{
+
+constexpr int kHeadDim = 128;
+constexpr int kRotaryDim = 64;
+constexpr int kElemsPerThread = kHeadDim / 32;
+
+// Match the established vLLM contract exactly: the normalized/RoPE result is
+// first materialized as BF16 and then cast, without an external FP8 scale.
+__device__ __forceinline__ __nv_fp8_e4m3 bf16RoundedToFp8(float value)
+{
+    return __nv_fp8_e4m3(__bfloat162float(__float2bfloat16_rn(value)));
+}
+
+__global__ void minimaxM3Fp8IndexerQKNormRopeKernel(__nv_bfloat16 const* qk, __nv_fp8_e4m3* q_out,
+    __nv_fp8_e4m3* k_cache, int const* out_cache_loc, int64_t page_stride, int64_t token_stride, int page_size,
+    int64_t num_pages, int num_tokens, int num_heads_q, float eps, __nv_bfloat16 const* q_weight,
+    __nv_bfloat16 const* k_weight, float base, int const* position_ids)
+{
+    int const warps_per_block = blockDim.x / 32;
+    int const warp_id = threadIdx.x / 32;
+    int const lane_id = threadIdx.x % 32;
+    int const global_warp = blockIdx.x * warps_per_block + warp_id;
+    int const total_heads = num_heads_q + 1;
+    int const token_idx = global_warp / total_heads;
+    int const local_head = global_warp % total_heads;
+    if (token_idx >= num_tokens)
+    {
+        return;
+    }
+
+    bool const is_q = local_head < num_heads_q;
+    int64_t const input_offset
+        = (static_cast<int64_t>(token_idx) * total_heads + local_head) * kHeadDim + lane_id * kElemsPerThread;
+
+    uint2 const packed_input = *reinterpret_cast<uint2 const*>(qk + input_offset);
+    float elements[kElemsPerThread];
+    float sum_squares = 0.0F;
+#pragma unroll
+    for (int pair = 0; pair < 2; ++pair)
+    {
+        auto const values = __bfloat1622float2(reinterpret_cast<__nv_bfloat162 const*>(&packed_input)[pair]);
+        elements[pair * 2] = values.x;
+        elements[pair * 2 + 1] = values.y;
+        // Preserve the accumulation order used by fusedQKNormRopeKernel.  A
+        // reassociated pair sum can move a final BF16 value across an FP8 bin.
+        sum_squares += values.x * values.x;
+        sum_squares += values.y * values.y;
+    }
+
+    sum_squares = tensorrt_llm::common::warpReduceSum(sum_squares);
+    float const rms_rcp = rsqrtf(sum_squares / static_cast<float>(kHeadDim) + eps);
+    auto const* weight = is_q ? q_weight : k_weight;
+#pragma unroll
+    for (int i = 0; i < kElemsPerThread; ++i)
+    {
+        int const dim = lane_id * kElemsPerThread + i;
+        elements[i] *= rms_rcp * (1.0F + __bfloat162float(weight[dim]));
+    }
+
+    // MiniMax-M3 uses NeoX partial RoPE: rotate the first 64 of 128 channels.
+    // Four elements per lane means the matching half is eight lanes away.
+    __syncwarp();
+    constexpr int kPairOffset = (kRotaryDim / 2) / kElemsPerThread;
+    // Keep the frequency calculation bitwise aligned with the shared fused
+    // QK-norm/RoPE kernel.  That kernel uses the fast base-2 intrinsics rather
+    // than powf, and the BF16-to-FP8 contract depends on the resulting rounding.
+    float const neg2_log2base_over_rd = -2.0F * __log2f(base) / static_cast<float>(kRotaryDim);
+#pragma unroll
+    for (int i = 0; i < kElemsPerThread; ++i)
+    {
+        int const dim = lane_id * kElemsPerThread + i;
+        float paired = __shfl_xor_sync(0xffffffff, elements[i], kPairOffset);
+        if (dim < kRotaryDim)
+        {
+            if (lane_id < kPairOffset)
+            {
+                paired = -paired;
+            }
+            int const dim_idx = (dim * 2) % kRotaryDim;
+            int const half_dim = dim_idx / 2;
+            float const frequency = exp2f(static_cast<float>(half_dim) * neg2_log2base_over_rd);
+            float sine;
+            float cosine;
+            __sincosf(static_cast<float>(position_ids[token_idx]) * frequency, &sine, &cosine);
+            elements[i] = elements[i] * cosine + paired * sine;
+        }
+    }
+    __syncwarp();
+
+    uint32_t packed_output = 0;
+    auto* fp8_values = reinterpret_cast<__nv_fp8_e4m3*>(&packed_output);
+#pragma unroll
+    for (int i = 0; i < kElemsPerThread; ++i)
+    {
+        fp8_values[i] = bf16RoundedToFp8(elements[i]);
+    }
+
+    __nv_fp8_e4m3* output;
+    if (is_q)
+    {
+        int64_t const output_offset
+            = (static_cast<int64_t>(token_idx) * num_heads_q + local_head) * kHeadDim + lane_id * kElemsPerThread;
+        output = q_out + output_offset;
+    }
+    else
+    {
+        int const slot = out_cache_loc[token_idx];
+        if (slot < 0)
+        {
+            return;
+        }
+        int const page = slot / page_size;
+        if (page >= num_pages)
+        {
+            return;
+        }
+        int const within_page = slot % page_size;
+        output = k_cache + static_cast<int64_t>(page) * page_stride + static_cast<int64_t>(within_page) * token_stride
+            + lane_id * kElemsPerThread;
+    }
+    *reinterpret_cast<uint32_t*>(output) = packed_output;
+}
+
+} // namespace
+
+void launchMinimaxM3Fp8IndexerQKNormRope(void const* qk, void* q_out, void* k_cache, int const* out_cache_loc,
+    int64_t page_stride, int64_t token_stride, int page_size, int64_t num_pages, int num_tokens, int num_heads_q,
+    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)
+{
+    TLLM_CHECK_WITH_INFO(head_dim == kHeadDim, "MiniMax-M3 FP8 indexer requires head_dim=128");
+    TLLM_CHECK_WITH_INFO(rotary_dim == kRotaryDim, "MiniMax-M3 FP8 indexer requires rotary_dim=64");
+    TLLM_CHECK_WITH_INFO(num_heads_q > 0, "MiniMax-M3 FP8 indexer requires at least one query head");
+
+    constexpr int kBlockSize = 256;
+    constexpr int kWarpsPerBlock = kBlockSize / 32;
+    int const total_warps = num_tokens * (num_heads_q + 1);
+    int const grid_size = common::divUp(total_warps, kWarpsPerBlock);
+    minimaxM3Fp8IndexerQKNormRopeKernel<<<grid_size, kBlockSize, 0, stream>>>(static_cast<__nv_bfloat16 const*>(qk),
+        static_cast<__nv_fp8_e4m3*>(q_out), static_cast<__nv_fp8_e4m3*>(k_cache), out_cache_loc, page_stride,
+        token_stride, page_size, num_pages, num_tokens, num_heads_q, eps, static_cast<__nv_bfloat16 const*>(q_weight),
+        static_cast<__nv_bfloat16 const*>(k_weight), base, position_ids);
+}
+
+} // namespace kernels
+
+TRTLLM_NAMESPACE_END

cpp/tensorrt_llm/kernels/minimaxM3Fp8IndexerKernel.h

@@ -0,0 +1,41 @@
+/*
+ * Copyright (c) 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.
+ * 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
+{
+
+// MiniMax-M3-specific index-branch producer. It applies Gemma RMSNorm and
+// NeoX partial RoPE to a packed BF16 [index-Q | index-K] projection, writes
+// index-Q as unscaled E4M3, and inserts index-K directly into the paged E4M3
+// HND cache. The direct cache store removes the standalone cast/scatter launch
+// from the decode graph.
+void launchMinimaxM3Fp8IndexerQKNormRope(void const* qk, void* q_out, void* k_cache, int const* out_cache_loc,
+    int64_t page_stride, int64_t token_stride, int page_size, int64_t num_pages, int num_tokens, int num_heads_q,
+    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);
+
+} // namespace kernels
+
+TRTLLM_NAMESPACE_END

cpp/tensorrt_llm/thop/CMakeLists.txt

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

cpp/tensorrt_llm/thop/minimaxM3Fp8IndexerOp.cpp

@@ -0,0 +1,89 @@
+/*
+ * Copyright (c) 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.
+ * 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/minimaxM3Fp8IndexerKernel.h"
+#include "tensorrt_llm/thop/thUtils.h"
+
+#include <ATen/cuda/CUDAContext.h>
+#include <torch/extension.h>
+
+TRTLLM_NAMESPACE_BEGIN
+
+namespace torch_ext
+{
+
+torch::Tensor minimaxM3Fp8IndexerQKNormRope(torch::Tensor const& qk, torch::Tensor& indexKCache,
+    torch::Tensor const& outCacheLoc, int64_t numHeadsQ, int64_t headDim, int64_t rotaryDim, double eps,
+    torch::Tensor const& qWeight, torch::Tensor const& kWeight, double base, torch::Tensor const& positionIds)
+{
+    TORCH_CHECK(qk.dim() == 2, "Index QK must be [num_tokens, (num_heads_q + 1) * head_dim]");
+    TORCH_CHECK(indexKCache.dim() == 4 && indexKCache.size(1) == 1,
+        "Index-K cache must be HND [num_pages, 1, page_size, head_dim]");
+    TORCH_CHECK(outCacheLoc.dim() == 1, "out_cache_loc must be one-dimensional");
+    TORCH_CHECK(positionIds.dim() == 1, "position_ids must be one-dimensional");
+    TORCH_CHECK(qWeight.dim() == 1 && kWeight.dim() == 1, "Q/K norm weights must be one-dimensional");
+
+    CHECK_INPUT(qk, torch::kBFloat16);
+    CHECK_INPUT(outCacheLoc, torch::kInt32);
+    CHECK_INPUT(positionIds, torch::kInt32);
+    CHECK_INPUT(qWeight, torch::kBFloat16);
+    CHECK_INPUT(kWeight, torch::kBFloat16);
+    TORCH_CHECK(indexKCache.is_cuda(), "Index-K cache must be on CUDA");
+    TORCH_CHECK(
+        indexKCache.scalar_type() == at::ScalarType::Float8_e4m3fn, "Index-K cache must use torch.float8_e4m3fn");
+
+    int64_t const numTokens = qk.size(0);
+    TORCH_CHECK(qk.size(1) == (numHeadsQ + 1) * headDim, "Index QK width must equal (num_heads_q + 1) * head_dim");
+    TORCH_CHECK(indexKCache.size(3) == headDim, "Index-K cache head dimension mismatch");
+    TORCH_CHECK(indexKCache.stride(3) == 1 && indexKCache.stride(2) == headDim,
+        "Index-K cache must have contiguous token rows in HND layout");
+    TORCH_CHECK(outCacheLoc.numel() >= numTokens, "out_cache_loc is shorter than num_tokens");
+    TORCH_CHECK(positionIds.numel() == numTokens, "position_ids length must equal num_tokens");
+    TORCH_CHECK(qWeight.numel() == headDim && kWeight.numel() == headDim, "Q/K norm weight width must equal head_dim");
+    TORCH_CHECK(qk.get_device() == indexKCache.get_device() && qk.get_device() == outCacheLoc.get_device()
+            && qk.get_device() == positionIds.get_device() && qk.get_device() == qWeight.get_device()
+            && qk.get_device() == kWeight.get_device(),
+        "All MiniMax-M3 FP8 indexer tensors must be on the same CUDA device");
+
+    auto const qOut = torch::empty({numTokens, numHeadsQ, headDim}, qk.options().dtype(at::ScalarType::Float8_e4m3fn));
+    if (numTokens == 0)
+    {
+        return qOut;
+    }
+    auto const stream = at::cuda::getCurrentCUDAStream(qk.get_device());
+    tensorrt_llm::kernels::launchMinimaxM3Fp8IndexerQKNormRope(qk.data_ptr(), qOut.data_ptr(), indexKCache.data_ptr(),
+        outCacheLoc.data_ptr<int>(), indexKCache.stride(0), indexKCache.stride(2), indexKCache.size(2),
+        indexKCache.size(0), numTokens, numHeadsQ, headDim, rotaryDim, static_cast<float>(eps), qWeight.data_ptr(),
+        kWeight.data_ptr(), static_cast<float>(base), positionIds.data_ptr<int>(), stream);
+    return qOut;
+}
+
+TORCH_LIBRARY_FRAGMENT(trtllm, m)
+{
+    m.def(
+        "minimax_m3_fp8_indexer_qk_norm_rope(Tensor qk, Tensor(a!) index_k_cache, Tensor out_cache_loc, int "
+        "num_heads_q, int head_dim, int rotary_dim, float eps, Tensor q_weight, Tensor k_weight, float base, Tensor "
+        "position_ids) -> Tensor");
+}
+
+TORCH_LIBRARY_IMPL(trtllm, CUDA, m)
+{
+    m.impl("minimax_m3_fp8_indexer_qk_norm_rope", &minimaxM3Fp8IndexerQKNormRope);
+}
+
+} // namespace torch_ext
+
+TRTLLM_NAMESPACE_END

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

@@ -162,7 +162,9 @@ def __init__(
         # then from ``sparse_attn_config``, then from the M3 checkpoint
         # convention (layers 0..2 dense, 3..N-1 sparse,
         # disable_index_value=True, sparse_index_dim=128).
-        sparse_attn_config = kwargs.get("sparse_attn_config")
+        sparse_attn_config = kwargs.get("sparse_attn_config") or kwargs.get(
+            "sparse_attention_config"
+        )
         num_layers = kwargs.get("num_layers")
 
         if sparse_index_dim is None:
@@ -181,6 +183,19 @@ def __init__(
         self.sparse_layer_ids = sorted(int(i) for i in sparse_layer_ids)
         self.disable_index_value_layer_ids = set(int(i) for i in disable_index_value_layer_ids)
         self.sparse_index_dim = int(sparse_index_dim)
+        self.indexer_kv_dtype = str(getattr(sparse_attn_config, "indexer_kv_dtype", "bf16"))
+        if self.indexer_kv_dtype not in ("bf16", "fp8"):
+            raise ValueError(
+                "MiniMax M3 indexer_kv_dtype must be 'bf16' or 'fp8', got "
+                f"{self.indexer_kv_dtype!r}."
+            )
+        if self.indexer_kv_dtype == "fp8" and (
+            set(self.sparse_layer_ids) - self.disable_index_value_layer_ids
+        ):
+            raise ValueError(
+                "MiniMax M3 FP8 index cache requires disable_index_value=True "
+                "for every sparse layer."
+            )
 
         super().__init__(*args, **kwargs)
 
@@ -251,7 +266,9 @@ def _compute_num_total_slots(self) -> int:
         return int((page_upper // kv_factor) * self.tokens_per_block)
 
     def _torch_dtype_for_index_cache(self) -> torch.dtype:
-        """Match the main cache dtype where possible, fall back to bf16."""
+        """Return the independently configured index-cache storage dtype."""
+        if self.indexer_kv_dtype == "fp8":
+            return torch.float8_e4m3fn
         if self.dtype == DataType.HALF:
             return torch.float16
         if self.dtype == DataType.FLOAT:

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

@@ -40,6 +40,7 @@ class MiniMaxM3SparseParams(SparseParams):
     score_type: str = "max"
     disable_index_value: bool = True
     implementation: Literal["triton", "msa"] = "triton"
+    indexer_kv_dtype: Literal["bf16", "fp8"] = "bf16"
 
     @property
     def indices_block_size(self) -> int: