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