[Triton/Gluon] Add config-aware repr to the quant kernels - #5100
Open
Boss2002n wants to merge 2 commits into
Open
[Triton/Gluon] Add config-aware repr to the quant kernels#5100Boss2002n wants to merge 2 commits into
Boss2002n wants to merge 2 commits into
Conversation
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.
Contributor
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
PR title tags: |
github-actions
Bot
changed the title
Contributor
There was a problem hiding this comment.
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_reprand 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.
| import triton | ||
| import triton.language as tl | ||
|
|
||
| from aiter.ops.triton.utils._triton.kernel_repr import make_kernel_repr |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
No description provided.