[Triton/Gluon] Add config-aware repr to the quant kernels by Boss2002n · Pull Request #5100 · ROCm/aiter

Boss2002n · 2026-08-29T04:43:54Z

No description provided.

42 launchable quant kernels had no repr, so they appeared in traces under their bare name with no way to match a row back to its config. Covers _triton_kernels/quant/{quant,sage_attention_quant,fused_fp8_quant, fused_mxfp4_quant,fused_mxfp8_quant}.py, the two quant wrapper modules that define kernels inline, and the gfx1250 gluon copy. 15 jit device helpers in these files are left alone -- they are called from inside other kernels, never launched with a grid, and the rule exempts them. Two of them are worth a look on their own: _rotate_quantize_qk_kernel has no references anywhere in the tree, and _general_quant_kernel is only ever called as a device function despite the _kernel suffix. sage_quant_v_fp4_colmajor_kernel and sage_quant_v_mxfp4_colmajor_kernel are launchable but have no constexpr parameters at all, so a repr could only render the bare name and would strip Triton's own specialization suffix. Left without one; say the word if you want them included. mxfp6_fmha_pack.py defines its kernels inside "if _HAVE_TRITON:" blocks, so its make_kernel_repr import goes inside the existing try/except -- the module is designed to import cleanly with no triton installed. Key lists follow the neighbouring kernels in each subtree: tuned block/tile sizes and the meaningful compile-time flags, with runtime pointers, strides and grid-extent constexprs excluded (the launch site was checked before deciding on each one). A key that is not a real parameter renders as NONE in every trace name, so every list was re-parsed from the file and cross-checked against the kernel signature, then rendered through the real make_kernel_repr to confirm no NONE appears. Bare @triton.jit(repr=...) is used throughout, matching the majority of the reprs already in the tree. Full-tree ruff (0.16.0, the CI pin) and black report exactly the same findings as main.

github-actions · 2026-08-29T04:44:02Z

🏷️ CI Guide

Runs automatically on every PR:

✅ Pre-checks (submodule verification, code formatting)
✅ Aiter op tests (gfx942 + gfx950)
✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
`ci:gfx1250-ffm-triton`	Run the five-shard gfx1250 FFM Triton test suite
`ci:triton-300x`	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
`multigpu`	Aiter multi-GPU tests on the 8-GPU runner
`ci:sglang`	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
`ci:atom`	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
`ci:atom_full`	ATOM accuracy suite for PR and main models from ATOM `models_accuracy.json`
`ci:vllm`	vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
`ci:all`	All standard extended tests (excludes `ci:atom_full`)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 5100 --add-label <label>

PR title tags:
Component tags ([Triton/Gluon], [HIP], [CK], [ASM], ...) are added to the PR title automatically from the changed files and re-synced on every push — change-type tags like [fix]/[Perf] and op tags like [MLA] are left untouched. Add the no-auto-title label to opt this PR out of title tagging.

Copilot

Pull request overview

This PR adds config-aware repr functions (via make_kernel_repr) to a set of Triton and Gluon quantization kernels so compiled artifacts / traces include key constexpr specialization parameters in their names.

Changes:

Import make_kernel_repr and define per-kernel repr builders listing relevant constexpr / heuristic keys.
Wire the repr builders into @triton.jit(repr=...) and @gluon.jit(repr=...) decorators across multiple quant kernel modules.

Reviewed changes

