[Triton/Gluon] Add MoE weight-gradient kernel

aiter/ops/triton/_triton_kernels/moe/moe_wgrad.py

@@ -0,0 +1,156 @@
+# SPDX-License-Identifier: MIT
+# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.
+
+"""Triton kernel for MoE weight-gradient computation.
+
+Computes dW[e] = sum_{tokens t assigned to expert e} grad[t].T @ input[t]
+using sorted token infrastructure from moe_align_block_size.
+"""
+
+import triton
+import triton.language as tl
+
+from aiter.ops.triton.utils._triton.pid_preprocessing import remap_xcd
+
+
+def _get_autotune_configs():
+    return [
+        triton.Config(
+            {"BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 64},
+            num_warps=4,
+            num_stages=2,
+        ),
+        triton.Config(
+            {"BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 128},
+            num_warps=4,
+            num_stages=2,
+        ),
+        triton.Config(
+            {"BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 64},
+            num_warps=8,
+            num_stages=2,
+        ),
+        triton.Config(
+            {"BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 128},
+            num_warps=8,
+            num_stages=2,
+        ),
+        triton.Config(
+            {"BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 64},
+            num_warps=8,
+            num_stages=2,
+        ),
+        triton.Config(
+            {"BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 128},
+            num_warps=4,
+            num_stages=2,
+        ),
+    ]
+
+
+@triton.autotune(
+    configs=_get_autotune_configs(),
+    key=["N", "K"],
+    # dW is accumulated via tl.atomic_add; reset it to zero before each
+    # benchmark trial so timing runs don't corrupt the final result.
+    reset_to_zero=["dw_ptr"],
+)
+# NOTE: repr= is intentionally omitted here. Combining @triton.autotune with
+# @triton.jit(repr=...) corrupts the kernel execution on current Triton versions.
+# The autotune key ["N", "K"] already encodes the compiled variant identity.
+@triton.jit
+def _moe_wgrad_kernel(
+    grad_ptr,
+    input_ptr,
+    dw_ptr,
+    sorted_token_ids_ptr,
+    expert_ids_ptr,
+    num_tokens_post_padded_ptr,
+    N,
+    K,
+    num_valid_tokens,
+    stride_gm,
+    stride_gn,
+    stride_im,
+    stride_ik,
+    stride_dwe,
+    stride_dwn,
+    stride_dwk,
+    num_sorted_tokens,
+    top_k: tl.constexpr,
+    BLOCK_SIZE_M: tl.constexpr,
+    NUM_XCDS: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+):
+    """Weight gradient kernel for MoE.
+
+    Computes dW[e] = sum_{tokens t assigned to expert e} grad[t].T @ input[t]
+
+    BLOCK_SIZE_M must match the block_size used in moe_align_block_size to
+    ensure expert_ids[pid_m] correctly describes all tokens in the block.
+    """
+    pid_m = tl.program_id(0)
+    pid_nk = tl.program_id(1)
+
+    num_tokens_post_padded = tl.load(num_tokens_post_padded_ptr)
+    if pid_m * BLOCK_SIZE_M >= num_tokens_post_padded:
+        return
+
+    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
+    num_pid_k = tl.cdiv(K, BLOCK_SIZE_K)
+
+    total_nk = num_pid_n * num_pid_k
+    if pid_nk >= total_nk:
+        return
+
+    pid_nk_remapped = remap_xcd(pid_nk, total_nk, NUM_XCDS)
+    pid_n = pid_nk_remapped // num_pid_k
+    pid_k = pid_nk_remapped % num_pid_k
+
+    off_expert = tl.load(expert_ids_ptr + pid_m).to(tl.int64)
+    if off_expert == -1:
+        return
+
+    offs_token_id = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M).to(tl.int64)
+    offs_token = tl.load(sorted_token_ids_ptr + offs_token_id)
+    token_mask = offs_token < num_valid_tokens
+
+    offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
+    offs_k = pid_k * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K)
+    n_mask = offs_n < N
+    k_mask = offs_k < K
+
+    grad_ptrs = (
+        grad_ptr
+        + (offs_token // top_k)[:, None] * stride_gm
+        + offs_n[None, :] * stride_gn
+    )
+    grad_block = tl.load(
+        grad_ptrs,
+        mask=token_mask[:, None] & n_mask[None, :],
+        other=0.0,
+    )
+
+    input_ptrs = (
+        input_ptr
+        + (offs_token // top_k)[:, None] * stride_im
+        + offs_k[None, :] * stride_ik
+    )
+    input_block = tl.load(
+        input_ptrs,
+        mask=token_mask[:, None] & k_mask[None, :],
+        other=0.0,
+    )
+
+    # grad.T @ input = [BLOCK_N, BLOCK_M] @ [BLOCK_M, BLOCK_K] = [BLOCK_N, BLOCK_K]
+    acc = tl.dot(tl.trans(grad_block), input_block)
+
+    dw_ptrs = (
+        dw_ptr
+        + off_expert * stride_dwe
+        + offs_n[:, None].to(tl.int64) * stride_dwn
+        + offs_k[None, :].to(tl.int64) * stride_dwk
+    )
+    dw_mask = n_mask[:, None] & k_mask[None, :]
+    tl.atomic_add(dw_ptrs, acc.to(dw_ptr.dtype.element_ty), mask=dw_mask, sem="relaxed")

aiter/ops/triton/moe/__init__.py

@@ -1,4 +1,5 @@
 from aiter.ops.triton.moe.moe_gemm_mxfp8 import moe_gemm_mxfp8
 from aiter.ops.triton.moe.moe_gemm_per_token import moe_gemm_per_token
+from aiter.ops.triton.moe.moe_wgrad import moe_wgrad
 
-__all__ = ["moe_gemm_mxfp8", "moe_gemm_per_token"]
+__all__ = ["moe_gemm_mxfp8", "moe_gemm_per_token", "moe_wgrad"]

aiter/ops/triton/moe/moe_wgrad.py

@@ -0,0 +1,86 @@
+# SPDX-License-Identifier: MIT
+# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.
+
+"""MoE weight-gradient wrapper (group_sizes interface)."""
+
+import torch
+import triton
+
+from aiter.ops.triton._triton_kernels.moe.moe_wgrad import _moe_wgrad_kernel
+from aiter.ops.triton.utils.device_info import get_num_xcds
+
+__all__ = ["moe_wgrad"]
+
+
+def moe_wgrad(
+    grad: torch.Tensor,
+    input: torch.Tensor,
+    sorted_token_ids: torch.Tensor,
+    expert_ids: torch.Tensor,
+    num_tokens_post_padded: torch.Tensor,
+    num_experts: int,
+    top_k: int,
+    weight_shape: tuple,
+    block_size_m: int = 64,
+) -> torch.Tensor:
+    """Compute MoE weight gradients using sorted token infrastructure.
+
+    Uses ``@triton.autotune`` to search optimal ``(BLOCK_SIZE_N, BLOCK_SIZE_K)``
+    on first launch, keyed on ``(N, K)`` dimensions.
+
+    Args:
+        grad: Gradient tensor ``[num_tokens, N]``.  When ``top_k > 1`` this
+            tensor must already incorporate per-route gating weights; the kernel
+            accumulates it as-is without applying gating internally.
+        input: Input activation tensor ``[num_tokens, K]``.
+        sorted_token_ids: From ``moe_align_block_size``.  Padding slots must
+            contain a value >= ``num_valid_tokens`` (i.e. >= ``num_tokens *
+            top_k``) so the token mask treats them as invalid; the kernel relies
+            on this sentinel rather than reading ``token_nums`` per block.
+        expert_ids: From ``moe_align_block_size``.
+        num_tokens_post_padded: From ``moe_align_block_size``.
+        num_experts: Number of experts.
+        top_k: Number of experts per token.
+        weight_shape: Shape of weight tensor ``(E, N, K)``.
+        block_size_m: Must match the ``block_size`` used in
+            ``moe_align_block_size``.
+
+    Returns:
+        ``dW`` tensor with shape ``weight_shape``.
+    """
+    _E, N, K = weight_shape
+
+    dW = torch.zeros(weight_shape, dtype=grad.dtype, device=grad.device)
+
+    num_sorted = sorted_token_ids.shape[0]
+    if num_sorted == 0:
+        return dW
+
+    grid = lambda META: (
+        num_sorted // block_size_m,
+        triton.cdiv(N, META["BLOCK_SIZE_N"]) * triton.cdiv(K, META["BLOCK_SIZE_K"]),
+    )
+
+    _moe_wgrad_kernel[grid](
+        grad,
+        input,
+        dW,
+        sorted_token_ids,
+        expert_ids,
+        num_tokens_post_padded,
+        N,
+        K,
+        grad.shape[0] * top_k,
+        grad.stride(0),
+        grad.stride(1),
+        input.stride(0),
+        input.stride(1),
+        dW.stride(0),
+        dW.stride(1),
+        dW.stride(2),
+        num_sorted,
+        top_k=top_k,
+        BLOCK_SIZE_M=block_size_m,
+        NUM_XCDS=get_num_xcds(),
+    )
+    return dW

op_tests/op_benchmarks/triton/bench_moe_wgrad.py

@@ -0,0 +1,109 @@
+# SPDX-License-Identifier: MIT
+# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.
+"""Benchmark for moe_wgrad.
+
+Usage:
+    python bench_moe_wgrad.py
+    python bench_moe_wgrad.py -metric bandwidth
+"""
+
+import argparse
+
+import torch
+import triton
+
+from aiter.ops.moe_op import moe_align_block_size
+from aiter.ops.triton.moe.moe_wgrad import moe_wgrad
+from op_tests.op_benchmarks.triton.utils.benchmark_utils import get_caller_name_no_ext
+
+# (num_tokens, E, N, K, top_k, label)
+_SHAPES = [
+    (4096, 8, 7168, 2048, 2, "DSv4-decode"),
+    (16384, 8, 7168, 2048, 2, "DSv4-prefill"),
+    (4096, 64, 2048, 7168, 6, "DSv4-E64-topk6"),
+]
+_BLOCK_SIZE_M = 64
+
+
+def _setup(num_tokens, E, N, K, top_k, device="cuda"):
+    torch.manual_seed(0)
+    grad = torch.randn(num_tokens, N, dtype=torch.bfloat16, device=device) * 0.1
+    inp = torch.randn(num_tokens, K, dtype=torch.bfloat16, device=device) * 0.1
+    scores = torch.randn(num_tokens, E, device=device)
+    topk_ids = scores.topk(top_k, dim=1).indices  # [num_tokens, top_k]
+
+    # moe_align_block_size requires pre-allocated output buffers.
+    max_blocks = triton.cdiv(num_tokens * top_k, _BLOCK_SIZE_M) + E
+    sorted_ids = torch.empty(
+        max_blocks * _BLOCK_SIZE_M, dtype=torch.int32, device=device
+    )
+    expert_ids = torch.empty(max_blocks, dtype=torch.int32, device=device)
+    token_nums = torch.empty(E, dtype=torch.int32, device=device)
+    ntpp = torch.empty(1, dtype=torch.int32, device=device)
+    moe_align_block_size(
+        topk_ids, E, _BLOCK_SIZE_M, sorted_ids, expert_ids, token_nums, ntpp
+    )
+    return grad, inp, sorted_ids, expert_ids, ntpp
+
+
+def benchmark(args):
+    unit = "ms" if args.metric == "time" else "GB/s"
+    x_vals = list(_SHAPES)
+
+    config = triton.testing.Benchmark(
+        x_names=["num_tokens", "E", "N", "K", "top_k", "label"],
+        x_vals=x_vals,
+        line_arg="provider",
+        line_vals=["moe_wgrad"],
+        line_names=[f"moe_wgrad ({unit})"],
+        styles=[("blue", "-")],
+        ylabel=unit,
+        plot_name=get_caller_name_no_ext(),
+        args={},
+    )
+
+    @triton.testing.perf_report([config])
+    def _run(num_tokens, E, N, K, top_k, label, provider):
+        grad, inp, sorted_ids, expert_ids, ntpp = _setup(num_tokens, E, N, K, top_k)
+        fn = lambda: moe_wgrad(
+            grad,
+            inp,
+            sorted_ids,
+            expert_ids,
+            ntpp,
+            num_experts=E,
+            top_k=top_k,
+            weight_shape=(E, N, K),
+            block_size_m=_BLOCK_SIZE_M,
+        )
+        # reads: grad(T*N) + inp(T*K); writes: dW(E*N*K); bf16 = 2 bytes
+        mem = (num_tokens * N + num_tokens * K + E * N * K) * 2
+        ms = triton.testing.do_bench(fn, warmup=25, rep=100)
+        if args.metric == "time":
+            return ms
+        return mem * 1e-9 / (ms * 1e-3)
+
+    _run.run(save_path="." if args.o else None, print_data=True, show_plots=False)
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(prog="Benchmark moe_wgrad", allow_abbrev=False)
+    parser.add_argument(
+        "-metric",
+        nargs="?",
+        const="time",
+        choices=["time", "bandwidth"],
+        default="time",
+    )
+    parser.add_argument("-o", action="store_true", default=False)
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+    torch.manual_seed(0)
+    benchmark(args)
+
+
+if __name__ == "__main__":
+    main()