Skip to content

[Triton/Gluon] Avoid the peeled sparse MLA spill on 64-bit gathers below triton 3.8 #6042

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
azaidy merged 3 commits into from
Oct 13, 2026
Merged

[Triton/Gluon] Avoid the peeled sparse MLA spill on 64-bit gathers below triton 3.8 #6042

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
042a289
Avoid the peeled sparse MLA spill on 64-bit gathers below triton 3.8
1am9trash Oct 1, 2026
1cbfb85
Fix version comparison for triton version check
1am9trash Oct 1, 2026
8672515
Merge branch 'main' into v41-decode-triton-3.7
1am9trash Oct 13, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 11 additions & 3 deletions aiter/ops/triton/attention/pa_decode_sparse.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@

import torch
import triton
from packaging.version import Version

from aiter.ops.triton._gluon_kernels.gfx950.attention.sparse_mla import (
_sparse_mla as _sparse_mla_gfx950,
Expand Down Expand Up @@ -46,6 +47,9 @@
except ImportError:
_HAS_SCALED_UPCAST = False

# Below Triton 3.8 the peeled 64-bit split-K kernel spills ~110 VGPRs to scratch.
_TRITON_GE_38 = Version(triton.__version__) >= Version("3.8.0")

_LOGGER = AiterTritonLogger()


Expand Down Expand Up @@ -788,10 +792,14 @@ def _pa_decode_sparse_gfx950_gluon(
if one_wg_per_cu:
waves_per_eu = 1

# Unpeeled is faster at prefill and, on the 64-bit gathers, unless split-K
# leaves each program a few tiles. Decode on buffer loads is faster peeled.
# Decode on buffer loads is faster peeled; prefill is faster unpeeled.
# On the 64-bit gathers unpeeled is faster unless split-K leaves each program
# a few tiles, and below Triton 3.8 the peeled kernel spills, so always unpeel.
short_splits = num_splits > 1 and row_tiles <= 4 * num_splits
unpeel = num_queries >= _PREFILL_MIN_ROWS if use_buffer_load else not short_splits
if use_buffer_load:
unpeel = num_queries >= _PREFILL_MIN_ROWS
else:
unpeel = not (_TRITON_GE_38 and short_splits)

main_splits = num_splits
if has_extra and avg_main > 0:
Expand Down