[sgl-kernel][1/2] Fused qk_norm_rope for GLM4.6 (#15141)

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);
 
   /*

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<float>(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<float>(low),
       static_cast<float>(high),
       static_cast<float>(attention_factor),
+      static_cast<int>(rotary_dim),
       stream);
 }

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,

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,
     )
 
 

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