KernelSwift: Optimize gated RMSNorm execution

scripts/run_gated_rms_norm_smoke.py

@@ -0,0 +1,77 @@
+import argparse
+
+import torch
+import torch.nn.functional as F
+
+import ntops
+
+
+def reference(
+    input, gate, weight, eps, group_size, norm_before_gate, activation
+):
+    input_float = input.float()
+    gate_float = gate.float()
+    activation_fn = torch.sigmoid if activation == "sigmoid" else F.silu
+
+    if not norm_before_gate:
+        input_float = input_float * activation_fn(gate_float)
+
+    grouped = input_float.reshape(*input.shape[:-1], -1, group_size)
+    variance = grouped.square().mean(dim=-1, keepdim=True)
+    output = grouped * torch.rsqrt(variance + eps)
+    output = output.reshape_as(input_float) * weight.float()
+
+    if norm_before_gate:
+        output = output * activation_fn(gate_float)
+
+    return output.to(input.dtype)
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--device", default="cuda")
+    args = parser.parse_args()
+
+    torch.manual_seed(42)
+    cases = (
+        ((3, 64), torch.float16, 64, False, "swish"),
+        ((2, 5, 128), torch.float16, 32, True, "swish"),
+        ((2, 3, 256), torch.float16, 64, False, "sigmoid"),
+        ((2, 7, 96), torch.float16, 48, False, "swish"),
+        ((4, 128), torch.bfloat16, 128, True, "sigmoid"),
+    )
+
+    for shape, dtype, group_size, norm_before_gate, activation in cases:
+        input = torch.randn(shape, device=args.device, dtype=dtype)
+        gate = torch.randn_like(input)
+        weight = torch.randn(shape[-1], device=args.device, dtype=torch.float16)
+        actual = ntops.torch.gated_rms_norm(
+            input,
+            gate,
+            weight,
+            group_size=group_size,
+            norm_before_gate=norm_before_gate,
+            activation=activation,
+        )
+        expected = reference(
+            input,
+            gate,
+            weight,
+            1e-5,
+            group_size,
+            norm_before_gate,
+            activation,
+        )
+        torch.testing.assert_close(actual, expected, rtol=2e-2, atol=2e-2)
+        print(
+            "PASS",
+            f"shape={shape}",
+            f"dtype={dtype}",
+            f"group_size={group_size}",
+            f"norm_before_gate={norm_before_gate}",
+            f"activation={activation}",
+        )
+
+
+if __name__ == "__main__":
+    main()

src/ntops/kernels/__init__.py

@@ -19,6 +19,7 @@
     eq,
     exp,
     ge,
+    gated_rms_norm,
     gelu,
     gt,
     instance_norm,
@@ -97,6 +98,7 @@
     "eq",
     "exp",
     "ge",
+    "gated_rms_norm",
     "gelu",
     "gt",
     "instance_norm",

src/ntops/kernels/gated_rms_norm.py

@@ -0,0 +1,123 @@
+import enum
+import functools
+
+import ninetoothed
+import ninetoothed.language as ntl
+from ninetoothed import Tensor
+
+
+class ActivationVariant(enum.IntEnum):
+    SILU = enum.auto()
+    SIGMOID = enum.auto()
+
+
+def arrangement(*tensors, block_size=None):
+    """Arrange grouped rows as one row-vector program per reduction group.
+
+    The wrapper presents tensors as ``prefix + (num_groups, group_size)``.
+    Flattening every dimension except the final group dimension gives a
+    logical ``(R, G)`` matrix.  Tiling it as ``(1, block_size)`` leaves the
+    row dimension as the program grid and gives each application a
+    ``(1, block_size)`` vector, with ``other=0`` supplying the tail padding.
+    """
+    if block_size is None:
+        block_size = ninetoothed.block_size()
+
+    def _arrange(tensor):
+        if tensor.ndim == 0:
+            return tensor
+        return tensor.flatten(end_dim=-1).tile((1, block_size))
+
+    return tuple(_arrange(tensor) for tensor in tensors)
+
+
+def application_norm_before_silu(
+    input, gate, weight, eps, output, num_normalized_elements
+):
+    value = input.to(ntl.float32)
+    rms = ntl.sqrt(ntl.sum(value * value) / num_normalized_elements + eps)
+    gate_float = gate.to(ntl.float32)
+    activated_gate = gate_float / (1 + ntl.exp(-gate_float))
+    output = value / rms * weight.to(ntl.float32) * activated_gate
+
+
+def application_norm_before_sigmoid(
+    input, gate, weight, eps, output, num_normalized_elements
+):
+    value = input.to(ntl.float32)
+    rms = ntl.sqrt(ntl.sum(value * value) / num_normalized_elements + eps)
+    gate_float = gate.to(ntl.float32)
+    activated_gate = 1 / (1 + ntl.exp(-gate_float))
+    output = value / rms * weight.to(ntl.float32) * activated_gate
+
+
+def application_norm_after_silu(
+    input, gate, weight, eps, output, num_normalized_elements
+):
+    value = input.to(ntl.float32)
+    gate_float = gate.to(ntl.float32)
+    value *= gate_float / (1 + ntl.exp(-gate_float))
+    rms = ntl.sqrt(ntl.sum(value * value) / num_normalized_elements + eps)
+    output = value / rms * weight.to(ntl.float32)
+
+
+def application_norm_after_sigmoid(
+    input, gate, weight, eps, output, num_normalized_elements
+):
+    value = input.to(ntl.float32)
+    gate_float = gate.to(ntl.float32)
+    value *= 1 / (1 + ntl.exp(-gate_float))
+    rms = ntl.sqrt(ntl.sum(value * value) / num_normalized_elements + eps)
+    output = value / rms * weight.to(ntl.float32)
+
+
+def application_norm_only(input, weight, eps, output, num_normalized_elements):
+    value = input.to(ntl.float32)
+    rms = ntl.sqrt(ntl.sum(value * value) / num_normalized_elements + eps)
+    output = value / rms * weight.to(ntl.float32)
+
+
+def premake(
+    ndim,
+    norm_before_gate=False,
+    activation=ActivationVariant.SILU,
+    input_dtype=None,
+    gate_dtype=None,
+    weight_dtype=None,
+    output_dtype=None,
+    block_size=None,
+    has_gate=True,
+):
+    arrangement_ = functools.partial(arrangement, block_size=block_size)
+
+    if not has_gate:
+        application = application_norm_only
+    elif norm_before_gate:
+        if activation == ActivationVariant.SIGMOID:
+            application = application_norm_before_sigmoid
+        else:
+            application = application_norm_before_silu
+    elif activation == ActivationVariant.SIGMOID:
+        application = application_norm_after_sigmoid
+    else:
+        application = application_norm_after_silu
+
+    if has_gate:
+        tensors = (
+            Tensor(ndim, other=0, dtype=input_dtype),
+            Tensor(ndim, other=0, dtype=gate_dtype),
+            Tensor(ndim, other=0, dtype=weight_dtype),
+            Tensor(0, dtype=ninetoothed.float64),
+            Tensor(ndim, dtype=output_dtype),
+            Tensor(0, constexpr=True),
+        )
+    else:
+        tensors = (
+            Tensor(ndim, other=0, dtype=input_dtype),
+            Tensor(ndim, other=0, dtype=weight_dtype),
+            Tensor(0, dtype=ninetoothed.float64),
+            Tensor(ndim, dtype=output_dtype),
+            Tensor(0, constexpr=True),
+        )
+
+    return arrangement_, application, tensors

