diff --git a/sgl-kernel/csrc/common_extension.cc b/sgl-kernel/csrc/common_extension.cc index 0145676e6..8c36e9a8f 100644 --- a/sgl-kernel/csrc/common_extension.cc +++ b/sgl-kernel/csrc/common_extension.cc @@ -278,7 +278,8 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) { "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) -> ()"); + "bool is_neox, Tensor position_ids, float factor, float low, float high, float attention_factor, int rotary_dim) " + "-> ()"); m.impl("fused_qk_norm_rope", torch::kCUDA, &fused_qk_norm_rope); /* diff --git a/sgl-kernel/csrc/moe/fused_qknorm_rope_kernel.cu b/sgl-kernel/csrc/moe/fused_qknorm_rope_kernel.cu index 0b6e52fc4..643fe65ff 100644 --- a/sgl-kernel/csrc/moe/fused_qknorm_rope_kernel.cu +++ b/sgl-kernel/csrc/moe/fused_qknorm_rope_kernel.cu @@ -120,8 +120,8 @@ __global__ void fusedQKNormRopeKernel( 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 -) { + float attention_factor, // attention_factor applied on cos and sin + int const rotary_dim) { int const warpsPerBlock = blockDim.x / 32; int const warpId = threadIdx.x / 32; int const laneId = threadIdx.x % 32; @@ -195,43 +195,51 @@ __global__ void fusedQKNormRopeKernel( float cos_vals[numElemsPerThread]; float sin_vals[numElemsPerThread]; float pos_id = static_cast(position_ids[tokenIdx]); + int const rotary_lanes = rotary_dim / numElemsPerThread; // rotary range + bool const applyRotary = (laneId < rotary_lanes); + if (applyRotary) { + if constexpr (interleave) { + // Perform interleaving. Fill cos_vals and sin_vals. + for (int i = 0; i < numElemsPerThread; i++) { + elements2[i] = (i % 2 == 0) ? -elements[i + 1] : elements[i - 1]; - 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, rotary_dim, half_dim, factor, low, high); + float theta = pos_id * freq; + __sincosf(theta, &sin_vals[i], &cos_vals[i]); } - 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]; + + } else { + // Neox style + // Before data exchange with in warp, we need to sync. + __syncwarp(); + int const half_rotary_lanes = rotary_lanes / 2; + unsigned int active_mask = (1u << rotary_lanes) - 1; + // Limitation: The operation below requires half_rotary_lanes to be a power of 2. + // because it relies on __shfl_xor_sync to exchange data within a warp. + for (int i = 0; i < numElemsPerThread; i++) { + elements2[i] = __shfl_xor_sync(active_mask, elements[i], half_rotary_lanes); + if (laneId < half_rotary_lanes) { + elements2[i] = -elements2[i]; + } + + int dim_idx = laneId * numElemsPerThread + i; + dim_idx = (dim_idx * 2) % rotary_dim; + int half_dim = dim_idx / 2; + float freq = compute_freq_yarn(base, rotary_dim, half_dim, factor, low, high); + float theta = pos_id * freq; + __sincosf(theta, &sin_vals[i], &cos_vals[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; } - // __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; @@ -270,6 +278,7 @@ void launchFusedQKNormRope( float low, float high, float attention_factor, + int const rotary_dim, cudaStream_t stream) { constexpr int blockSize = 256; int const warpsPerBlock = blockSize / 32; @@ -297,7 +306,8 @@ void launchFusedQKNormRope( factor, low, high, - attention_factor); + attention_factor, + rotary_dim); }); break; case 128: @@ -316,7 +326,8 @@ void launchFusedQKNormRope( factor, low, high, - attention_factor); + attention_factor, + rotary_dim); }); break; case 256: @@ -335,7 +346,8 @@ void launchFusedQKNormRope( factor, low, high, - attention_factor); + attention_factor, + rotary_dim); }); break; default: @@ -363,8 +375,8 @@ void fused_qk_norm_rope( 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 -) { + double attention_factor, // attention_factor applied on cos and sin + int64_t rotary_dim) { // 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]"); @@ -372,7 +384,14 @@ void fused_qk_norm_rope( 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"); - + TORCH_CHECK(rotary_dim % (head_dim / 32) == 0, "rotary_dim must be divisible by numElemsPerThread"); + if (is_neox) { + int64_t half_rotary_lanes = rotary_dim / (head_dim / 32) / 2; + TORCH_CHECK( + half_rotary_lanes >= 1 && (half_rotary_lanes & (half_rotary_lanes - 1)) == 0, + "half_rotary_lanes must be a power of 2 for neox style, got ", + half_rotary_lanes); + } CHECK_INPUT(qkv, torch::kBFloat16); CHECK_INPUT(position_ids, torch::kInt32); CHECK_INPUT(q_weight, torch::kBFloat16); @@ -404,5 +423,6 @@ void fused_qk_norm_rope( static_cast(low), static_cast(high), static_cast(attention_factor), + static_cast(rotary_dim), stream); } diff --git a/sgl-kernel/include/sgl_kernel_ops.h b/sgl-kernel/include/sgl_kernel_ops.h index e171ab316..2eb0856aa 100644 --- a/sgl-kernel/include/sgl_kernel_ops.h +++ b/sgl-kernel/include/sgl_kernel_ops.h @@ -401,7 +401,8 @@ void fused_qk_norm_rope( double factor, double low, double high, - double attention_factor); + double attention_factor, + int64_t rotary_dim); void cutlass_fp4_group_mm( torch::Tensor& output, diff --git a/sgl-kernel/python/sgl_kernel/moe.py b/sgl-kernel/python/sgl_kernel/moe.py index 42c20ca53..d85e4b602 100755 --- a/sgl-kernel/python/sgl_kernel/moe.py +++ b/sgl-kernel/python/sgl_kernel/moe.py @@ -267,6 +267,7 @@ def fused_qk_norm_rope( low: float, high: float, attention_factor: float, + rotary_dim: Optional[int] = None, ) -> None: torch.ops.sgl_kernel.fused_qk_norm_rope( qkv, @@ -284,6 +285,7 @@ def fused_qk_norm_rope( low, high, attention_factor, + rotary_dim if rotary_dim is not None else head_dim, ) diff --git a/sgl-kernel/tests/test_fused_qk_norm_rope.py b/sgl-kernel/tests/test_fused_qk_norm_rope.py index 05a74facf..7f1daf5aa 100644 --- a/sgl-kernel/tests/test_fused_qk_norm_rope.py +++ b/sgl-kernel/tests/test_fused_qk_norm_rope.py @@ -41,6 +41,7 @@ def torch_ref_rms_norm_rope( base, is_neox, position_ids, + partial_rotary_factor, ): """ PyTorch reference implementation of RMSNorm+RoPE for verification. @@ -60,6 +61,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] + partial_rotary_factor: Partial rotary factor Returns: Combined tensor with Q and K parts normalized and RoPE applied @@ -91,6 +93,7 @@ def torch_ref_rms_norm_rope( is_neox_style=is_neox, rope_scaling=None, dual_chunk_attention_config=None, + partial_rotary_factor=partial_rotary_factor, ) rotary_emb = rotary_emb.to(qkv.device) @@ -127,10 +130,12 @@ num_heads_groups = [ (32, 8, 8), # Qwen3-4B, Qwen3-8B, Qwen3-30B-A3B (40, 8, 8), # Qwen3-14B (64, 8, 8), # Qwen3-32B, Qwen3-235B-A22B + (12, 1, 1), # GLM4.6 TP8 ] num_tokens_list = [1, 3, 8, 32, 256] is_neox_list = [False, True] dtypes = [torch.bfloat16] +partial_rotary_factor_list = [1.0, 0.5] @pytest.mark.skipif(not _is_cuda, reason="Skipping CUDA/ROCm only tests.") @@ -139,7 +144,10 @@ dtypes = [torch.bfloat16] @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): +@pytest.mark.parametrize("partial_rotary_factor", partial_rotary_factor_list) +def test_fused_qk_norm_rope( + head_dim, num_heads_group, num_tokens, is_neox, dtype, partial_rotary_factor +): """ Test the fused QK RMSNorm + RoPE operation with various configurations. @@ -198,6 +206,7 @@ def test_fused_qk_norm_rope(head_dim, num_heads_group, num_tokens, is_neox, dtyp low, high, attention_factor, + int(head_dim * partial_rotary_factor), ) output = qkv # This op is inplace @@ -214,6 +223,7 @@ def test_fused_qk_norm_rope(head_dim, num_heads_group, num_tokens, is_neox, dtyp base, is_neox, position_ids, + partial_rotary_factor, ) # Compare outputs from custom kernel vs reference implementation