From 2ba44374374467f9301055b337edd74c6201c46b 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 1/3] [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 535de64c049b..ff605cb90ffc 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 858348f914ba..a4770d21ec30 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 @@ -824,8 +824,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, + ) From 2aa9032f130d91aa3591bf876c8701193bdbd25d Mon Sep 17 00:00:00 2001 From: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> Date: Wed, 5 Aug 2026 18:42:09 -0700 Subject: [PATCH 2/3] [None][fix] Address review feedback on fused QK norm + RoPE FP8 path Bound rotary_dim to 1..head_dim before dispatch; rotary_dim == 0 passed the evenness check and then divided by zero deriving the RoPE frequencies. This matches the predicate the sibling Triton path already gates on. Replace launchFusedQKNormRopeOut with an FP8-only launchFusedQKNormRopeToFp8. The BF16 out-of-place branch had no caller and no test, and left V uninitialized for any genuinely out-of-place caller. Guard the remaining process_v=false case on input and output aliasing so that combination cannot be reached. Tighten the FP8 test tolerance from rtol=0.2 (3x looser than E4M3 needs, so it would not catch swapped norm weights) to 0.07, assert V matches torch's E4M3 cast bit-exactly, and parametrize use_gemma so the Gemma norm MiniMax-M3 actually ships is covered. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> --- .../kernels/fusedQKNormRopeKernel.cu | 51 ++++++++++--------- .../kernels/fusedQKNormRopeKernel.h | 18 +++---- cpp/tensorrt_llm/thop/fusedQKNormRopeOp.cpp | 12 ++--- .../test_fused_qk_norm_rope.py | 47 +++++++++++++---- 4 files changed, 75 insertions(+), 53 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.cu b/cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.cu index 1e45a0775258..bcd56342f931 100644 --- a/cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.cu +++ b/cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.cu @@ -59,8 +59,10 @@ __device__ __forceinline__ float selectMRopePosId(int const* position_ids, int t // 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)). +// FP8 E4M3 packs pairs via __nv_fp8x2_e4m3, i.e. round-to-nearest-even with +// __NV_SATFINITE clamping to +/-448. That matches torch's +// .to(torch.float8_e4m3fn) for in-range values; out-of-range values clamp to +// +/-448 here whereas torch produces NaN. template __device__ __forceinline__ void storeHeadElements( OutT* out, int offsetThread, float const (&elements)[numElemsPerThread]) @@ -69,6 +71,7 @@ __device__ __forceinline__ void storeHeadElements( if constexpr (std::is_same_v) { vec_T vec; +#pragma unroll for (int i = 0; i < vecSize; i++) { __nv_bfloat162 vals = __float22bfloat162_rn(make_float2(elements[2 * i], elements[2 * i + 1])); @@ -186,6 +189,7 @@ __global__ void fusedQKNormRopeKernel( // Load from the BF16 input. { vec_T vec = *reinterpret_cast(&qkv_in[offsetThread]); +#pragma unroll for (int i = 0; i < vecSize; i++) { float2 vals = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162*>(reinterpret_cast(&vec) + i)); @@ -381,7 +385,15 @@ static void launchFusedQKNormRopeImpl(__nv_bfloat16 const* qkv_in, OutT* qkv_out TLLM_CHECK(attention_factor == 1.0f); } - TLLM_CHECK_WITH_INFO(rotary_dim % 2 == 0, "rotary_dim must be even"); + // rotary_dim == 0 would divide by zero when deriving the RoPE frequencies, and + // rotary_dim > head_dim breaks the warp-level pairing assumptions. + TLLM_CHECK_WITH_INFO(rotary_dim > 0 && rotary_dim <= head_dim && rotary_dim % 2 == 0, + "rotary_dim must be positive, even and no greater than head_dim (got rotary_dim=%d, head_dim=%d)", rotary_dim, + head_dim); + // Skipping V is only well-defined in-place, where the V slots of the output + // already hold their final values; out-of-place they would stay uninitialized. + TLLM_CHECK_WITH_INFO(process_v || static_cast(qkv_in) == static_cast(qkv_out), + "process_v=false requires qkv_in and qkv_out to alias"); if (!interleave) { // To allow warp-level pairing for partial rope @@ -446,29 +458,18 @@ void launchFusedQKNormRope(void* qkv, int const num_tokens, int const num_heads_ 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) +void launchFusedQKNormRopeToFp8(void const* qkv_in, void* 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) { - 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); - } + // Out-of-place FP8: qkv_out is a fresh buffer, so V must be copy-cast too. + launchFusedQKNormRopeImpl<__nv_fp8_e4m3>(reinterpret_cast<__nv_bfloat16 const*>(qkv_in), + reinterpret_cast<__nv_fp8_e4m3*>(qkv_out), /*process_v=*/true, 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); } } // namespace kernels diff --git a/cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.h b/cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.h index 4eaaf77cd9f0..98701348479a 100644 --- a/cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.h +++ b/cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.h @@ -51,21 +51,15 @@ 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. +// Out-of-place FP8 variant of launchFusedQKNormRope that reads a BF16 qkv input +// and writes __nv_fp8_e4m3 to a separate output buffer. // // 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: +// not need separate cast kernels for Q/K/V. Q and K get RMSNorm + RoPE; V 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 +void launchFusedQKNormRopeToFp8(void const* qkv_in, // BF16 input [num_tokens, total_heads*head_dim] + void* qkv_out, // FP8 E4M3 output buffer, same layout as input 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, diff --git a/cpp/tensorrt_llm/thop/fusedQKNormRopeOp.cpp b/cpp/tensorrt_llm/thop/fusedQKNormRopeOp.cpp index 5d7d84f43ce6..72b912b07b23 100644 --- a/cpp/tensorrt_llm/thop/fusedQKNormRopeOp.cpp +++ b/cpp/tensorrt_llm/thop/fusedQKNormRopeOp.cpp @@ -126,12 +126,12 @@ torch::Tensor fused_qk_norm_rope_to_fp8(torch::Tensor const& qkv, // [num_tokens 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, + tensorrt_llm::kernels::launchFusedQKNormRopeToFp8(qkv.data_ptr(), out.data_ptr(), 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; 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 cbf0aa02a428..1dcfd8317ba4 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 @@ -20,6 +20,7 @@ def torch_ref_rms_norm_rope( base, is_neox, position_ids, + use_gemma=False, ): """ PyTorch reference implementation of RMSNorm+RoPE for verification. @@ -40,6 +41,7 @@ def torch_ref_rms_norm_rope( base: Base value for RoPE calculations is_neox: Whether to use NeoX style RoPE position_ids: Position IDs for RoPE of shape [num_tokens] + use_gemma: Whether QK norm is Gemma-style RMSNorm (scale by (1 + weight)) Returns: Combined tensor with Q and K parts normalized and RoPE applied @@ -64,8 +66,12 @@ def torch_ref_rms_norm_rope( v = qkv[:, q_size + k_size :] # Create and apply RMSNorm modules with custom weights - q_norm = RMSNorm(hidden_size=head_dim, eps=eps).to(qkv.device).to(qkv.dtype) - k_norm = RMSNorm(hidden_size=head_dim, eps=eps).to(qkv.device).to(qkv.dtype) + q_norm = ( + RMSNorm(hidden_size=head_dim, eps=eps, use_gemma=use_gemma).to(qkv.device).to(qkv.dtype) + ) + k_norm = ( + RMSNorm(hidden_size=head_dim, eps=eps, use_gemma=use_gemma).to(qkv.device).to(qkv.dtype) + ) # Set the weights to the provided weights q_norm.weight.data.copy_(q_weight) @@ -356,8 +362,11 @@ def test_fused_qk_norm_rope_gemma_mrope( @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]) +# MiniMax-M3 stores every RMSNorm as a Gemma norm, so use_gemma=True is the +# configuration that actually ships through this op. +@pytest.mark.parametrize("use_gemma", [False, True]) def test_fused_qk_norm_rope_to_fp8( - head_dim, num_heads_group, num_tokens, partial_rotary_factor, is_neox + head_dim, num_heads_group, num_tokens, partial_rotary_factor, is_neox, use_gemma ): """Test the FP8 out-variant of fused QK RMSNorm + RoPE. @@ -366,8 +375,10 @@ def test_fused_qk_norm_rope_to_fp8( 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. + 3. V is bit-exactly the E4M3 cast of the input V (pure copy-cast, so this + pins down the kernel's rounding against torch's). + 4. The dequantized Q/K output matches the BF16 fused reference within + FP8-appropriate tolerance. """ device = "cuda" dtype = torch.bfloat16 @@ -405,7 +416,7 @@ def test_fused_qk_norm_rope_to_fp8( high, attention_factor, True, # is_qk_norm - False, # use_gemma (standard RMSNorm reference below) + use_gemma, False, # use_mrope (plain RoPE) 0, # mrope_section1 0, # mrope_section2 @@ -416,6 +427,18 @@ def test_fused_qk_norm_rope_to_fp8( # 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) + # V is copy-cast only (no norm, no RoPE), so it must match torch's E4M3 cast + # exactly. This pins down the kernel's round-to-nearest-even FP8 store, which + # nothing else here verifies. The inputs are standard normal, so they stay + # well inside E4M3 range and the two casts' clamp/NaN behaviors don't diverge. + v_start = (num_heads_q + num_heads_k) * head_dim + torch.testing.assert_close( + out_fp8[:, v_start:].float(), + qkv_ref[:, v_start:].to(torch.float8_e4m3fn).float(), + rtol=0.0, + atol=0.0, + ) + ref_output = torch_ref_rms_norm_rope( qkv_ref, num_heads_q, @@ -429,14 +452,18 @@ def test_fused_qk_norm_rope_to_fp8( base, is_neox, position_ids, + use_gemma=use_gemma, ) - # 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). + # The op folds the E4M3 cast into the epilogue, so the dominant error is the + # single E4M3 rounding: 3 mantissa bits give a worst-case relative error of + # 2^-4 = 6.25%, plus the ~2^-9 the BF16 reference itself carries. The kernel + # accumulates in fp32 and casts straight to FP8, i.e. one rounding step fewer + # than the reference's fp32 -> bf16 -> fp8, so it should clear this + # comfortably. torch.testing.assert_close( out_fp8.float(), ref_output.float(), - rtol=0.2, + rtol=0.07, atol=0.1, ) From 31c4853b84f5cf5a7ecf441936e9f72876b0a683 Mon Sep 17 00:00:00 2001 From: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> Date: Wed, 5 Aug 2026 20:21:39 -0700 Subject: [PATCH 3/3] [None][chore] Trim review comments and dedupe validation Drop the .reshape() and dtype-guard changes in msa_sparse_gqa.py and msa_backend.py: .view() already handles the strided column-views the model now hands over, and .to() on an already-FP8 tensor is a no-op, so both files return to their original form. Cut the explanatory comments added across the kernel, op, model and test down to what is not already obvious from the code. Fold the duplicated input validation in fusedQKNormRopeOp.cpp into one helper shared by both operators, and reject negative head counts and out-of-int-range dimensions there before any shape arithmetic. Use static_cast rather than reinterpret_cast for the void* launcher conversions, tuple[...] instead of Tuple[...] in the new annotations, and rename the test's head-group constant to _FP8_NUM_HEADS_GROUPS. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com> --- .../kernels/fusedQKNormRopeKernel.cu | 52 ++++------ .../kernels/fusedQKNormRopeKernel.h | 10 +- cpp/tensorrt_llm/thop/fusedQKNormRopeOp.cpp | 95 ++++++++++--------- .../attention_backend/fmha/msa_sparse_gqa.py | 14 +-- .../sparse/minimax_m3/msa_backend.py | 7 +- .../_torch/models/modeling_minimaxm3.py | 61 ++++-------- .../test_fused_qk_norm_rope.py | 38 +++----- 7 files changed, 107 insertions(+), 170 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.cu b/cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.cu index bcd56342f931..2eaf4ba15fa6 100644 --- a/cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.cu +++ b/cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.cu @@ -57,12 +57,9 @@ __device__ __forceinline__ float selectMRopePosId(int const* position_ids, int t return static_cast(position_ids[sec * num_tokens + tokenIdx]); } -// 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, i.e. round-to-nearest-even with -// __NV_SATFINITE clamping to +/-448. That matches torch's -// .to(torch.float8_e4m3fn) for in-range values; out-of-range values clamp to -// +/-448 here whereas torch produces NaN. +// Store a per-thread run of `numElemsPerThread` float elements, converting to +// OutT. The FP8 path saturates to +/-448, whereas torch's .to(float8_e4m3fn) +// produces NaN for out-of-range values. template __device__ __forceinline__ void storeHeadElements( OutT* out, int offsetThread, float const (&elements)[numElemsPerThread]) @@ -92,12 +89,10 @@ __device__ __forceinline__ void storeHeadElements( } // 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. +// writing to a (possibly different-dtype) output buffer. // head_dim: the dimension of each head // interleave: interleave=!is_neox. -// 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. +// OutT: output element type (__nv_bfloat16 or __nv_fp8_e4m3). template __global__ void fusedQKNormRopeKernel( __nv_bfloat16 const* qkv_in, // Combined QKV input [num_tokens, (num_heads_q+num_heads_k+num_heads_v)*head_dim] @@ -136,7 +131,6 @@ __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 this warp processes @@ -149,8 +143,7 @@ __global__ void fusedQKNormRopeKernel( bool const isQ = 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 headIdx; // index within the head's own Q/K/V segment int segStart; // element offset of the segment start within a token row if (isQ) { @@ -186,7 +179,7 @@ __global__ void fusedQKNormRopeKernel( // Sum of squares for RMSNorm float sumOfSquares = 0.0f; - // Load from the BF16 input. + // Load. { vec_T vec = *reinterpret_cast(&qkv_in[offsetThread]); #pragma unroll @@ -201,7 +194,7 @@ __global__ void fusedQKNormRopeKernel( } } - // V heads are copy-cast only: skip norm and RoPE and store the raw values. + // V heads are copy-cast only: no norm, no RoPE. if (isV) { storeHeadElements(qkv_out, offsetThread, elements); @@ -354,7 +347,7 @@ __global__ void fusedQKNormRopeKernel( } } - // Store to the (templated) output. + // Store. storeHeadElements(qkv_out, offsetThread, elements); } @@ -385,13 +378,10 @@ static void launchFusedQKNormRopeImpl(__nv_bfloat16 const* qkv_in, OutT* qkv_out TLLM_CHECK(attention_factor == 1.0f); } - // rotary_dim == 0 would divide by zero when deriving the RoPE frequencies, and - // rotary_dim > head_dim breaks the warp-level pairing assumptions. TLLM_CHECK_WITH_INFO(rotary_dim > 0 && rotary_dim <= head_dim && rotary_dim % 2 == 0, "rotary_dim must be positive, even and no greater than head_dim (got rotary_dim=%d, head_dim=%d)", rotary_dim, head_dim); - // Skipping V is only well-defined in-place, where the V slots of the output - // already hold their final values; out-of-place they would stay uninitialized. + // Skipping V leaves the output's V slots untouched, which is only meaningful in place. TLLM_CHECK_WITH_INFO(process_v || static_cast(qkv_in) == static_cast(qkv_out), "process_v=false requires qkv_in and qkv_out to alias"); if (!interleave) @@ -404,7 +394,6 @@ static void launchFusedQKNormRopeImpl(__nv_bfloat16 const* qkv_in, OutT* qkv_out constexpr int blockSize = 256; int const warpsPerBlock = blockSize / 32; - // 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; @@ -450,12 +439,11 @@ void launchFusedQKNormRope(void* qkv, int const num_tokens, int const num_heads_ 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); + launchFusedQKNormRopeImpl<__nv_bfloat16>(static_cast<__nv_bfloat16 const*>(qkv), static_cast<__nv_bfloat16*>(qkv), + /*process_v=*/false, num_tokens, num_heads_q, num_heads_k, num_heads_v, head_dim, rotary_dim, eps, + static_cast<__nv_bfloat16 const*>(q_weight), static_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 launchFusedQKNormRopeToFp8(void const* qkv_in, void* qkv_out, int const num_tokens, int const num_heads_q, @@ -464,11 +452,11 @@ void launchFusedQKNormRopeToFp8(void const* qkv_in, void* qkv_out, int const num 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) { - // Out-of-place FP8: qkv_out is a fresh buffer, so V must be copy-cast too. - launchFusedQKNormRopeImpl<__nv_fp8_e4m3>(reinterpret_cast<__nv_bfloat16 const*>(qkv_in), - reinterpret_cast<__nv_fp8_e4m3*>(qkv_out), /*process_v=*/true, 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, + // Out-of-place, so V has to be copy-cast rather than left untouched. + launchFusedQKNormRopeImpl<__nv_fp8_e4m3>(static_cast<__nv_bfloat16 const*>(qkv_in), + static_cast<__nv_fp8_e4m3*>(qkv_out), /*process_v=*/true, num_tokens, num_heads_q, num_heads_k, num_heads_v, + head_dim, rotary_dim, eps, static_cast<__nv_bfloat16 const*>(q_weight), + static_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); } } // namespace kernels diff --git a/cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.h b/cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.h index 98701348479a..fd2401f592f1 100644 --- a/cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.h +++ b/cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.h @@ -51,13 +51,9 @@ void launchFusedQKNormRope( int mrope_section1, // mrope_section[1] (height) int mrope_section2); // mrope_section[2] (width) -// Out-of-place FP8 variant of launchFusedQKNormRope that reads a BF16 qkv input -// and writes __nv_fp8_e4m3 to a separate output buffer. -// -// 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 is -// copy-cast only. The output layout matches the input: -// [num_tokens, (num_heads_q+num_heads_k+num_heads_v)*head_dim]. +// Out-of-place FP8 variant of launchFusedQKNormRope, folding the FP8 +// activation-quant into the norm+RoPE epilogue. Q and K get RMSNorm + RoPE; V is +// copy-cast only. void launchFusedQKNormRopeToFp8(void const* qkv_in, // BF16 input [num_tokens, total_heads*head_dim] void* qkv_out, // FP8 E4M3 output buffer, same layout as input int const num_tokens, int const num_heads_q, int const num_heads_k, int const num_heads_v, int const head_dim, diff --git a/cpp/tensorrt_llm/thop/fusedQKNormRopeOp.cpp b/cpp/tensorrt_llm/thop/fusedQKNormRopeOp.cpp index 72b912b07b23..12c124b85aaf 100644 --- a/cpp/tensorrt_llm/thop/fusedQKNormRopeOp.cpp +++ b/cpp/tensorrt_llm/thop/fusedQKNormRopeOp.cpp @@ -18,6 +18,7 @@ #include "tensorrt_llm/thop/thUtils.h" #include +#include #include TRTLLM_NAMESPACE_BEGIN @@ -25,6 +26,47 @@ TRTLLM_NAMESPACE_BEGIN namespace torch_ext { +namespace +{ + +// Shared input validation for the in-place and out-of-place operators. Returns num_tokens. +int64_t validateFusedQKNormRopeInputs(torch::Tensor const& qkv, torch::Tensor const& position_ids, + torch::Tensor const& q_weight, torch::Tensor const& k_weight, int64_t num_heads_q, int64_t num_heads_k, + int64_t num_heads_v, int64_t head_dim, bool use_mrope) +{ + TORCH_CHECK(qkv.dim() == 2, "QKV tensor must be 2D: [num_tokens, (num_heads_q+num_heads_k+num_heads_v)*head_dim]"); + // Plain RoPE: position_ids is 1D [num_tokens]. Interleaved mRoPE: 2D [3, num_tokens]. + 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"); + + // The kernel narrows these to int, so reject anything that would not survive it. + TORCH_CHECK(num_heads_q >= 0 && num_heads_k >= 0 && num_heads_v >= 0 && head_dim > 0, + "Head counts must be non-negative and head_dim must be positive"); + int64_t total_heads = num_heads_q + num_heads_k + num_heads_v; + TORCH_CHECK( + num_tokens <= std::numeric_limits::max() && total_heads * head_dim <= std::numeric_limits::max(), + "QKV dimensions exceed the kernel's supported range"); + TORCH_CHECK( + qkv.size(1) == total_heads * head_dim, "QKV tensor size must match total number of heads and head dimension"); + + return num_tokens; +} + +} // namespace + // Function for fused QK Norm and RoPE // This operator applies RMS normalization and RoPE to Q and K tensors in a single CUDA kernel. // The OP performs operations in-place on the input qkv tensor. @@ -53,28 +95,8 @@ void fused_qk_norm_rope( int64_t mrope_section2 // mrope_section[2] (width) ) { - // Input validation - TORCH_CHECK(qkv.dim() == 2, "QKV tensor must be 2D: [num_tokens, (num_heads_q+num_heads_k+num_heads_v)*head_dim]"); - // Plain RoPE: position_ids is 1D [num_tokens]. Interleaved mRoPE: 2D [3, num_tokens]. - 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"); + int64_t num_tokens = validateFusedQKNormRopeInputs( + qkv, position_ids, q_weight, k_weight, num_heads_q, num_heads_k, num_heads_v, head_dim, use_mrope); auto stream = at::cuda::getCurrentCUDAStream(qkv.get_device()); @@ -89,39 +111,18 @@ 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. +// Out-of-place FP8 variant of fused_qk_norm_rope: applies RMSNorm + RoPE to Q/K, +// copy-casts V, and returns a new FP8 (E4M3) tensor of the same shape. 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 num_tokens = validateFusedQKNormRopeInputs( + qkv, position_ids, q_weight, k_weight, num_heads_q, num_heads_k, num_heads_v, head_dim, use_mrope); 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()); 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 ff605cb90ffc..535de64c049b 100644 --- a/tensorrt_llm/_torch/attention_backend/fmha/msa_sparse_gqa.py +++ b/tensorrt_llm/_torch/attention_backend/fmha/msa_sparse_gqa.py @@ -139,24 +139,16 @@ def run_msa_paged_gqa( kv_cache_manager, layer_idx, metadata.msa_out_cache_loc[:num_tokens], k, v ) - # 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. + q_view = q.view(num_tokens, attn.num_heads, head_dim) 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. 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. + # KV-cache scales, so the scale is 1.0 and this is a plain E4M3 cast. use_fp8 = k_paged.dtype == torch.float8_e4m3fn - if use_fp8 and q_view.dtype != torch.float8_e4m3fn: + if use_fp8: 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 a4770d21ec30..858348f914ba 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 @@ -824,11 +824,8 @@ 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 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) + 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) 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 0db6ea8c1675..b1ca4789dd6e 100644 --- a/tensorrt_llm/_torch/models/modeling_minimaxm3.py +++ b/tensorrt_llm/_torch/models/modeling_minimaxm3.py @@ -714,10 +714,7 @@ 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. + # Whether the main K/V cache is stored as FP8 E4M3. quant_config = getattr(model_config, "quant_config", None) self.main_kv_is_fp8 = bool( quant_config is not None @@ -886,15 +883,13 @@ def _fused_qk_norm_rope( channels, matching M3's whole-head norm with front partial RoPE, and leaves the V heads untouched. - 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. + When out_fp8 is True, an out-of-place variant runs instead and returns a + fresh FP8 E4M3 tensor with Q/K normed and roped and V copy-cast, leaving + qkv 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. + 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. """ if position_ids is None or qkv.dtype != torch.bfloat16: return None @@ -911,7 +906,6 @@ def _fused_qk_norm_rope( 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, @@ -970,52 +964,36 @@ 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. - """ + """Whether the MSA fmha_sm100 backend handles this layer's attention.""" 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. + Only the MSA backend consumes an FP8 paged K/V cache directly; the + Triton/SDPA reference paths and the index branch stay 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]: + ) -> 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). + The MSA kernels read q with its real strides and scatter k/v into the + paged cache with an indexed copy, so the split column-views can be + handed over as-is. Other backends keep the previous contiguity. """ 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]: + 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. + The index analogue of _split_main_qkv; the index cache is bf16, so there + is no FP8 variant here. """ idx_q, idx_k = fused_idx.split([self.index_q_size, self.index_k_size], dim=-1) if self._msa_backend_active(): @@ -1333,9 +1311,8 @@ def _forward_attention_core( idx_k: Optional[torch.Tensor], attn_metadata: AttentionMetadata, ) -> torch.Tensor: - # 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. + # q may be FP8 on the MSA FP8-KV path, so pin the output to the compute + # 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 ) 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 1dcfd8317ba4..f85e7b922bd9 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 @@ -348,37 +348,28 @@ 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 = [ +_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_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]) -# MiniMax-M3 stores every RMSNorm as a Gemma norm, so use_gemma=True is the -# configuration that actually ships through this op. +# MiniMax-M3 stores every RMSNorm as a Gemma norm, so use_gemma=True is what ships. @pytest.mark.parametrize("use_gemma", [False, True]) def test_fused_qk_norm_rope_to_fp8( head_dim, num_heads_group, num_tokens, partial_rotary_factor, is_neox, use_gemma ): """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. V is bit-exactly the E4M3 cast of the input V (pure copy-cast, so this - pins down the kernel's rounding against torch's). - 4. The dequantized Q/K output matches the BF16 fused reference within - FP8-appropriate tolerance. + Checks that the input is left untouched, the output is FP8 E4M3 of the same + shape, V is a bit-exact E4M3 cast of the input V, and Q/K match the BF16 + reference after dequantization. """ device = "cuda" dtype = torch.bfloat16 @@ -427,10 +418,9 @@ def test_fused_qk_norm_rope_to_fp8( # 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) - # V is copy-cast only (no norm, no RoPE), so it must match torch's E4M3 cast - # exactly. This pins down the kernel's round-to-nearest-even FP8 store, which - # nothing else here verifies. The inputs are standard normal, so they stay - # well inside E4M3 range and the two casts' clamp/NaN behaviors don't diverge. + # V is copy-cast only, so it pins down the kernel's FP8 rounding against + # torch's. The inputs stay well inside E4M3 range, so the two casts' + # clamp/NaN behaviors don't diverge. v_start = (num_heads_q + num_heads_k) * head_dim torch.testing.assert_close( out_fp8[:, v_start:].float(), @@ -455,12 +445,8 @@ def test_fused_qk_norm_rope_to_fp8( use_gemma=use_gemma, ) - # The op folds the E4M3 cast into the epilogue, so the dominant error is the - # single E4M3 rounding: 3 mantissa bits give a worst-case relative error of - # 2^-4 = 6.25%, plus the ~2^-9 the BF16 reference itself carries. The kernel - # accumulates in fp32 and casts straight to FP8, i.e. one rounding step fewer - # than the reference's fp32 -> bf16 -> fp8, so it should clear this - # comfortably. + # Dominated by the single E4M3 rounding: 3 mantissa bits give a worst-case + # relative error of 2^-4 = 6.25%, plus the ~2^-9 the BF16 reference carries. torch.testing.assert_close( out_fp8.float(), ref_output.float(),