Copilot reviewed 8 out of 8 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
aiter/ops/triton/quant/mxfp6_fmha_pack.py	Adds config-aware repr to FP6 packing kernels.
aiter/ops/triton/quant/fused_mxfp4_quant.py	Adds config-aware repr to the `_fused_quant_fp8_sort_kernel` wrapper-local Triton kernel.
aiter/ops/triton/_triton_kernels/quant/sage_attention_quant.py	Adds config-aware repr to multiple Sage quant kernels in this module.
aiter/ops/triton/_triton_kernels/quant/quant.py	Adds config-aware repr to several per-tensor/token quant kernels.
aiter/ops/triton/_triton_kernels/quant/fused_mxfp8_quant.py	Adds config-aware repr to fused RMS/dual-RMS/flatten MXFP8 quant kernels.
aiter/ops/triton/_triton_kernels/quant/fused_mxfp4_quant.py	Adds config-aware repr to several fused MXFP4 quant kernels.
aiter/ops/triton/_triton_kernels/quant/fused_fp8_quant.py	Adds config-aware repr to multiple fused FP8 quant kernels.
aiter/ops/triton/_gluon_kernels/gfx1250/quant/fused_mxfp4_quant.py	Adds config-aware repr to Gluon gfx1250 fused MXFP4 quant kernels.

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

Copilot

Pull request overview

Copilot reviewed 8 out of 8 changed files in this pull request and generated no new comments.

github-actions · 2026-08-31T05:37:15Z

🏷️ CI Guide

Runs automatically on every PR:

✅ Pre-checks (submodule verification, code formatting)
✅ Aiter op tests (gfx942 + gfx950)
✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
`ci:gfx1250-ffm-triton`	Run the five-shard gfx1250 FFM Triton test suite
`ci:triton-300x`	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
`multigpu`	Aiter multi-GPU tests on the 8-GPU runner
`ci:sglang`	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
`ci:atom`	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
`ci:atom_full`	ATOM accuracy suite for PR and main models from ATOM `models_accuracy.json`
`ci:vllm`	vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
`ci:all`	All standard extended tests (excludes `ci:atom_full`)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 5100 --add-label <label>

PR title tags & labels:
Component tags ([Triton/Gluon], [HIP], [CK], [ASM], ...) are added to the PR title and as PR labels automatically from the changed files and re-synced on every push — change-type tags like [fix]/[Perf], op tags like [MLA], and human labels (ci:*) are left untouched. Add the no-auto-title label to opt this PR out.

Copilot

Pull request overview

Copilot reviewed 8 out of 8 changed files in this pull request and generated 2 comments.

make_kernel_repr renders each key's value into the kernel name, and Triton uses that name as a function identifier. A non-integral float becomes "..._2.5" and breaks at the dot; a negative becomes "..._-448" and breaks at the minus. Either way the compile fails. That is what the 513 failures in test_fused_fp8_quant.py and test_fused_rms_gated_fp8_group_quant.py were: invalid function identifier: 359654017857143e-06_ACTIVATION_SILU FP8_MIN_SCALING_FACTOR sat immediately before ACTIVATION in _fused_rms_fp8_group_quant_kernel's key list, and the test passes 1.0 / (fp8_max * 512) = 4.3596540178571429e-06. Triton parsed the name up to the first dot and was left with the fragment above. Removed, all for the same reason: DTYPE_MAX, DTYPE_MIN, FP8_MAX, FP8_MIN torch.finfo bounds; the MIN pair is negated at the call site FP8_MIN_SCALING_FACTOR caller-supplied float dtype_max mha_v4_per_tensor_scale_kernel takes `dtype_max / clip`, a true division Audited exhaustively rather than by name: every one of the 91 remaining keys was resolved to the expressions its launch sites actually pass, and none is a true division, a negation, a non-integral float literal or a finfo bound. The kernels keep their block sizes and compile-time flags, which is what a trace row needs to identify the specialization. Two things a reviewer may want to follow up on, both pre-existing and untouched here: the reprs on main in _triton_kernels/attention/{extend_attention, mla_decode_rope}.py list logit_cap, which is fine only because every current caller passes 0.0; and eleven repr keys elsewhere in the tree are not constexpr parameters of the kernel they name, so they render as NONE in every trace name (harmless, but the segment carries no information).

Copilot

Pull request overview

Copilot reviewed 8 out of 8 changed files in this pull request and generated no new comments.

Suppressed comments (2)

aiter/ops/triton/_triton_kernels/quant/fused_mxfp4_quant.py:324

ACTIVATION is passed as a Triton JIT function (see wrapper uses _get_activation_from_str(...)), and make_kernel_repr falls back to str(value) for unknown types. For JIT functions this typically includes a memory address, making the compiled kernel name non-deterministic across runs and harming compile cache reuse / trace readability. Prefer excluding callable constexprs from config_keys (or switch to passing a stable string activation name and dispatch inside the kernel).

_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",
        "ACTIVATION",
        "scaleN",
        "scaleM_pad",
        "scaleN_pad",
        "SHUFFLE",
        "X_HAS_SPLITK",
        "X_NUM_KSPLIT",
        "X_NUM_KSPLIT_POW2",
    ],
)

aiter/ops/triton/_triton_kernels/quant/fused_fp8_quant.py:502

ACTIVATION here is passed as a Triton JIT function in the wrapper path, and make_kernel_repr will stringify it (often including a memory address). That makes the kernel repr non-deterministic across runs and can defeat compile caching / make traces noisy. Exclude callable constexprs from config_keys (or pass an activation name string and dispatch inside the kernel).

_fused_reduce_act_mul_fp8_group_quant_repr = make_kernel_repr(
    "_fused_reduce_act_mul_fp8_group_quant",
    [
        "ACTIVATION",
        "BLOCK_SIZE_M2",
        "BLOCK_SIZE_N1",
        "BLOCK_SIZE_N2",
        "QUANT_BLOCK_SIZE",
        "X_HAS_SPLITK",
        "X_NUM_KSPLIT",
        "X_NUM_KSPLIT_POW2",
        "X_MASK",
    ],
)

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

Copilot

Pull request overview

Copilot reviewed 8 out of 8 changed files in this pull request and generated no new comments.

Suppressed comments (1)

Previously missed (1) — in code that hasn't changed since the last review.

aiter/ops/triton/_triton_kernels/quant/fused_mxfp4_quant.py:317

The kernel repr for _fused_reduce_act_mul_and_dynamic_mxfp4_quant_kernel omits the ACTIVATION constexpr, which can make compiled artifact names ambiguous across specializations that differ only by activation behavior. Include ACTIVATION in the config_keys list so the repr fully reflects compile-time specialization.

        "MXFP4_QUANT_BLOCK_SIZE",
        "EVEN_M_N",
        "SCALING_MODE",
        "scaleN",
        "scaleM_pad",

One conflict, in _gluon_kernels/gfx1250/quant/fused_mxfp4_quant.py, and unlike the other two branches it is an add/add rather than a revert: main landed a new kernel, _gluon_fused_reduce_rms_mxfp4_quant_kernel, immediately above _gluon_fused_dynamic_mxfp4_quant_moe_sort_kernel, which is exactly where this branch had put that kernel's repr definition. Kept both -- main's new kernel byte for byte, then the repr, then @gluon.jit(repr=...) in place of the bare decorator main's hunk carried for the kernel below it. All 38 reprs the branch adds survive, and every @triton.jit/@gluon.jit repr=... in the tree resolves to a definition. Note for a follow-up: the kernel main just added has no repr. It is launched from quant/fused_mxfp4_quant.py:560, so it is a candidate for the same treatment, but adding it inside a merge commit would hide it from review.

Copilot

Pull request overview

Copilot reviewed 8 out of 8 changed files in this pull request and generated 3 comments.

Suppressed comments (3)

aiter/ops/triton/quant/mxfp6_fmha_pack.py:450

