From f008ad0841311c9a9222155b666f25472b6881e8 Mon Sep 17 00:00:00 2001 From: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> Date: Tue, 21 Jul 2026 22:48:31 +0000 Subject: [PATCH] [None][perf] Fold q/k/v quantization into qknorm_rope_fused kernel & remove contiguous calls Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> --- .../kernels/fusedQKNormRopeKernel.cu | 191 +++++++++++++----- .../kernels/fusedQKNormRopeKernel.h | 20 ++ cpp/tensorrt_llm/thop/fusedQKNormRopeOp.cpp | 72 +++++++ .../attention_backend/fmha/msa_sparse_gqa.py | 14 +- .../sparse/minimax_m3/msa_backend.py | 7 +- .../_torch/models/modeling_minimaxm3.py | 123 +++++++++-- .../test_fused_qk_norm_rope.py | 100 +++++++++ 7 files changed, 457 insertions(+), 70 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.cu b/cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.cu index e06b0f200e4b..8987234060fb 100644 --- a/cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.cu +++ b/cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.cu @@ -24,6 +24,7 @@ #include #include #include +#include TRTLLM_NAMESPACE_BEGIN @@ -56,15 +57,52 @@ __device__ __forceinline__ float selectMRopePosId(int const* position_ids, int t return static_cast(position_ids[sec * num_tokens + tokenIdx]); } -// Perform per-head QK Norm and RoPE in a single kernel. +// Store a per-thread run of `numElemsPerThread` float elements to the output +// head, converting to the output dtype. BF16 uses the packed uint vector store; +// FP8 E4M3 packs pairs via __nv_fp8x2_e4m3 (saturating round-to-nearest, matching +// torch's .to(torch.float8_e4m3fn)). +template +__device__ __forceinline__ void storeHeadElements( + OutT* out, int offsetThread, float const (&elements)[numElemsPerThread]) +{ + using vec_T = typename tensorrt_llm::common::packed_as::type; + if constexpr (std::is_same_v) + { + vec_T vec; + for (int i = 0; i < vecSize; i++) + { + __nv_bfloat162 vals = __float22bfloat162_rn(make_float2(elements[2 * i], elements[2 * i + 1])); + reinterpret_cast<__nv_bfloat162&>(*(reinterpret_cast(&vec) + i)) = vals; + } + *reinterpret_cast(&out[offsetThread]) = vec; + } + else // __nv_fp8_e4m3 + { + static_assert(numElemsPerThread % 2 == 0, "FP8 store expects an even element count per thread"); +#pragma unroll + for (int i = 0; i < numElemsPerThread; i += 2) + { + __nv_fp8x2_e4m3 packed(make_float2(elements[i], elements[i + 1])); + reinterpret_cast<__nv_fp8x2_storage_t*>(&out[offsetThread])[i / 2] = packed.__x; + } + } +} + +// Perform per-head QK Norm and RoPE in a single kernel, reading a BF16 input and +// writing the result to a (possibly different-dtype) output buffer. // head_dim: the dimension of each head // interleave: interleave=!is_neox. -template +// OutT: output element type (__nv_bfloat16 for in-place/BF16, __nv_fp8_e4m3 for FP8). +// When process_v is true, V heads are copy-cast into the output (no norm/RoPE); +// otherwise only Q/K heads are processed and V output slots are left untouched. +template __global__ void fusedQKNormRopeKernel( - __nv_bfloat16* qkv, // Combined QKV tensor [num_tokens, (num_heads_q+num_heads_k+num_heads_v)*head_dim] + __nv_bfloat16 const* qkv_in, // Combined QKV input [num_tokens, (num_heads_q+num_heads_k+num_heads_v)*head_dim] + OutT* qkv_out, // Output buffer, same layout as qkv_in int const num_heads_q, // Number of query heads int const num_heads_k, // Number of key heads int const num_heads_v, // Number of value heads + bool const process_v, // Whether to copy-cast V heads into qkv_out int const rotary_dim, // Dimension for RoPE float const eps, // Epsilon for RMS normalization __nv_bfloat16 const* q_weight, // RMSNorm weights for query @@ -95,17 +133,37 @@ __global__ void fusedQKNormRopeKernel( // Total number of attention heads (Q and K) int const total_qk_heads = num_heads_q + num_heads_k; + // Heads actually processed by this launch: Q + K, plus V when copy-casting. + int const total_proc_heads = total_qk_heads + (process_v ? num_heads_v : 0); - // Determine which token and head type (Q or K) this warp processes - int const tokenIdx = globalWarpIdx / total_qk_heads; - int const localHeadIdx = globalWarpIdx % total_qk_heads; + // Determine which token and head this warp processes + int const tokenIdx = globalWarpIdx / total_proc_heads; + int const localHeadIdx = globalWarpIdx % total_proc_heads; // Skip if this warp is assigned beyond the number of tokens if (tokenIdx >= num_tokens) return; bool const isQ = localHeadIdx < num_heads_q; - int const headIdx = isQ ? localHeadIdx : localHeadIdx - num_heads_q; + bool const isV = localHeadIdx >= total_qk_heads; + // headIdx is the head's index within its own (Q/K/V) segment. + int headIdx; + int segStart; // element offset of the segment start within a token row + if (isQ) + { + headIdx = localHeadIdx; + segStart = 0; + } + else if (!isV) + { + headIdx = localHeadIdx - num_heads_q; + segStart = num_heads_q * head_dim; + } + else + { + headIdx = localHeadIdx - total_qk_heads; + segStart = total_qk_heads * head_dim; + } int const num_heads = num_heads_q + num_heads_k + num_heads_v; @@ -119,25 +177,15 @@ __global__ void fusedQKNormRopeKernel( constexpr int vecSize = elemSizeBytes / 4; // Use packed_as to perform loading/saving. using vec_T = typename tensorrt_llm::common::packed_as::type; - int offsetWarp; // Offset for the warp - if (isQ) - { - // Q segment: token offset + head offset within Q segment - offsetWarp = tokenIdx * num_heads * head_dim + headIdx * head_dim; - } - else - { - // K segment: token offset + entire Q segment + head offset within K segment - offsetWarp = tokenIdx * num_heads * head_dim + num_heads_q * head_dim + headIdx * head_dim; - } + int const offsetWarp = tokenIdx * num_heads * head_dim + segStart + headIdx * head_dim; int offsetThread = offsetWarp + laneId * numElemsPerThread; // Sum of squares for RMSNorm float sumOfSquares = 0.0f; - // Load. + // Load from the BF16 input. { - vec_T vec = *reinterpret_cast(&qkv[offsetThread]); + vec_T vec = *reinterpret_cast(&qkv_in[offsetThread]); for (int i = 0; i < vecSize; i++) { float2 vals = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162*>(reinterpret_cast(&vec) + i)); @@ -149,6 +197,13 @@ __global__ void fusedQKNormRopeKernel( } } + // V heads are copy-cast only: skip norm and RoPE and store the raw values. + if (isV) + { + storeHeadElements(qkv_out, offsetThread, elements); + return; + } + if (is_qk_norm) { // Reduce sum across warp using the utility function @@ -284,17 +339,8 @@ __global__ void fusedQKNormRopeKernel( } } - // Store. - { - vec_T vec; - for (int i = 0; i < vecSize; i++) - { - __nv_bfloat162 vals = __float22bfloat162_rn(make_float2(elements[2 * i], elements[2 * i + 1])); - reinterpret_cast<__nv_bfloat162&>(*(reinterpret_cast(&vec) + i)) = vals; - } - vec_T* outputPtr = reinterpret_cast(&qkv[offsetThread]); - *outputPtr = vec; - } + // Store to the (templated) output. + storeHeadElements(qkv_out, offsetThread, elements); } // Borrowed from @@ -311,11 +357,13 @@ __global__ void fusedQKNormRopeKernel( __VA_ARGS__ \ } -void launchFusedQKNormRope(void* qkv, int const num_tokens, int const num_heads_q, int const num_heads_k, - int const num_heads_v, int const head_dim, int const rotary_dim, float const eps, void const* q_weight, - void const* k_weight, float const base, bool const interleave, int const* position_ids, float factor, float low, - float high, float attention_factor, cudaStream_t stream, bool is_qk_norm, bool use_gemma, bool use_mrope, - int mrope_section1, int mrope_section2) +template +static void launchFusedQKNormRopeImpl(__nv_bfloat16 const* qkv_in, OutT* qkv_out, bool const process_v, + int const num_tokens, int const num_heads_q, int const num_heads_k, int const num_heads_v, int const head_dim, + int const rotary_dim, float const eps, __nv_bfloat16 const* q_weight, __nv_bfloat16 const* k_weight, + float const base, bool const interleave, int const* position_ids, float factor, float low, float high, + float attention_factor, cudaStream_t stream, bool is_qk_norm, bool use_gemma, bool use_mrope, int mrope_section1, + int mrope_section2) { if (factor == 1.0f) { @@ -333,8 +381,9 @@ void launchFusedQKNormRope(void* qkv, int const num_tokens, int const num_heads_ constexpr int blockSize = 256; int const warpsPerBlock = blockSize / 32; - int const totalQKHeads = num_heads_q + num_heads_k; - int const totalWarps = num_tokens * totalQKHeads; + // Q + K heads, plus V heads when copy-casting them into the output. + int const totalProcHeads = num_heads_q + num_heads_k + (process_v ? num_heads_v : 0); + int const totalWarps = num_tokens * totalProcHeads; int const gridSize = common::divUp(totalWarps, warpsPerBlock); dim3 gridDim(gridSize); @@ -346,34 +395,70 @@ void launchFusedQKNormRope(void* qkv, int const num_tokens, int const num_heads_ { case 64: DISPATCH_INTERLEAVE(interleave, INTERLEAVE, { - fusedQKNormRopeKernel<64, INTERLEAVE> - <<>>(reinterpret_cast<__nv_bfloat16*>(qkv), num_heads_q, num_heads_k, - num_heads_v, rotary_dim, eps, reinterpret_cast<__nv_bfloat16 const*>(q_weight), - reinterpret_cast<__nv_bfloat16 const*>(k_weight), base, position_ids, num_tokens, factor, low, high, - attention_factor, is_qk_norm, use_gemma, use_mrope, mrope_section1, mrope_section2); + fusedQKNormRopeKernel<64, INTERLEAVE, OutT><<>>(qkv_in, qkv_out, num_heads_q, + num_heads_k, num_heads_v, process_v, rotary_dim, eps, q_weight, k_weight, base, position_ids, + num_tokens, factor, low, high, attention_factor, is_qk_norm, use_gemma, use_mrope, mrope_section1, + mrope_section2); }); break; case 128: DISPATCH_INTERLEAVE(interleave, INTERLEAVE, { - fusedQKNormRopeKernel<128, INTERLEAVE> - <<>>(reinterpret_cast<__nv_bfloat16*>(qkv), num_heads_q, num_heads_k, - num_heads_v, rotary_dim, eps, reinterpret_cast<__nv_bfloat16 const*>(q_weight), - reinterpret_cast<__nv_bfloat16 const*>(k_weight), base, position_ids, num_tokens, factor, low, high, - attention_factor, is_qk_norm, use_gemma, use_mrope, mrope_section1, mrope_section2); + fusedQKNormRopeKernel<128, INTERLEAVE, OutT><<>>(qkv_in, qkv_out, num_heads_q, + num_heads_k, num_heads_v, process_v, rotary_dim, eps, q_weight, k_weight, base, position_ids, + num_tokens, factor, low, high, attention_factor, is_qk_norm, use_gemma, use_mrope, mrope_section1, + mrope_section2); }); break; case 256: DISPATCH_INTERLEAVE(interleave, INTERLEAVE, { - fusedQKNormRopeKernel<256, INTERLEAVE> - <<>>(reinterpret_cast<__nv_bfloat16*>(qkv), num_heads_q, num_heads_k, - num_heads_v, rotary_dim, eps, reinterpret_cast<__nv_bfloat16 const*>(q_weight), - reinterpret_cast<__nv_bfloat16 const*>(k_weight), base, position_ids, num_tokens, factor, low, high, - attention_factor, is_qk_norm, use_gemma, use_mrope, mrope_section1, mrope_section2); + fusedQKNormRopeKernel<256, INTERLEAVE, OutT><<>>(qkv_in, qkv_out, num_heads_q, + num_heads_k, num_heads_v, process_v, rotary_dim, eps, q_weight, k_weight, base, position_ids, + num_tokens, factor, low, high, attention_factor, is_qk_norm, use_gemma, use_mrope, mrope_section1, + mrope_section2); }); break; default: TLLM_THROW("Unsupported head dimension for fusedQKNormRope: %d", head_dim); } } + +void launchFusedQKNormRope(void* qkv, int const num_tokens, int const num_heads_q, int const num_heads_k, + int const num_heads_v, int const head_dim, int const rotary_dim, float const eps, void const* q_weight, + void const* k_weight, float const base, bool const interleave, int const* position_ids, float factor, float low, + float high, float attention_factor, cudaStream_t stream, bool is_qk_norm, bool use_gemma, bool use_mrope, + int mrope_section1, int mrope_section2) +{ + // In-place BF16: input and output alias the same buffer; V is left untouched. + launchFusedQKNormRopeImpl<__nv_bfloat16>(reinterpret_cast<__nv_bfloat16 const*>(qkv), + reinterpret_cast<__nv_bfloat16*>(qkv), /*process_v=*/false, num_tokens, num_heads_q, num_heads_k, num_heads_v, + head_dim, rotary_dim, eps, reinterpret_cast<__nv_bfloat16 const*>(q_weight), + reinterpret_cast<__nv_bfloat16 const*>(k_weight), base, interleave, position_ids, factor, low, high, + attention_factor, stream, is_qk_norm, use_gemma, use_mrope, mrope_section1, mrope_section2); +} + +void launchFusedQKNormRopeOut(void const* qkv_in, void* qkv_out, bool out_fp8, bool process_v, int const num_tokens, + int const num_heads_q, int const num_heads_k, int const num_heads_v, int const head_dim, int const rotary_dim, + float const eps, void const* q_weight, void const* k_weight, float const base, bool const interleave, + int const* position_ids, float factor, float low, float high, float attention_factor, cudaStream_t stream, + bool is_qk_norm, bool use_gemma, bool use_mrope, int mrope_section1, int mrope_section2) +{ + auto const* in = reinterpret_cast<__nv_bfloat16 const*>(qkv_in); + auto const* qw = reinterpret_cast<__nv_bfloat16 const*>(q_weight); + auto const* kw = reinterpret_cast<__nv_bfloat16 const*>(k_weight); + if (out_fp8) + { + launchFusedQKNormRopeImpl<__nv_fp8_e4m3>(in, reinterpret_cast<__nv_fp8_e4m3*>(qkv_out), process_v, num_tokens, + num_heads_q, num_heads_k, num_heads_v, head_dim, rotary_dim, eps, qw, kw, base, interleave, position_ids, + factor, low, high, attention_factor, stream, is_qk_norm, use_gemma, use_mrope, mrope_section1, + mrope_section2); + } + else + { + launchFusedQKNormRopeImpl<__nv_bfloat16>(in, reinterpret_cast<__nv_bfloat16*>(qkv_out), process_v, num_tokens, + num_heads_q, num_heads_k, num_heads_v, head_dim, rotary_dim, eps, qw, kw, base, interleave, position_ids, + factor, low, high, attention_factor, stream, is_qk_norm, use_gemma, use_mrope, mrope_section1, + mrope_section2); + } +} } // namespace kernels TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.h b/cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.h index 4e2421cb57a2..4eaaf77cd9f0 100644 --- a/cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.h +++ b/cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.h @@ -51,6 +51,26 @@ void launchFusedQKNormRope( int mrope_section1, // mrope_section[1] (height) int mrope_section2); // mrope_section[2] (width) +// Out-of-place variant of launchFusedQKNormRope that reads a BF16 qkv input and +// writes the result to a separate output buffer, optionally as FP8 E4M3. +// +// This folds the FP8 activation-quant into the norm+RoPE epilogue so callers do +// not need separate cast kernels for Q/K/V. Q and K get RMSNorm + RoPE; V (when +// process_v is true) is copy-cast only. The output layout matches the input: +// [num_tokens, (num_heads_q+num_heads_k+num_heads_v)*head_dim]. +// +// out_fp8=false writes BF16 output (a plain out-of-place variant); out_fp8=true +// writes __nv_fp8_e4m3. When out_fp8 is false, process_v must be true only if the +// caller wants V copied into the output (otherwise V slots are left untouched). +void launchFusedQKNormRopeOut(void const* qkv_in, // BF16 input [num_tokens, total_heads*head_dim] + void* qkv_out, // Output buffer (BF16 or FP8 E4M3), same layout as input + bool out_fp8, // Whether qkv_out is FP8 E4M3 (else BF16) + bool process_v, // Whether to copy-cast the V heads into qkv_out + int const num_tokens, int const num_heads_q, int const num_heads_k, int const num_heads_v, int const head_dim, + int const rotary_dim, float const eps, void const* q_weight, void const* k_weight, float const base, + bool const interleave, int const* position_ids, float factor, float low, float high, float attention_factor, + cudaStream_t stream, bool is_qk_norm, bool use_gemma, bool use_mrope, int mrope_section1, int mrope_section2); + } // namespace kernels TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/thop/fusedQKNormRopeOp.cpp b/cpp/tensorrt_llm/thop/fusedQKNormRopeOp.cpp index 4ff4cff6d3ba..5d7d84f43ce6 100644 --- a/cpp/tensorrt_llm/thop/fusedQKNormRopeOp.cpp +++ b/cpp/tensorrt_llm/thop/fusedQKNormRopeOp.cpp @@ -89,6 +89,66 @@ void fused_qk_norm_rope( static_cast(mrope_section1), static_cast(mrope_section2)); } +// Out-of-place FP8 variant of fused_qk_norm_rope. +// +// Reads a BF16 qkv tensor, applies RMSNorm + RoPE to Q/K and copy-casts V, and +// returns a new FP8 (E4M3) tensor of the same shape. This folds the FP8 +// activation-quant into the norm+RoPE epilogue so callers (e.g. the MiniMax-M3 +// MSA path with an FP8 KV cache) do not need separate q/k/v cast kernels. +torch::Tensor fused_qk_norm_rope_to_fp8(torch::Tensor const& qkv, // [num_tokens, (num_q+num_k+num_v)*head_dim] BF16 + int64_t num_heads_q, int64_t num_heads_k, int64_t num_heads_v, int64_t head_dim, int64_t rotary_dim, double eps, + torch::Tensor const& q_weight, torch::Tensor const& k_weight, double base, bool is_neox, + torch::Tensor const& position_ids, double factor, double low, double high, double attention_factor, bool is_qk_norm, + bool use_gemma, bool use_mrope, int64_t mrope_section1, int64_t mrope_section2) +{ + TORCH_CHECK(qkv.dim() == 2, "QKV tensor must be 2D: [num_tokens, (num_heads_q+num_heads_k+num_heads_v)*head_dim]"); + TORCH_CHECK(position_ids.dim() == 1 || (position_ids.dim() == 2 && position_ids.size(0) == 3), + "Position IDs must be 1D [num_tokens] (plain RoPE) or 2D [3, num_tokens] (mRoPE)"); + TORCH_CHECK(!use_mrope || position_ids.dim() == 2, "use_mrope requires 2D [3, num_tokens] position_ids"); + TORCH_CHECK(q_weight.dim() == 1, "Query weights must be 1D: [head_dim]"); + TORCH_CHECK(k_weight.dim() == 1, "Key weights must be 1D: [head_dim]"); + TORCH_CHECK(q_weight.size(0) == head_dim, "Query weights size must match head dimension"); + TORCH_CHECK(k_weight.size(0) == head_dim, "Key weights size must match head dimension"); + + CHECK_INPUT(qkv, torch::kBFloat16); + CHECK_INPUT(position_ids, torch::kInt32); + CHECK_INPUT(q_weight, torch::kBFloat16); + CHECK_INPUT(k_weight, torch::kBFloat16); + + int64_t num_tokens = qkv.size(0); + TORCH_CHECK(position_ids.size(-1) == num_tokens, "Number of tokens in position_ids must match QKV"); + + int64_t total_heads = num_heads_q + num_heads_k + num_heads_v; + TORCH_CHECK( + qkv.size(1) == total_heads * head_dim, "QKV tensor size must match total number of heads and head dimension"); + + auto out = torch::empty({num_tokens, total_heads * head_dim}, qkv.options().dtype(torch::kFloat8_e4m3fn)); + + auto stream = at::cuda::getCurrentCUDAStream(qkv.get_device()); + + tensorrt_llm::kernels::launchFusedQKNormRopeOut(qkv.data_ptr(), out.data_ptr(), /*out_fp8=*/true, + /*process_v=*/true, static_cast(num_tokens), static_cast(num_heads_q), static_cast(num_heads_k), + static_cast(num_heads_v), static_cast(head_dim), static_cast(rotary_dim), + static_cast(eps), q_weight.data_ptr(), k_weight.data_ptr(), static_cast(base), !is_neox, + reinterpret_cast(position_ids.data_ptr()), static_cast(factor), static_cast(low), + static_cast(high), static_cast(attention_factor), stream, is_qk_norm, use_gemma, use_mrope, + static_cast(mrope_section1), static_cast(mrope_section2)); + + return out; +} + +// Meta (fake) implementation for torch.compile / tracing: only shape+dtype. +torch::Tensor fused_qk_norm_rope_to_fp8_meta(torch::Tensor const& qkv, int64_t num_heads_q, int64_t num_heads_k, + int64_t num_heads_v, int64_t head_dim, int64_t /*rotary_dim*/, double /*eps*/, torch::Tensor const& /*q_weight*/, + torch::Tensor const& /*k_weight*/, double /*base*/, bool /*is_neox*/, torch::Tensor const& /*position_ids*/, + double /*factor*/, double /*low*/, double /*high*/, double /*attention_factor*/, bool /*is_qk_norm*/, + bool /*use_gemma*/, bool /*use_mrope*/, int64_t /*mrope_section1*/, int64_t /*mrope_section2*/) +{ + int64_t num_tokens = qkv.size(0); + int64_t total_heads = num_heads_q + num_heads_k + num_heads_v; + return torch::empty({num_tokens, total_heads * head_dim}, qkv.options().dtype(torch::kFloat8_e4m3fn)); +} + // Register the PyTorch operators TORCH_LIBRARY_FRAGMENT(trtllm, m) { @@ -98,12 +158,24 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m) "eps, Tensor q_weight, Tensor k_weight, float base, bool is_neox, Tensor position_ids, float factor, float " "low, float high, float attention_factor, bool is_qk_norm, bool use_gemma, bool use_mrope, int " "mrope_section1, int mrope_section2) -> ()"); + m.def( + "fused_qk_norm_rope_to_fp8(Tensor qkv, int num_heads_q, int num_heads_k, int num_heads_v, int head_dim, int " + "rotary_dim, float eps, Tensor q_weight, Tensor k_weight, float base, bool is_neox, Tensor position_ids, float " + "factor, float low, float high, float attention_factor, bool is_qk_norm, bool use_gemma, bool use_mrope, int " + "mrope_section1, int mrope_section2) -> Tensor"); } // Register the CUDA implementation TORCH_LIBRARY_IMPL(trtllm, CUDA, m) { m.impl("fused_qk_norm_rope", &fused_qk_norm_rope); + m.impl("fused_qk_norm_rope_to_fp8", &fused_qk_norm_rope_to_fp8); +} + +// Register the Meta implementation (shape/dtype inference for torch.compile). +TORCH_LIBRARY_IMPL(trtllm, Meta, m) +{ + m.impl("fused_qk_norm_rope_to_fp8", &fused_qk_norm_rope_to_fp8_meta); } } // namespace torch_ext diff --git a/tensorrt_llm/_torch/attention_backend/fmha/msa_sparse_gqa.py b/tensorrt_llm/_torch/attention_backend/fmha/msa_sparse_gqa.py index 5578f36e3ca8..7917f96c71dc 100644 --- a/tensorrt_llm/_torch/attention_backend/fmha/msa_sparse_gqa.py +++ b/tensorrt_llm/_torch/attention_backend/fmha/msa_sparse_gqa.py @@ -139,16 +139,24 @@ def run_msa_paged_gqa( kv_cache_manager, layer_idx, metadata.msa_out_cache_loc[:num_tokens], k, v ) - q_view = q.view(num_tokens, attn.num_heads, head_dim) + # q may be a strided column-view of a fused [q|k|v] buffer (the model skips + # the split contiguous copy on this path). fmha_sm100 reads q's real strides + # through the TMA descriptor, so reshape here is a zero-copy view for that + # layout; it only falls back to a copy for an otherwise non-viewable q. + q_view = q.reshape(num_tokens, attn.num_heads, head_dim) + # output is freshly allocated and contiguous; view keeps out_view aliasing it + # so the kernel's in-place write lands in the caller's buffer. out_view = output.view(num_tokens, attn.num_heads, head_dim) k_paged, v_paged = msa_paged_kv(kv_cache_manager, layer_idx) sm_scale = (head_dim**-0.5) / float(attn.q_scaling) # The fmha_sm100 variant is chosen from q.dtype and shares one dtype across # q/k/v, so q must be FP8 to match an FP8 paged K/V. MiniMax-M3 has no - # KV-cache scales, so the scale is 1.0 and this is a plain E4M3 cast. + # KV-cache scales, so the scale is 1.0 and this is a plain E4M3 cast. When the + # model's fused QK-norm+RoPE already emitted FP8 q/k/v (the FP8-KV fast path), + # this .to() is a no-op; it stays as a safety net for callers that pass bf16 q. use_fp8 = k_paged.dtype == torch.float8_e4m3fn - if use_fp8: + if use_fp8 and q_view.dtype != torch.float8_e4m3fn: q_view = q_view.to(torch.float8_e4m3fn) run_msa_sparse_gqa( 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 a7313417b010..683ee67372f0 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 @@ -997,8 +997,11 @@ 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]) - 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) + # idx_q and idx_k may be strided column-views of a fused buffer, so + # 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) diff --git a/tensorrt_llm/_torch/models/modeling_minimaxm3.py b/tensorrt_llm/_torch/models/modeling_minimaxm3.py index cd69ec4bba0b..a993e426abac 100644 --- a/tensorrt_llm/_torch/models/modeling_minimaxm3.py +++ b/tensorrt_llm/_torch/models/modeling_minimaxm3.py @@ -712,6 +712,17 @@ def __init__( # (fp8/fp4) only changes cache storage, not the activation dtype. self.attn_activation_dtype = config.torch_dtype + # Whether the main K/V cache is stored as FP8 E4M3. When true and the MSA + # backend is active, the fused QK-norm+RoPE kernel emits FP8 q/k/v + # directly, so the separate q-cast and cache-write casts collapse into one + # kernel. This only moves where the E4M3 conversion happens, not its value. + quant_config = getattr(model_config, "quant_config", None) + self.main_kv_is_fp8 = bool( + quant_config is not None + and quant_config.quant_mode is not None + and quant_config.quant_mode.has_fp8_kv_cache() + ) + # Per-head Gemma RMSNorm — one set of weights shared across heads. self.q_norm = RMSNorm( hidden_size=self.head_dim_value, @@ -867,6 +878,7 @@ def _fused_qk_norm_rope( head_dim: int, q_norm: RMSNorm, k_norm: RMSNorm, + out_fp8: bool = False, ) -> Optional[torch.Tensor]: """Fuse per-head Gemma RMSNorm and partial RoPE into one kernel. @@ -877,6 +889,11 @@ def _fused_qk_norm_rope( whole-head norm with front partial RoPE, and applies the same Gemma (1 + weight) scaling as apply_qk_norm. + When out_fp8 is True, an out-of-place FP8 variant runs instead: it reads + the bf16 fused qkv and returns a fresh FP8 E4M3 tensor with Q/K normed + and roped and V copy-cast, folding the FP8 activation quant into the + norm+RoPE epilogue. The input qkv is left untouched. + Returns None, leaving qkv untouched, when the fused path does not apply so callers fall back to separate norm and RoPE. This happens when activations are not bf16 (the kernel is bf16-only), when RoPE has no @@ -895,6 +912,32 @@ def _fused_qk_norm_rope( rotary_dim = int(self.pos_embd_params.rope.dim) # The kernel assumes a contiguous [num_tokens, total_heads * head_dim]. qkv = qkv.contiguous() + position_ids_i32 = position_ids.reshape(-1).contiguous().to(torch.int32) + if out_fp8: + # Out-of-place FP8 variant: returns a fresh E4M3 [q|k|v] tensor. + return torch.ops.trtllm.fused_qk_norm_rope_to_fp8( + qkv, + num_heads_q, + num_heads_k, + num_heads_v, + head_dim, + rotary_dim, + q_norm.variance_epsilon, + q_norm.weight, + k_norm.weight, + self.pos_embd_params.rope.theta, + self.pos_embd_params.is_neox, + position_ids_i32, + 1.0, # factor: no YARN (M3 has no rope_scaling) + 0.0, # low + 0.0, # high + 1.0, # attention_factor + True, # is_qk_norm + self.use_gemma_norm, # use_gemma + False, # use_mrope + 0, # mrope_section1 + 0, # mrope_section2 + ) torch.ops.trtllm.fused_qk_norm_rope( qkv, num_heads_q, @@ -907,7 +950,7 @@ def _fused_qk_norm_rope( k_norm.weight, self.pos_embd_params.rope.theta, self.pos_embd_params.is_neox, - position_ids.reshape(-1).contiguous().to(torch.int32), + position_ids_i32, 1.0, # factor: no YARN (M3 has no rope_scaling) 0.0, # low 0.0, # high @@ -929,6 +972,62 @@ def _expect_fused_qk_norm_rope(self, position_ids: Optional[torch.Tensor]) -> bo """ return self.attn_activation_dtype == torch.bfloat16 and position_ids is not None + def _msa_backend_active(self) -> bool: + """Whether the MSA fmha_sm100 backend handles this layer's attention. + + Used to gate MSA-only main-branch optimizations (FP8 q/k/v emission and + skipping the split q/k contiguous copies). The Triton/SDPA reference + backends are left on the conservative contiguous/bf16 path. + """ + return isinstance(self.attn, MiniMaxM3MsaSparseAttention) + + def _emit_fp8_main_qkv(self) -> bool: + """Whether the main-branch fused QK-norm+RoPE should emit FP8 q/k/v. + + Only the MSA backend consumes an FP8 paged K/V cache directly (the + kernel variant shares one dtype across q/k/v). The Triton/SDPA reference + paths keep bf16, so gate on the MSA backend being active in addition to + the cache being FP8. The index branch always stays bf16. + """ + return self.main_kv_is_fp8 and self._msa_backend_active() + + def _split_main_qkv( + self, fused_qkv: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Split the fused [q|k|v] buffer into per-tensor q/k/v. + + The MSA fmha_sm100 kernel reads q with its real strides through the TMA + descriptor (both the dense and the packed-decode Q load paths address + global memory with q.stride()), and k/v are scattered into the paged + cache by an indexed copy that tolerates a strided source. So on the MSA + backend the split column-views can be handed over directly with no + contiguous copy. Other backends keep the previous contiguity (q/k made + contiguous, v a column-slice view). + """ + q, k, v = fused_qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + if self._msa_backend_active(): + return q, k, v + return q.contiguous(), k.contiguous(), v + + def _split_index_qk(self, fused_idx: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + """Split the fused [idx_q|idx_k] buffer into per-tensor idx_q/idx_k. + + The index analogue of _split_main_qkv. The index cache is bf16, so there + is no FP8 output variant here. On the MSA backend the split is stride + safe: idx_q feeds the fmha_sm100 proxy, which loads Q via TMA using its + real strides, and idx_k is scattered into the paged index-K cache by an + indexed copy that tolerates a strided source. The split column-views are + handed over directly with no contiguous copy. Other backends fall back to + contiguous idx_q and idx_k. + """ + idx_q, idx_k = fused_idx.split( + [self.sparse_num_index_heads * self.sparse_index_dim, self.sparse_index_dim], + dim=-1, + ) + if self._msa_backend_active(): + return idx_q, idx_k + return idx_q.contiguous(), idx_k.contiguous() + def forward( self, position_ids: Optional[torch.IntTensor] = None, @@ -1033,11 +1132,10 @@ def _dense_forward( head_dim=self.head_dim, q_norm=self.q_norm, k_norm=self.k_norm, + out_fp8=self._emit_fp8_main_qkv(), ) if fused_qkv is not None: - q, k, v = fused_qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) - # Match the fallback contiguity; V stays a column-slice view. - q, k = q.contiguous(), k.contiguous() + q, k, v = self._split_main_qkv(fused_qkv) else: assert not self._expect_fused_qk_norm_rope(position_ids), ( f"MiniMax-M3 dense attention (layer {self.layer_idx}) expected the " @@ -1268,7 +1366,12 @@ def _forward_attention_core( idx_k: Optional[torch.Tensor], attn_metadata: AttentionMetadata, ) -> torch.Tensor: - output = q.new_empty((q.shape[0], self.num_heads * self.head_dim)) + # Attention output is always the compute dtype (bf16); q may be FP8 when + # the MSA FP8-KV path emits FP8 q/k/v, so pin the dtype rather than + # inheriting it from q. + output = q.new_empty( + (q.shape[0], self.num_heads * self.head_dim), dtype=self.attn_activation_dtype + ) if self.register_to_config and is_torch_compiling(): minimax_m3_attn_custom_op_inplace( q, @@ -1406,10 +1509,10 @@ def _main_norm_rope(): head_dim=self.head_dim, q_norm=self.q_norm, k_norm=self.k_norm, + out_fp8=self._emit_fp8_main_qkv(), ) if fused_qkv is not None: - q, k, v = fused_qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) - return q.contiguous(), k.contiguous(), v + return self._split_main_qkv(fused_qkv) assert not self._expect_fused_qk_norm_rope(position_ids), ( f"MiniMax-M3 sparse attention (layer {self.layer_idx}) expected the " f"fused QK-norm+RoPE kernel (bf16 activations, head_dim=" @@ -1435,11 +1538,7 @@ def _index_norm_rope(): k_norm=self.index_k_norm, ) if fused_idx is not None: - idx_q, idx_k = fused_idx.split( - [self.sparse_num_index_heads * self.sparse_index_dim, self.sparse_index_dim], - dim=-1, - ) - return idx_q.contiguous(), idx_k.contiguous() + return self._split_index_qk(fused_idx) assert not self._expect_fused_qk_norm_rope(position_ids), ( f"MiniMax-M3 sparse index branch (layer {self.layer_idx}) expected the " f"fused QK-norm+RoPE kernel (bf16 activations, index_dim=" diff --git a/tests/unittest/_torch/thop/parallel_hw_agnostic/test_fused_qk_norm_rope.py b/tests/unittest/_torch/thop/parallel_hw_agnostic/test_fused_qk_norm_rope.py index b33886a1c233..cbf0aa02a428 100644 --- a/tests/unittest/_torch/thop/parallel_hw_agnostic/test_fused_qk_norm_rope.py +++ b/tests/unittest/_torch/thop/parallel_hw_agnostic/test_fused_qk_norm_rope.py @@ -340,3 +340,103 @@ def test_fused_qk_norm_rope_gemma_mrope( ) torch.testing.assert_close(output, ref_output, rtol=5e-2, atol=1e-1) + + +# FP8 out-variant coverage. Includes an M3-like GQA shape (8 Q / 1 KV, head_dim +# 128) so the MiniMax-M3 FP8-KV path geometry is exercised directly. +fp8_num_heads_groups = [ + (16, 8, 8), + (32, 8, 8), + (8, 1, 1), # MiniMax-M3 sharded GQA (num_heads=8, num_kv_heads=1) +] + + +@pytest.mark.parametrize("head_dim", [64, 128]) +@pytest.mark.parametrize("num_heads_group", fp8_num_heads_groups) +@pytest.mark.parametrize("num_tokens", [1, 3, 8, 256]) +@pytest.mark.parametrize("is_neox", [False, True]) +@pytest.mark.parametrize("partial_rotary_factor", [1.0, 0.5]) +def test_fused_qk_norm_rope_to_fp8( + head_dim, num_heads_group, num_tokens, partial_rotary_factor, is_neox +): + """Test the FP8 out-variant of fused QK RMSNorm + RoPE. + + The op reads a BF16 qkv, applies RMSNorm + RoPE to Q/K and copy-casts V, and + returns a new FP8 (E4M3) tensor, folding the FP8 activation-quant into the + norm+RoPE epilogue. Verifies: + 1. The input qkv is left untouched (out-of-place). + 2. Output dtype is FP8 E4M3 with the same shape. + 3. The dequantized output matches the BF16 fused reference (Q/K normed+roped, + V unchanged) within FP8-appropriate tolerance. + """ + device = "cuda" + dtype = torch.bfloat16 + num_heads_q, num_heads_k, num_heads_v = num_heads_group + hidden_size = (num_heads_q + num_heads_k + num_heads_v) * head_dim + + torch.random.manual_seed(0) + qkv = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device) + qkv_ref = qkv.clone() + + position_ids = torch.arange(num_tokens, dtype=torch.int32, device=device) + 100 + q_weight = torch.randn(head_dim, dtype=dtype, device=device) * 5.0 + k_weight = torch.randn(head_dim, dtype=dtype, device=device) * 5.0 + + eps = 1e-5 + base = 10000.0 + factor, low, high, attention_factor = 1.0, 0, 0, 1.0 + rotary_dim = int(head_dim * partial_rotary_factor) + + out_fp8 = torch.ops.trtllm.fused_qk_norm_rope_to_fp8( + qkv, + num_heads_q, + num_heads_k, + num_heads_v, + head_dim, + rotary_dim, + eps, + q_weight, + k_weight, + base, + is_neox, + position_ids, + factor, + low, + high, + attention_factor, + True, # is_qk_norm + False, # use_gemma (standard RMSNorm reference below) + False, # use_mrope (plain RoPE) + 0, # mrope_section1 + 0, # mrope_section2 + ) + + assert out_fp8.dtype == torch.float8_e4m3fn + assert tuple(out_fp8.shape) == (num_tokens, hidden_size) + # Out-of-place: the BF16 input must be left byte-for-byte unchanged. + torch.testing.assert_close(qkv, qkv_ref, rtol=0.0, atol=0.0) + + ref_output = torch_ref_rms_norm_rope( + qkv_ref, + num_heads_q, + num_heads_k, + num_heads_v, + head_dim, + rotary_dim, + eps, + q_weight, + k_weight, + base, + is_neox, + position_ids, + ) + + # The op folds the E4M3 cast into the epilogue; compare the dequantized + # result against the BF16 reference with FP8-appropriate tolerance (E4M3 has + # 3 mantissa bits, so ~1/8 relative resolution). + torch.testing.assert_close( + out_fp8.float(), + ref_output.float(), + rtol=0.2, + atol=0.1, + )