[Triton/Gluon] Move the sage-attention launch params into the config tree - #5106
Open
Boss2002n wants to merge 3 commits into
Open
[Triton/Gluon] Move the sage-attention launch params into the config tree#5106Boss2002n wants to merge 3 commits into
Boss2002n wants to merge 3 commits into
Conversation
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
Moves Sage-attention (fav3_sage) forward launch parameters out of Python and into the nested Triton config tree, so launch tuning is data-driven and consistent with the rest of the config-loading machinery.
Changes:
- Add nested-layout
DEFAULT.jsonconfig tables forfav3_sageon gfx942 and gfx950, and forfav3_sage_mxfp4on gfx950. - Update
fav3_sageto load per-arch configs viaresolve_config_dir()/load_config_json(), with a deterministic gfx942 fallback when the running arch has no table. - Update the MXFP4 wrapper to load its forward launch config from the config tree (gfx950-only).
Reviewed changes
Copilot reviewed 5 out of 5 changed files in this pull request and generated no comments.
Show a summary per file
| File | Description |
|---|---|
| aiter/ops/triton/configs/gfx950/triton/attention/fav3_sage/DEFAULT.json | Adds gfx950 default + optional fwd_a3_tuned preset for fav3_sage forward launch params. |
| aiter/ops/triton/configs/gfx950/triton/attention/fav3_sage_mxfp4/DEFAULT.json | Adds gfx950 default forward launch params for the MXFP4 variant. |
| aiter/ops/triton/configs/gfx942/triton/attention/fav3_sage/DEFAULT.json | Adds gfx942 default + optional fwd_a3_tuned preset for fav3_sage forward launch params. |
| aiter/ops/triton/attention/fav3_sage.py | Replaces hardcoded arch-specific dicts with config-tree loading + gfx942 fallback. |
| aiter/ops/triton/attention/fav3_sage_attention_mxfp4_wrapper.py | Replaces hardcoded launch dict with config-tree loading for gfx950. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
"Tuning values live in JSON, never in Python" -- two sage config selectors
still returned hardcoded dicts:
attention/fav3_sage.py three arch branches, one dict each
attention/fav3_sage_attention_mxfp4_wrapper.py one dict behind an arch guard
Both now read <arch>/triton/attention/{fav3_sage,fav3_sage_mxfp4}/DEFAULT.json
through resolve_config_dir() + load_config_json(). The a3_tuned preset moves
into the JSON as a `fwd_a3_tuned` overlay on `fwd`; the `if a3_tuned:` branch
and the mxfp4 wrapper's non-gfx950 guard are control flow and stay in Python.
The JSON carries exactly the numbers the Python returned. Verified by
extracting both versions of the function and comparing their output for every
architecture across the full block_m x block_n x a3_tuned matrix -- equal in
value, key order and value type (PRE_LOAD_V stays a bool, not 0).
Two things a reviewer should weigh in on:
- The old `else:` branch was byte-identical to gfx942, so arches with no tuned
table of their own read the gfx942 file through an explicit
`arch=_FALLBACK_ARCH` retry. That preserves today's behaviour exactly, but it
is a second instance of the arch-fallback pattern that so far only MHC uses.
Tuning those arches and dropping the fallback is the better end state.
- Both functions copy before returning. load_config_json hands back the shared
cached dict and get_sage_fwd_configs mutates what it returns, so without the
copy an a3_tuned=True call would poison the cache for every later caller.
Boss2002n
force-pushed
the
satya/triton-sage-config-to-json
branch
from
August 31, 2026 06:05
60081ea to
9c09646
Compare
Removed comments regarding import behavior and tracing paths.
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 5 out of 5 changed files in this pull request and generated no new comments.
Suppressed comments (1)
aiter/ops/triton/attention/fav3_sage.py:44
- os.scandir() returns an iterator that holds an open directory handle; relying on implicit closure after sorted() can leave the fd open longer than intended (especially on non-CPython runtimes). Use a context manager to deterministically close the scandir handle once the entries are consumed.
for entry in sorted(os.scandir(AITER_TRITON_CONFIGS_PATH), key=lambda e: e.name):
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.