[Triton/Gluon] [GFX1250] fused_add_rmsnorm_pad() gluon equivalent function

aiter/ops/triton/_gluon_kernels/gfx1250/norm/fused_add_rmsnorm_pad.py

@@ -0,0 +1,163 @@
+# SPDX-License-Identifier: MIT
+# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.
+
+import triton
+import triton.language as tl
+from triton.experimental import gluon
+from triton.experimental.gluon import language as gl
+
+from aiter.ops.triton.utils._triton.kernel_repr import make_kernel_repr
+
+
+@triton.jit
+def _rmsmorm_op(row, weight, n_cols, epsilon):
+    row_norm = row * row
+    row_norm = tl.sum(row_norm, axis=-1)
+    norm_factor = tl.math.rsqrt((row_norm / n_cols) + epsilon)
+
+    rms_norm = row * norm_factor * weight
+    return rms_norm
+
+
+_gluon_fused_add_rmsnorm_pad_kernel_repr = make_kernel_repr(
+    "_gluon_fused_add_rmsnorm_pad_kernel",
+    ["BLOCK_SIZE_N", "HAS_RES", "num_warps"],
+)
+
+
+@gluon.jit(repr=_gluon_fused_add_rmsnorm_pad_kernel_repr)
+def _gluon_fused_add_rmsnorm_pad_kernel(
+    x_ptr,
+    res_ptr,
+    out_ptr,
+    res_out_ptr,
+    weight_ptr,
+    eps,
+    M,
+    N,
+    N_out,
+    x_stride_m,
+    x_stride_n,
+    res_stride_m,
+    res_stride_n,
+    out_stride_m,
+    out_stride_n,
+    res_out_stride_m,
+    res_out_stride_n,
+    HAS_RES: gl.constexpr,
+    BLOCK_SIZE_N: gl.constexpr,
+    num_warps: gl.constexpr,
+):
+    start_pid = gl.program_id(0)
+    NUM_WARPS: gl.constexpr = gl.num_warps()
+    WARP_SIZE: gl.constexpr = 32
+    CTA_NUM_THREADS: gl.constexpr = NUM_WARPS * WARP_SIZE
+
+    # create 1d layout for n_offs
+    n_offs_layout: gl.constexpr = gl.BlockedLayout(
+        [(BLOCK_SIZE_N + CTA_NUM_THREADS - 1) // CTA_NUM_THREADS],  # size per thread
+        [WARP_SIZE],  # threads per warp
+        [NUM_WARPS],  # warps per cta
+        [0],  # order
+    )
+
+    # create shared layouts for x, res, weights, out
+    sharedLayout1D: gl.constexpr = gl.SwizzledSharedLayout(1, 1, 1, order=[1, 0])
+
+    # create tensor descriptors for x, res, weights, out
+    x_desc = gl.amd.gfx1250.tdm.make_tensor_descriptor(
+        x_ptr,
+        [M, N],
+        [x_stride_m, x_stride_n],
+        [1, BLOCK_SIZE_N],
+        sharedLayout1D,
+    )
+    weights_desc = gl.amd.gfx1250.tdm.make_tensor_descriptor(
+        weight_ptr,
+        [1, N],
+        [0, 1],
+        [1, BLOCK_SIZE_N],
+        sharedLayout1D,
+    )
+    out_desc = gl.amd.gfx1250.tdm.make_tensor_descriptor(
+        out_ptr,
+        [M, N_out],
+        [out_stride_m, out_stride_n],
+        [1, BLOCK_SIZE_N],
+        sharedLayout1D,
+    )
+
+    # create tensor descriptors for res, res_out if applicable
+    if HAS_RES:
+        res_desc = gl.amd.gfx1250.tdm.make_tensor_descriptor(
+            res_ptr,
+            [M, N],
+            [res_stride_m, res_stride_n],
+            [1, BLOCK_SIZE_N],
+            sharedLayout1D,
+        )
+        res_out_desc = gl.amd.gfx1250.tdm.make_tensor_descriptor(
+            res_out_ptr,
+            [M, N],
+            [res_out_stride_m, res_out_stride_n],
+            [1, BLOCK_SIZE_N],
+            sharedLayout1D,
+        )
+
+    # allocate shared memory for x, weights, out
+    smemX = gl.allocate_shared_memory(
+        x_ptr.dtype.element_ty, [1, BLOCK_SIZE_N], sharedLayout1D
+    )
+    smemWeights = gl.allocate_shared_memory(
+        weight_ptr.dtype.element_ty, [1, BLOCK_SIZE_N], sharedLayout1D
+    )
+    smemOut = gl.allocate_shared_memory(
+        out_ptr.dtype.element_ty, [1, BLOCK_SIZE_N], sharedLayout1D
+    )
+
+    # allocate shared memory for res, res_out if applicable
+    if HAS_RES:
+        smemRes = gl.allocate_shared_memory(
+            res_ptr.dtype.element_ty, [1, BLOCK_SIZE_N], sharedLayout1D
+        )
+        smemResOut = gl.allocate_shared_memory(
+            res_out_ptr.dtype.element_ty, [1, BLOCK_SIZE_N], sharedLayout1D
+        )
+
+    # load x, res (if applicable), and weights
+    gl.amd.gfx1250.tdm.async_load(x_desc, [start_pid, 0], smemX)
+    if HAS_RES:
+        gl.amd.gfx1250.tdm.async_load(res_desc, [start_pid, 0], smemRes)
+    gl.amd.gfx1250.tdm.async_load(weights_desc, [0, 0], smemWeights)
+
+    if HAS_RES:
+        gl.amd.gfx1250.tdm.async_wait(2)
+    else:
+        gl.amd.gfx1250.tdm.async_wait(1)
+
+    smemX_1d = smemX.reshape([BLOCK_SIZE_N])
+    x = smemX_1d.load(n_offs_layout).to(gl.float32)
+
+    # reshape res and add to x if applicable
+    if HAS_RES:
+        gl.amd.gfx1250.tdm.async_wait(1)
+        smemRes_1d = smemRes.reshape([BLOCK_SIZE_N])
+        res = smemRes_1d.load(n_offs_layout).to(gl.float32)
+        x = x + res
+        smemResOut.store(x.reshape([1, BLOCK_SIZE_N]).to(res_out_ptr.dtype.element_ty))
+        gl.amd.gfx1250.tdm.async_store(res_out_desc, [start_pid, 0], smemResOut)
+
+    if HAS_RES:
+        gl.amd.gfx1250.tdm.async_wait(1)
+    else:
+        gl.amd.gfx1250.tdm.async_wait(0)
+
+    smemWeights_1d = smemWeights.reshape([BLOCK_SIZE_N])
+    w = smemWeights_1d.load(n_offs_layout).to(gl.float32)
+    # compute rmsnorm
+    out = _rmsmorm_op(x, w, N, eps).to(out_ptr.dtype.element_ty)
+    # write out to LDS
+    smemOut.store(out.reshape([1, BLOCK_SIZE_N]))
+    gl.amd.gfx1250.tdm.async_store(out_desc, [start_pid, 0], smemOut)
+
+    gl.amd.gfx1250.tdm.async_wait(0)

aiter/ops/triton/configs/gfx1250/gluon/normalization/fused_add_rmsnorm_pad/DEFAULT.json

@@ -0,0 +1,7 @@
+{
+    "N_LEQ_32": { "num_warps": 1 },
+    "N_LEQ_64": { "num_warps": 2 },
+    "N_LEQ_128": { "num_warps": 4 },
+    "N_LEQ_256": { "num_warps": 8 },
+    "any": { "num_warps": 8 }
+}

aiter/ops/triton/normalization/fused_add_rmsnorm_pad.py

@@ -1,9 +1,28 @@
 import torch
 import triton
 
+from aiter.ops.triton._gluon_kernels.gfx1250.norm.fused_add_rmsnorm_pad import (
+    _gluon_fused_add_rmsnorm_pad_kernel,
+)
 from aiter.ops.triton._triton_kernels.normalization.fused_add_rmsnorm_pad import (
     _fused_add_rmsnorm_pad,
 )
+from aiter.ops.triton.utils._triton.arch_info import get_arch
+from aiter.ops.triton.utils.config_utils import load_config_json, resolve_config_dir
+from aiter.ops.triton.utils.logger import AiterTritonLogger
+
+_LOGGER = AiterTritonLogger()
+
+
+def _get_config(block_size_n: int, backend: str) -> dict:
+    config_dir = resolve_config_dir(
+        "normalization", "FUSED_ADD_RMSNORM_PAD", backend=backend
+    )
+    raw = load_config_json(f"{config_dir}/DEFAULT.json")
+    for bound in sorted(int(k[len("N_LEQ_") :]) for k in raw if k.startswith("N_LEQ_")):
+        if block_size_n <= bound:
+            return dict(raw[f"N_LEQ_{bound}"])
+    return dict(raw["any"])
 
 
 def fused_add_rmsnorm_pad(
@@ -12,9 +31,69 @@ def fused_add_rmsnorm_pad(
     epsilon: float,
     res: torch.Tensor = None,
     x_pad_to_multiple: int = 0,
+    backend: str | None = None,
 ):
     M, N = x.shape
 
+    if backend is None:
+        backend = "gluon" if get_arch() == "gfx1250" else "triton"
+
+    backend = backend.lower()
+    assert backend in (
+        "triton",
+        "gluon",
+    ), f"Unknown backend '{backend}', must be 'triton' or 'gluon'"
+
+    if backend == "gluon":
+        if get_arch() == "gfx1250":
+
+            if x_pad_to_multiple > 0:
+                N_out = triton.cdiv(N, x_pad_to_multiple) * x_pad_to_multiple
+            else:
+                N_out = N
+            out = torch.empty((M, N_out), dtype=x.dtype, device=x.device)
+
+            res_out = None
+            if res is not None:
+                M2, N2 = res.shape
+                assert M == M2, "Shape error!"
+                assert N == N2, "Shape error!"
+                res_out = torch.empty((M, N), dtype=res.dtype, device=res.device)
+            BLOCK_SIZE_N = triton.next_power_of_2(N_out)
+            config = _get_config(BLOCK_SIZE_N, "gluon")
+            NUM_WARPS = config["num_warps"]
+            _gluon_fused_add_rmsnorm_pad_kernel[(M,)](
+                x,
+                res,
+                out,
+                res_out,
+                weight,
+                epsilon,
+                M,
+                N,
+                N_out,
+                x.stride(0),
+                x.stride(1),
+                res.stride(0) if res is not None else 0,
+                res.stride(1) if res is not None else 0,
+                out.stride(0),
+                out.stride(1),
+                res_out.stride(0) if res is not None else 0,
+                res_out.stride(1) if res is not None else 0,
+                HAS_RES=(res is not None),
+                BLOCK_SIZE_N=BLOCK_SIZE_N,
+                num_warps=NUM_WARPS,
+            )
+
+            if res is not None:
+                return out, res_out
+            return out
+        else:
+            _LOGGER.warning(
+                "Gluon kernel is only supported on gfx1250, using triton kernel instead"
+            )
+            backend = "triton"
+
     if x_pad_to_multiple > 0:
         N_out = triton.cdiv(N, x_pad_to_multiple) * x_pad_to_multiple
     else:
@@ -28,6 +107,7 @@ def fused_add_rmsnorm_pad(
         assert N == N2, "Shape error!"
         res_out = torch.empty((M, N), dtype=res.dtype, device=res.device)
     BLOCK_SIZE_N = triton.next_power_of_2(N_out)
+
     _fused_add_rmsnorm_pad[(M,)](
         x,
         res,

op_tests/triton_tests/normalization/test_fused_add_rmsnorm_pad.py

@@ -3,6 +3,7 @@
 import torch.nn.functional as F
 
 from aiter.ops.triton.normalization.fused_add_rmsnorm_pad import fused_add_rmsnorm_pad
+from aiter.ops.triton.utils._triton.arch_info import get_arch
 
 
 def generate_inputs(M, N, has_res, dtype):
@@ -37,18 +38,29 @@ def run_torch(x, weight, eps=1e-6, res=None, pad_to_multiple=0):
 @pytest.mark.parametrize("has_res", [False, True])
 @pytest.mark.parametrize("pad_to_multiple", [0, 256])
 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
-def test_mul_add(M: int, N: int, has_res: bool, pad_to_multiple: int, dtype):
+@pytest.mark.parametrize("backend", [None, "gluon", "triton"])
+def test_mul_add(
+    M: int,
+    N: int,
+    has_res: bool,
+    pad_to_multiple: int,
+    dtype,
+    backend: str | None,
+):
+
+    if backend == "gluon" and get_arch() != "gfx1250":
+        pytest.skip("Gluon kernel is only supported on gfx1250")
 
     x, weight, res = generate_inputs(M, N, has_res, dtype)
 
     if has_res:
         x_torch, res_torch = run_torch(x, weight, 1e-6, res, pad_to_multiple)
         x_triton, res_triton = fused_add_rmsnorm_pad(
-            x, weight, 1e-6, res, pad_to_multiple
+            x, weight, 1e-6, res, pad_to_multiple, backend
         )
     else:
         x_torch = run_torch(x, weight, 1e-6, res, pad_to_multiple)
-        x_triton = fused_add_rmsnorm_pad(x, weight, 1e-6, res, pad_to_multiple)
+        x_triton = fused_add_rmsnorm_pad(x, weight, 1e-6, res, pad_to_multiple, backend)
 
     torch.testing.assert_close(x_torch.to(dtype), x_triton)
     if has_res: