diff --git a/cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.cu b/cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.cu index 6e26dfce65bf..2eaf4ba15fa6 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,50 @@ __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, 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]) +{ + using vec_T = typename tensorrt_llm::common::packed_as::type; + 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])); + 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 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 or __nv_fp8_e4m3). +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 +131,35 @@ __global__ void fusedQKNormRopeKernel( // Total number of attention heads (Q and K) int const total_qk_heads = num_heads_q + num_heads_k; + 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; + 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) + { + 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,17 +173,7 @@ __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 @@ -137,7 +181,8 @@ __global__ void fusedQKNormRopeKernel( // Load. { - vec_T vec = *reinterpret_cast(&qkv[offsetThread]); + 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)); @@ -149,6 +194,13 @@ __global__ void fusedQKNormRopeKernel( } } + // V heads are copy-cast only: no norm, no RoPE. + if (isV) + { + storeHeadElements(qkv_out, offsetThread, elements); + return; + } + if (is_qk_norm) { // Reduce sum across warp using the utility function @@ -296,16 +348,7 @@ __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; - } + storeHeadElements(qkv_out, offsetThread, elements); } // Borrowed from @@ -322,18 +365,25 @@ __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) { TLLM_CHECK(attention_factor == 1.0f); } - TLLM_CHECK_WITH_INFO(rotary_dim % 2 == 0, "rotary_dim must be even"); + 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 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) { // To allow warp-level pairing for partial rope @@ -344,8 +394,8 @@ 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; + 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 +407,58 @@ 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) +{ + 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, + 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) +{ + // 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 TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.h b/cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.h index 4e2421cb57a2..fd2401f592f1 100644 --- a/cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.h +++ b/cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.h @@ -51,6 +51,16 @@ void launchFusedQKNormRope( int mrope_section1, // mrope_section[1] (height) int mrope_section2); // mrope_section[2] (width) +// 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, + 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..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,6 +111,45 @@ void fused_qk_norm_rope( static_cast(mrope_section1), static_cast(mrope_section2)); } +// 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) +{ + 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; + 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::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; +} + +// 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 +159,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/models/modeling_minimaxm3.py b/tensorrt_llm/_torch/models/modeling_minimaxm3.py index 8d3e0ddcac1e..b1ca4789dd6e 100644 --- a/tensorrt_llm/_torch/models/modeling_minimaxm3.py +++ b/tensorrt_llm/_torch/models/modeling_minimaxm3.py @@ -714,6 +714,14 @@ 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. + 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 +872,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,6 +883,10 @@ 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 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 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. @@ -891,6 +904,31 @@ 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: + 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 +941,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 +963,43 @@ 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.""" + 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 + 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]: + """Split the fused [q|k|v] buffer into per-tensor q/k/v. + + 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]: + """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 variant here. + """ + 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 +1100,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 +1311,11 @@ 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)) + # 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 + ) if self.register_to_config and is_torch_compiling(): minimax_m3_attn_custom_op_inplace( q, @@ -1370,10 +1448,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 +1477,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..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 @@ -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) @@ -340,3 +346,110 @@ def test_fused_qk_norm_rope_gemma_mrope( ) torch.testing.assert_close(output, ref_output, rtol=5e-2, atol=1e-1) + + +_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]) +# 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. + + 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 + 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 + use_gemma, + 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) + + # 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(), + 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, + num_heads_k, + num_heads_v, + head_dim, + rotary_dim, + eps, + q_weight, + k_weight, + base, + is_neox, + position_ids, + use_gemma=use_gemma, + ) + + # 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(), + rtol=0.07, + atol=0.1, + )