From ff548443311e22061770b13616a99516bb0d631a Mon Sep 17 00:00:00 2001 From: CearX <56916338+CearX@users.noreply.github.com> Date: Mon, 31 Aug 2026 23:22:23 +0800 Subject: [PATCH 1/6] Add gated RMSNorm operator --- scripts/run_gated_rms_norm_smoke.py | 77 +++++++++++++++++++++++++ src/ntops/kernels/__init__.py | 2 + src/ntops/kernels/gated_rms_norm.py | 87 +++++++++++++++++++++++++++++ src/ntops/torch/__init__.py | 2 + src/ntops/torch/gated_rms_norm.py | 66 ++++++++++++++++++++++ tests/test_gated_rms_norm.py | 83 +++++++++++++++++++++++++++ 6 files changed, 317 insertions(+) create mode 100644 scripts/run_gated_rms_norm_smoke.py create mode 100644 src/ntops/kernels/gated_rms_norm.py create mode 100644 src/ntops/torch/gated_rms_norm.py create mode 100644 tests/test_gated_rms_norm.py diff --git a/scripts/run_gated_rms_norm_smoke.py b/scripts/run_gated_rms_norm_smoke.py new file mode 100644 index 0000000..42d2116 --- /dev/null +++ b/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() diff --git a/src/ntops/kernels/__init__.py b/src/ntops/kernels/__init__.py index 12d337b..7a99880 100644 --- a/src/ntops/kernels/__init__.py +++ b/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", diff --git a/src/ntops/kernels/gated_rms_norm.py b/src/ntops/kernels/gated_rms_norm.py new file mode 100644 index 0000000..433b2ab --- /dev/null +++ b/src/ntops/kernels/gated_rms_norm.py @@ -0,0 +1,87 @@ +import enum +import functools + +import ninetoothed +import ninetoothed.language as ntl +from ninetoothed import Tensor + +from ntops.kernels.reduction import arrangement + + +class ActivationVariant(enum.IntEnum): + SILU = enum.auto() + SIGMOID = enum.auto() + + +def application_norm_before_silu( + input, gate, weight, eps, output, num_normalized_elements +): + value = input[0].to(ntl.float32) + rms = ntl.sqrt(ntl.sum(value * value, axis=0) / num_normalized_elements + eps) + gate_float = gate[0].to(ntl.float32) + activated_gate = gate_float / (1 + ntl.exp(-gate_float)) + output[0] = value / rms * weight[0].to(ntl.float32) * activated_gate + + +def application_norm_before_sigmoid( + input, gate, weight, eps, output, num_normalized_elements +): + value = input[0].to(ntl.float32) + rms = ntl.sqrt(ntl.sum(value * value, axis=0) / num_normalized_elements + eps) + gate_float = gate[0].to(ntl.float32) + activated_gate = 1 / (1 + ntl.exp(-gate_float)) + output[0] = value / rms * weight[0].to(ntl.float32) * activated_gate + + +def application_norm_after_silu( + input, gate, weight, eps, output, num_normalized_elements +): + value = input[0].to(ntl.float32) + gate_float = gate[0].to(ntl.float32) + value *= gate_float / (1 + ntl.exp(-gate_float)) + rms = ntl.sqrt(ntl.sum(value * value, axis=0) / num_normalized_elements + eps) + output[0] = value / rms * weight[0].to(ntl.float32) + + +def application_norm_after_sigmoid( + input, gate, weight, eps, output, num_normalized_elements +): + value = input[0].to(ntl.float32) + gate_float = gate[0].to(ntl.float32) + value *= 1 / (1 + ntl.exp(-gate_float)) + rms = ntl.sqrt(ntl.sum(value * value, axis=0) / num_normalized_elements + eps) + output[0] = value / rms * weight[0].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, +): + arrangement_ = functools.partial(arrangement, dim=-1, block_size=block_size) + + if 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 + + tensors = ( + Tensor(ndim, other=0, dtype=input_dtype), + Tensor(ndim, other=0, dtype=gate_dtype), + Tensor(ndim, dtype=weight_dtype), + Tensor(0, dtype=ninetoothed.float64), + Tensor(ndim, dtype=output_dtype), + Tensor(0, constexpr=True), + ) + + return arrangement_, application, tensors diff --git a/src/ntops/torch/__init__.py b/src/ntops/torch/__init__.py index ad6fd4c..8ea5d36 100644 --- a/src/ntops/torch/__init__.py +++ b/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", diff --git a/src/ntops/torch/gated_rms_norm.py b/src/ntops/torch/gated_rms_norm.py new file mode 100644 index 0000000..30d846c --- /dev/null +++ b/src/ntops/torch/gated_rms_norm.py @@ -0,0 +1,66 @@ +import torch + +import ntops +from ntops.kernels.gated_rms_norm import ActivationVariant +from ntops.torch.utils import _cached_make + + +def gated_rms_norm( + input, + gate, + weight, + eps=1e-5, + group_size=None, + norm_before_gate=False, + activation="swish", +): + if input.shape != gate.shape: + raise ValueError("`input` and `gate` must have the same shape.") + + 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.") + + activation_variants = { + "silu": ActivationVariant.SILU, + "swish": ActivationVariant.SILU, + "sigmoid": ActivationVariant.SIGMOID, + } + + try: + activation_variant = activation_variants[activation] + except KeyError as error: + raise ValueError( + "`activation` must be one of `silu`, `swish`, or `sigmoid`." + ) from error + + grouped_shape = input.shape[:-1] + (hidden_size // group_size, group_size) + grouped_input = input.reshape(grouped_shape) + grouped_gate = gate.reshape(grouped_shape) + grouped_weight = weight.expand_as(input).reshape(grouped_shape) + grouped_output = torch.empty_like(grouped_input) + + kernel = _cached_make( + ntops.kernels.gated_rms_norm.premake, + grouped_input.ndim, + norm_before_gate=norm_before_gate, + activation=activation_variant, + block_size=1 << (group_size - 1).bit_length(), + ) + kernel( + grouped_input, + grouped_gate, + grouped_weight, + eps, + grouped_output, + group_size, + ) + + return grouped_output.reshape(input.shape) diff --git a/tests/test_gated_rms_norm.py b/tests/test_gated_rms_norm.py new file mode 100644 index 0000000..f34ac42 --- /dev/null +++ b/tests/test_gated_rms_norm.py @@ -0,0 +1,83 @@ +import pytest +import torch +import torch.nn.functional as F + +import ntops +from tests.skippers import skip_if_cuda_not_available + + +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) + + effective_group_size = group_size or input.shape[-1] + grouped = input_float.reshape(*input.shape[:-1], -1, effective_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) + + +@skip_if_cuda_not_available +@pytest.mark.parametrize("shape", ((3, 64), (2, 5, 128), (2, 3, 7, 256))) +@pytest.mark.parametrize("dtype", (torch.float16, torch.bfloat16)) +@pytest.mark.parametrize("group_size", (None, 32, 64)) +@pytest.mark.parametrize("norm_before_gate", (False, True)) +@pytest.mark.parametrize("activation", ("swish", "sigmoid")) +def test_gated_rms_norm( + shape, dtype, group_size, norm_before_gate, activation +): + input = torch.randn(shape, device="cuda", dtype=dtype) + gate = torch.randn_like(input) + weight = torch.randn(shape[-1], device="cuda", dtype=dtype) + + actual = ntops.torch.gated_rms_norm( + input, + gate, + weight, + eps=1e-5, + 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) + + +@skip_if_cuda_not_available +def test_gated_rms_norm_non_contiguous_input(): + input = torch.randn((4, 3, 128), device="cuda", dtype=torch.float16).transpose( + 0, 1 + ) + gate = torch.randn_like(input) + weight = torch.randn(128, device="cuda", dtype=torch.float16) + + actual = ntops.torch.gated_rms_norm(input, gate, weight, group_size=64) + expected = _reference(input, gate, weight, 1e-5, 64, False, "swish") + + torch.testing.assert_close(actual, expected, rtol=2e-2, atol=2e-2) + + +@skip_if_cuda_not_available +def test_gated_rms_norm_non_power_of_two_group(): + input = torch.randn((2, 7, 96), device="cuda", dtype=torch.float16) + gate = torch.randn_like(input) + weight = torch.randn(96, device="cuda", dtype=torch.float16) + + actual = ntops.torch.gated_rms_norm(input, gate, weight, group_size=48) + expected = _reference(input, gate, weight, 1e-5, 48, False, "swish") + + torch.testing.assert_close(actual, expected, rtol=2e-2, atol=2e-2) From e7502ff58107967c6c5275058f052b0bcd54d971 Mon Sep 17 00:00:00 2001 From: CearX <56916338+CearX@users.noreply.github.com> Date: Mon, 31 Aug 2026 23:22:23 +0800 Subject: [PATCH 2/6] Align gated RMSNorm with vLLM contract --- src/ntops/kernels/gated_rms_norm.py | 36 ++++++++++++----- src/ntops/torch/gated_rms_norm.py | 31 +++++++++------ tests/test_gated_rms_norm.py | 60 +++++++++++++++++++++++++---- 3 files changed, 99 insertions(+), 28 deletions(-) diff --git a/src/ntops/kernels/gated_rms_norm.py b/src/ntops/kernels/gated_rms_norm.py index 433b2ab..0464a95 100644 --- a/src/ntops/kernels/gated_rms_norm.py +++ b/src/ntops/kernels/gated_rms_norm.py @@ -53,6 +53,12 @@ def application_norm_after_sigmoid( output[0] = value / rms * weight[0].to(ntl.float32) +def application_norm_only(input, weight, eps, output, num_normalized_elements): + value = input[0].to(ntl.float32) + rms = ntl.sqrt(ntl.sum(value * value, axis=0) / num_normalized_elements + eps) + output[0] = value / rms * weight[0].to(ntl.float32) + + def premake( ndim, norm_before_gate=False, @@ -62,10 +68,13 @@ def premake( weight_dtype=None, output_dtype=None, block_size=None, + has_gate=True, ): arrangement_ = functools.partial(arrangement, dim=-1, block_size=block_size) - if norm_before_gate: + if not has_gate: + application = application_norm_only + elif norm_before_gate: if activation == ActivationVariant.SIGMOID: application = application_norm_before_sigmoid else: @@ -75,13 +84,22 @@ def premake( else: application = application_norm_after_silu - tensors = ( - Tensor(ndim, other=0, dtype=input_dtype), - Tensor(ndim, other=0, dtype=gate_dtype), - Tensor(ndim, dtype=weight_dtype), - Tensor(0, dtype=ninetoothed.float64), - Tensor(ndim, dtype=output_dtype), - Tensor(0, constexpr=True), - ) + if has_gate: + tensors = ( + Tensor(ndim, other=0, dtype=input_dtype), + Tensor(ndim, other=0, dtype=gate_dtype), + Tensor(ndim, 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, dtype=weight_dtype), + Tensor(0, dtype=ninetoothed.float64), + Tensor(ndim, dtype=output_dtype), + Tensor(0, constexpr=True), + ) return arrangement_, application, tensors diff --git a/src/ntops/torch/gated_rms_norm.py b/src/ntops/torch/gated_rms_norm.py index 30d846c..1787980 100644 --- a/src/ntops/torch/gated_rms_norm.py +++ b/src/ntops/torch/gated_rms_norm.py @@ -7,14 +7,17 @@ def gated_rms_norm( input, - gate, - weight, + gate=None, + weight=None, eps=1e-5, group_size=None, norm_before_gate=False, activation="swish", ): - if input.shape != gate.shape: + 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.") hidden_size = input.shape[-1] @@ -43,24 +46,28 @@ def gated_rms_norm( grouped_shape = input.shape[:-1] + (hidden_size // group_size, group_size) grouped_input = input.reshape(grouped_shape) - grouped_gate = gate.reshape(grouped_shape) + grouped_gate = None if gate is None else gate.reshape(grouped_shape) grouped_weight = weight.expand_as(input).reshape(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(), ) - kernel( - grouped_input, - grouped_gate, - grouped_weight, - eps, - grouped_output, - group_size, - ) + 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) diff --git a/tests/test_gated_rms_norm.py b/tests/test_gated_rms_norm.py index f34ac42..b9d1dad 100644 --- a/tests/test_gated_rms_norm.py +++ b/tests/test_gated_rms_norm.py @@ -10,10 +10,10 @@ def _reference( input, gate, weight, eps, group_size, norm_before_gate, activation ): input_float = input.float() - gate_float = gate.float() + gate_float = None if gate is None else gate.float() activation_fn = torch.sigmoid if activation == "sigmoid" else F.silu - if not norm_before_gate: + if gate is not None and not norm_before_gate: input_float = input_float * activation_fn(gate_float) effective_group_size = group_size or input.shape[-1] @@ -22,7 +22,7 @@ def _reference( output = grouped * torch.rsqrt(variance + eps) output = output.reshape_as(input_float) * weight.float() - if norm_before_gate: + if gate is not None and norm_before_gate: output = output * activation_fn(gate_float) return output.to(input.dtype) @@ -30,7 +30,7 @@ def _reference( @skip_if_cuda_not_available @pytest.mark.parametrize("shape", ((3, 64), (2, 5, 128), (2, 3, 7, 256))) -@pytest.mark.parametrize("dtype", (torch.float16, torch.bfloat16)) +@pytest.mark.parametrize("dtype", (torch.float16, torch.bfloat16, torch.float32)) @pytest.mark.parametrize("group_size", (None, 32, 64)) @pytest.mark.parametrize("norm_before_gate", (False, True)) @pytest.mark.parametrize("activation", ("swish", "sigmoid")) @@ -54,7 +54,35 @@ def test_gated_rms_norm( input, gate, weight, 1e-5, group_size, norm_before_gate, activation ) - torch.testing.assert_close(actual, expected, rtol=2e-2, atol=2e-2) + torch.testing.assert_close(actual, expected, rtol=1e-2, atol=1e-2) + + +@skip_if_cuda_not_available +@pytest.mark.parametrize("dtype", (torch.float16, torch.bfloat16, torch.float32)) +@pytest.mark.parametrize( + ("shape", "group_size"), + ( + ((3, 64), None), + ((3, 64), 32), + ((3, 64), 64), + ((2, 5, 128), None), + ((2, 5, 128), 32), + ((2, 5, 128), 64), + ((2, 7, 96), None), + ((2, 7, 96), 48), + ), +) +def test_gated_rms_norm_without_gate(shape, dtype, group_size): + + input = torch.randn(shape, device="cuda", dtype=dtype) + weight = torch.randn(shape[-1], device="cuda", dtype=dtype) + + actual = ntops.torch.gated_rms_norm( + input, None, weight, eps=1e-5, group_size=group_size + ) + expected = _reference(input, None, weight, 1e-5, group_size, False, "swish") + + torch.testing.assert_close(actual, expected, rtol=1e-2, atol=1e-2) @skip_if_cuda_not_available @@ -68,7 +96,7 @@ def test_gated_rms_norm_non_contiguous_input(): actual = ntops.torch.gated_rms_norm(input, gate, weight, group_size=64) expected = _reference(input, gate, weight, 1e-5, 64, False, "swish") - torch.testing.assert_close(actual, expected, rtol=2e-2, atol=2e-2) + torch.testing.assert_close(actual, expected, rtol=1e-2, atol=1e-2) @skip_if_cuda_not_available @@ -80,4 +108,22 @@ def test_gated_rms_norm_non_power_of_two_group(): actual = ntops.torch.gated_rms_norm(input, gate, weight, group_size=48) expected = _reference(input, gate, weight, 1e-5, 48, False, "swish") - torch.testing.assert_close(actual, expected, rtol=2e-2, atol=2e-2) + torch.testing.assert_close(actual, expected, rtol=1e-2, atol=1e-2) + + +def test_gated_rms_norm_rejects_invalid_contract_on_cpu(): + input = torch.randn(2, 8) + gate = torch.randn(2, 7) + weight = torch.randn(8) + + with pytest.raises(ValueError, match="input.*gate"): + ntops.torch.gated_rms_norm(input, gate, weight) + + with pytest.raises(ValueError, match="weight"): + ntops.torch.gated_rms_norm(input, None, torch.randn(7)) + + with pytest.raises(ValueError, match="group_size"): + ntops.torch.gated_rms_norm(input, None, weight, group_size=3) + + with pytest.raises(ValueError, match="activation"): + ntops.torch.gated_rms_norm(input, None, weight, activation="gelu") From 053d00440f48081dbc19c66ac96635a2ce305417 Mon Sep 17 00:00:00 2001 From: CearX <56916338+CearX@users.noreply.github.com> Date: Mon, 31 Aug 2026 23:22:23 +0800 Subject: [PATCH 3/6] Add gated RMSNorm benchmark harness --- scripts/benchmark_gated_rms_norm.py | 890 ++++++++++++++++++++++++++++ 1 file changed, 890 insertions(+) create mode 100644 scripts/benchmark_gated_rms_norm.py diff --git a/scripts/benchmark_gated_rms_norm.py b/scripts/benchmark_gated_rms_norm.py new file mode 100644 index 0000000..d79be73 --- /dev/null +++ b/scripts/benchmark_gated_rms_norm.py @@ -0,0 +1,890 @@ +"""Fixed-configuration benchmark for :func:`ntops.torch.gated_rms_norm`. + +The script intentionally keeps the benchmark protocol in a fresh process. In +particular, ``max_num_configs`` is set to one before any tensor or operator +call. The gated RMSNorm wrapper derives and passes a next-power-of-two +``block_size`` from ``group_size``; this benchmark records that derived value +but does not tune or override it. + +The default case list is a small, representative vLLM-style matrix rather +than a Cartesian product. It covers token counts 128/1024/4096, hidden sizes +128/256/1024, full/64/128 groups, FP16/BF16, both gate orders, and both gated +and ungated calls. Use ``--quick`` or the filters when a shorter run is +needed. Output is JSONL by default and can be changed to CSV with +``--format csv``. +""" + +from __future__ import annotations + +import argparse +import csv +import dataclasses +import datetime as _datetime +import importlib.metadata +import json +import statistics +import subprocess +import sys +import time +from pathlib import Path +from typing import Callable, Iterable, Optional, Sequence + +import torch +import torch.nn.functional as F + +import ninetoothed +import ntops +import ntops.torch.utils + + +_SCRIPT_REPO_ROOT = Path(__file__).resolve().parents[1] +_CORRECTNESS_RTOL = 1e-2 +_CORRECTNESS_ATOL = 1e-2 +_DEFAULT_EPS = 1e-5 + + +@dataclasses.dataclass(frozen=True) +class Case: + """One fixed benchmark configuration.""" + + case_id: str + tokens: int + hidden_size: int + group_size: Optional[int] + dtype_name: str + norm_before_gate: bool + activation: str + has_gate: bool + quick: bool = False + + @property + def dtype(self) -> torch.dtype: + return { + "float16": torch.float16, + "bfloat16": torch.bfloat16, + }[self.dtype_name] + + @property + def effective_group_size(self) -> int: + return self.hidden_size if self.group_size is None else self.group_size + + @property + def block_size(self) -> int: + # Keep this in sync with ntops.torch.gated_rms_norm: the wrapper + # explicitly passes this next-power-of-two value to premake(). + group_size = self.effective_group_size + return 1 << (group_size - 1).bit_length() + + +# A curated matrix keeps the default protocol useful without running every +# Cartesian combination. Every listed dimension is represented in the set. +CASES: tuple[Case, ...] = ( + Case( + "t128_h128_gfull_f16_after_silu", + 128, + 128, + None, + "float16", + False, + "swish", + True, + True, + ), + Case( + "t128_h256_g64_bf16_before_silu", + 128, + 256, + 64, + "bfloat16", + True, + "swish", + True, + False, + ), + Case( + "t128_h1024_g128_f16_before_sigmoid", + 128, + 1024, + 128, + "float16", + True, + "sigmoid", + True, + False, + ), + Case( + "t1024_h128_g64_bf16_after_sigmoid", + 1024, + 128, + 64, + "bfloat16", + False, + "sigmoid", + True, + True, + ), + Case( + "t1024_h256_gfull_f16_before_silu", + 1024, + 256, + None, + "float16", + True, + "swish", + True, + False, + ), + Case( + "t1024_h1024_g128_bf16_after_silu", + 1024, + 1024, + 128, + "bfloat16", + False, + "swish", + True, + False, + ), + Case( + "t4096_h128_g128_f16_after_silu", + 4096, + 128, + 128, + "float16", + False, + "swish", + True, + True, + ), + Case( + "t4096_h256_g64_bf16_before_sigmoid", + 4096, + 256, + 64, + "bfloat16", + True, + "sigmoid", + True, + ), + Case( + "t4096_h1024_gfull_f16_before_silu", + 4096, + 1024, + None, + "float16", + True, + "swish", + True, + False, + ), + Case( + "t128_h256_gfull_bf16_nogate", + 128, + 256, + None, + "bfloat16", + False, + "swish", + False, + True, + ), + Case( + "t1024_h1024_g64_f16_nogate", + 1024, + 1024, + 64, + "float16", + False, + "swish", + False, + ), + Case( + "t4096_h128_gfull_bf16_nogate", + 4096, + 128, + None, + "bfloat16", + False, + "swish", + False, + ), +) + + +def _positive_int(value: str) -> int: + parsed = int(value) + if parsed <= 0: + raise argparse.ArgumentTypeError("must be a positive integer") + return parsed + + +def _nonnegative_int(value: str) -> int: + parsed = int(value) + if parsed < 0: + raise argparse.ArgumentTypeError("must be a non-negative integer") + return parsed + + +def _parse_int_filter(value: Optional[str], option_name: str) -> Optional[set[int]]: + if value is None: + return None + + result: set[int] = set() + for item in value.split(","): + item = item.strip() + if not item: + continue + try: + result.add(int(item)) + except ValueError as error: + raise argparse.ArgumentTypeError( + f"{option_name} expects comma-separated integers" + ) from error + if not result: + raise argparse.ArgumentTypeError(f"{option_name} must not be empty") + return result + + +def _parse_group_filter(value: Optional[str]) -> Optional[set[Optional[int]]]: + if value is None: + return None + + result: set[Optional[int]] = set() + for item in value.split(","): + item = item.strip().lower() + if not item: + continue + if item in {"none", "full"}: + result.add(None) + continue + try: + group_size = int(item) + except ValueError as error: + raise argparse.ArgumentTypeError( + "--group-size expects comma-separated integers or none" + ) from error + if group_size <= 0: + raise argparse.ArgumentTypeError("group sizes must be positive") + result.add(group_size) + if not result: + raise argparse.ArgumentTypeError("--group-size must not be empty") + return result + + +def _parse_bool_filter(value: str) -> Optional[bool]: + normalized = value.lower() + if normalized == "all": + return None + if normalized in {"true", "yes", "1"}: + return True + if normalized in {"false", "no", "0"}: + return False + raise argparse.ArgumentTypeError("expected all, true, or false") + + +def _parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--device", default="cuda", help="CUDA device, e.g. cuda:0") + parser.add_argument( + "--seed", type=int, default=20240828, help="base seed (case seeds are derived)" + ) + parser.add_argument( + "--warmup", + type=_nonnegative_int, + default=20, + help="untimed warmup calls per implementation (default: 20)", + ) + parser.add_argument( + "--repeat", + type=_positive_int, + default=100, + help="calls per timed trial (default: 100)", + ) + parser.add_argument( + "--trials", + type=_positive_int, + default=3, + help="timed trials per implementation (default: 3)", + ) + parser.add_argument("--eps", type=float, default=_DEFAULT_EPS) + parser.add_argument( + "--format", choices=("jsonl", "csv"), default="jsonl", dest="output_format" + ) + parser.add_argument( + "--output", + type=Path, + help="write records to this file instead of stdout", + ) + parser.add_argument( + "--variant", + default="ntops", + help="reserved result label; this script currently invokes ntops only", + ) + parser.add_argument( + "--quick", + action="store_true", + help="run only the small quick subset of the case library", + ) + parser.add_argument( + "--case", + action="append", + dest="case_ids", + help="case id (repeat or use comma-separated ids); default runs all cases", + ) + parser.add_argument( + "--dtype", + choices=("all", "float16", "bfloat16"), + default="all", + help="dtype filter", + ) + parser.add_argument("--tokens", help="comma-separated token-count filter") + parser.add_argument("--hidden-size", help="comma-separated hidden-size filter") + parser.add_argument("--group-size", help="comma-separated group-size filter") + parser.add_argument( + "--has-gate", + type=_parse_bool_filter, + default=None, + metavar="{all,true,false}", + help="has_gate filter", + ) + parser.add_argument( + "--norm-before-gate", + type=_parse_bool_filter, + default=None, + metavar="{all,true,false}", + help="gate-order filter", + ) + return parser + + +def _select_cases(args: argparse.Namespace) -> list[Case]: + selected: Iterable[Case] = CASES + if args.quick: + selected = (case for case in selected if case.quick) + + requested_case_ids: Optional[set[str]] = None + if args.case_ids: + requested_case_ids = { + case_id.strip() + for group in args.case_ids + for case_id in group.split(",") + if case_id.strip() + } + known_case_ids = {case.case_id for case in CASES} + unknown_case_ids = requested_case_ids - known_case_ids + if unknown_case_ids: + raise ValueError( + "unknown case id(s): " + ", ".join(sorted(unknown_case_ids)) + ) + selected = (case for case in selected if case.case_id in requested_case_ids) + + token_filter = _parse_int_filter(args.tokens, "--tokens") + hidden_filter = _parse_int_filter(args.hidden_size, "--hidden-size") + group_filter = _parse_group_filter(args.group_size) + + def matches(case: Case) -> bool: + return ( + (args.dtype == "all" or case.dtype_name == args.dtype) + and (token_filter is None or case.tokens in token_filter) + and (hidden_filter is None or case.hidden_size in hidden_filter) + and (group_filter is None or case.group_size in group_filter) + and (args.has_gate is None or case.has_gate == args.has_gate) + and ( + args.norm_before_gate is None + or case.norm_before_gate == args.norm_before_gate + ) + ) + + result = [case for case in selected if matches(case)] + if not result: + raise ValueError("case filters selected no cases") + return result + + +def _reference( + input: torch.Tensor, + gate: Optional[torch.Tensor], + weight: torch.Tensor, + eps: float, + group_size: Optional[int], + norm_before_gate: bool, + activation: str, +) -> torch.Tensor: + """Pure PyTorch reference, intentionally independent of ntops.""" + + input_float = input.float() + gate_float = None if gate is None else gate.float() + activation_fn: Callable[[torch.Tensor], torch.Tensor] = ( + torch.sigmoid if activation == "sigmoid" else F.silu + ) + + if gate_float is not None and not norm_before_gate: + input_float = input_float * activation_fn(gate_float) + + effective_group_size = group_size or input.shape[-1] + grouped = input_float.reshape(*input.shape[:-1], -1, effective_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 gate_float is not None and norm_before_gate: + output = output * activation_fn(gate_float) + + return output.to(input.dtype) + + +def _case_seed(base_seed: int, case_index: int) -> int: + # Avoid Python's randomized hash so case data remains reproducible across + # fresh processes. + return (base_seed + 1009 * (case_index + 1)) % (2**63 - 1) + + +def _make_inputs( + case: Case, device: torch.device, seed: int +) -> tuple[torch.Tensor, Optional[torch.Tensor], torch.Tensor]: + torch.manual_seed(seed) + input = torch.randn( + (case.tokens, case.hidden_size), device=device, dtype=case.dtype + ) + gate = ( + torch.randn_like(input) + if case.has_gate + else None + ) + weight = torch.randn((case.hidden_size,), device=device, dtype=case.dtype) + return input, gate, weight + + +def _call_ntops( + case: Case, + input: torch.Tensor, + gate: Optional[torch.Tensor], + weight: torch.Tensor, + eps: float, +) -> torch.Tensor: + # Do not pass a competing block size here: the wrapper explicitly computes + # next_power_of_two(group_size) and supplies it to the NineToothed premake. + return ntops.torch.gated_rms_norm( + input, + gate, + weight, + eps=eps, + group_size=case.group_size, + norm_before_gate=case.norm_before_gate, + activation=case.activation, + ) + + +def _cuda_sync(device: torch.device) -> None: + torch.cuda.synchronize(device) + + +def _time_trials_cuda_events( + fn: Callable[[], torch.Tensor], + device: torch.device, + warmup: int, + repeat: int, + trials: int, +) -> list[float]: + for _ in range(warmup): + fn() + _cuda_sync(device) + + results: list[float] = [] + with torch.cuda.device(device): + for _ in range(trials): + try: + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + except (NotImplementedError, RuntimeError) as error: + raise _EventTimingUnavailable(error) from error + for _ in range(repeat): + fn() + try: + end.record() + except (NotImplementedError, RuntimeError) as error: + raise _EventTimingUnavailable(error) from error + _cuda_sync(device) + try: + elapsed_ms = start.elapsed_time(end) + except (NotImplementedError, RuntimeError) as error: + raise _EventTimingUnavailable(error) from error + results.append(elapsed_ms * 1000.0 / repeat) + return results + + +class _EventTimingUnavailable(Exception): + """Internal signal used only to select the labelled host-timer fallback.""" + + +def _time_trials_perf_counter( + fn: Callable[[], torch.Tensor], + device: torch.device, + warmup: int, + repeat: int, + trials: int, +) -> list[float]: + """Portable fallback for backends whose CUDA Event lacks timing support.""" + + for _ in range(warmup): + fn() + _cuda_sync(device) + + results: list[float] = [] + for _ in range(trials): + _cuda_sync(device) + start = time.perf_counter() + for _ in range(repeat): + fn() + _cuda_sync(device) + results.append((time.perf_counter() - start) * 1e6 / repeat) + return results + + +def _time_trials( + fn: Callable[[], torch.Tensor], + device: torch.device, + warmup: int, + repeat: int, + trials: int, +) -> tuple[list[float], str]: + try: + return ( + _time_trials_cuda_events(fn, device, warmup, repeat, trials), + "cuda_event", + ) + except _EventTimingUnavailable as error: + # Event timing is preferred, but some non-NVIDIA CUDA backends expose + # only synchronization. Do not fabricate a result: the fallback is + # explicitly labelled in every record. + print( + f"warning: CUDA Event timing unavailable; using perf_counter ({error})", + file=sys.stderr, + ) + return ( + _time_trials_perf_counter(fn, device, warmup, repeat, trials), + "perf_counter", + ) + + +def _package_version(distribution: str, module: object) -> str: + try: + return importlib.metadata.version(distribution) + except importlib.metadata.PackageNotFoundError: + return str(getattr(module, "__version__", "unknown")) + + +def _git_description(repo_root: Path) -> str: + try: + result = subprocess.run( + ["git", "describe", "--always", "--dirty"], + cwd=repo_root, + check=False, + capture_output=True, + text=True, + timeout=2, + ) + except (OSError, subprocess.TimeoutExpired): + return "unknown" + if result.returncode != 0: + return "unknown" + return result.stdout.strip() or "unknown" + + +def _utc_now() -> str: + return _datetime.datetime.now(_datetime.timezone.utc).isoformat() + + +def _max_error( + actual: torch.Tensor, expected: torch.Tensor +) -> tuple[float, float]: + actual_float = actual.float() + expected_float = expected.float() + difference = (actual_float - expected_float).abs() + denominator = expected_float.abs().clamp_min(1e-12) + relative = difference / denominator + return float(difference.max().item()), float(relative.max().item()) + + +_OUTPUT_FIELDS: tuple[str, ...] = ( + "schema_version", + "record_type", + "status", + "variant", + "started_at_utc", + "finished_at_utc", + "device", + "device_name", + "torch_version", + "ntops_version", + "ninetoothed_version", + "ntops_git", + "seed", + "case_index", + "case_id", + "tokens", + "hidden_size", + "shape", + "dtype", + "group_size", + "effective_group_size", + "block_size", + "norm_before_gate", + "activation", + "has_gate", + "eps", + "correctness_rtol", + "correctness_atol", + "correctness_max_abs", + "correctness_max_rel", + "warmup", + "repeat", + "trials", + "timer", + "ntops_trials_us", + "ntops_median_us", + "reference_trials_us", + "reference_median_us", + "speedup_reference_over_ntops", + "error", +) + + +def _base_record( + case: Case, + case_index: int, + case_seed: int, + args: argparse.Namespace, + device: torch.device, + device_name: str, + run_started_at: str, +) -> dict[str, object]: + return { + "schema_version": 1, + "record_type": "gated_rms_norm_benchmark", + "status": "error", + "variant": args.variant, + "started_at_utc": run_started_at, + "finished_at_utc": None, + "device": str(device), + "device_name": device_name, + "torch_version": torch.__version__, + "ntops_version": _package_version("ntops", ntops), + "ninetoothed_version": _package_version("ninetoothed", ninetoothed), + "ntops_git": _git_description(_SCRIPT_REPO_ROOT), + "seed": case_seed, + "case_index": case_index, + "case_id": case.case_id, + "tokens": case.tokens, + "hidden_size": case.hidden_size, + "shape": [case.tokens, case.hidden_size], + "dtype": case.dtype_name, + "group_size": case.group_size, + "effective_group_size": case.effective_group_size, + "block_size": case.block_size, + "norm_before_gate": case.norm_before_gate, + "activation": case.activation, + "has_gate": case.has_gate, + "eps": args.eps, + "correctness_rtol": _CORRECTNESS_RTOL, + "correctness_atol": _CORRECTNESS_ATOL, + "correctness_max_abs": None, + "correctness_max_rel": None, + "warmup": args.warmup, + "repeat": args.repeat, + "trials": args.trials, + "timer": None, + "ntops_trials_us": None, + "ntops_median_us": None, + "reference_trials_us": None, + "reference_median_us": None, + "speedup_reference_over_ntops": None, + "error": None, + } + + +def _benchmark_case( + case: Case, + case_index: int, + args: argparse.Namespace, + device: torch.device, + device_name: str, + run_started_at: str, +) -> dict[str, object]: + case_seed = _case_seed(args.seed, case_index) + record = _base_record( + case, case_index, case_seed, args, device, device_name, run_started_at + ) + + try: + input, gate, weight = _make_inputs(case, device, case_seed) + with torch.no_grad(): + expected = _reference( + input, + gate, + weight, + args.eps, + case.group_size, + case.norm_before_gate, + case.activation, + ) + actual = _call_ntops(case, input, gate, weight, args.eps) + _cuda_sync(device) + + max_abs, max_rel = _max_error(actual, expected) + record["correctness_max_abs"] = max_abs + record["correctness_max_rel"] = max_rel + if not torch.allclose( + actual, + expected, + rtol=_CORRECTNESS_RTOL, + atol=_CORRECTNESS_ATOL, + equal_nan=False, + ): + record["status"] = "correctness_failed" + record["error"] = ( + "ntops output did not meet torch reference tolerance " + f"(rtol={_CORRECTNESS_RTOL}, atol={_CORRECTNESS_ATOL})" + ) + return record + + ntops_fn = lambda: _call_ntops(case, input, gate, weight, args.eps) + reference_fn = lambda: _reference( + input, + gate, + weight, + args.eps, + case.group_size, + case.norm_before_gate, + case.activation, + ) + + # Each implementation is warmed up and timed independently. A + # correctness failure returns above, so failed cases are never timed. + ntops_trials, timer = _time_trials( + ntops_fn, device, args.warmup, args.repeat, args.trials + ) + reference_trials, reference_timer = _time_trials( + reference_fn, device, args.warmup, args.repeat, args.trials + ) + if timer != reference_timer: + timer = f"{timer}+{reference_timer}" + + ntops_median = float(statistics.median(ntops_trials)) + reference_median = float(statistics.median(reference_trials)) + record.update( + { + "status": "passed", + "timer": timer, + "ntops_trials_us": ntops_trials, + "ntops_median_us": ntops_median, + "reference_trials_us": reference_trials, + "reference_median_us": reference_median, + "speedup_reference_over_ntops": reference_median / ntops_median, + } + ) + except Exception as error: # preserve one machine-readable record per case + record["status"] = "error" + record["error"] = f"{type(error).__name__}: {error}" + finally: + record["finished_at_utc"] = _utc_now() + return record + + +def _csv_value(value: object) -> object: + if isinstance(value, (list, tuple, dict)): + return json.dumps(value, sort_keys=True, separators=(",", ":")) + if value is None: + return "" + return value + + +def _write_records(records: Sequence[dict[str, object]], args: argparse.Namespace) -> None: + output = sys.stdout + should_close = False + if args.output is not None: + output = args.output.open("w", newline="", encoding="utf-8") + should_close = True + + try: + if args.output_format == "csv": + writer = csv.DictWriter( + output, + fieldnames=_OUTPUT_FIELDS, + extrasaction="ignore", + lineterminator="\n", + ) + writer.writeheader() + for record in records: + writer.writerow({key: _csv_value(record.get(key)) for key in _OUTPUT_FIELDS}) + else: + for record in records: + output.write( + json.dumps( + record, + sort_keys=True, + separators=(",", ":"), + ensure_ascii=False, + ) + + "\n" + ) + output.flush() + finally: + if should_close: + output.close() + + +def main(argv: Optional[Sequence[str]] = None) -> int: + # This must remain the first runtime action. In a fresh process it fixes + # the NineToothed search space before any ntops operator is called. + ntops.torch.utils.set_default_max_num_configs(1) + + args = _parser().parse_args(argv) + try: + cases = _select_cases(args) + except (argparse.ArgumentTypeError, ValueError) as error: + print(f"error: {error}", file=sys.stderr) + return 2 + + if args.eps < 0: + print("error: --eps must be non-negative", file=sys.stderr) + return 2 + + device = torch.device(args.device) + if device.type != "cuda": + print("error: gated RMSNorm benchmark requires a CUDA device", file=sys.stderr) + return 2 + if not torch.cuda.is_available(): + print("error: torch.cuda is unavailable", file=sys.stderr) + return 2 + + device_index = torch.cuda.current_device() if device.index is None else device.index + device = torch.device("cuda", device_index) + device_name = torch.cuda.get_device_name(device) + run_started_at = _utc_now() + + records: list[dict[str, object]] = [] + for case_index, case in enumerate(cases): + print( + f"[{case_index + 1}/{len(cases)}] {case.case_id}", + file=sys.stderr, + flush=True, + ) + records.append( + _benchmark_case( + case, + case_index, + args, + device, + device_name, + run_started_at, + ) + ) + + _write_records(records, args) + return 0 if all(record["status"] == "passed" for record in records) else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) From db9092223788d0b654af250a6bac9181a586ea1c Mon Sep 17 00:00:00 2001 From: CearX <56916338+CearX@users.noreply.github.com> Date: Mon, 31 Aug 2026 23:22:23 +0800 Subject: [PATCH 4/6] Optimize gated RMSNorm row reductions --- src/ntops/kernels/gated_rms_norm.py | 66 ++++++++++++++++++----------- src/ntops/torch/gated_rms_norm.py | 5 ++- 2 files changed, 46 insertions(+), 25 deletions(-) diff --git a/src/ntops/kernels/gated_rms_norm.py b/src/ntops/kernels/gated_rms_norm.py index 0464a95..815f54a 100644 --- a/src/ntops/kernels/gated_rms_norm.py +++ b/src/ntops/kernels/gated_rms_norm.py @@ -5,58 +5,76 @@ import ninetoothed.language as ntl from ninetoothed import Tensor -from ntops.kernels.reduction import arrangement - 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[0].to(ntl.float32) - rms = ntl.sqrt(ntl.sum(value * value, axis=0) / num_normalized_elements + eps) - gate_float = gate[0].to(ntl.float32) + 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[0] = value / rms * weight[0].to(ntl.float32) * activated_gate + output = value / rms * weight.to(ntl.float32) * activated_gate def application_norm_before_sigmoid( input, gate, weight, eps, output, num_normalized_elements ): - value = input[0].to(ntl.float32) - rms = ntl.sqrt(ntl.sum(value * value, axis=0) / num_normalized_elements + eps) - gate_float = gate[0].to(ntl.float32) + 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[0] = value / rms * weight[0].to(ntl.float32) * activated_gate + output = value / rms * weight.to(ntl.float32) * activated_gate def application_norm_after_silu( input, gate, weight, eps, output, num_normalized_elements ): - value = input[0].to(ntl.float32) - gate_float = gate[0].to(ntl.float32) + 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, axis=0) / num_normalized_elements + eps) - output[0] = value / rms * weight[0].to(ntl.float32) + 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[0].to(ntl.float32) - gate_float = gate[0].to(ntl.float32) + 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, axis=0) / num_normalized_elements + eps) - output[0] = value / rms * weight[0].to(ntl.float32) + 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[0].to(ntl.float32) - rms = ntl.sqrt(ntl.sum(value * value, axis=0) / num_normalized_elements + eps) - output[0] = value / rms * weight[0].to(ntl.float32) + 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( @@ -70,7 +88,7 @@ def premake( block_size=None, has_gate=True, ): - arrangement_ = functools.partial(arrangement, dim=-1, block_size=block_size) + arrangement_ = functools.partial(arrangement, block_size=block_size) if not has_gate: application = application_norm_only @@ -88,7 +106,7 @@ def premake( tensors = ( Tensor(ndim, other=0, dtype=input_dtype), Tensor(ndim, other=0, dtype=gate_dtype), - Tensor(ndim, dtype=weight_dtype), + Tensor(ndim, other=0, dtype=weight_dtype), Tensor(0, dtype=ninetoothed.float64), Tensor(ndim, dtype=output_dtype), Tensor(0, constexpr=True), @@ -96,7 +114,7 @@ def premake( else: tensors = ( Tensor(ndim, other=0, dtype=input_dtype), - Tensor(ndim, dtype=weight_dtype), + Tensor(ndim, other=0, dtype=weight_dtype), Tensor(0, dtype=ninetoothed.float64), Tensor(ndim, dtype=output_dtype), Tensor(0, constexpr=True), diff --git a/src/ntops/torch/gated_rms_norm.py b/src/ntops/torch/gated_rms_norm.py index 1787980..85e5360 100644 --- a/src/ntops/torch/gated_rms_norm.py +++ b/src/ntops/torch/gated_rms_norm.py @@ -47,7 +47,10 @@ def gated_rms_norm( grouped_shape = input.shape[:-1] + (hidden_size // group_size, group_size) grouped_input = input.reshape(grouped_shape) grouped_gate = None if gate is None else gate.reshape(grouped_shape) - grouped_weight = weight.expand_as(input).reshape(grouped_shape) + prefix = input.shape[:-1] + num_groups = hidden_size // group_size + 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( From 08bda486af74a792e0b01f397c8c7de70bca3937 Mon Sep 17 00:00:00 2001 From: CearX <56916338+CearX@users.noreply.github.com> Date: Mon, 31 Aug 2026 23:22:24 +0800 Subject: [PATCH 5/6] Add gated RMSNorm launch path profiler --- scripts/profile_gated_rms_norm_launch.py | 439 +++++++++++++++++++++++ 1 file changed, 439 insertions(+) create mode 100644 scripts/profile_gated_rms_norm_launch.py diff --git a/scripts/profile_gated_rms_norm_launch.py b/scripts/profile_gated_rms_norm_launch.py new file mode 100644 index 0000000..3c9e0ee --- /dev/null +++ b/scripts/profile_gated_rms_norm_launch.py @@ -0,0 +1,439 @@ +"""Diagnose gated RMSNorm launch, preparation, and allocation overhead. + +The default case is deliberately fixed to the representative configuration +used by the Track3 investigation: ``tokens=1024``, ``hidden_size=128``, +``group_size=64``, BF16 input/gate/weight, sigmoid gate after normalization. +All five paths use the same input storage and, for the direct paths, the same +grouped tensor objects and one ``_cached_make`` kernel: + +``A`` public wrapper (new output and grouped views on every call) +``B`` direct kernel with one persistent output (prepared identity should hit) +``C`` direct kernel with a new ``empty_like`` output on every call +``D`` only ``empty_like`` allocation +``E`` independent PyTorch reference + +This is a diagnostic, not a submission benchmark. CUDA Events are required +for the measurements. If the backend cannot provide event timing, the +script exits with a labelled error instead of mixing a host-timer result into +the record. +""" + +from __future__ import annotations + +import argparse +import datetime as _datetime +import importlib.metadata +import json +import statistics +from pathlib import Path +from typing import Callable, Optional, Sequence + + +_TOKENS = 1024 +_HIDDEN_SIZE = 128 +_GROUP_SIZE = 64 +_DTYPE_NAME = "bfloat16" +_WARMUP = 20 +_REPEAT = 100 +_TRIALS = 3 +_EPS = 1e-5 +_RTOL = 1e-2 +_ATOL = 1e-2 + + +class _EventTimingUnavailable(RuntimeError): + """The selected CUDA backend does not implement timed CUDA Events.""" + + +def _parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--device", default="cuda", help="CUDA device, e.g. cuda:0") + parser.add_argument("--seed", type=int, default=20240828) + parser.add_argument("--warmup", type=_nonnegative_int, default=_WARMUP) + parser.add_argument("--repeat", type=_positive_int, default=_REPEAT) + parser.add_argument("--trials", type=_positive_int, default=_TRIALS) + parser.add_argument("--eps", type=float, default=_EPS) + parser.add_argument( + "--output", type=Path, help="write the single JSON record here instead of stdout" + ) + return parser + + +def _positive_int(value: str) -> int: + parsed = int(value) + if parsed <= 0: + raise argparse.ArgumentTypeError("must be a positive integer") + return parsed + + +def _nonnegative_int(value: str) -> int: + parsed = int(value) + if parsed < 0: + raise argparse.ArgumentTypeError("must be non-negative") + return parsed + + +def _utc_now() -> str: + return _datetime.datetime.now(_datetime.timezone.utc).isoformat() + + +def _version(distribution: str, module: object) -> str: + try: + return importlib.metadata.version(distribution) + except importlib.metadata.PackageNotFoundError: + return str(getattr(module, "__version__", "unknown")) + + +def _sync(torch, device) -> None: + torch.cuda.synchronize(device) + + +def _event_error(stage: str, error: BaseException) -> _EventTimingUnavailable: + return _EventTimingUnavailable( + "CUDA Event timing unavailable on this backend " + f"while {stage}; no host-timer fallback is used: {type(error).__name__}: {error}" + ) + + +def _time_events( + torch, + device, + fn: Callable[[], object], + warmup: int, + repeat: int, + trials: int, +) -> list[float]: + """Return microseconds per call, using only CUDA Event timing.""" + + for _ in range(warmup): + fn() + _sync(torch, device) + + results: list[float] = [] + with torch.cuda.device(device): + for _ in range(trials): + try: + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + except (AttributeError, NotImplementedError, RuntimeError) as error: + raise _event_error("creating timing events", error) from error + + try: + start.record() + except (AttributeError, NotImplementedError, RuntimeError) as error: + raise _event_error("recording the start event", error) from error + + for _ in range(repeat): + fn() + + try: + end.record() + except (AttributeError, NotImplementedError, RuntimeError) as error: + raise _event_error("recording the end event", error) from error + _sync(torch, device) + try: + elapsed_ms = start.elapsed_time(end) + except (AttributeError, NotImplementedError, RuntimeError) as error: + raise _event_error("reading event elapsed time", error) from error + results.append(float(elapsed_ms) * 1000.0 / repeat) + return results + + +def _reference(torch, input, gate, weight, eps: float, group_size: int): + value = input.float() * torch.sigmoid(gate.float()) + rms = torch.sqrt( + torch.sum(value * value, dim=-1, keepdim=True) / group_size + eps + ) + return (value / rms * weight.float()).to(input.dtype) + + +def _error(torch, actual, expected) -> tuple[float, float]: + difference = (actual.float() - expected.float()).abs() + denominator = expected.float().abs().clamp_min(1e-12) + relative = difference / denominator + return float(difference.max().item()), float(relative.max().item()) + + +def _check(torch, actual, expected) -> dict[str, object]: + result: dict[str, object] = { + "status": "passed", + "max_abs": None, + "max_rel": None, + "shape": list(actual.shape), + "dtype": str(actual.dtype), + } + if actual.shape != expected.shape or actual.dtype != expected.dtype: + result["status"] = "failed_shape_or_dtype" + result["expected_shape"] = list(expected.shape) + result["expected_dtype"] = str(expected.dtype) + return result + + max_abs, max_rel = _error(torch, actual, expected) + result["max_abs"] = max_abs + result["max_rel"] = max_rel + if not torch.allclose(actual, expected, rtol=_RTOL, atol=_ATOL, equal_nan=False): + result["status"] = "failed_tolerance" + return result + + +def _measure_variant( + torch, + device, + name: str, + fn: Callable[[], object], + warmup: int, + repeat: int, + trials: int, +) -> dict[str, object]: + trials_us = _time_events(torch, device, fn, warmup, repeat, trials) + return { + "description": name, + "trials_us_per_call": trials_us, + "median_us_per_call": float(statistics.median(trials_us)), + } + + +def _write_record(record: dict[str, object], output: Optional[Path]) -> None: + line = json.dumps(record, sort_keys=True, separators=(",", ":"), ensure_ascii=False) + if output is None: + print(line) + return + output.write_text(line + "\n", encoding="utf-8") + + +def _run(args, torch, ntops, ninetoothed) -> dict[str, object]: + if args.eps < 0: + raise ValueError("--eps must be non-negative") + requested_device = torch.device(args.device) + if requested_device.type != "cuda": + raise ValueError("gated RMSNorm launch diagnosis requires a CUDA device") + if not torch.cuda.is_available(): + raise RuntimeError("torch.cuda is unavailable") + + device_index = ( + torch.cuda.current_device() + if requested_device.index is None + else requested_device.index + ) + device = torch.device("cuda", device_index) + device_name = torch.cuda.get_device_name(device) + + torch.manual_seed(args.seed) + input = torch.randn( + (_TOKENS, _HIDDEN_SIZE), device=device, dtype=torch.bfloat16 + ) + gate = torch.randn_like(input) + weight = torch.randn((_HIDDEN_SIZE,), device=device, dtype=torch.bfloat16) + + grouped_shape = (_TOKENS, _HIDDEN_SIZE // _GROUP_SIZE, _GROUP_SIZE) + grouped_input = input.reshape(grouped_shape) + grouped_gate = gate.reshape(grouped_shape) + grouped_weight = weight.reshape((1, _HIDDEN_SIZE // _GROUP_SIZE, _GROUP_SIZE)) + grouped_weight = grouped_weight.expand(grouped_shape) + fixed_output = torch.empty_like(grouped_input) + + activation = ntops.kernels.gated_rms_norm.ActivationVariant.SIGMOID + kernel = ntops.torch.utils._cached_make( + ntops.kernels.gated_rms_norm.premake, + grouped_input.ndim, + has_gate=True, + norm_before_gate=False, + activation=activation, + block_size=1 << (_GROUP_SIZE - 1).bit_length(), + ) + same_kernel = kernel is ntops.torch.utils._cached_make( + ntops.kernels.gated_rms_norm.premake, + grouped_input.ndim, + has_gate=True, + norm_before_gate=False, + activation=activation, + block_size=1 << (_GROUP_SIZE - 1).bit_length(), + ) + if not same_kernel: + raise RuntimeError("_cached_make did not return one shared kernel object") + + def public_wrapper(): + return ntops.torch.gated_rms_norm( + input, + gate, + weight, + eps=args.eps, + group_size=_GROUP_SIZE, + norm_before_gate=False, + activation="sigmoid", + ) + + def direct_fixed(): + kernel( + grouped_input, + grouped_gate, + grouped_weight, + args.eps, + fixed_output, + _GROUP_SIZE, + ) + return fixed_output + + def direct_new_output(): + output = torch.empty_like(grouped_input) + kernel( + grouped_input, + grouped_gate, + grouped_weight, + args.eps, + output, + _GROUP_SIZE, + ) + return output + + def empty_only(): + return torch.empty_like(grouped_input) + + def reference(): + return _reference(torch, grouped_input, grouped_gate, grouped_weight, args.eps, _GROUP_SIZE) + + with torch.no_grad(): + expected = reference() + public_result = public_wrapper() + direct_fixed_result = direct_fixed() + direct_new_result = direct_new_output() + _sync(torch, device) + + correctness = { + "A_public_wrapper": _check( + torch, public_result.reshape(grouped_shape), expected + ), + "B_direct_fixed_output": _check(torch, direct_fixed_result, expected), + "C_direct_new_output": _check(torch, direct_new_result, expected), + "D_empty_only": {"status": "not_applicable"}, + "E_pytorch_reference": {"status": "reference"}, + } + if any( + item["status"] not in {"passed", "reference", "not_applicable"} + for item in correctness.values() + ): + raise RuntimeError("one or more NineToothed paths failed correctness") + + variants = { + "A_public_wrapper": _measure_variant( + torch, + device, + "new wrapper output and grouped views per call", + public_wrapper, + args.warmup, + args.repeat, + args.trials, + ), + "B_direct_fixed_output": _measure_variant( + torch, + device, + "direct kernel with persistent output and tensor identities", + direct_fixed, + args.warmup, + args.repeat, + args.trials, + ), + "C_direct_new_output": _measure_variant( + torch, + device, + "direct kernel with new empty_like output per call", + direct_new_output, + args.warmup, + args.repeat, + args.trials, + ), + "D_empty_only": _measure_variant( + torch, + device, + "only torch.empty_like(grouped_input)", + empty_only, + args.warmup, + args.repeat, + args.trials, + ), + "E_pytorch_reference": _measure_variant( + torch, + device, + "independent PyTorch reference", + reference, + args.warmup, + args.repeat, + args.trials, + ), + } + + return { + "schema_version": 1, + "record_type": "gated_rms_norm_launch_diagnosis", + "status": "passed", + "started_at_utc": _utc_now(), + "finished_at_utc": _utc_now(), + "case": { + "tokens": _TOKENS, + "hidden_size": _HIDDEN_SIZE, + "group_size": _GROUP_SIZE, + "dtype": _DTYPE_NAME, + "norm_before_gate": False, + "activation": "sigmoid", + "gate_position": "after", + "eps": args.eps, + "shape": list(grouped_shape), + "block_size": 1 << (_GROUP_SIZE - 1).bit_length(), + }, + "device": { + "requested": args.device, + "resolved": str(device), + "name": device_name, + }, + "versions": { + "torch": torch.__version__, + "ntops": _version("ntops", ntops), + "ninetoothed": _version("ninetoothed", ninetoothed), + }, + "seed": args.seed, + "protocol": { + "timer": "cuda_event", + "warmup": args.warmup, + "repeat": args.repeat, + "trials": args.trials, + "correctness_rtol": _RTOL, + "correctness_atol": _ATOL, + "kernel_shared_by": ["B_direct_fixed_output", "C_direct_new_output"], + "max_num_configs": 1, + }, + "correctness": correctness, + "variants": variants, + } + + +def main(argv: Optional[Sequence[str]] = None) -> int: + # Keep this before argument parsing, tensor creation, and every operator + # call. The process running this script must set the search bound first. + import ntops + + ntops.torch.utils.set_default_max_num_configs(1) + + import ninetoothed + import torch + + try: + args = _parser().parse_args(argv) + record = _run(args, torch, ntops, ninetoothed) + except (argparse.ArgumentError, ValueError, RuntimeError) as error: + record = { + "schema_version": 1, + "record_type": "gated_rms_norm_launch_diagnosis", + "status": "error", + "finished_at_utc": _utc_now(), + "error": f"{type(error).__name__}: {error}", + } + output = getattr(locals().get("args", None), "output", None) + _write_record(record, output) + return 1 + + _write_record(record, args.output) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) From be43fc55e100f8e3f30602fcf2042db6b6bfc560 Mon Sep 17 00:00:00 2001 From: CearX <56916338+CearX@users.noreply.github.com> Date: Mon, 31 Aug 2026 23:22:24 +0800 Subject: [PATCH 6/6] Reduce gated RMSNorm wrapper overhead --- src/ntops/torch/gated_rms_norm.py | 25 +++++++++++++------------ 1 file changed, 13 insertions(+), 12 deletions(-) diff --git a/src/ntops/torch/gated_rms_norm.py b/src/ntops/torch/gated_rms_norm.py index 85e5360..d536044 100644 --- a/src/ntops/torch/gated_rms_norm.py +++ b/src/ntops/torch/gated_rms_norm.py @@ -4,6 +4,12 @@ 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, @@ -20,7 +26,9 @@ def gated_rms_norm( if gate is not None and input.shape != gate.shape: raise ValueError("`input` and `gate` must have the same shape.") - hidden_size = input.shape[-1] + 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],)`.") @@ -31,24 +39,17 @@ def gated_rms_norm( if group_size <= 0 or hidden_size % group_size != 0: raise ValueError("`group_size` must be a positive divisor of the hidden size.") - activation_variants = { - "silu": ActivationVariant.SILU, - "swish": ActivationVariant.SILU, - "sigmoid": ActivationVariant.SIGMOID, - } - try: - activation_variant = activation_variants[activation] + activation_variant = _ACTIVATION_VARIANTS[activation] except KeyError as error: raise ValueError( "`activation` must be one of `silu`, `swish`, or `sigmoid`." ) from error - grouped_shape = input.shape[:-1] + (hidden_size // group_size, group_size) + 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) - prefix = input.shape[:-1] - num_groups = hidden_size // group_size 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) @@ -73,4 +74,4 @@ def gated_rms_norm( group_size, ) - return grouped_output.reshape(input.shape) + return grouped_output.reshape(input_shape)