[Triton/Gluon] Add config-aware repr to the rope and normalization kernels - #5099
Open
Boss2002n wants to merge 1 commit into
Open
[Triton/Gluon] Add config-aware repr to the rope and normalization kernels#5099Boss2002n wants to merge 1 commit into
Boss2002n wants to merge 1 commit into
Conversation
32 launchable kernels across _triton_kernels/rope/{rope,fused_qkv_split_qk_rope,
fused_qkv_split_qk_norm_rope_cache}.py and _triton_kernels/normalization/
{norm,rmsnorm,fused_add_rmsnorm_pad,fused_rmsnorm_add}.py had no repr.
12 jit device helpers are left alone (the _get_neox_rotated_x / _get_gptj_
rotated_x family, _rms_norm, _per_token_quant, _rmsnorm_op and friends) -- they
are inlined into other kernels, never launched.
NUM_PRGMS is excluded from the five rmsnorm kernels that declare it: every
launch site is literally grid = (NUM_PRGMS,) with NUM_PRGMS = min(rows,
get_num_sms()), so it is device occupancy rather than a tuned parameter and
would give a different trace name per input size. num_stages is kept for the
four rope kernels that take it as a real constexpr -- it is tuned at the launch
site, and attention/mla.py and unified_attention.py already list it.
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
Adds config-aware Triton repr functions (via make_kernel_repr) to RoPE and normalization kernels under aiter/ops/triton/_triton_kernels/, improving kernel naming/traceability of compiled specializations.
Changes:
- Introduces
make_kernel_reprusage and attaches@triton.jit(repr=...)to RoPE kernels (including fused QKV+RoPE variants). - Introduces
make_kernel_reprusage and attaches@triton.jit(repr=...)to RMSNorm/LayerNorm kernels (including fused variants and some bwd kernels). - Adds per-kernel repr key lists intended to encode relevant
tl.constexprspecialization parameters into artifact names.
Reviewed changes
Copilot reviewed 7 out of 7 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
| aiter/ops/triton/_triton_kernels/rope/rope.py | Adds make_kernel_repr and config-aware repr= to multiple RoPE kernels. |
| aiter/ops/triton/_triton_kernels/rope/fused_qkv_split_qk_rope.py | Adds config-aware repr= to the fused QKV split + RoPE kernel. |
| aiter/ops/triton/_triton_kernels/rope/fused_qkv_split_qk_norm_rope_cache.py | Adds config-aware repr= to the fused QKV split + norm + RoPE + cache kernel. |
| aiter/ops/triton/_triton_kernels/normalization/rmsnorm.py | Adds config-aware repr= to RMSNorm forward/fused/quant and bwd kernels. |
| aiter/ops/triton/_triton_kernels/normalization/norm.py | Adds config-aware repr= to LayerNorm forward/fused/quant and bwd kernels. |
| aiter/ops/triton/_triton_kernels/normalization/fused_rmsnorm_add.py | Adds config-aware repr= to the fused RMSNorm (+ optional residual) kernel. |
| aiter/ops/triton/_triton_kernels/normalization/fused_add_rmsnorm_pad.py | Adds config-aware repr= to the fused add + RMSNorm + pad kernel. |
Suppressed comments (4)
aiter/ops/triton/_triton_kernels/normalization/rmsnorm.py:183
- This kernel takes NUM_PRGMS as a tl.constexpr and uses it in tl.range(..., NUM_PRGMS, ...), but NUM_PRGMS is missing from the repr keys, reducing trace-name fidelity across specializations.
"DUMP_INTERMEDIATE",
"BLOCK_SIZE",
"USE_BLOCKED",
],
)
aiter/ops/triton/_triton_kernels/normalization/rmsnorm.py:422
- NUM_PRGMS is a tl.constexpr parameter of this kernel and affects the persistent-loop traversal, but it’s not included in the repr keys; add it so specialization names remain config-aware.
[
"BLOCK_SIZE",
"USE_BLOCKED",
],
aiter/ops/triton/_triton_kernels/normalization/rmsnorm.py:582
- The kernel signature includes NUM_PRGMS (constexpr) but the repr key list omits it, so different NUM_PRGMS specializations will share the same compiled name suffix.
"IS_SMOOTH",
"BLOCK_SIZE",
"USE_BLOCKED",
],
aiter/ops/triton/_triton_kernels/normalization/rmsnorm.py:820
- NUM_PRGMS is a constexpr that changes which rows each program processes (tl.range stride), but it’s missing from the repr keys. Include it so the repr stays aligned with specialization-relevant constants.
[
"BLOCK_SIZE",
"USE_BLOCKED",
],
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
+47
to
+50
| [ | ||
| "BLOCK_SIZE", | ||
| "USE_BLOCKED", | ||
| ], |
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.