[Diffsuion & JIT_kernel] QKNorm cross heads kernel (#18073)

python/sglang/jit_kernel/tests/test_qknorm_across_heads.py

@@ -0,0 +1,75 @@
+import itertools
+
+import pytest
+import torch
+import triton
+
+
+def sglang_jit_qknorm_across_heads(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    q_weight: torch.Tensor,
+    k_weight: torch.Tensor,
+) -> None:
+    from sglang.jit_kernel.norm import fused_inplace_qknorm_across_heads
+
+    fused_inplace_qknorm_across_heads(q, k, q_weight, k_weight)
+
+
+def sglang_aot_qknorm_across_heads(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    q_weight: torch.Tensor,
+    k_weight: torch.Tensor,
+) -> None:
+    from sgl_kernel import rmsnorm
+
+    rmsnorm(q, q_weight, out=q)
+    rmsnorm(k, k_weight, out=k)
+
+
+@torch.compile()
+def torch_impl_qknorm_across_heads(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    q_weight: torch.Tensor,
+    k_weight: torch.Tensor,
+    eps: float = 1e-6,
+) -> None:
+    q_mean = q.float().pow(2).mean(dim=-1, keepdim=True)
+    k_mean = k.float().pow(2).mean(dim=-1, keepdim=True)
+    q_norm = (q_mean + eps).rsqrt()
+    k_norm = (k_mean + eps).rsqrt()
+    q.copy_(q.float() * q_norm * q_weight.float())
+    k.copy_(k.float() * k_norm * k_weight.float())
+
+
+BS_LIST = [2**n for n in range(0, 14)]
+BS_LIST += [x + 1 + i for i, x in enumerate(BS_LIST)]
+HIDDEN_DIM_LIST = [512, 1024, 2048, 4096]
+DEVICE = "cuda"
+DTYPE = torch.bfloat16
+
+
+@pytest.mark.parametrize(
+    "batch_size,hidden_dim",
+    list(itertools.product(BS_LIST, HIDDEN_DIM_LIST)),
+)
+def test_qknorm_across_heads(batch_size: int, hidden_dim: int) -> None:
+    q = torch.randn(batch_size, hidden_dim, device=DEVICE, dtype=DTYPE)
+    k = torch.randn(batch_size, hidden_dim, device=DEVICE, dtype=DTYPE)
+    q_weight = torch.randn(hidden_dim, device=DEVICE, dtype=DTYPE)
+    k_weight = torch.randn(hidden_dim, device=DEVICE, dtype=DTYPE)
+
+    q_k_jit = (q.clone(), k.clone())
+    q_k_aot = (q.clone(), k.clone())
+
+    sglang_jit_qknorm_across_heads(q_k_jit[0], q_k_jit[1], q_weight, k_weight)
+    sglang_aot_qknorm_across_heads(q_k_aot[0], q_k_aot[1], q_weight, k_weight)
+
+    triton.testing.assert_close(q_k_jit[0], q_k_aot[0], atol=1e-2, rtol=1e-2)
+    triton.testing.assert_close(q_k_jit[1], q_k_aot[1], atol=1e-2, rtol=1e-2)
+
+
+if __name__ == "__main__":
+    pytest.main([__file__])