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

Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>

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,

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,