Skip to content

[Triton] Clean up MoE elementwise kernels #5259

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
vgokhale wants to merge 3 commits into
base: main
Choose a base branch
from
Draft

[Triton] Clean up MoE elementwise kernels #5259

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
a8968e2
[Triton] Deduplicate the MoE quant kernels against aiter/ops/triton/q…
vgokhale Sep 2, 2026
f9a812d
[Triton] Trim the comments added by the MoE quant dedup
vgokhale Sep 3, 2026
7965462
[Triton] Fold the static per-tensor quant launcher into its wrapper
vgokhale Sep 3, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
74 changes: 0 additions & 74 deletions aiter/ops/triton/_triton_kernels/moe/moe_op_gemm_a4w4.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,6 @@
import triton.language as tl

from aiter.ops.triton._triton_kernels.moe.activations import _swiglu
from aiter.ops.triton._triton_kernels.quant.quant import _mxfp4_quant_op
from aiter.ops.triton.utils._triton.kernel_repr import make_kernel_repr
from aiter.ops.triton.utils._triton.pid_preprocessing import pid_grid

Expand Down Expand Up @@ -127,79 +126,6 @@ def unswizzle_mx_scale_gfx1250(
return scale_buffer_slice


_mxfp4_quant_kernel_repr = make_kernel_repr(
"_mxfp4_quant_kernel",
[
"BLOCK_SIZE_M",
"BLOCK_SIZE_N",
"MXFP4_QUANT_BLOCK_SIZE",
"EVEN_M_N",
],
)


@triton.jit(repr=_mxfp4_quant_kernel_repr)
def _mxfp4_quant_kernel(
x_ptr,
x_fp4_ptr,
bs_ptr,
stride_x_m,
stride_x_n,
stride_fp4_m,
stride_fp4_n,
stride_bs_m,
stride_bs_n,
M,
N,
BLOCK_SIZE_M: tl.constexpr,
BLOCK_SIZE_N: tl.constexpr,
MXFP4_QUANT_BLOCK_SIZE: tl.constexpr,
EVEN_M_N: tl.constexpr,
):
pid_m = tl.program_id(0)
pid_n = tl.program_id(1)

offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
mask_m = offs_m < M
mask_n = offs_n < N
mask = mask_m[:, None] & mask_n[None, :]
x_offs = offs_m[:, None] * stride_x_m + offs_n[None, :] * stride_x_n
x = tl.load(x_ptr + x_offs, mask=mask, other=0).to(tl.float32)

out_tensor, bs_e8m0 = _mxfp4_quant_op(
x, BLOCK_SIZE_N, BLOCK_SIZE_M, MXFP4_QUANT_BLOCK_SIZE
)

# Store quantized x blocks
out_offs_n = pid_n * BLOCK_SIZE_N // 2 + tl.arange(0, BLOCK_SIZE_N // 2)
out_offs = offs_m[:, None] * stride_fp4_m + out_offs_n[None, :] * stride_fp4_n

if EVEN_M_N:
tl.store(x_fp4_ptr + out_offs, out_tensor)
else:
out_mask = (offs_m < M)[:, None] & (out_offs_n < (N // 2))[None, :]
tl.store(x_fp4_ptr + out_offs, out_tensor, mask=out_mask)

# Store scale blocks
NUM_QUANT_BLOCKS: tl.constexpr = BLOCK_SIZE_N // MXFP4_QUANT_BLOCK_SIZE
num_blocks_total = N // MXFP4_QUANT_BLOCK_SIZE
bs_offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
bs_offs_n = pid_n * NUM_QUANT_BLOCKS + tl.arange(0, NUM_QUANT_BLOCKS)

bs_offs = bs_offs_m[:, None] * stride_bs_m + bs_offs_n[None, :] * stride_bs_n
bs_mask = (bs_offs_m < M)[:, None] & (bs_offs_n < num_blocks_total)[None, :]

if EVEN_M_N:
tl.store(bs_ptr + bs_offs, bs_e8m0)
else:
tl.store(
bs_ptr + bs_offs,
bs_e8m0,
mask=bs_mask,
)


_moe_gemm_a4w4_repr = make_kernel_repr(
"_moe_gemm_a4w4",
[
Expand Down
Loading