Skip to content

[None][perf] Fold q/k/v quantization into qknorm_rope_fused kernel & remove contiguous #17093

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
brb-nv wants to merge 1 commit into
base: main
Choose a base branch
from
Open

[None][perf] Fold q/k/v quantization into qknorm_rope_fused kernel & remove contiguous #17093

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
191 changes: 138 additions & 53 deletions cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@
#include <cuda_fp16.h>
#include <cuda_fp8.h>
#include <cuda_runtime.h>
#include <type_traits>

TRTLLM_NAMESPACE_BEGIN

Expand Down Expand Up @@ -56,15 +57,52 @@ __device__ __forceinline__ float selectMRopePosId(int const* position_ids, int t
return static_cast<float>(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 <typename OutT, int numElemsPerThread, int vecSize>
__device__ __forceinline__ void storeHeadElements(
OutT* out, int offsetThread, float const (&elements)[numElemsPerThread])
{
using vec_T = typename tensorrt_llm::common::packed_as<uint, vecSize>::type;
if constexpr (std::is_same_v<OutT, __nv_bfloat16>)
{
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<uint*>(&vec) + i)) = vals;
}
*reinterpret_cast<vec_T*>(&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 <int head_dim, bool interleave>
// 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 <int head_dim, bool interleave, typename OutT>
__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
Expand Down Expand Up @@ -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;

Expand All @@ -119,25 +177,15 @@ __global__ void fusedQKNormRopeKernel(
constexpr int vecSize = elemSizeBytes / 4; // Use packed_as<uint, vecSize> to perform loading/saving.
using vec_T = typename tensorrt_llm::common::packed_as<uint, vecSize>::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<vec_T const*>(&qkv[offsetThread]);
vec_T vec = *reinterpret_cast<vec_T const*>(&qkv_in[offsetThread]);
for (int i = 0; i < vecSize; i++)

@yunruis yunruis Aug 4, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why not #pragma unroll here?

{
float2 vals = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162*>(reinterpret_cast<uint*>(&vec) + i));
Expand All @@ -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<OutT, numElemsPerThread, vecSize>(qkv_out, offsetThread, elements);
return;
}

if (is_qk_norm)
{
// Reduce sum across warp using the utility function
Expand Down Expand Up @@ -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<uint*>(&vec) + i)) = vals;
}
vec_T* outputPtr = reinterpret_cast<vec_T*>(&qkv[offsetThread]);
*outputPtr = vec;
}
// Store to the (templated) output.
storeHeadElements<OutT, numElemsPerThread, vecSize>(qkv_out, offsetThread, elements);
}

// Borrowed from
Expand All @@ -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 <typename OutT>
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)
Comment on lines +371 to +377

@coderabbitai coderabbitai Bot Aug 3, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🎯 Functional Correctness | 🟠 Major | ⚡ Quick win

Validate rotary_dim before dispatch.

The FP8 operator checks tensor shapes but does not constrain rotary_dim. A value of 0 passes the existing evenness check and then causes division by zero during RoPE frequency calculation. A value greater than head_dim also violates the kernel pairing assumptions. Reject values outside 1..head_dim before calculating launch dimensions.

Proposed fix 
-    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, no greater than head_dim, and even");
🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.cu` around lines 371 - 377,
Validate rotary_dim in launchFusedQKNormRopeImpl before calculating launch
dimensions or dispatching the kernel, rejecting values less than 1 or greater
than head_dim while preserving the existing evenness validation. Ensure invalid
values cannot reach RoPE frequency calculation or kernel launch.

{
if (factor == 1.0f)
{
Expand All @@ -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);
Expand All @@ -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>
<<<gridDim, blockDim, 0, stream>>>(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><<<gridDim, blockDim, 0, stream>>>(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>
<<<gridDim, blockDim, 0, stream>>>(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><<<gridDim, blockDim, 0, stream>>>(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>
<<<gridDim, blockDim, 0, stream>>>(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><<<gridDim, blockDim, 0, stream>>>(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);
}
Comment on lines +465 to +471

@yunruis yunruis Aug 4, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This out_fp8=false branch is unreachable. The only caller of launchFusedQKNormRopeOut is fused_qk_norm_rope_to_fp8 in thop/fusedQKNormRopeOp.cpp, which hardcodes /*out_fp8=*/true, /*process_v=*/true, so the BF16 out-of-place path ships with no caller and no test.

It also carries the footgun the header comment acknowledges: with process_v=false, the V region of qkv_out is never written. That is harmless for the in-place BF16 entry point (input and output alias, so V keeps its original values), but any genuinely out-of-place caller passes a freshly allocated buffer, so V would be left uninitialized.

Suggest either dropping the out_fp8 parameter and keeping only the FP8 specialization until a real BF16 caller exists, or adding a TLLM_CHECK(process_v || qkv_in == qkv_out) at the top of launchFusedQKNormRopeOut so the undefined-V combination cannot be reached.

}
} // namespace kernels

TRTLLM_NAMESPACE_END
20 changes: 20 additions & 0 deletions cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Loading
Loading