From 2a0c68e7e4e32d9459948f5e9677a16ce3630299 Mon Sep 17 00:00:00 2001 From: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> Date: Wed, 22 Jul 2026 12:54:36 -0700 Subject: [PATCH] [None][perf] Fold q/k/v quantization into qknorm_rope_fused kernel & remove contiguous (#16699) 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 | 124 ++++++++++-- .../test_fused_qk_norm_rope.py | 100 +++++++++ 7 files changed, 458 insertions(+), 70 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.cu b/cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.cu index 6e26dfce65bf..1e45a0775258 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 @@ -295,17 +350,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 @@ -322,11 +368,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) { @@ -344,8 +392,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); @@ -357,34 +406,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 0bbb1ad1509a..ead436b72b6c 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 3b0bbb0d2b77..d319c12c722d 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 @@ -728,8 +728,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 8d3e0ddcac1e..0db6ea8c1675 100644 --- a/tensorrt_llm/_torch/models/modeling_minimaxm3.py +++ b/tensorrt_llm/_torch/models/modeling_minimaxm3.py @@ -714,6 +714,17 @@ def __init__( # quantization changes cache storage only, so this stays the compute 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, @@ -864,6 +875,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. @@ -874,9 +886,15 @@ def _fused_qk_norm_rope( channels, matching M3's whole-head norm with front partial RoPE, and leaves the V heads untouched. - Returns None with qkv unmodified when the kernel does not apply, so the - caller runs norm and RoPE separately: non-bf16 activations (the kernel - is bf16-only), missing position_ids, or no rotary embedding. + 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 + position_ids, or when no partial-RoPE rotary_emb exists. """ if position_ids is None or qkv.dtype != torch.bfloat16: return None @@ -891,6 +909,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, @@ -903,7 +947,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 @@ -925,6 +969,59 @@ 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.index_q_size, self.index_k_size], 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, @@ -1025,11 +1122,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 contiguity of the separate path; 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 " @@ -1237,7 +1333,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, @@ -1370,10 +1471,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=" @@ -1399,8 +1500,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.index_q_size, self.index_k_size], 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, + )