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

Merged
brb-nv merged 3 commits into from
Aug 8, 2026
Merged

[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
3 commits
Select commit Hold shift + click to select a range
2ba4437
[None][perf] Fold q/k/v quantization into qknorm_rope_fused kernel & …
brb-nv Jul 22, 2026
2aa9032
[None][fix] Address review feedback on fused QK norm + RoPE FP8 path
brb-nv Aug 6, 2026
31c4853
[None][chore] Trim review comments and dedupe validation
brb-nv Aug 6, 2026
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
178 changes: 126 additions & 52 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,50 @@ __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, converting to
// OutT. The FP8 path saturates to +/-448, whereas torch's .to(float8_e4m3fn)
// produces NaN for out-of-range values.
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;
#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<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 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 or __nv_fp8_e4m3).
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 +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;

Expand All @@ -119,25 +173,16 @@ __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.
{
vec_T vec = *reinterpret_cast<vec_T const*>(&qkv[offsetThread]);
vec_T vec = *reinterpret_cast<vec_T const*>(&qkv_in[offsetThread]);
#pragma unroll
for (int i = 0; i < vecSize; i++)
Comment thread
brb-nv marked this conversation as resolved.
{
float2 vals = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162*>(reinterpret_cast<uint*>(&vec) + i));
Expand All @@ -149,6 +194,13 @@ __global__ void fusedQKNormRopeKernel(
}
}

// V heads are copy-cast only: no norm, no RoPE.
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 @@ -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<uint*>(&vec) + i)) = vals;
}
vec_T* outputPtr = reinterpret_cast<vec_T*>(&qkv[offsetThread]);
*outputPtr = vec;
}
storeHeadElements<OutT, numElemsPerThread, vecSize>(qkv_out, offsetThread, elements);
}

// Borrowed from
Expand All @@ -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 <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 thread
coderabbitai[bot] marked this conversation as resolved.
{
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<void const*>(qkv_in) == static_cast<void const*>(qkv_out),
"process_v=false requires qkv_in and qkv_out to alias");
if (!interleave)
{
// To allow warp-level pairing for partial rope
Expand All @@ -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);
Expand All @@ -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>
<<<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)
{
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
10 changes: 10 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,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
Loading
Loading