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

2025-12-18 17:07:04 +08:00
parent fea2d5211d
commit 4792d1f452
5 changed files with 77 additions and 43 deletions
@@ -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);

  /*