chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[sgl-kernel][1/2] Fused qk_norm_rope for Qwen3-MoE (#14036)

Browse Source 
Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>
This commit is contained in:
Yuan Luo
2025-11-28 12:25:15 +08:00
committed by GitHub
parent bce40fa217
commit e12c78aab6
7 changed files with 696 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
sgl-kernel/CMakeLists.txt
View File

@@ -318,6 +318,7 @@ set(SOURCES
"csrc/moe/marlin_moe_wna16/ops.cu"
"csrc/moe/moe_align_kernel.cu"
"csrc/moe/moe_fused_gate.cu"
"csrc/moe/fused_qknorm_rope_kernel.cu"
"csrc/moe/kimi_k2_moe_fused_gate.cu"
"csrc/moe/moe_sum.cu"
"csrc/moe/moe_sum_reduce.cu"

8
sgl-kernel/csrc/common_extension.cc
View File

@@ -264,9 +264,17 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
m.def("shuffle_rows(Tensor input, Tensor dst2src_map, Tensor output) -> ()");
m.impl("shuffle_rows", torch::kCUDA, &shuffle_rows);
m.def("apply_shuffle_mul_sum(Tensor input, Tensor output, Tensor permutation, Tensor? factors) -> ()");
m.impl("apply_shuffle_mul_sum", torch::kCUDA, &apply_shuffle_mul_sum);
m.def(
"fused_qk_norm_rope(Tensor! qkv, int num_heads_q, "
"int num_heads_k, int num_heads_v, int head_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) -> ()");
m.impl("fused_qk_norm_rope", torch::kCUDA, &fused_qk_norm_rope);
/*
* From csrc/moe/cutlass_moe/w4a8
*/

408
sgl-kernel/csrc/moe/fused_qknorm_rope_kernel.cu Normal file
View File

@@ -0,0 +1,408 @@
/*
* Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
// Adapted from
// https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.cu
#include <ATen/cuda/Exceptions.h>
#include <c10/cuda/CUDAGuard.h>
#include <c10/cuda/CUDAStream.h>
#include <cuda_bf16.h>
#include <cuda_fp16.h>
#include <cuda_fp8.h>
#include <cuda_runtime.h>
#include <torch/all.h>
#include <torch/cuda.h>
#include <cmath>
#define CHECK_TYPE(x, st) \
TORCH_CHECK(x.scalar_type() == st, #x " dtype is ", x.scalar_type(), ", while ", st, " is expected")
#define CHECK_TH_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
#define CHECK_INPUT(x, st) \
CHECK_TH_CUDA(x); \
CHECK_CONTIGUOUS(x); \
CHECK_TYPE(x, st)
#define FINAL_MASK 0xffffffff
namespace tensorrt_llm::common {
template <typename T, int num>
struct packed_as;
// Specialization for packed_as used in this kernel.
template <>
struct packed_as<uint, 1> {
using type = uint;
};
template <>
struct packed_as<uint, 2> {
using type = uint2;
};
template <>
struct packed_as<uint, 4> {
using type = uint4;
};
template <typename T>
__inline__ __device__ T warpReduceSum(T val) {
#pragma unroll
for (int mask = 16; mask > 0; mask >>= 1)
val += __shfl_xor_sync(FINAL_MASK, val, mask,
32); //__shfl_sync bf16 return float when sm < 80
return val;
}
template <typename T>
inline __device__ __host__ T divUp(T m, T n) {
return (m + n - 1) / n;
}
} // namespace tensorrt_llm::common
namespace tensorrt_llm::kernels {
__device__ inline float compute_freq_yarn(float base, int head_dim, int half_dim, float factor, float low, float high) {
float freq = powf(base, -2.0f * half_dim / static_cast<float>(head_dim));
if (factor != 1.0f) {
float inv_freq_extrapolation = freq;
float inv_freq_interpolation = freq / factor;
float high_adj = high;
if (fabsf(low - high_adj) <= 1e-6f) {
high_adj += 0.001f;
}
float linear_func = (static_cast<float>(half_dim) - low) / (high_adj - low);
float ramp_func = fminf(fmaxf(linear_func, 0.0f), 1.0f);
float inv_freq_extrapolation_factor = 1.0f - ramp_func;
freq = inv_freq_interpolation * (1.0f - inv_freq_extrapolation_factor) +
inv_freq_extrapolation * inv_freq_extrapolation_factor;
}
return freq;
}
////////////////////////////////////////////////////////////////////////////////////////////////////
// Perform per-head QK Norm and RoPE in a single kernel.
// head_dim: the dimension of each head
// interleave: interleave=!is_neox.
template <int head_dim, bool interleave>
__global__ void fusedQKNormRopeKernel(
__nv_bfloat16* qkv, // Combined QKV tensor [num_tokens, (num_heads_q+num_heads_k+num_heads_v)*head_dim]
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
float const eps, // Epsilon for RMS normalization
__nv_bfloat16 const* q_weight, // RMSNorm weights for query
__nv_bfloat16 const* k_weight, // RMSNorm weights for key
float const base, // Base for RoPE computation
int const* position_ids, // Position IDs for RoPE
int const num_tokens, // Number of tokens
// parameters for yarn
float factor, // factor in rope_scaling in config.json. When it is not 1.0, it means the model is using yarn.
float low, // threshold for high frequency
float high, // threshold for low frequency
float attention_factor // attention_factor applied on cos and sin
) {
int const warpsPerBlock = blockDim.x / 32;
int const warpId = threadIdx.x / 32;
int const laneId = threadIdx.x % 32;
// Calculate global warp index to determine which head/token this warp processes
int const globalWarpIdx = blockIdx.x * warpsPerBlock + warpId;
// Total number of attention heads (Q and K)
int const total_qk_heads = num_heads_q + num_heads_k;
// 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;
// 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;
int const num_heads = num_heads_q + num_heads_k + num_heads_v;
static_assert(
head_dim % (32 * 2) == 0,
"head_dim must be divisible by 64 (each warp processes one head, and each thread gets even number of "
"elements)");
constexpr int numElemsPerThread = head_dim / 32;
float elements[numElemsPerThread];
constexpr int elemSizeBytes = numElemsPerThread * sizeof(__nv_bfloat16);
static_assert(elemSizeBytes % 4 == 0, "numSizeBytes must be a multiple of 4");
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 offsetThread = offsetWarp + laneId * numElemsPerThread;
// Sum of squares for RMSNorm
float sumOfSquares = 0.0f;
// Load.
{
vec_T vec = *reinterpret_cast<vec_T const*>(&qkv[offsetThread]);
for (int i = 0; i < vecSize; i++) {
float2 vals = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162*>(reinterpret_cast<uint*>(&vec) + i));
sumOfSquares += vals.x * vals.x;
sumOfSquares += vals.y * vals.y;
elements[2 * i] = vals.x;
elements[2 * i + 1] = vals.y;
}
}
// Reduce sum across warp using the utility function
sumOfSquares = tensorrt_llm::common::warpReduceSum(sumOfSquares);
// Compute RMS normalization factor
float rms_rcp = rsqrtf(sumOfSquares / static_cast<float>(head_dim) + eps);
// Normalize elements
for (int i = 0; i < numElemsPerThread; i++) {
int dim = laneId * numElemsPerThread + i;
float weight = isQ ? __bfloat162float(q_weight[dim]) : __bfloat162float(k_weight[dim]);
elements[i] *= rms_rcp * weight;
}
// Apply RoPE to normalized elements
float elements2[numElemsPerThread]; // Additional buffer required for RoPE.
float cos_vals[numElemsPerThread];
float sin_vals[numElemsPerThread];
float pos_id = static_cast<float>(position_ids[tokenIdx]);
if constexpr (interleave) {
// Perform interleaving. Fill cos_vals and sin_vals.
for (int i = 0; i < numElemsPerThread; i++) {
if (i % 2 == 0) {
elements2[i] = -elements[i + 1];
} else {
elements2[i] = elements[i - 1];
}
int dim_idx = laneId * numElemsPerThread + i;
int half_dim = dim_idx / 2;
float freq = compute_freq_yarn(base, head_dim, half_dim, factor, low, high);
float theta = pos_id * freq;
__sincosf(theta, &sin_vals[i], &cos_vals[i]);
}
} else {
// Before data exchange with in warp, we need to sync.
__syncwarp();
// Get the data from the other half of the warp. Fill cos_vals and sin_vals.
for (int i = 0; i < numElemsPerThread; i++) {
elements2[i] = __shfl_xor_sync(0xffffffff, elements[i], 16);
if (laneId < 16) {
elements2[i] = -elements2[i];
}
int dim_idx = laneId * numElemsPerThread + i;
dim_idx = (dim_idx * 2) % head_dim;
int half_dim = dim_idx / 2;
float freq = compute_freq_yarn(base, head_dim, half_dim, factor, low, high);
float theta = pos_id * freq;
__sincosf(theta, &sin_vals[i], &cos_vals[i]);
}
// __shfl_xor_sync does not provide memfence. Need to sync again.
__syncwarp();
}
for (int i = 0; i < numElemsPerThread; i++) {
elements[i] = (elements[i] * cos_vals[i] + elements2[i] * sin_vals[i]) * attention_factor;
}
// 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;
}
}
// Borrowed from
// https://github.com/flashinfer-ai/flashinfer/blob/8125d079a43e9a0ba463a4ed1b639cefd084cec9/include/flashinfer/pos_enc.cuh#L568
#define DISPATCH_INTERLEAVE(interleave, INTERLEAVE, ...) \
if (interleave) { \
const bool INTERLEAVE = true; \
__VA_ARGS__ \
} else { \
const bool INTERLEAVE = false; \
__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,
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) {
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 gridSize = common::divUp(totalWarps, warpsPerBlock);
dim3 gridDim(gridSize);
dim3 blockDim(blockSize);
// Head dimensions should be a multiple of 64
// Add more cases as needed
switch (head_dim) {
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,
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);
});
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,
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);
});
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,
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);
});
break;
default:
TORCH_CHECK(false, "Unsupported head dimension for fusedQKNormRope: ", head_dim);
}
}
} // namespace tensorrt_llm::kernels
// 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.
void fused_qk_norm_rope(
torch::Tensor& qkv, // Combined QKV tensor [num_tokens, (num_heads_q+num_heads_k+num_heads_v)*head_dim]
int64_t num_heads_q, // Number of query heads
int64_t num_heads_k, // Number of key heads
int64_t num_heads_v, // Number of value heads
int64_t head_dim, // Dimension per head
double eps, // Epsilon for RMS normalization
torch::Tensor& q_weight, // RMSNorm weights for query [head_dim]
torch::Tensor& k_weight, // RMSNorm weights for key [head_dim]
double base, // Base for RoPE computation
bool is_neox, // Whether RoPE is applied in Neox style
torch::Tensor& position_ids, // Position IDs for RoPE [num_tokens]
// parameters for yarn
double factor, // factor in rope_scaling in config.json. When it is not 1.0, it means the model is using yarn.
double low, // threshold for high frequency
double high, // threshold for low frequency
double attention_factor // attention_factor applied on cos and sin
) {
// 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]");
TORCH_CHECK(position_ids.dim() == 1, "Position IDs must be 1D: [num_tokens]");
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(0) == 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 stream = at::cuda::getCurrentCUDAStream(qkv.get_device());
tensorrt_llm::kernels::launchFusedQKNormRope(
reinterpret_cast<__nv_bfloat16*>(qkv.data_ptr()),
static_cast<int>(num_tokens),
static_cast<int>(num_heads_q),
static_cast<int>(num_heads_k),
static_cast<int>(num_heads_v),
static_cast<int>(head_dim),
static_cast<float>(eps),
reinterpret_cast<__nv_bfloat16*>(q_weight.data_ptr()),
reinterpret_cast<__nv_bfloat16*>(k_weight.data_ptr()),
static_cast<float>(base),
!is_neox, // interleave
reinterpret_cast<int const*>(position_ids.data_ptr()),
static_cast<float>(factor),
static_cast<float>(low),
static_cast<float>(high),
static_cast<float>(attention_factor),
stream);
}

17
sgl-kernel/include/sgl_kernel_ops.h
View File

@@ -378,6 +378,23 @@ void apply_shuffle_mul_sum(
const torch::Tensor& permutation,
const std::optional<torch::Tensor>& factors);
void fused_qk_norm_rope(
torch::Tensor& qkv,
int64_t num_heads_q,
int64_t num_heads_k,
int64_t num_heads_v,
int64_t head_dim,
double eps,
torch::Tensor& q_weight,
torch::Tensor& k_weight,
double base,
bool is_neox,
torch::Tensor& position_ids,
double factor,
double low,
double high,
double attention_factor);
void cutlass_fp4_group_mm(
torch::Tensor& output,
const torch::Tensor& a,

1
sgl-kernel/python/sgl_kernel/__init__.py
View File

@@ -84,6 +84,7 @@ from sgl_kernel.moe import (
apply_shuffle_mul_sum,
cutlass_fp4_group_mm,
fp8_blockwise_scaled_grouped_mm,
fused_qk_norm_rope,
kimi_k2_moe_fused_gate,
moe_align_block_size,
moe_fused_gate,

36
sgl-kernel/python/sgl_kernel/moe.py
View File

@@ -251,6 +251,42 @@ def apply_shuffle_mul_sum(
)
def fused_qk_norm_rope(
qkv: torch.Tensor,
num_heads_q: int,
num_heads_k: int,
num_heads_v: int,
head_dim: int,
eps: float,
q_weight: torch.Tensor,
k_weight: torch.Tensor,
base: float,
is_neox: bool,
position_ids: torch.Tensor,
factor: float,
low: float,
high: float,
attention_factor: float,
) -> None:
torch.ops.sgl_kernel.fused_qk_norm_rope(
qkv,
num_heads_q,
num_heads_k,
num_heads_v,
head_dim,
eps,
q_weight,
k_weight,
base,
is_neox,
position_ids,
factor,
low,
high,
attention_factor,
)
def cutlass_fp4_group_mm(
a_fp4,
b_fp4,

225
sgl-kernel/tests/test_fused_qk_norm_rope.py Normal file
View File

@@ -0,0 +1,225 @@
import pytest
import torch
from sgl_kernel import fused_qk_norm_rope as sgl_fused_qk_norm_rope
from sglang.srt.layers.layernorm import RMSNorm
from sglang.srt.layers.rotary_embedding import get_rope
from sglang.srt.server_args import (
ServerArgs,
get_global_server_args,
set_global_server_args_for_scheduler,
)
from sglang.srt.utils import (
cpu_has_amx_support,
is_cpu,
is_cuda,
is_hip,
is_npu,
is_xpu,
)
_is_cuda = is_cuda()
_is_hip = is_hip()
_is_cpu = is_cpu()
_is_cpu_amx_available = cpu_has_amx_support()
_is_npu = is_npu()
_is_xpu = is_xpu()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@torch.inference_mode()
def torch_ref_rms_norm_rope(
qkv,
num_heads_q,
num_heads_k,
num_heads_v,
head_dim,
eps,
q_weight,
k_weight,
base,
is_neox,
position_ids,
):
"""
PyTorch reference implementation of RMSNorm+RoPE for verification.
Uses SGLang's own RMSNorm and RotaryEmbedding modules to ensure consistency
with the expected behavior of the fused kernel.
Args:
qkv: Combined QKV tensor of shape [num_tokens, hidden_size]
num_heads_q: Number of query heads
num_heads_k: Number of key heads
num_heads_v: Number of value heads (unused for normalization/RoPE but needed for tensor splitting)
head_dim: Dimension of each head
eps: Epsilon value for RMS normalization
q_weight: RMSNorm weights for query [head_dim]
k_weight: RMSNorm weights for key [head_dim]
base: Base value for RoPE calculations
is_neox: Whether to use NeoX style RoPE
position_ids: Position IDs for RoPE of shape [num_tokens]
Returns:
Combined tensor with Q and K parts normalized and RoPE applied
"""
# Get input shape information
num_tokens = qkv.shape[0]
hidden_size = qkv.shape[1]
# Calculate dimensions for Q, K, V segments
q_size = num_heads_q * head_dim
k_size = num_heads_k * head_dim
v_size = num_heads_v * head_dim
# Verify dimensions match
assert (
hidden_size == q_size + k_size + v_size
), f"Hidden size {hidden_size} doesn't match Q+K+V dimensions {q_size + k_size + v_size}"
# Split the tensor into Q, K, V parts
q = qkv[:, :q_size]
k = qkv[:, q_size : q_size + k_size]
v = qkv[:, q_size + k_size :]
rotary_emb = get_rope(
head_dim,
rotary_dim=head_dim,
max_position=8192,
base=10000,
is_neox_style=is_neox,
rope_scaling=None,
dual_chunk_attention_config=None,
)
rotary_emb = rotary_emb.to(qkv.device)
# Create and apply RMSNorm modules with custom weights
q_norm = RMSNorm(hidden_size=head_dim, eps=eps).to(qkv.device).to(qkv.dtype)
q_norm.weight.data.copy_(q_weight)
k_norm = RMSNorm(hidden_size=head_dim, eps=eps).to(qkv.device).to(qkv.dtype)
k_norm.weight.data.copy_(k_weight)
q_by_head = q.reshape(-1, head_dim)
q_by_head = q_norm(q_by_head)
k_by_head = k.reshape(-1, head_dim)
k_by_head = k_norm(k_by_head)
q = q_by_head.view(q.shape)
k = k_by_head.view(k.shape)
[q_rope, k_rope] = rotary_emb(
position_ids,
q,
k,
fused_set_kv_buffer_arg=None,
)
# Combine Q, K, V back together
result = torch.cat([q_rope, k_rope, v], dim=1)
return result
head_dims = [64, 128]
# (Q heads, K heads, V heads)
num_heads_groups = [
(16, 8, 8), # Qwen3-0.6B, Qwen3-1.7B
(32, 8, 8), # Qwen3-4B, Qwen3-8B, Qwen3-30B-A3B
(40, 8, 8), # Qwen3-14B
(64, 8, 8), # Qwen3-32B, Qwen3-235B-A22B
]
num_tokens_list = [1, 3, 8, 32, 256]
is_neox_list = [False, True]
dtypes = [torch.bfloat16]
@pytest.mark.skipif(not _is_cuda, reason="Skipping CUDA/ROCm only tests.")
@pytest.mark.parametrize("head_dim", head_dims)
@pytest.mark.parametrize("num_heads_group", num_heads_groups)
@pytest.mark.parametrize("num_tokens", num_tokens_list)
@pytest.mark.parametrize("is_neox", is_neox_list)
@pytest.mark.parametrize("dtype", dtypes)
def test_fused_qk_norm_rope(head_dim, num_heads_group, num_tokens, is_neox, dtype):
"""
Test the fused QK RMSNorm + RoPE operation with various configurations.
This test verifies that the fused kernel correctly applies:
1. RMSNorm to both query (Q) and key (K) portions of the QKV tensor
2. Rotary Position Embeddings (RoPE) to the normalized Q and K
3. Leaves the value (V) portion unchanged
Args:
head_dim: Dimension of each attention head
num_heads_group: Tuple of (num_heads_q, num_heads_k, num_heads_v)
num_tokens: Number of tokens to process
dtype: Data type (float16 or bfloat16)
"""
set_global_server_args_for_scheduler(ServerArgs(model_path="dummy"))
device = "cuda"
torch_dtype = dtype
# Unpack head counts
num_heads_q, num_heads_k, num_heads_v = num_heads_group
# Calculate total hidden dimension
hidden_size = (num_heads_q + num_heads_k + num_heads_v) * head_dim
# Generate random inputs directly as 2D [num_tokens, hidden_size]
torch.random.manual_seed(0)
qkv = torch.randn(num_tokens, hidden_size, dtype=torch_dtype, device=device)
qkv_copy = qkv.clone()
# Generate position IDs with +100 offset to test decoding scenarios
position_ids = torch.arange(num_tokens, dtype=torch.int32, device=device) + 100
# Generate random weights for RMSNorm
q_weight = torch.randn(head_dim, dtype=torch_dtype, device=device) * 5.0
k_weight = torch.randn(head_dim, dtype=torch_dtype, device=device) * 5.0
# Set RMSNorm and RoPE parameters
eps = 1e-5
base = 10000.0
factor, low, high, attention_factor = 1.0, 0, 0, 1.0
# Run the custom fusedQKNormRope operation
sgl_fused_qk_norm_rope(
qkv,
num_heads_q,
num_heads_k,
num_heads_v,
head_dim,
eps,
q_weight,
k_weight,
base,
is_neox,
position_ids,
factor,
low,
high,
attention_factor,
)
output = qkv # This op is inplace
# Compute reference output using TensorRT-LLM modules
ref_output = torch_ref_rms_norm_rope(
qkv_copy,
num_heads_q,
num_heads_k,
num_heads_v,
head_dim,
eps,
q_weight,
k_weight,
base,
is_neox,
position_ids,
)
# Compare outputs from custom kernel vs reference implementation
torch.testing.assert_close(
output,
ref_output,
rtol=5e-2,
atol=1e-1,
)