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

cpp/tensorrt_llm/kernels/minimaxM3Fp8IndexerKernel.cu

@@ -0,0 +1,190 @@
+/*
+ * 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>
+#include <limits>
+
+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];
+        // Production msa_out_cache_loc contains only allocated live-token
+        // slots: KVCacheManagerV2 canonicalizes padded BAD_PAGE_INDEX entries
+        // before build_paged_kv_slot_mapping selects the live ranges. Keep
+        // these guards so direct custom-op callers cannot corrupt the cache
+        // when they supply a sentinel or stale out-of-range slot.
+        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* qOut, void* kCache, int const* outCacheLoc,
+    int64_t pageStride, int64_t tokenStride, int pageSize, int64_t numPages, int numTokens, int numHeadsQ, int headDim,
+    int rotaryDim, float eps, void const* qWeight, void const* kWeight, float base, int const* positionIds,
+    cudaStream_t stream)
+{
+    TLLM_CHECK_WITH_INFO(headDim == kHeadDim, "MiniMax-M3 FP8 indexer requires head_dim=128");
+    TLLM_CHECK_WITH_INFO(rotaryDim == kRotaryDim, "MiniMax-M3 FP8 indexer requires rotary_dim=64");
+    TLLM_CHECK_WITH_INFO(numHeadsQ > 0, "MiniMax-M3 FP8 indexer requires at least one query head");
+
+    constexpr int kBlockSize = 256;
+    constexpr int kWarpsPerBlock = kBlockSize / 32;
+    int64_t const totalWarps = static_cast<int64_t>(numTokens) * (static_cast<int64_t>(numHeadsQ) + 1);
+    int64_t const gridSize64 = common::divUp(totalWarps, static_cast<int64_t>(kWarpsPerBlock));
+    TLLM_CHECK_WITH_INFO(gridSize64 <= std::numeric_limits<int>::max(), "MiniMax-M3 FP8 indexer grid is too large");
+    int const gridSize = static_cast<int>(gridSize64);
+    minimaxM3Fp8IndexerQKNormRopeKernel<<<gridSize, kBlockSize, 0, stream>>>(static_cast<__nv_bfloat16 const*>(qk),
+        static_cast<__nv_fp8_e4m3*>(qOut), static_cast<__nv_fp8_e4m3*>(kCache), outCacheLoc, pageStride, tokenStride,
+        pageSize, numPages, numTokens, numHeadsQ, eps, static_cast<__nv_bfloat16 const*>(qWeight),
+        static_cast<__nv_bfloat16 const*>(kWeight), base, positionIds);
+    sync_check_cuda_error(stream);
+}
+
+} // namespace kernels
+
+TRTLLM_NAMESPACE_END

cpp/tensorrt_llm/kernels/minimaxM3Fp8IndexerKernel.h

@@ -0,0 +1,55 @@
+/*
+ * 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.
+//!
+//! 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.
+//!
+//! `qk` must be a contiguous BF16 `[num_tokens, (num_heads_q + 1) *
+//! head_dim]` tensor whose base address is 8-byte aligned. `qOut` is a
+//! contiguous E4M3 `[num_tokens, num_heads_q, head_dim]` output. `kCache` is
+//! an E4M3 HND cache `[num_pages, 1, page_size, head_dim]`; its base address
+//! and every page start must be 4-byte aligned. `outCacheLoc` and
+//! `positionIds` contain one int32 value per token. The norm weights are BF16
+//! vectors of `head_dim` elements.
+//!
+//! \param pageStride Distance in E4M3 elements between cache pages.
+//! \param tokenStride Distance in E4M3 elements between tokens in a page.
+//! \param pageSize Number of token slots per cache page.
+//! \param numPages Number of addressable pages in `kCache`.
+void launchMinimaxM3Fp8IndexerQKNormRope(void const* qk, void* qOut, void* kCache, int const* outCacheLoc,
+    int64_t pageStride, int64_t tokenStride, int pageSize, int64_t numPages, int numTokens, int numHeadsQ, int headDim,
+    int rotaryDim, float eps, void const* qWeight, void const* kWeight, float base, int const* positionIds,
+    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
   minimaxM3SelectBlocksOp.cpp
   mlaRopeInplaceOp.cpp
   ncclCommunicatorOp.cpp

cpp/tensorrt_llm/thop/minimaxM3Fp8IndexerOp.cpp

@@ -0,0 +1,137 @@
+/*
+ * 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>
+
+#include <cmath>
+#include <cstdint>
+#include <limits>
+
+TRTLLM_NAMESPACE_BEGIN
+
+namespace torch_ext
+{
+
+namespace
+{
+constexpr int64_t kMinimaxM3IndexKHeads = 1;
+constexpr int64_t kMinimaxM3HeadDim = 128;
+constexpr int64_t kMinimaxM3RotaryDim = 64;
+} // namespace
+
+//! Normalize and rotate MiniMax-M3 index Q/K, returning Q and caching K as E4M3.
+//!
+//! `qk` is contiguous CUDA BF16 `[num_tokens, (numHeadsQ + 1) * headDim]`.
+//! `indexKCache` is a mutable CUDA E4M3 HND cache
+//! `[num_pages, 1, page_size, headDim]`; valid `outCacheLoc` entries address
+//! `page * page_size + token`, while malformed negative or out-of-range entries
+//! are defensively skipped. `qWeight` and `kWeight` are contiguous CUDA BF16
+//! vectors of `headDim` elements, and `positionIds` is contiguous CUDA int32
+//! with one entry per token.
+//!
+//! `numHeadsQ` must be positive, `headDim` must be 128, `rotaryDim` must be 64,
+//! and both `eps` and `base` must be finite and positive. The function mutates
+//! `indexKCache` in place and returns contiguous E4M3
+//! `[num_tokens, numHeadsQ, headDim]` index Q.
+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) == kMinimaxM3IndexKHeads,
+        "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(numHeadsQ > 0, "num_heads_q must be greater than zero");
+    TORCH_CHECK(headDim == kMinimaxM3HeadDim, "MiniMax-M3 FP8 indexer requires head_dim=128");
+    TORCH_CHECK(rotaryDim == kMinimaxM3RotaryDim, "MiniMax-M3 FP8 indexer requires rotary_dim=64");
+    TORCH_CHECK(std::isfinite(eps) && eps > 0.0, "eps must be finite and greater than zero");
+    TORCH_CHECK(std::isfinite(base) && base > 0.0, "RoPE base must be finite and greater than zero");
+    auto const epsFloat = static_cast<float>(eps);
+    auto const baseFloat = static_cast<float>(base);
+    TORCH_CHECK(std::isfinite(epsFloat) && epsFloat > 0.0F, "eps must remain finite and positive in float32");
+    TORCH_CHECK(std::isfinite(baseFloat) && baseFloat > 0.0F, "RoPE base must remain finite and positive in float32");
+    TORCH_CHECK(indexKCache.size(0) > 0, "Index-K cache must contain at least one page");
+    TORCH_CHECK(indexKCache.size(2) > 0, "Index-K cache page_size must be greater than zero");
+    TORCH_CHECK(numTokens <= std::numeric_limits<int>::max(), "num_tokens exceeds the CUDA kernel's int range");
+    TORCH_CHECK(numHeadsQ <= std::numeric_limits<int>::max(), "num_heads_q exceeds the CUDA kernel's int range");
+    TORCH_CHECK(indexKCache.size(2) <= std::numeric_limits<int>::max(),
+        "Index-K cache page_size exceeds the CUDA kernel's int range");
+    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");
+    constexpr uintptr_t kQkAlignment = 8;
+    constexpr uintptr_t kCacheAlignment = 4;
+    TORCH_CHECK(reinterpret_cast<uintptr_t>(qk.data_ptr()) % kQkAlignment == 0,
+        "Index QK must start at an 8-byte-aligned address for vectorized BF16 loads");
+    TORCH_CHECK(reinterpret_cast<uintptr_t>(indexKCache.data_ptr()) % kCacheAlignment == 0,
+        "Index-K cache must start at a 4-byte-aligned address for packed E4M3 stores");
+    TORCH_CHECK(indexKCache.stride(0) % static_cast<int64_t>(kCacheAlignment) == 0,
+        "Index-K cache page stride must be a multiple of 4 E4M3 elements for packed stores");
+    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, epsFloat, qWeight.data_ptr(), kWeight.data_ptr(),
+        baseFloat, 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