From 6842ede4a3774c7459e03aa324537c39d4c4b87b Mon Sep 17 00:00:00 2001 From: peihengh <259410613+peihu-nv@users.noreply.github.com> Date: Tue, 21 Jul 2026 19:37:14 -0700 Subject: [PATCH] [None][perf] Use FP8 MiniMax-M3 MSA indexer QK Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> --- .../kernels/minimaxM3Fp8IndexerKernel.cu | 166 ++++++++++++++++++ .../kernels/minimaxM3Fp8IndexerKernel.h | 41 +++++ cpp/tensorrt_llm/thop/CMakeLists.txt | 1 + .../thop/minimaxM3Fp8IndexerOp.cpp | 90 ++++++++++ .../sparse/minimax_m3/cache_manager.py | 21 ++- .../sparse/minimax_m3/common.py | 1 + .../sparse/minimax_m3/msa_backend.py | 23 ++- .../_torch/custom_ops/cpp_custom_ops.py | 8 + .../_torch/models/modeling_minimaxm3.py | 61 ++++++- tensorrt_llm/llmapi/llm_args.py | 16 ++ .../usage/llm_args_golden_manifest.json | 10 ++ .../sparse/test_minimax_m3_msa_backend.py | 82 +++++++++ .../test_minimax_m3_fp8_indexer.py | 126 +++++++++++++ 13 files changed, 639 insertions(+), 7 deletions(-) create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Fp8IndexerKernel.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Fp8IndexerKernel.h create mode 100644 cpp/tensorrt_llm/thop/minimaxM3Fp8IndexerOp.cpp create mode 100644 tests/unittest/_torch/thop/parallel_hw_agnostic/test_minimax_m3_fp8_indexer.py diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Fp8IndexerKernel.cu b/cpp/tensorrt_llm/kernels/minimaxM3Fp8IndexerKernel.cu new file mode 100644 index 000000000000..0f255ef3b7ae --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Fp8IndexerKernel.cu @@ -0,0 +1,166 @@ +/* + * 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 +#include +#include + +#include + +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, + 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(token_idx) * total_heads + local_head) * kHeadDim + lane_id * kElemsPerThread; + + uint2 const packed_input = *reinterpret_cast(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; + sum_squares += values.x * values.x + values.y * values.y; + } + + sum_squares = tensorrt_llm::common::warpReduceSum(sum_squares); + float const rms_rcp = rsqrtf(sum_squares / static_cast(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; +#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 = powf(base, -2.0F * half_dim / static_cast(kRotaryDim)); + float sine; + float cosine; + __sincosf(static_cast(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(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]; + int const page = slot / page_size; + int const within_page = slot % page_size; + output = k_cache + static_cast(page) * page_stride + static_cast(within_page) * token_stride + + lane_id * kElemsPerThread; + } + *reinterpret_cast(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, 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<<>>(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_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 diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Fp8IndexerKernel.h b/cpp/tensorrt_llm/kernels/minimaxM3Fp8IndexerKernel.h new file mode 100644 index 000000000000..0d23eed8cf53 --- /dev/null +++ b/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 +#include + +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, 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 diff --git a/cpp/tensorrt_llm/thop/CMakeLists.txt b/cpp/tensorrt_llm/thop/CMakeLists.txt index 855147742214..722748b335ac 100644 --- a/cpp/tensorrt_llm/thop/CMakeLists.txt +++ b/cpp/tensorrt_llm/thop/CMakeLists.txt @@ -111,6 +111,7 @@ add_library( IndexerKCacheScatterOp.cpp IndexerTopKOp.cpp minimaxM3SelectBlocksOp.cpp + minimaxM3Fp8IndexerOp.cpp mlaRopeInplaceOp.cpp ncclCommunicatorOp.cpp allocateOutput.cpp diff --git a/cpp/tensorrt_llm/thop/minimaxM3Fp8IndexerOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3Fp8IndexerOp.cpp new file mode 100644 index 000000000000..211de47e799f --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3Fp8IndexerOp.cpp @@ -0,0 +1,90 @@ +/* + * 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 +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +torch::Tensor minimax_m3_fp8_indexer_qk_norm_rope(torch::Tensor const& qk, torch::Tensor& index_k_cache, + torch::Tensor const& out_cache_loc, int64_t num_heads_q, int64_t head_dim, int64_t rotary_dim, double eps, + torch::Tensor const& q_weight, torch::Tensor const& k_weight, double base, torch::Tensor const& position_ids) +{ + TORCH_CHECK(qk.dim() == 2, "Index QK must be [num_tokens, (num_heads_q + 1) * head_dim]"); + TORCH_CHECK(index_k_cache.dim() == 4 && index_k_cache.size(1) == 1, + "Index-K cache must be HND [num_pages, 1, page_size, head_dim]"); + TORCH_CHECK(out_cache_loc.dim() == 1, "out_cache_loc must be one-dimensional"); + TORCH_CHECK(position_ids.dim() == 1, "position_ids must be one-dimensional"); + TORCH_CHECK(q_weight.dim() == 1 && k_weight.dim() == 1, "Q/K norm weights must be one-dimensional"); + + CHECK_INPUT(qk, torch::kBFloat16); + CHECK_INPUT(out_cache_loc, torch::kInt32); + CHECK_INPUT(position_ids, torch::kInt32); + CHECK_INPUT(q_weight, torch::kBFloat16); + CHECK_INPUT(k_weight, torch::kBFloat16); + TORCH_CHECK(index_k_cache.is_cuda(), "Index-K cache must be on CUDA"); + TORCH_CHECK( + index_k_cache.scalar_type() == at::ScalarType::Float8_e4m3fn, "Index-K cache must use torch.float8_e4m3fn"); + + int64_t const num_tokens = qk.size(0); + TORCH_CHECK(qk.size(1) == (num_heads_q + 1) * head_dim, "Index QK width must equal (num_heads_q + 1) * head_dim"); + TORCH_CHECK(index_k_cache.size(3) == head_dim, "Index-K cache head dimension mismatch"); + TORCH_CHECK(index_k_cache.stride(3) == 1 && index_k_cache.stride(2) == head_dim, + "Index-K cache must have contiguous token rows in HND layout"); + TORCH_CHECK(out_cache_loc.numel() >= num_tokens, "out_cache_loc is shorter than num_tokens"); + TORCH_CHECK(position_ids.numel() == num_tokens, "position_ids length must equal num_tokens"); + TORCH_CHECK( + q_weight.numel() == head_dim && k_weight.numel() == head_dim, "Q/K norm weight width must equal head_dim"); + TORCH_CHECK(qk.get_device() == index_k_cache.get_device() && qk.get_device() == out_cache_loc.get_device() + && qk.get_device() == position_ids.get_device() && qk.get_device() == q_weight.get_device() + && qk.get_device() == k_weight.get_device(), + "All MiniMax-M3 FP8 indexer tensors must be on the same CUDA device"); + + auto q_out = torch::empty({num_tokens, num_heads_q, head_dim}, qk.options().dtype(at::ScalarType::Float8_e4m3fn)); + if (num_tokens == 0) + { + return q_out; + } + auto stream = at::cuda::getCurrentCUDAStream(qk.get_device()); + tensorrt_llm::kernels::launchMinimaxM3Fp8IndexerQKNormRope(qk.data_ptr(), q_out.data_ptr(), + index_k_cache.data_ptr(), out_cache_loc.data_ptr(), index_k_cache.stride(0), index_k_cache.stride(2), + index_k_cache.size(2), num_tokens, num_heads_q, head_dim, rotary_dim, static_cast(eps), + q_weight.data_ptr(), k_weight.data_ptr(), static_cast(base), position_ids.data_ptr(), stream); + return q_out; +} + +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", &minimax_m3_fp8_indexer_qk_norm_rope); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END diff --git a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/cache_manager.py b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/cache_manager.py index fd0f72159f1f..2baefd5cc6c1 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/cache_manager.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/cache_manager.py @@ -176,7 +176,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: @@ -195,6 +197,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) @@ -265,7 +280,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: diff --git a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/common.py b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/common.py index 9f2863846df2..c508563bb720 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/common.py +++ b/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: diff --git a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.py b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.py index 5cf0512b2509..e80c76857c6a 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.py @@ -976,6 +976,7 @@ def __init__( head_dim=head_dim, ) self.disable_index_value = bool(sparse_params.disable_index_value) + self.indexer_kv_dtype = str(sparse_params.indexer_kv_dtype) self._validate_msa_preconditions() self.indexer = MsaIndexer(self.m3_config) @@ -1009,7 +1010,7 @@ def support_fused_rope(cls) -> bool: def run_indexer( self, idx_q: torch.Tensor, - idx_k: torch.Tensor, + idx_k: Optional[torch.Tensor], metadata, *, idx_sm_scale: Optional[float] = None, @@ -1028,10 +1029,24 @@ def run_indexer( # reshape to keep them zero-copy. The proxy fmha_sm100 and the index-K # scatter below both honor the source strides. idx_q_view = idx_q.reshape(num_tokens, config.num_index_heads, config.sparse_index_dim) - idx_k_view = idx_k.reshape(num_tokens, 1, config.sparse_index_dim) - - metadata.msa_write_idx_k(self.layer_idx, idx_k_view) idx_k_cache = metadata.msa_idx_k_cache(self.layer_idx) + if idx_k is not None: + idx_k_view = idx_k.reshape(num_tokens, 1, config.sparse_index_dim) + metadata.msa_write_idx_k(self.layer_idx, idx_k_view) + elif idx_k_cache.dtype != torch.float8_e4m3fn or idx_q_view.dtype != torch.float8_e4m3fn: + raise ValueError( + "A missing live index-K is valid only when the fused MiniMax-M3 " + "producer already emitted FP8 index-Q and inserted FP8 index-K." + ) + # The FP8 indexer mirrors vLLM's unscaled E4M3 contract: normalized + # index Q/K are cast directly and the proxy accumulates their QK scores + # in FP32. Block ordering is invariant to the omitted positive scale. + # The fused production path arrives here with E4M3 Q and an already + # populated cache. The explicit conversion is retained for standalone + # callers that supply BF16 Q/K to an E4M3-configured backend. + if idx_k_cache.dtype == torch.float8_e4m3fn: + if idx_q_view.dtype != torch.float8_e4m3fn: + idx_q_view = idx_q_view.to(torch.float8_e4m3fn) # One selection path. Decode passes the graph-safe proxy plan plus the # proxy scratch shaped to the live query count. Prefill and mixed batches diff --git a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py index 851a55c32f3f..ce3e410c2d8d 100644 --- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py @@ -282,6 +282,14 @@ def _(scores, n_valid_blocks, topk, init_blocks, local_blocks): return scores.new_empty((scores.shape[2], scores.shape[0], topk), dtype=torch.int32) + @torch.library.register_fake("trtllm::minimax_m3_fp8_indexer_qk_norm_rope") + def _(qk, index_k_cache, out_cache_loc, num_heads_q, head_dim, rotary_dim, + eps, q_weight, k_weight, base, position_ids): + del index_k_cache, out_cache_loc, rotary_dim, eps, q_weight, k_weight + del base, position_ids + return qk.new_empty((qk.shape[0], num_heads_q, head_dim), + dtype=torch.float8_e4m3fn) + @torch.library.register_fake("trtllm::userbuffers_allreduce_finalize") def _(input, force_applying_finalize): return torch.empty_like(input) diff --git a/tensorrt_llm/_torch/models/modeling_minimaxm3.py b/tensorrt_llm/_torch/models/modeling_minimaxm3.py index b7c6820053bd..60b7c70f9368 100644 --- a/tensorrt_llm/_torch/models/modeling_minimaxm3.py +++ b/tensorrt_llm/_torch/models/modeling_minimaxm3.py @@ -964,6 +964,61 @@ def _fused_qk_norm_rope( ) return qkv + def _fused_fp8_index_qk_norm_rope( + self, + idx_qk: torch.Tensor, + position_ids: Optional[torch.Tensor], + attn_metadata: AttentionMetadata, + ) -> Optional[torch.Tensor]: + """Produce raw-E4M3 index-Q and insert index-K in one M3-only kernel. + + This path is deliberately narrower than the general fused QK kernel: + it is enabled only for the MSA FP8-indexer configuration and the exact + M3 Gemma-RMSNorm + NeoX partial-RoPE geometry. The kernel rounds the + normalized/RoPE values through BF16 before E4M3 conversion, matching + the former fused-BF16-kernel followed by ``Tensor.to(E4M3)`` contract. + """ + if not isinstance(self.attn, MiniMaxM3MsaSparseAttention): + return None + if self.attn.indexer_kv_dtype != "fp8": + return None + if position_ids is None or idx_qk.dtype != torch.bfloat16: + raise NotImplementedError( + "MiniMax-M3 fused FP8 indexer requires BF16 activations and position_ids." + ) + if ( + self.rotary_emb is None + or self.pos_embd_params is None + or self.pos_embd_params.rope is None + ): + raise NotImplementedError("MiniMax-M3 fused FP8 indexer requires partial RoPE.") + if not self.pos_embd_params.is_neox: + raise NotImplementedError("MiniMax-M3 fused FP8 indexer requires NeoX RoPE.") + if not self.use_gemma_norm: + raise NotImplementedError("MiniMax-M3 fused FP8 indexer requires Gemma RMSNorm.") + + rotary_dim = int(self.pos_embd_params.rope.dim) + index_k_cache = attn_metadata.msa_idx_k_cache(self.layer_idx) + if index_k_cache.dtype != torch.float8_e4m3fn: + raise ValueError( + "MiniMax-M3 fused FP8 indexer requires an E4M3 index-K cache, " + f"got {index_k_cache.dtype}." + ) + num_tokens = int(idx_qk.shape[0]) + return torch.ops.trtllm.minimax_m3_fp8_indexer_qk_norm_rope( + idx_qk.contiguous(), + index_k_cache, + attn_metadata.msa_out_cache_loc[:num_tokens], + self.sparse_num_index_heads, + self.sparse_index_dim, + rotary_dim, + self.index_q_norm.variance_epsilon, + self.index_q_norm.weight, + self.index_k_norm.weight, + self.pos_embd_params.rope.theta, + position_ids.reshape(-1).contiguous().to(torch.int32), + ).flatten(1) + def _expect_fused_qk_norm_rope(self, position_ids: Optional[torch.Tensor]) -> bool: """Whether the fused path must run rather than fall back. @@ -1433,7 +1488,7 @@ def _msa_attention_core( only) and builds the ``forward_args`` the FMHA reads. """ if self.is_sparse_attention_layer: - assert idx_q is not None and idx_k is not None + assert idx_q is not None # Publish the selected blocks so the FMHA runs the sparse path. kv_block_indexes = self.attn.run_indexer(idx_q, idx_k, attn_metadata) forward_args = AttentionForwardArgs(output=output, topk_indices=kv_block_indexes) @@ -1528,6 +1583,10 @@ def _main_norm_rope(): def _index_norm_rope(): idx_qk = self.index_qk_proj(hidden_states) + fp8_idx_q = self._fused_fp8_index_qk_norm_rope(idx_qk, position_ids, attn_metadata) + if fp8_idx_q is not None: + # Index-K was inserted directly into the paged side cache. + return fp8_idx_q, None fused_idx = self._fused_qk_norm_rope( idx_qk, position_ids, diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 317272c9e2c8..e60e923edd51 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -728,6 +728,14 @@ class MiniMaxM3SparseAttentionConfig(BaseSparseAttentionConfig): default=True, description="If True, skip the index V branch (M3 checkpoint default).", ) + indexer_kv_dtype: Literal["bf16", "fp8"] = Field( + default="bf16", + description= + "Storage and score-compute dtype for normalized index Q/K. 'fp8' uses " + "unscaled E4M3 values with FP32 score accumulation and is supported only " + "by the MSA implementation.", + status="prototype", + ) num_attention_heads: Optional[int] = Field( default=None, description= @@ -755,6 +763,13 @@ def _validate_msa_block_size(self): raise ValueError( "MiniMax-M3 'msa' implementation requires sparse_block_size == " f"128, got {self.sparse_block_size}.") + if self.indexer_kv_dtype == "fp8" and self.implementation != "msa": + raise ValueError( + "MiniMax-M3 indexer_kv_dtype='fp8' currently requires the " + "'msa' implementation.") + if self.indexer_kv_dtype == "fp8" and not self.sparse_disable_index_value: + raise ValueError("MiniMax-M3 indexer_kv_dtype='fp8' requires " + "sparse_disable_index_value=True.") return self def supports_backend(self, backend: str) -> bool: @@ -777,6 +792,7 @@ def to_sparse_params(self, **kwargs): score_type=self.sparse_score_type, disable_index_value=self.sparse_disable_index_value, implementation=self.implementation, + indexer_kv_dtype=self.indexer_kv_dtype, ) def to_sparse_metadata_params(self, **kwargs): diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index 066389d7be27..bbbaefff258c 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -1430,6 +1430,16 @@ "kind": "categorical", "path": "sparse_attention_config.indexer_k_dtype" }, + { + "allowed_values": [ + "bf16", + "fp8" + ], + "annotation": "Literal['bf16', 'fp8']", + "converter": "", + "kind": "categorical", + "path": "sparse_attention_config.indexer_kv_dtype" + }, { "allowed_values": [], "annotation": "Optional[int]", diff --git a/tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py b/tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py index 1c777d593ae9..3ce0bdd9033a 100644 --- a/tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py +++ b/tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py @@ -35,6 +35,20 @@ def test_msa_requires_block_size_128(): assert cfg.sparse_block_size == 64 +def test_msa_fp8_indexer_config_is_explicit_and_lowered(): + cfg = MiniMaxM3SparseAttentionConfig(implementation="msa", indexer_kv_dtype="fp8") + assert cfg.to_sparse_params().indexer_kv_dtype == "fp8" + + with pytest.raises(ValueError, match=r"requires the 'msa' implementation"): + MiniMaxM3SparseAttentionConfig(implementation="triton", indexer_kv_dtype="fp8") + with pytest.raises(ValueError, match=r"sparse_disable_index_value=True"): + MiniMaxM3SparseAttentionConfig( + implementation="msa", + indexer_kv_dtype="fp8", + sparse_disable_index_value=False, + ) + + def test_msa_metadata_rejects_undersized_max_score_buffer(): metadata_cls = MiniMaxM3MsaSparseAttention.Metadata metadata = metadata_cls.__new__(metadata_cls) @@ -165,6 +179,74 @@ def fake_select_blocks_from_maxscore(*args, **kwargs): assert result is expected +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_msa_fp8_cache_converts_live_index_query_before_scoring(): + from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.common import MiniMaxM3SparseConfig + + config = MiniMaxM3SparseConfig( + num_q_heads=4, + num_kv_heads=4, + head_dim=128, + num_index_heads=4, + sparse_index_dim=128, + block_size=128, + topk=16, + ) + attention = MiniMaxM3MsaSparseAttention.__new__(MiniMaxM3MsaSparseAttention) + attention.m3_config = config + attention.layer_idx = 3 + captured = {} + + class FakeIndexer: + def select_blocks(self, idx_q, idx_k, **kwargs): + captured["idx_q"] = idx_q + captured["idx_k"] = idx_k + captured["kwargs"] = kwargs + return torch.zeros(2, 4, 16, dtype=torch.int32, device="cuda") + + attention.indexer = FakeIndexer() + + class FakeMetadata: + msa_decode_proxy_plan = None + msa_eager_proxy_plan = (False, 0, 2, {}, None) + msa_eager_all_blocks_empty = False + msa_eager_n_valid_blocks = torch.ones(2, dtype=torch.int32, device="cuda") + msa_kv_indices = torch.arange(2, dtype=torch.int32, device="cuda") + msa_qo_lens_cpu = torch.ones(2, dtype=torch.int32) + msa_kv_lens_cpu = torch.full((2,), 128, dtype=torch.int32) + msa_qo_offset_cpu = torch.full((2,), 127, dtype=torch.int32) + + def __init__(self): + self.cache = torch.empty(2, 1, 128, 128, dtype=torch.float8_e4m3fn, device="cuda") + + def msa_write_idx_k(self, layer_idx, idx_k): + captured["write"] = (layer_idx, idx_k) + + def msa_idx_k_cache(self, layer_idx): + captured["read_layer"] = layer_idx + return self.cache + + idx_q = torch.randn(2, 4 * 128, dtype=torch.bfloat16, device="cuda") + idx_k = torch.randn(2, 128, dtype=torch.bfloat16, device="cuda") + result = attention.run_indexer(idx_q, idx_k, FakeMetadata()) + + assert result.shape == (2, 4, 16) + assert captured["idx_q"].dtype == torch.float8_e4m3fn + assert captured["idx_k"].dtype == torch.float8_e4m3fn + assert captured["idx_k"].stride(0) > captured["idx_k"].shape[-1] + assert captured["write"][0] == 3 + assert captured["write"][1].data_ptr() == idx_k.data_ptr() + + # The production fused producer has already inserted K and passes no live + # K tensor; E4M3 Q must flow to the scorer without a duplicate cache write. + captured.pop("write") + fused_q = idx_q.to(torch.float8_e4m3fn) + result = attention.run_indexer(fused_q, None, FakeMetadata()) + assert result.shape == (2, 4, 16) + assert captured["idx_q"].data_ptr() == fused_q.data_ptr() + assert "write" not in captured + + def test_msa_proxy_max_score_strided_index_k_matches_packed(): if not torch.cuda.is_available(): pytest.skip("CUDA required") diff --git a/tests/unittest/_torch/thop/parallel_hw_agnostic/test_minimax_m3_fp8_indexer.py b/tests/unittest/_torch/thop/parallel_hw_agnostic/test_minimax_m3_fp8_indexer.py new file mode 100644 index 000000000000..f577507b3831 --- /dev/null +++ b/tests/unittest/_torch/thop/parallel_hw_agnostic/test_minimax_m3_fp8_indexer.py @@ -0,0 +1,126 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +import torch + + +def _reference(qk, num_heads_q, q_weight, k_weight, position_ids): + reference = qk.clone() + torch.ops.trtllm.fused_qk_norm_rope( + reference, + num_heads_q, + 1, + 0, + 128, + 64, + 1e-5, + q_weight, + k_weight, + 10000.0, + True, # is_neox + position_ids, + 1.0, + 0.0, + 0.0, + 1.0, + True, # is_qk_norm + True, # use_gemma + False, # use_mrope + 0, + 0, + ) + q, k = reference.split([num_heads_q * 128, 128], dim=-1) + return q.view(q.shape[0], num_heads_q, 128).to(torch.float8_e4m3fn), k.to(torch.float8_e4m3fn) + + +def _strided_cache(num_pages, page_size=128, stride_scale=7): + backing = torch.zeros( + num_pages * stride_scale, + 1, + page_size, + 128, + dtype=torch.float8_e4m3fn, + device="cuda", + ) + return backing[::stride_scale] + + +def _run(qk, cache, slots, q_weight, k_weight, position_ids, num_heads_q=4): + return torch.ops.trtllm.minimax_m3_fp8_indexer_qk_norm_rope( + qk, + cache, + slots, + num_heads_q, + 128, + 64, + 1e-5, + q_weight, + k_weight, + 10000.0, + position_ids, + ) + + +@pytest.mark.parametrize("num_tokens", [1, 16, 129]) +def test_minimax_m3_fp8_indexer_matches_bf16_then_cast(num_tokens): + torch.manual_seed(1234) + num_heads_q = 4 + page_size = 128 + qk = torch.randn( + num_tokens, + (num_heads_q + 1) * 128, + dtype=torch.bfloat16, + device="cuda", + ) + q_weight = torch.randn(128, dtype=torch.bfloat16, device="cuda") + k_weight = torch.randn(128, dtype=torch.bfloat16, device="cuda") + position_ids = torch.arange(num_tokens, dtype=torch.int32, device="cuda") + 8192 + within = torch.arange(num_tokens, dtype=torch.int32, device="cuda") % page_size + pages = torch.arange(num_tokens, dtype=torch.int32, device="cuda") + slots = pages * page_size + within + cache = _strided_cache(num_tokens) + + q_out = _run(qk, cache, slots, q_weight, k_weight, position_ids, num_heads_q) + q_ref, k_ref = _reference(qk, num_heads_q, q_weight, k_weight, position_ids) + k_out = cache[pages.long(), 0, within.long()] + + assert torch.equal(q_out.view(torch.uint8), q_ref.contiguous().view(torch.uint8)) + assert torch.equal(k_out.view(torch.uint8), k_ref.contiguous().view(torch.uint8)) + + +def test_minimax_m3_fp8_indexer_cuda_graph_replay_updates_outputs(): + torch.manual_seed(5678) + num_tokens = 16 + num_heads_q = 4 + qk = torch.randn( + num_tokens, + (num_heads_q + 1) * 128, + dtype=torch.bfloat16, + device="cuda", + ) + q_weight = torch.randn(128, dtype=torch.bfloat16, device="cuda") + k_weight = torch.randn(128, dtype=torch.bfloat16, device="cuda") + position_ids = torch.arange(num_tokens, dtype=torch.int32, device="cuda") + 4096 + pages = torch.arange(num_tokens, dtype=torch.int32, device="cuda") + within = (pages * 11) % 128 + slots = pages * 128 + within + cache = _strided_cache(num_tokens) + + for _ in range(3): + _run(qk, cache, slots, q_weight, k_weight, position_ids, num_heads_q) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + q_out = _run(qk, cache, slots, q_weight, k_weight, position_ids, num_heads_q) + + first_q = q_out.clone() + qk.copy_(torch.randn_like(qk)) + graph.replay() + torch.cuda.synchronize() + q_ref, k_ref = _reference(qk, num_heads_q, q_weight, k_weight, position_ids) + k_out = cache[pages.long(), 0, within.long()] + + assert not torch.equal(q_out.view(torch.uint8), first_q.view(torch.uint8)) + assert torch.equal(q_out.view(torch.uint8), q_ref.contiguous().view(torch.uint8)) + assert torch.equal(k_out.view(torch.uint8), k_ref.contiguous().view(torch.uint8))