diff --git a/sgl-kernel/CMakeLists.txt b/sgl-kernel/CMakeLists.txt index b025ee0bc..7b5d857a4 100644 --- a/sgl-kernel/CMakeLists.txt +++ b/sgl-kernel/CMakeLists.txt @@ -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" diff --git a/sgl-kernel/csrc/common_extension.cc b/sgl-kernel/csrc/common_extension.cc index 55ce98697..f01342fa5 100644 --- a/sgl-kernel/csrc/common_extension.cc +++ b/sgl-kernel/csrc/common_extension.cc @@ -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 */ diff --git a/sgl-kernel/csrc/moe/fused_qknorm_rope_kernel.cu b/sgl-kernel/csrc/moe/fused_qknorm_rope_kernel.cu new file mode 100644 index 000000000..0b6e52fc4 --- /dev/null +++ b/sgl-kernel/csrc/moe/fused_qknorm_rope_kernel.cu @@ -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 +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#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 +struct packed_as; + +// Specialization for packed_as used in this kernel. +template <> +struct packed_as { + using type = uint; +}; + +template <> +struct packed_as { + using type = uint2; +}; + +template <> +struct packed_as { + using type = uint4; +}; + +template +__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 +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(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(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 +__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 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 offsetThread = offsetWarp + laneId * numElemsPerThread; + + // Sum of squares for RMSNorm + float sumOfSquares = 0.0f; + + // Load. + { + vec_T vec = *reinterpret_cast(&qkv[offsetThread]); + for (int i = 0; i < vecSize; i++) { + float2 vals = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162*>(reinterpret_cast(&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(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(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(&vec) + i)) = vals; + } + vec_T* outputPtr = reinterpret_cast(&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><<>>( + 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><<>>( + 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><<>>( + 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(num_tokens), + static_cast(num_heads_q), + static_cast(num_heads_k), + static_cast(num_heads_v), + static_cast(head_dim), + static_cast(eps), + reinterpret_cast<__nv_bfloat16*>(q_weight.data_ptr()), + reinterpret_cast<__nv_bfloat16*>(k_weight.data_ptr()), + static_cast(base), + !is_neox, // interleave + reinterpret_cast(position_ids.data_ptr()), + static_cast(factor), + static_cast(low), + static_cast(high), + static_cast(attention_factor), + stream); +} diff --git a/sgl-kernel/include/sgl_kernel_ops.h b/sgl-kernel/include/sgl_kernel_ops.h index c80536894..312756c1c 100644 --- a/sgl-kernel/include/sgl_kernel_ops.h +++ b/sgl-kernel/include/sgl_kernel_ops.h @@ -378,6 +378,23 @@ void apply_shuffle_mul_sum( const torch::Tensor& permutation, const std::optional& 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, diff --git a/sgl-kernel/python/sgl_kernel/__init__.py b/sgl-kernel/python/sgl_kernel/__init__.py index 5dc22ee62..40cfd37fe 100644 --- a/sgl-kernel/python/sgl_kernel/__init__.py +++ b/sgl-kernel/python/sgl_kernel/__init__.py @@ -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, diff --git a/sgl-kernel/python/sgl_kernel/moe.py b/sgl-kernel/python/sgl_kernel/moe.py index e14eebcb2..42c20ca53 100755 --- a/sgl-kernel/python/sgl_kernel/moe.py +++ b/sgl-kernel/python/sgl_kernel/moe.py @@ -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, diff --git a/sgl-kernel/tests/test_fused_qk_norm_rope.py b/sgl-kernel/tests/test_fused_qk_norm_rope.py new file mode 100644 index 000000000..05a74facf --- /dev/null +++ b/sgl-kernel/tests/test_fused_qk_norm_rope.py @@ -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, + )