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

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

No description provided.

github-actions · 2026-08-29T04:43:43Z

🏷️ 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 5097 --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

Adds config-aware repr strings (via make_kernel_repr) to a set of Triton/Gluon attention-related kernels so compiled artifacts and traces encode key compile-time specialization parameters.

Changes:

Imports make_kernel_repr and defines per-kernel _..._repr callables capturing relevant tl.constexpr / compile-time keys.
Applies repr= to multiple @triton.jit / @gluon.jit kernel entrypoints across attention kernels.
Extends kernel naming coverage to decode, unified attention, sparse attention, FP8 MQA logits, and related utilities.

Reviewed changes

Copilot reviewed 12 out of 12 changed files in this pull request and generated 7 comments.

Show a summary per file

File	Description
aiter/ops/triton/attention/mla_decode.py	Adds `repr=` for MLA decode kernels to improve specialization-aware naming.
aiter/ops/triton/_triton_kernels/attention/unified_attention.py	Adds config-aware `repr=` for unified attention 2D kernel.
aiter/ops/triton/_triton_kernels/attention/unified_attention_sparse_mla.py	Adds config-aware `repr=` for sparse MLA unified attention kernel.
aiter/ops/triton/_triton_kernels/attention/sparse_attention_dsv4.py	Adds `repr=` for the autotuned sparse attention prefill kernel.
aiter/ops/triton/_triton_kernels/attention/pa_mqa_logits.py	Adds `repr=` for paged MQA logits Triton kernels.
aiter/ops/triton/_triton_kernels/attention/lean_atten_paged.py	Adds `repr=` for Lean Attention paged decode kernel.
aiter/ops/triton/_triton_kernels/attention/fp8_mqa_logits.py	Adds `repr=` for FP8 MQA logits Triton kernel.
aiter/ops/triton/_triton_kernels/attention/fav3_sage_attention.py	Adds `repr=` for SAGE attention forward kernel.
aiter/ops/triton/_triton_kernels/attention/fav3_sage_attention_mxfp4.py	Adds `repr=` for MXFP4 SAGE attention forward kernel.
aiter/ops/triton/_triton_kernels/attention/block_lut.py	Adds `repr=` for block attention mask → LUT kernel.
aiter/ops/triton/_gluon_kernels/gfx950/attention/fp8_mqa_logits.py	Adds `repr=` for gfx950 Gluon FP8 MQA logits kernel.
aiter/ops/triton/_gluon_kernels/gfx1250/attention/fp8_mqa_logits.py	Adds `repr=` for gfx1250 Gluon FP8 MQA logits kernel.

Suppressed comments (2)

aiter/ops/triton/attention/mla_decode.py:304

_fwd_grouped_kernel_stage1 is launched with tuning kwargs (num_warps/num_stages and, on ROCm, waves_per_eu/matrix_instr_nonkdim/kpack), but the repr key list omits them. Add these keys so different tuned variants don’t share the same repr name.

