80 lines
2.6 KiB
Python
80 lines
2.6 KiB
Python
"""RMSNorm with HF LlamaRMSNorm semantics (cast to dtype before weight multiply)."""
|
|
|
|
from __future__ import annotations
|
|
|
|
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,
|
|
)
|
|
|
|
if TYPE_CHECKING:
|
|
from tvm_ffi.module import Module
|
|
|
|
_CTA_BLOCK_SIZE = 512
|
|
_WARP_SIZE = 32
|
|
|
|
|
|
def is_supported_rmsnorm_hf_hidden_size(hidden_size: int) -> bool:
|
|
"""Return True iff the JIT rmsnorm_hf kernel supports this hidden size.
|
|
|
|
Two launch configs cover the practical range:
|
|
- Warp kernel: ``[32, 512)`` in multiples of 32 (q/k RMSNorm head dims).
|
|
- CTA kernel: ``>= 512`` in multiples of 512 (token RMSNorms).
|
|
"""
|
|
if _WARP_SIZE <= hidden_size < _CTA_BLOCK_SIZE and hidden_size % _WARP_SIZE == 0:
|
|
return True
|
|
return hidden_size >= _CTA_BLOCK_SIZE and hidden_size % _CTA_BLOCK_SIZE == 0
|
|
|
|
|
|
@cache_once
|
|
def _jit_rmsnorm_hf_module(hidden_size: int, dtype: torch.dtype) -> Module:
|
|
args = make_cpp_args(hidden_size, is_arch_support_pdl(), dtype)
|
|
kernel_cls = (
|
|
"HFRMSNormWarpKernel" if hidden_size < _CTA_BLOCK_SIZE else "HFRMSNormKernel"
|
|
)
|
|
return load_jit(
|
|
"rmsnorm_hf",
|
|
*args,
|
|
cuda_files=["elementwise/rmsnorm_hf.cuh"],
|
|
cuda_wrappers=[("rmsnorm_hf", f"{kernel_cls}<{args}>::run")],
|
|
)
|
|
|
|
|
|
def rmsnorm_hf(
|
|
input: torch.Tensor,
|
|
weight: torch.Tensor,
|
|
eps: float = 1e-6,
|
|
out: Optional[torch.Tensor] = None,
|
|
) -> torch.Tensor:
|
|
"""RMSNorm: ``out = weight * cast_dtype(rsqrt(mean(x^2) + eps) * x)``.
|
|
|
|
``input`` must be 2D ``(num_tokens, hidden_size)``; callers with
|
|
higher-rank tensors should reshape first. ``hidden_size`` must satisfy
|
|
:func:`is_supported_rmsnorm_hf_hidden_size`. Empty inputs return an empty
|
|
output without launching the kernel.
|
|
"""
|
|
if input.dtype not in (torch.float16, torch.bfloat16):
|
|
raise RuntimeError(f"rmsnorm_hf: input must be fp16 or bf16, got {input.dtype}")
|
|
if input.dim() != 2:
|
|
raise RuntimeError(f"rmsnorm_hf: input must be 2D, got {input.dim()}D")
|
|
hidden_size = input.size(-1)
|
|
if not is_supported_rmsnorm_hf_hidden_size(hidden_size):
|
|
raise RuntimeError(
|
|
f"rmsnorm_hf: unsupported hidden_size={hidden_size} "
|
|
f"(must be a multiple of {_WARP_SIZE} in [{_WARP_SIZE}, {_CTA_BLOCK_SIZE}) "
|
|
f"or a multiple of {_CTA_BLOCK_SIZE})"
|
|
)
|
|
if out is None:
|
|
out = torch.empty_like(input)
|
|
if input.numel() == 0:
|
|
return out
|
|
module = _jit_rmsnorm_hf_module(hidden_size, input.dtype)
|
|
module.rmsnorm_hf(input, weight, out, eps)
|
|
return out
|