Skip to content

[Triton/Gluon] Add config-aware repr to the quant kernels #5100

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

Merged
Boss2002n merged 11 commits into from
Sep 2, 2026
Merged

[Triton/Gluon] Add config-aware repr to the quant kernels #5100

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
b6e4a55
[Triton] Add config-aware repr to the quant kernels
Boss2002n Aug 29, 2026
0b09adb
Merge branch 'main' into satya/triton-kernel-repr-quant
Boss2002n Aug 29, 2026
f585dbd
Merge branch 'main' into satya/triton-kernel-repr-quant
Boss2002n Aug 31, 2026
b1dd1fd
[Triton] Drop repr keys that render an invalid Triton kernel name
Boss2002n Aug 31, 2026
7466cf4
Merge branch 'main' into satya/triton-kernel-repr-quant
Boss2002n Aug 31, 2026
0840d94
Potential fix for pull request finding
Boss2002n Sep 1, 2026
c3c9e5d
Potential fix for pull request finding
Boss2002n Sep 1, 2026
d492d31
Merge branch 'main' into satya/triton-kernel-repr-quant
Boss2002n Sep 1, 2026
fe68f78
Potential fix for pull request finding
Boss2002n Sep 1, 2026
370fa4b
Potential fix for pull request finding
Boss2002n Sep 1, 2026
293a366
Potential fix for pull request finding
Boss2002n Sep 1, 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
39 changes: 37 additions & 2 deletions aiter/ops/triton/_gluon_kernels/gfx1250/quant/fused_mxfp4_quant.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,8 @@
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


# rms norm op copied from triton
@triton.jit
Expand Down Expand Up @@ -124,13 +126,34 @@ def _mxfp4_quant_op(
return x_fp4, bs_e8m0.reshape(BLOCK_SIZE_M, NUM_QUANT_BLOCKS)


_gluon_fused_rms_mxfp4_quant_repr = make_kernel_repr(
"_gluon_fused_rms_mxfp4_quant_kernel",
[
"BLOCK_SIZE_M",
"BLOCK_SIZE_N",
"BLOCK_SIZE_N2",
"MXFP4_QUANT_BLOCK_SIZE",
"HAS_SECOND_INPUT",
"FIRST_INPUT_RES",
"FIRST_INPUT_OUT",
"SCALE_N",
"SCALE_M_PAD",
"SCALE_N_PAD",
"SHUFFLE",
"SHUFFLE_PAD",
"EVEN_M_N",
"ROWS_PER_CTA",
],
)


@triton.heuristics(
{
"EVEN_M_N": lambda args: args["M"] % args["ROWS_PER_CTA"] == 0
and args["N1"] % (args["BLOCK_SIZE_N"]) == 0,
}
)
@gluon.jit
@gluon.jit(repr=_gluon_fused_rms_mxfp4_quant_repr)
def _gluon_fused_rms_mxfp4_quant_kernel(
x1_ptr,
w1_ptr,
Expand Down Expand Up @@ -879,7 +902,19 @@ def _gluon_fused_reduce_rms_mxfp4_quant_kernel(
gl.amd.gfx1250.tdm.async_wait(0)


@gluon.jit
_gluon_fused_dynamic_mxfp4_quant_moe_sort_repr = make_kernel_repr(
"_gluon_fused_dynamic_mxfp4_quant_moe_sort_kernel",
[
"MXFP4_QUANT_BLOCK_SIZE",
"BLOCK_SIZE_Mx",
"BLOCK_SIZE_M",
"BLOCK_SIZE_N",
"TOPK",
],
)


@gluon.jit(repr=_gluon_fused_dynamic_mxfp4_quant_moe_sort_repr)
def _gluon_fused_dynamic_mxfp4_quant_moe_sort_kernel(
x_ptr,
x_fp4_ptr,
Expand Down
102 changes: 96 additions & 6 deletions aiter/ops/triton/_triton_kernels/quant/fused_fp8_quant.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
import triton
import triton.language as tl

from aiter.ops.triton.utils._triton.kernel_repr import make_kernel_repr

try:
from triton.language.extra.libdevice import fast_dividef, fast_expf
except ImportError:
Expand Down Expand Up @@ -39,7 +41,19 @@ def _fp8_quant_op(
return x, scale_out


@triton.jit
_fused_rms_fp8_per_tensor_static_quant_repr = make_kernel_repr(
"_fused_rms_fp8_per_tensor_static_quant_kernel",
[
"BLOCK_SIZE_N",
"HAVE_SECOND_INPUT",
"FIRST_INPUT_RES",
"FIRST_INPUT_OUT",
"RMSNORM_CONVERT_TO_INP1_TYPE",
],
)


@triton.jit(repr=_fused_rms_fp8_per_tensor_static_quant_repr)
def _fused_rms_fp8_per_tensor_static_quant_kernel(
inp1_ptr,
weight1_ptr,
Expand Down Expand Up @@ -148,7 +162,30 @@ def _fused_rms_fp8_per_tensor_static_quant_kernel(
)


@triton.jit
_fused_rms_fp8_group_quant_repr = make_kernel_repr(
"_fused_rms_fp8_group_quant_kernel",
[
"BLOCK_SIZE_N",
"QUANT_BLOCK_SIZE",
"HAVE_SECOND_INPUT",
"FIRST_INPUT_RES",
"FIRST_INPUT_OUT",
"GATED_RMS_FP8",
"RMS_TILE",
"ROWS_PER_BLOCK",
"GROUP_SIZE_GATED",
"NUM_GROUPS_GATED",
"BLOCK_G",
"HAS_BIAS_GATED",
"HAS_Z_GATED",
"NORM_BEFORE_GATE",
"USE_UE8M0",
"ACTIVATION",
],
)


@triton.jit(repr=_fused_rms_fp8_group_quant_repr)
def _fused_rms_fp8_group_quant_kernel(
inp1_ptr,
weight1_ptr,
Expand Down Expand Up @@ -386,7 +423,16 @@ def _fused_rms_fp8_group_quant_kernel(
)


@triton.jit
_fused_flatten_fp8_group_quant_repr = make_kernel_repr(
"_fused_flatten_fp8_group_quant_kernel",
[
"BLOCK_SIZE_N2",
"QUANT_BLOCK_SIZE",
],
)


@triton.jit(repr=_fused_flatten_fp8_group_quant_repr)
def _fused_flatten_fp8_group_quant_kernel(
x_ptr,
out_ptr,
Expand Down Expand Up @@ -440,7 +486,22 @@ def _fused_flatten_fp8_group_quant_kernel(
)


@triton.jit
_fused_reduce_act_mul_fp8_group_quant_repr = make_kernel_repr(
"_fused_reduce_act_mul_fp8_group_quant",
[
"BLOCK_SIZE_M2",
"BLOCK_SIZE_N1",
"BLOCK_SIZE_N2",
"QUANT_BLOCK_SIZE",
"X_HAS_SPLITK",
"X_NUM_KSPLIT",
"X_NUM_KSPLIT_POW2",
"X_MASK",
],
)


@triton.jit(repr=_fused_reduce_act_mul_fp8_group_quant_repr)
def _fused_reduce_act_mul_fp8_group_quant(
x_ptr,
y_ptr,
Expand Down Expand Up @@ -589,7 +650,27 @@ def _fused_reduce_act_mul_fp8_group_quant(
)


@triton.jit
_fused_reduce_rms_fp8_group_quant_repr = make_kernel_repr(
"_fused_reduce_rms_fp8_group_quant_kernel",
[
"BLOCK_SIZE_N1",
"BLOCK_SIZE_N2",
"BLOCK_SIZE_N3",
"N_MASK1",
"N_MASK2",
"N_MASK3",
"QUANT_BLOCK_SIZE",
"HAVE_SECOND_INPUT",
"FIRST_INPUT_RES",
"FIRST_INPUT_OUT",
"HAS_SPLITK",
"NUM_SPLITK",
"NUM_SPLITK_POW2",
],
)


@triton.jit(repr=_fused_reduce_rms_fp8_group_quant_repr)
def _fused_reduce_rms_fp8_group_quant_kernel(
inp1_ptr,
weight1_ptr,
Expand Down Expand Up @@ -837,7 +918,16 @@ def _fused_reduce_rms_fp8_group_quant_kernel(
)


@triton.jit
_fused_silu_mul_fp8_per_tensor_static_quant_repr = make_kernel_repr(
"_fused_silu_mul_fp8_per_tensor_static_quant_kernel",
[
"BLOCK_SIZE_N",
"SILU_CONVERT_TO_INP_TYPE",
],
)


@triton.jit(repr=_fused_silu_mul_fp8_per_tensor_static_quant_repr)
def _fused_silu_mul_fp8_per_tensor_static_quant_kernel(
inp_ptr,
out_fp8_ptr,
Expand Down
105 changes: 100 additions & 5 deletions aiter/ops/triton/_triton_kernels/quant/fused_mxfp4_quant.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
import triton
import triton.language as tl

from aiter.ops.triton.utils._triton.kernel_repr import make_kernel_repr

from .quant import _mxfp4_quant_op


Expand All @@ -16,6 +18,27 @@ def _rmsmorm_op(row, weight, n_cols, epsilon):
return rms_norm


_fused_rms_mxfp4_quant_repr = make_kernel_repr(
"_fused_rms_mxfp4_quant_kernel",
[
"BLOCK_SIZE_M",
"BLOCK_SIZE_N",
"BLOCK_SIZE_N2",
"MXFP4_QUANT_BLOCK_SIZE",
"HAS_SECOND_INPUT",
"FIRST_INPUT_RES",
"FIRST_INPUT_OUT",
"SCALE_N",
"SCALE_M_PAD",
"SCALE_N_PAD",
"SHUFFLE",
"SHUFFLE_PAD",
"EVEN_M_N",
"EVEN_M_N2",
],
)


@triton.heuristics(
{
"EVEN_M_N": lambda args: args["M"] % args["BLOCK_SIZE_M"] == 0
Expand All @@ -24,7 +47,7 @@ def _rmsmorm_op(row, weight, n_cols, epsilon):
and args["N2"] % (args["BLOCK_SIZE_N2"]) == 0,
}
)
@triton.jit
@triton.jit(repr=_fused_rms_mxfp4_quant_repr)
def _fused_rms_mxfp4_quant_kernel(
x1_ptr,
w1_ptr,
Expand Down Expand Up @@ -222,7 +245,16 @@ def _fused_rms_mxfp4_quant_kernel(
)


@triton.jit
_fused_flatten_mxfp4_quant_repr = make_kernel_repr(
"_fused_flatten_mxfp4_quant",
[
"BLOCK_SIZE_N2",
"MXFP4_QUANT_BLOCK_SIZE",
],
)


@triton.jit(repr=_fused_flatten_mxfp4_quant_repr)
def _fused_flatten_mxfp4_quant(
x_ptr,
out_ptr,
Expand Down Expand Up @@ -268,13 +300,38 @@ def _fused_flatten_mxfp4_quant(
)


_fused_reduce_act_mul_and_dynamic_mxfp4_quant_repr = make_kernel_repr(
"_fused_reduce_act_mul_and_dynamic_mxfp4_quant_kernel",
[
"BLOCK_SIZE_M1",
"BLOCK_SIZE_N1",
"BLOCK_SIZE_M2",
"BLOCK_SIZE_N2",
"NUM_ITER",
"NUM_STAGES",
"MXFP4_QUANT_BLOCK_SIZE",
"EVEN_M_N",
"SCALING_MODE",
"scaleN",
"scaleM_pad",
"scaleN_pad",
"SHUFFLE",
"X_HAS_SPLITK",
"X_NUM_KSPLIT",
"X_NUM_KSPLIT_POW2",
"num_warps",
"num_stages",
],
)
Comment thread
Copilot marked this conversation as resolved.


@triton.heuristics(
{
"EVEN_M_N": lambda args: args["M"] % args["BLOCK_SIZE_M1"] == 0
and args["N1"] % (args["BLOCK_SIZE_N1"] * args["NUM_ITER"]) == 0,
}
)
@triton.jit
@triton.jit(repr=_fused_reduce_act_mul_and_dynamic_mxfp4_quant_repr)
def _fused_reduce_act_mul_and_dynamic_mxfp4_quant_kernel(
x_ptr,
y_ptr,
Expand Down Expand Up @@ -482,6 +539,32 @@ def _fused_reduce_act_mul_and_dynamic_mxfp4_quant_kernel(
)


_fused_reduce_rms_mxfp4_quant_repr = make_kernel_repr(
"_fused_reduce_rms_mxfp4_quant_kernel",
[
"BLOCK_SIZE_M",
"BLOCK_SIZE_N",
"BLOCK_SIZE_N2",
"BLOCK_SIZE_N3",
"MXFP4_QUANT_BLOCK_SIZE",
"HAS_SECOND_INPUT",
"FIRST_INPUT_RES",
"FIRST_INPUT_OUT",
"HAS_SPLITK",
"NUM_SPLITK",
"NUM_SPLITK_POW2",
"SCALE_N",
"SCALE_M_PAD",
"SCALE_N_PAD",
"SHUFFLE",
"SHUFFLE_PAD",
"EVEN_M_N",
"EVEN_M_N2",
"EVEN_M_N3",
],
)


@triton.heuristics(
{
"EVEN_M_N": lambda args: args["M"] % args["BLOCK_SIZE_M"] == 0
Expand All @@ -492,7 +575,7 @@ def _fused_reduce_act_mul_and_dynamic_mxfp4_quant_kernel(
and args["N3"] % (args["BLOCK_SIZE_N3"]) == 0,
}
)
@triton.jit
@triton.jit(repr=_fused_reduce_rms_mxfp4_quant_repr)
def _fused_reduce_rms_mxfp4_quant_kernel(
x1_ptr,
w1_ptr,
Expand Down Expand Up @@ -818,7 +901,19 @@ def _fused_reduce_rms_mxfp4_quant_kernel(
)


@triton.jit
_fused_dynamic_mxfp4_quant_moe_sort_repr = make_kernel_repr(
"_fused_dynamic_mxfp4_quant_moe_sort_kernel",
[
"MXFP4_QUANT_BLOCK_SIZE",
"BLOCK_SIZE_Mx",
"BLOCK_SIZE_M",
"BLOCK_SIZE_N",
"TOPK",
],
)


@triton.jit(repr=_fused_dynamic_mxfp4_quant_moe_sort_repr)
def _fused_dynamic_mxfp4_quant_moe_sort_kernel(
x_ptr,
x_fp4_ptr,
Expand Down
Loading
Loading