src/ntops/torch/__init__.py

@@ -18,6 +18,7 @@
 from ntops.torch.eq import eq
 from ntops.torch.exp import exp
 from ntops.torch.ge import ge
+from ntops.torch.gated_rms_norm import gated_rms_norm
 from ntops.torch.gelu import gelu
 from ntops.torch.gt import gt
 from ntops.torch.instance_norm import instance_norm
@@ -96,6 +97,7 @@
     "eq",
     "exp",
     "ge",
+    "gated_rms_norm",
     "gelu",
     "gt",
     "instance_norm",

src/ntops/torch/gated_rms_norm.py

@@ -0,0 +1,77 @@
+import torch
+
+import ntops
+from ntops.kernels.gated_rms_norm import ActivationVariant
+from ntops.torch.utils import _cached_make
+
+_ACTIVATION_VARIANTS = {
+    "silu": ActivationVariant.SILU,
+    "swish": ActivationVariant.SILU,
+    "sigmoid": ActivationVariant.SIGMOID,
+}
+
+
+def gated_rms_norm(
+    input,
+    gate=None,
+    weight=None,
+    eps=1e-5,
+    group_size=None,
+    norm_before_gate=False,
+    activation="swish",
+):
+    if weight is None:
+        raise ValueError("`weight` must be provided.")
+
+    if gate is not None and input.shape != gate.shape:
+        raise ValueError("`input` and `gate` must have the same shape.")
+
+    input_shape = input.shape
+    prefix = input_shape[:-1]
+    hidden_size = input_shape[-1]
+
+    if weight.shape != (hidden_size,):
+        raise ValueError("`weight` must have shape `(input.shape[-1],)`.")
+
+    if group_size is None:
+        group_size = hidden_size
+
+    if group_size <= 0 or hidden_size % group_size != 0:
+        raise ValueError("`group_size` must be a positive divisor of the hidden size.")
+
+    try:
+        activation_variant = _ACTIVATION_VARIANTS[activation]
+    except KeyError as error:
+        raise ValueError(
+            "`activation` must be one of `silu`, `swish`, or `sigmoid`."
+        ) from error
+
+    num_groups = hidden_size // group_size
+    grouped_shape = prefix + (num_groups, group_size)
+    grouped_input = input.reshape(grouped_shape)
+    grouped_gate = None if gate is None else gate.reshape(grouped_shape)
+    grouped_weight = weight.reshape((1,) * len(prefix) + (num_groups, group_size))
+    grouped_weight = grouped_weight.expand(grouped_shape)
+    grouped_output = torch.empty_like(grouped_input)
+
+    kernel = _cached_make(
+        ntops.kernels.gated_rms_norm.premake,
+        grouped_input.ndim,
+        has_gate=gate is not None,
+        norm_before_gate=norm_before_gate,
+        activation=activation_variant,
+        block_size=1 << (group_size - 1).bit_length(),
+    )
+    if gate is None:
+        kernel(grouped_input, grouped_weight, eps, grouped_output, group_size)
+    else:
+        kernel(
+            grouped_input,
+            grouped_gate,
+            grouped_weight,
+            eps,
+            grouped_output,
+            group_size,
+        )
+
+    return grouped_output.reshape(input_shape)