_fwd_grouped_kernel_stage1_repr = make_kernel_repr(
    "_fwd_grouped_kernel_stage1",
    [
        "kv_group_num",
        "q_head_num",
        "BLOCK_DMODEL",

aiter/ops/triton/attention/mla_decode.py:590

_fwd_kernel_stage2 is launched with num_warps/num_stages and (on ROCm) waves_per_eu/matrix_instr_nonkdim/kpack, but the repr key list omits those compile-time variants. Include the tuning keys so repr names differentiate the compiled specializations.

_fwd_kernel_stage2_repr = make_kernel_repr(
    "_fwd_kernel_stage2",
    [
        "NUM_KV_SPLITS",
        "BLOCK_DV",
        "Lv",
    ],

💡 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 12 out of 12 changed files in this pull request and generated 1 comment.

Suppressed comments (4)

aiter/ops/triton/_triton_kernels/attention/unified_attention.py:59

kernel_unified_attention_2d is launched with a per-shape config that includes num_warps, num_stages, and waves_per_eu (see aiter/ops/triton/attention/unified_attention.py:149-156), but those keys are not included in the make_kernel_repr(...) key list here. This can produce identical compiled artifact names for different kernel specializations, defeating the purpose of a config-aware repr (compare to kernel_unified_attention_3d_repr including these keys in this file).

_kernel_unified_attention_2d_repr = make_kernel_repr(
    "kernel_unified_attention_2d",
    [
        "num_query_heads",
        "num_queries_per_kv",

aiter/ops/triton/_triton_kernels/attention/unified_attention_sparse_mla.py:47

This kernel is launched with compile-time launch parameters num_warps and num_stages (see aiter/ops/triton/attention/unified_attention_sparse_mla.py:58-95), but the new make_kernel_repr(...) list here does not include them. That can cause different specializations to share the same repr name, making profiling/debugging ambiguous.

_kernel_unified_attention_sparse_mla_2d_repr = make_kernel_repr(
    "_kernel_unified_attention_sparse_mla_2d",
    [
        "num_query_heads",
        "num_queries_per_kv",

aiter/ops/triton/_triton_kernels/attention/sparse_attention_dsv4.py:270

_sparse_attn_prefill_kernel is autotuned over multiple configurations that vary waves_per_eu, matrix_instr_nonkdim, num_warps, and num_stages (see _get_prefill_autotune_configs() just above), but the new make_kernel_repr(...) list does not include those tuning keys. This can lead to identical repr names for different autotuned variants.

_sparse_attn_prefill_kernel_repr = make_kernel_repr(
    "_sparse_attn_prefill_kernel",
    [
        "HAS_ATTN_SINK",
        "BLOCK_H",

aiter/ops/triton/_triton_kernels/attention/pa_mqa_logits.py:232

_deepgemm_fp8_paged_mqa_logits_stage1 is launched with the compile-time kwarg waves_per_eu (see aiter/ops/triton/attention/pa_mqa_logits.py:253-256), but the new make_kernel_repr(...) key list here does not include it. If WavePerEU changes, different specializations will share the same repr name.

_deepgemm_fp8_paged_mqa_logits_stage1_repr = make_kernel_repr(
    "_deepgemm_fp8_paged_mqa_logits_stage1",
    [
        "ChunkQ",
        "ChunkK",

Copilot

Pull request overview

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

Copilot

Pull request overview

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

19 launchable attention kernels had no repr: the five real pa_mqa_logits kernels, the four defined inline in attention/mla_decode.py, block_lut, both fav3_sage variants, fp8_mqa_logits (triton plus both gluon arch copies), lean_atten_paged, sparse_attention_dsv4, unified_attention's 2d kernel, and unified_attention_sparse_mla. The three _gluon_deepgemm_* entries in _triton_kernels/attention/pa_mqa_logits.py are deliberately left alone: their bodies are just "pass" -- they exist only so an AOT load has a matching signature -- and both triton.compile call sites pass an options dict that already sets "name" to a literal, so the repr callback is never consulted on that path. Naming for that op is driven by options["name"], not by repr. kernel_unified_attention_2d does NOT get the 3d kernel's key list: six of those keys do not exist on the 2d kernel and would have rendered NONE. It uses its own constexpr set. MAX_SEQLENS_Q / MAX_SEQLENS_K are dropped from both sage reprs -- MAX_SEQLENS_Q is exactly the grid extent at the launch site, so including it would vary the trace name per sequence length while the tuned config stayed put. (Noting the other side of that argument: flash_attn_triton_amd lists both in its autotune keys for the same kernel family.) 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. Keys whose value can be negative or a non-integral float are deliberately excluded. make_kernel_repr renders the value into the kernel name, and Triton then rejects the name -- "CompilationError: invalid function identifier" -- so e.g. a 2.5 becomes "..._2.5" and the name breaks at the dot, and a -448 breaks at the minus. Every key list here was checked against the values its launch sites actually pass. Dropped on that basis: logit_cap on the two mla_decode stage-1 kernels. Note the pre-existing reprs in _triton_kernels/attention/extend_attention.py and mla_decode_rope.py still list logit_cap; those are untouched here, but they carry the same latent failure if a caller ever passes a non-integral cap.

Copilot

Pull request overview

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

Suppressed comments (2)

aiter/ops/triton/attention/mla_decode.py:312

The repr key list for _fwd_grouped_kernel_stage1 omits logit_cap even though it is a tl.constexpr argument. This makes the compiled artifact/trace name ambiguous across different logit_cap specializations.

        "NUM_KV_SPLITS",
        "PAGE_SIZE",
        "Lk",
        "Lv",
    ],

aiter/ops/triton/_triton_kernels/attention/sparse_attention_dsv4.py:273

This kernel is autotuned, but the repr key list does not include the autotune-selected meta-parameters (e.g., matrix_instr_nonkdim, waves_per_eu, num_warps, num_stages). Without these, trace/compiled artifact names won't identify which tuned config was selected.

        "HAS_ATTN_SINK",
        "BLOCK_H",
        "BLOCK_D",
        "BLOCK_K",
    ],

+        "NUM_KV_SPLITS",
+        "PAGE_SIZE",
+        "Lk",
+        "Lv",
+    ],


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 attention kernels~~ [Triton/Gluon] Add config-aware repr to the attention 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

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:12 View session

Copilot AI reviewed Aug 29, 2026

View reviewed changes

Comment thread aiter/ops/triton/_triton_kernels/attention/pa_mqa_logits.py

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

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

Copilot AI reviewed Aug 29, 2026

View reviewed changes

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

Copilot started reviewing on behalf of Boss2002n August 29, 2026 15:24 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:18

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

Boss2002n force-pushed the satya/triton-kernel-repr-attention branch from cb35f8b to a78cbdc Compare August 31, 2026 06:06

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

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

Copilot AI reviewed Aug 31, 2026

View reviewed changes

Comment thread aiter/ops/triton/attention/mla_decode.py

Comment on lines +51 to +55

"NUM_KV_SPLITS",

"PAGE_SIZE",

"Lk",

"Lv",

],

Navigation Menu

Uh oh!

[Triton/Gluon] Add config-aware repr to the attention kernels - #5097

[Triton/Gluon] Add config-aware repr to the attention kernels#5097
Boss2002n wants to merge 1 commit into
mainfrom
satya/triton-kernel-repr-attention

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!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

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!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 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!

Uh oh!

Uh oh!

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!

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!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants