[Diffusion] Add qknorm rope fuse kernel (#21440)

python/sglang/jit_kernel/diffusion/qknorm_rope.py

@@ -0,0 +1,97 @@
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING
+
+import torch
+
+from sglang.jit_kernel.utils import (
+    cache_once,
+    is_arch_support_pdl,
+    load_jit,
+    make_cpp_args,
+)
+from sglang.srt.utils.custom_op import register_custom_op
+
+if TYPE_CHECKING:
+    from tvm_ffi.module import Module
+
+
+logger = logging.getLogger(__name__)
+
+
+@cache_once
+def _jit_qknorm_rope_module(
+    head_dim: int,
+    rope_dim: int,
+    is_neox: bool,
+    dtype: torch.dtype,
+) -> Module:
+    args = make_cpp_args(head_dim, rope_dim, is_neox, is_arch_support_pdl(), dtype)
+    return load_jit(
+        "qknorm_rope",
+        *args,
+        cuda_files=["diffusion/qknorm_rope.cuh"],
+        cuda_wrappers=[("qknorm_rope", f"QKNormRopeKernel<{args}>::run")],
+    )
+
+
+@torch.compiler.assume_constant_result
+@cache_once
+def can_use_fused_inplace_qknorm_rope(
+    head_dim: int,
+    rope_dim: int,
+    is_neox: bool,
+    dtype: torch.dtype,
+) -> bool:
+    if head_dim not in (64, 128, 256):
+        logger.warning(f"Unsupported head_dim={head_dim} for JIT fused QKNorm+RoPE")
+        return False
+    if rope_dim <= 0 or rope_dim > head_dim:
+        logger.warning(
+            f"Unsupported rope_dim={rope_dim} for head_dim={head_dim} in fused QKNorm+RoPE"
+        )
+        return False
+    elems_per_thread = head_dim // 32
+    if rope_dim % elems_per_thread != 0:
+        logger.warning(
+            "rope_dim=%s must be divisible by per-thread width=%s for fused QKNorm+RoPE",
+            rope_dim,
+            elems_per_thread,
+        )
+        return False
+    if is_neox:
+        rotary_lanes = rope_dim // elems_per_thread
+        if rotary_lanes < 2 or rotary_lanes & (rotary_lanes - 1):
+            logger.warning(
+                "rope_dim=%s yields invalid rotary_lanes=%s for neox fused QKNorm+RoPE; rotary lane count must be a power of 2",
+                rope_dim,
+                rotary_lanes,
+            )
+            return False
+    try:
+        _jit_qknorm_rope_module(head_dim, rope_dim, is_neox, dtype)
+        return True
+    except Exception as e:
+        logger.warning(f"Failed to load JIT fused QKNorm+RoPE kernel: {e}")
+        return False
+
+
+@register_custom_op(mutates_args=["q", "k"])
+def fused_inplace_qknorm_rope(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    q_weight: torch.Tensor,
+    k_weight: torch.Tensor,
+    cos_sin_cache: torch.Tensor,
+    positions: torch.Tensor,
+    *,
+    is_neox: bool,
+    eps: float = 1e-6,
+    head_dim: int = 0,
+    rope_dim: int = 0,
+) -> None:
+    head_dim = head_dim or q.size(-1)
+    rope_dim = rope_dim or cos_sin_cache.size(-1)
+    module = _jit_qknorm_rope_module(head_dim, rope_dim, is_neox, q.dtype)
+    module.qknorm_rope(q, k, q_weight, k_weight, cos_sin_cache, positions, eps)