The config-aware repr doesn’t include the Triton meta-parameter num_warps, but this kernel is launched with num_warps=4 in reorder_fp6_k_lds_order_triton. Adding it helps distinguish compiled variants in traces (consistent with e.g. aiter/ops/triton/_triton_kernels/gemm/basic/gemm_afp4wfp4.py:11-24).

    _gather_k_lds_repr = make_kernel_repr(
        "_gather_k_lds_kernel",
        [
            "DATA_TILE_BYTES",
            "BLOCK",

aiter/ops/triton/quant/mxfp6_fmha_pack.py:492

The config-aware repr doesn’t include the Triton meta-parameter num_warps, but this kernel is launched with num_warps=4 in reorder_fp6_k_lds_order_triton. Including it keeps compiled artifact names aligned with the specialization (consistent with other kernels’ repr usage).

    _fill_k_scale_tail_repr = make_kernel_repr(
        "_fill_k_scale_tail_kernel",
        [
            "TILE_BYTES",
            "SCALE_TAIL_OFFSET",

aiter/ops/triton/_triton_kernels/quant/quant.py:672

This kernel is launched with explicit Triton meta-parameters num_warps/num_stages (see aiter/ops/triton/quant/quant.py), but the new repr doesn’t include them. Adding these keys helps ensure trace names differentiate compiled variants even when only meta-parameters change.

_dynamic_nvfp4_quant_repr = make_kernel_repr(
    "_dynamic_nvfp4_quant_kernel",
    [
        "BLOCK_SIZE_M",
        "BLOCK_SIZE_N",

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

Copilot

Pull request overview

Copilot reviewed 8 out of 8 changed files in this pull request and generated no new comments.

Suppressed comments (5)

Previously missed (3) — in code that hasn't changed since the last review.

aiter/ops/triton/_triton_kernels/quant/fused_fp8_quant.py:500

The repr key list omits ACTIVATION even though it is a tl.constexpr meta-parameter that changes the computation (the kernel calls ACTIVATION(x)). Include it in config_keys so the repr reflects the actual specialization.

        "X_HAS_SPLITK",
        "X_NUM_KSPLIT",
        "X_NUM_KSPLIT_POW2",
        "X_MASK",
    ],

aiter/ops/triton/_triton_kernels/quant/fused_mxfp8_quant.py:29

This kernel has a tl.constexpr NUM_PRGMS (controls the persistent row loop) but the repr key list doesn’t include it. Omitting it makes the generated kernel name less representative of the actual specialization used at runtime.

This issue also appears on line 121 of the same file.

    [
        "BLOCK_SIZE_K",
        "QUANT_BLOCK_SIZE",
    ],

aiter/ops/triton/_triton_kernels/quant/quant.py:445

The kernel repr is meant to be config-aware, but this kernel’s behavior depends on the tl.constexpr NUM_PRGMS (row loop stride) and the wrapper passes different NUM_PRGMS values based on M. Without including NUM_PRGMS in config_keys, compiled artifact names/traces can collide across specializations that only differ by NUM_PRGMS.

This issue also appears on line 510 of the same file.

    [
        "BLOCK_SIZE_N",
        "QUANT_BLOCK_SIZE",
    ],

aiter/ops/triton/_triton_kernels/quant/quant.py:513

Same as above: this kernel also has a tl.constexpr NUM_PRGMS that affects its persistent row loop, but the repr key list omits it, making the repr not fully specialization-aware.

    [
        "BLOCK_SIZE_N",
        "QUANT_BLOCK_SIZE",
    ],

aiter/ops/triton/_triton_kernels/quant/fused_mxfp8_quant.py:125

This kernel’s specialization includes tl.constexpr NUM_PRGMS (persistent row loop bound), but it’s not included in the repr config_keys. That makes traces/artifact names ambiguous across different NUM_PRGMS settings.

    [
        "BLOCK_SIZE_KQ",
        "BLOCK_SIZE_KK",
        "QUANT_BLOCK_SIZE",
    ],

Boss2002n requested review from a team and a lite review from Copilot August 29, 2026 04:43

github-actions Bot changed the title ~~[Triton] Add config-aware repr to the quant kernels~~ [Triton/Gluon] Add config-aware repr to the quant kernels Aug 29, 2026

Copilot started reviewing on behalf of Boss2002n August 29, 2026 04:44 View session

Copilot AI reviewed Aug 29, 2026

View reviewed changes

Comment thread aiter/ops/triton/_triton_kernels/quant/sage_attention_quant.py

Merge branch 'main' into satya/triton-kernel-repr-quant

0b09adb

Copilot AI review requested due to automatic review settings August 29, 2026 15:11

Copilot started reviewing on behalf of Boss2002n August 29, 2026 15:11 View session

Copilot AI reviewed Aug 29, 2026

View reviewed changes

zufayu requested review from a team and Dewei-Wang-sh August 31, 2026 01:17

Boss2002n closed this Aug 31, 2026

Boss2002n reopened this Aug 31, 2026

Boss2002n marked this pull request as draft August 31, 2026 05:37

Merge branch 'main' into satya/triton-kernel-repr-quant

f585dbd

Boss2002n removed request for a team and Dewei-Wang-sh August 31, 2026 05:49

Boss2002n marked this pull request as ready for review August 31, 2026 05:49

Copilot AI review requested due to automatic review settings August 31, 2026 05:49

github-actions Bot added the Triton/Gluon label Aug 31, 2026

Copilot started reviewing on behalf of Boss2002n August 31, 2026 05:50 View session

Copilot AI reviewed Aug 31, 2026

View reviewed changes

Comment thread aiter/ops/triton/_triton_kernels/quant/fused_mxfp4_quant.py

Comment thread aiter/ops/triton/_triton_kernels/quant/fused_fp8_quant.py Outdated

Copilot AI review requested due to automatic review settings August 31, 2026 12:52

Merge branch 'main' into satya/triton-kernel-repr-quant

7466cf4

Copilot started reviewing on behalf of Boss2002n August 31, 2026 12:53 View session

Copilot AI reviewed Aug 31, 2026

View reviewed changes

Boss2002n requested a review from brunomazzottiamd August 31, 2026 14:37

Potential fix for pull request finding

0840d94

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

Copilot AI review requested due to automatic review settings September 1, 2026 13:05

Potential fix for pull request finding

c3c9e5d

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

Copilot started reviewing on behalf of Boss2002n September 1, 2026 13:06 View session

Copilot AI reviewed Sep 1, 2026

View reviewed changes

Copilot AI review requested due to automatic review settings September 1, 2026 13:18

Copilot started reviewing on behalf of Boss2002n September 1, 2026 13:19 View session

Copilot AI reviewed Sep 1, 2026

View reviewed changes

Comment thread aiter/ops/triton/_triton_kernels/quant/fused_mxfp4_quant.py

Comment thread aiter/ops/triton/_triton_kernels/quant/quant.py

Comment thread aiter/ops/triton/quant/mxfp6_fmha_pack.py

Potential fix for pull request finding

fe68f78

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

Copilot AI review requested due to automatic review settings September 1, 2026 16:08

Boss2002n and others added 2 commits September 1, 2026 12:09

Potential fix for pull request finding

370fa4b

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

Potential fix for pull request finding

293a366

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

Copilot started reviewing on behalf of Boss2002n September 1, 2026 16:10 View session

Copilot AI reviewed Sep 1, 2026

View reviewed changes

brunomazzottiamd approved these changes Sep 1, 2026

View reviewed changes

Boss2002n merged commit 1c1d846 into main Sep 2, 2026
84 of 85 checks passed

Boss2002n deleted the satya/triton-kernel-repr-quant branch September 2, 2026 02:53

Navigation Menu

Uh oh!

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

[Triton/Gluon] Add config-aware repr to the quant kernels#5100
Boss2002n merged 11 commits into
mainfrom
satya/triton-kernel-repr-quant

Boss2002n commented Aug 29, 2026

Uh oh!

github-actions Bot commented Aug 29, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Copilot AI left a comment

Uh oh!

github-actions Bot commented Aug 31, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Uh oh!

Copilot AI left a comment

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

Boss2002n commented Aug 29, 2026

Uh oh!

github-actions Bot commented Aug 29, 2026

🏷️ CI Guide

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

github-actions Bot commented Aug 31, 2026

🏷️ CI Guide

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants