chore: import upstream snapshot with attribution

python/sglang/jit_kernel/norm.py

@@ -0,0 +1,179 @@
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Optional
+
+import torch
+
+from sglang.jit_kernel.utils import (
+    cache_once,
+    is_arch_support_pdl,
+    load_jit,
+    make_cpp_args,
+)
+from sglang.kernel_api_logging import debug_kernel_api
+
+if TYPE_CHECKING:
+    from tvm_ffi.module import Module
+
+
+logger = logging.getLogger(__name__)
+
+
+@cache_once
+def _jit_qknorm_module(head_dim: int, dtype: torch.dtype) -> Module:
+    args = make_cpp_args(head_dim, is_arch_support_pdl(), dtype)
+    return load_jit(
+        "qknorm",
+        *args,
+        cuda_files=["elementwise/qknorm.cuh"],
+        cuda_wrappers=[("qknorm", f"QKNormKernel<{args}>::run")],
+    )
+
+
+_RMSNORM_WARP_SIZES = frozenset({64, 128, 256})
+_RMSNORM_MAX_HIDDEN_SIZE = 16384
+_RMSNORM_HALF_BLOCK_MIN_SIZE = 2048
+
+
+def _is_supported_rmsnorm_hidden_size(d: int) -> bool:
+    return d in _RMSNORM_WARP_SIZES or (
+        (d > 256 and d % 256 == 0 and d <= 8192)
+        or (d >= 8192 and d % 512 == 0 and d <= 16384)
+    )
+
+
+def _rmsnorm_kernel_class(hidden_size: int) -> str:
+    if hidden_size in _RMSNORM_WARP_SIZES:
+        return "RMSNormWarpKernel"
+    if hidden_size == 512:
+        return "RMSNormHalfKernel"
+    if hidden_size >= _RMSNORM_HALF_BLOCK_MIN_SIZE:
+        if hidden_size % 512 == 0:
+            return "RMSNormHalfKernel"
+    return "RMSNormKernel"
+
+
+@cache_once
+def _jit_rmsnorm_module(hidden_size: int, dtype: torch.dtype) -> Module:
+    args = make_cpp_args(hidden_size, is_arch_support_pdl(), dtype)
+    kernel_class = f"{_rmsnorm_kernel_class(hidden_size)}<{args}>"
+    return load_jit(
+        "rmsnorm",
+        *args,
+        cuda_files=["elementwise/rmsnorm.cuh"],
+        cuda_wrappers=[("rmsnorm", f"{kernel_class}::run")],
+    )
+
+
+def is_supported_jit_fused_add_rmsnorm_hidden_size(hidden_size: int) -> bool:
+    return hidden_size > 0 and hidden_size % 16 == 0 and hidden_size <= 8192
+
+
+@cache_once
+def _jit_fused_add_rmsnorm_module(
+    dtype: torch.dtype, cast_x_before_out_mul: bool
+) -> Module:
+    args = make_cpp_args(cast_x_before_out_mul, dtype)
+    return load_jit(
+        "fused_add_rmsnorm",
+        *args,
+        cuda_files=["elementwise/fused_add_rmsnorm.cuh"],
+        cuda_wrappers=[("fused_add_rmsnorm", f"FusedAddRMSNormKernel<{args}>::run")],
+    )
+
+
+@cache_once
+def _jit_qknorm_across_heads_module(dtype: torch.dtype) -> Module:
+    args = make_cpp_args(dtype)
+    return load_jit(
+        "qknorm_across_heads",
+        *args,
+        cuda_files=["elementwise/qknorm_across_heads.cuh"],
+        cuda_wrappers=[
+            ("qknorm_across_heads", f"QKNormAcrossHeadsKernel<{args}>::run")
+        ],
+    )
+
+
+@torch.compiler.assume_constant_result
+@cache_once
+def can_use_fused_inplace_qknorm(head_dim: int, dtype: torch.dtype) -> bool:
+    if head_dim not in [64, 128, 256, 512, 1024]:
+        logger.warning(f"Unsupported head_dim={head_dim} for JIT QK-Norm kernel")
+        return False
+    try:
+        _jit_qknorm_module(head_dim, dtype)
+        return True
+    except Exception as e:
+        logger.warning(f"Failed to load JIT QK-Norm kernel: {e}")
+        return False
+
+
+@debug_kernel_api
+def fused_inplace_qknorm(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    q_weight: torch.Tensor,
+    k_weight: torch.Tensor,
+    eps: float = 1e-6,
+    *,
+    head_dim: int = 0,
+) -> None:
+    head_dim = head_dim or q.size(-1)
+    module = _jit_qknorm_module(head_dim, q.dtype)
+    module.qknorm(q, k, q_weight, k_weight, eps)
+
+
+@debug_kernel_api
+def rmsnorm(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    out: Optional[torch.Tensor] = None,
+    eps: float = 1e-6,
+) -> None:
+    out = out if out is not None else input
+    hidden_size = input.size(-1)
+    if not _is_supported_rmsnorm_hidden_size(hidden_size):
+        raise RuntimeError(
+            f"jit rmsnorm: unsupported hidden_size={hidden_size}. "
+            f"Supported: {sorted(_RMSNORM_WARP_SIZES)}, and multiples of 256 in "
+            f"(256, {_RMSNORM_MAX_HIDDEN_SIZE}]."
+        )
+    module = _jit_rmsnorm_module(hidden_size, input.dtype)
+    module.rmsnorm(input, weight, out, eps)
+
+
+@debug_kernel_api
+def fused_add_rmsnorm(
+    input: torch.Tensor,
+    residual: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float = 1e-6,
+    *,
+    cast_x_before_out_mul: bool = False,
+) -> None:
+    module = _jit_fused_add_rmsnorm_module(input.dtype, cast_x_before_out_mul)
+    module.fused_add_rmsnorm(input, residual, weight, eps)
+
+
+@debug_kernel_api
+def fused_inplace_qknorm_across_heads(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    q_weight: torch.Tensor,
+    k_weight: torch.Tensor,
+    eps: float = 1e-6,
+) -> None:
+    """
+    Fused inplace QK normalization across all heads.
+
+    Args:
+        q: Query tensor of shape [batch_size, num_heads * head_dim]
+        k: Key tensor of shape [batch_size, num_heads * head_dim]
+        q_weight: Query weight tensor of shape [num_heads * head_dim]
+        k_weight: Key weight tensor of shape [num_heads * head_dim]
+        eps: Epsilon for numerical stability
+    """
+    module = _jit_qknorm_across_heads_module(q.dtype)
+    module.qknorm_across_heads(q, k, q_weight, k_weight, eps)