Skip to content

[Triton/Gluon] Move attention configs to nested layout and unify their resolution #5019

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
brunomazzottiamd merged 2 commits into from
Aug 27, 2026
Merged

[Triton/Gluon] Move attention configs to nested layout and unify their resolution #5019

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
a8b8b08
[Triton] Move attention configs to nested layout and unify their reso…
Boss2002n Aug 26, 2026
98e0404
Merge branch 'main' into satya/migrate-configs-attention
Boss2002n Aug 26, 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
10 changes: 4 additions & 6 deletions aiter/ops/triton/_triton_kernels/attention/extend_attention.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,10 +24,10 @@
import triton.language as tl

from aiter.ops.triton._triton_kernels.activation import _tanh
from aiter.ops.triton.utils._triton import arch_info
from aiter.ops.triton.utils._triton.kernel_repr import make_kernel_repr
from aiter.ops.triton.utils._triton.pid_preprocessing import remap_xcd
from aiter.ops.triton.utils.core import AITER_TRITON_CONFIGS_PATH, load_config_json
from aiter.ops.triton.utils.core import load_config_json
from aiter.ops.triton.utils.gemm_config_utils import resolve_config_dir
Comment thread
Boss2002n marked this conversation as resolved.
Resolved

_fwd_kernel_extend_repr = make_kernel_repr(
"_fwd_kernel",
Expand Down Expand Up @@ -320,10 +320,8 @@ def _fwd_kernel(

@functools.lru_cache(maxsize=1024)
def _get_config(HEAD_SIZE, dtype):
dev = arch_info.get_arch()
config = load_config_json(
f"{AITER_TRITON_CONFIGS_PATH}/{dev}-EXTEND_ATTENTION.json"
)
cfg_dir, _ = resolve_config_dir("attention", "EXTEND_ATTENTION", backend="triton")
config = load_config_json(f"{cfg_dir}/DEFAULT.json")

# HEAD_SIZE 192 = 128 head and 64 pe head dim
if (HEAD_SIZE > 192) or dtype == torch.float32:
Expand Down
16 changes: 6 additions & 10 deletions aiter/ops/triton/_triton_kernels/attention/hstu_attention.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,9 +21,9 @@
# @manual=//triton:triton
import triton.language as tl

from aiter.ops.triton.utils._triton import arch_info
from aiter.ops.triton.utils._triton.kernel_repr import make_kernel_repr
from aiter.ops.triton.utils.core import AITER_TRITON_CONFIGS_PATH, load_config_json
from aiter.ops.triton.utils.core import load_config_json
from aiter.ops.triton.utils.gemm_config_utils import resolve_config_dir
Comment thread
brunomazzottiamd marked this conversation as resolved.
Resolved

try:
from triton.language.extra.libdevice import (
Expand Down Expand Up @@ -860,10 +860,8 @@ def _hstu_attn_bwd(
def _get_fwd_config(
AUTOTUNE_Z: int,
):
dev = arch_info.get_arch()
config = load_config_json(
f"{AITER_TRITON_CONFIGS_PATH}/hstu_attn/{dev}-HSTU_ATTN_FWD.json",
)
cfg_dir, _ = resolve_config_dir("attention", "HSTU_ATTN_FWD", backend="triton")
config = load_config_json(f"{cfg_dir}/DEFAULT.json")

if AUTOTUNE_Z < 512:
batch_key = "small_batch"
Expand All @@ -879,10 +877,8 @@ def _get_fwd_config(
def _get_bwd_config(
AUTOTUNE_Z: int,
):
dev = arch_info.get_arch()
config = load_config_json(
f"{AITER_TRITON_CONFIGS_PATH}/hstu_attn/{dev}-HSTU_ATTN_BWD.json",
)
cfg_dir, _ = resolve_config_dir("attention", "HSTU_ATTN_BWD", backend="triton")
config = load_config_json(f"{cfg_dir}/DEFAULT.json")

if AUTOTUNE_Z < 512:
batch_key = "small_batch"
Expand Down
8 changes: 4 additions & 4 deletions aiter/ops/triton/_triton_kernels/attention/mha.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,14 +7,14 @@
import triton
import triton.language as tl

from aiter.ops.triton.utils._triton import arch_info
from aiter.ops.triton.utils._triton.kernel_repr import make_kernel_repr
from aiter.ops.triton.utils._triton.mha_kernel_utils import _compute_fp8_scaling_factors
from aiter.ops.triton.utils._triton.pid_preprocessing import (
remap_workgroup_spatial,
remap_xcd,
)
from aiter.ops.triton.utils.core import AITER_TRITON_CONFIGS_PATH, load_config_json
from aiter.ops.triton.utils.core import load_config_json
from aiter.ops.triton.utils.gemm_config_utils import resolve_config_dir
Comment thread
brunomazzottiamd marked this conversation as resolved.
Resolved


@triton.jit
Expand Down Expand Up @@ -952,8 +952,8 @@ def _get_config(
has_pe: bool = False,
head_dim_v: int | None = None,
):
dev = arch_info.get_arch()
config = load_config_json(f"{AITER_TRITON_CONFIGS_PATH}/{dev}-MHA-DEFAULT.json")
cfg_dir, _ = resolve_config_dir("attention", "MHA", backend="triton")
config = load_config_json(f"{cfg_dir}/DEFAULT.json")
fwd_cfg = config["fwd"]
has_dropout_or_fp32 = enable_dropout or dtype == torch.float32
# TODO: pe + dropout is not tuned
Expand Down
8 changes: 4 additions & 4 deletions aiter/ops/triton/_triton_kernels/attention/mha_fused_bwd.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,11 +6,11 @@
import triton
import triton.language as tl

from aiter.ops.triton.utils._triton import arch_info
from aiter.ops.triton.utils._triton.kernel_repr import make_kernel_repr
from aiter.ops.triton.utils._triton.mha_kernel_utils import _compute_fp8_scaling_factors
from aiter.ops.triton.utils._triton.pid_preprocessing import remap_xcd
from aiter.ops.triton.utils.core import AITER_TRITON_CONFIGS_PATH, load_config_json
from aiter.ops.triton.utils.core import load_config_json
from aiter.ops.triton.utils.gemm_config_utils import resolve_config_dir
Comment thread
brunomazzottiamd marked this conversation as resolved.
Resolved

# This function computes delta given output Out and gradient DO
# Here is the I/O shape:
Expand Down Expand Up @@ -1061,6 +1061,6 @@ def _bwd_kernel_dkdvdq_noncausal(

@functools.lru_cache(maxsize=1024)
def _get_config():
dev = arch_info.get_arch()
config = load_config_json(f"{AITER_TRITON_CONFIGS_PATH}/{dev}-MHA-DEFAULT.json")
cfg_dir, _ = resolve_config_dir("attention", "MHA", backend="triton")
config = load_config_json(f"{cfg_dir}/DEFAULT.json")
return config["bkwd_fused"]
8 changes: 4 additions & 4 deletions aiter/ops/triton/_triton_kernels/attention/mha_onekernel_bwd.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,10 +6,10 @@
import triton # type: ignore
import triton.language as tl # type: ignore

from aiter.ops.triton.utils._triton import arch_info
from aiter.ops.triton.utils._triton.kernel_repr import make_kernel_repr
from aiter.ops.triton.utils._triton.mha_kernel_utils import _compute_fp8_scaling_factors
from aiter.ops.triton.utils.core import AITER_TRITON_CONFIGS_PATH, load_config_json
from aiter.ops.triton.utils.core import load_config_json
from aiter.ops.triton.utils.gemm_config_utils import resolve_config_dir
Comment thread
brunomazzottiamd marked this conversation as resolved.
Resolved

# NOTE: triton fails to import tl.constexprs so create them here for the file
DROPOUT_USE_PYTORCH = False
Expand Down Expand Up @@ -1768,6 +1768,6 @@ def bwd_kernel_noncausal(

@functools.lru_cache(maxsize=1024)
def _get_config():
dev = arch_info.get_arch()
config = load_config_json(f"{AITER_TRITON_CONFIGS_PATH}/{dev}-MHA-DEFAULT.json")
cfg_dir, _ = resolve_config_dir("attention", "MHA", backend="triton")
config = load_config_json(f"{cfg_dir}/DEFAULT.json")
return config["bkwd_onekernel"]
10 changes: 4 additions & 6 deletions aiter/ops/triton/_triton_kernels/attention/mla_decode_rope.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,10 +29,10 @@
import triton.language as tl

from aiter.ops.triton._triton_kernels.activation import _tanh
from aiter.ops.triton.utils._triton import arch_info
from aiter.ops.triton.utils._triton.kernel_repr import make_kernel_repr
from aiter.ops.triton.utils._triton.pid_preprocessing import remap_xcd
from aiter.ops.triton.utils.core import AITER_TRITON_CONFIGS_PATH, load_config_json
from aiter.ops.triton.utils.core import load_config_json
from aiter.ops.triton.utils.gemm_config_utils import resolve_config_dir
Comment thread
brunomazzottiamd marked this conversation as resolved.
Resolved

_fwd_grouped_kernel_stage1_rope_repr = make_kernel_repr(
"_fwd_grouped_kernel_stage1_rope",
Expand Down Expand Up @@ -403,7 +403,5 @@ def _fwd_kernel_stage2(

@functools.lru_cache(maxsize=1024)
def _get_config():
dev = arch_info.get_arch()
return load_config_json(
f"{AITER_TRITON_CONFIGS_PATH}/{dev}-MLA_DECODE_ROPE-DEFAULT.json",
)
cfg_dir, _ = resolve_config_dir("attention", "MLA_DECODE_ROPE", backend="triton")
Comment thread
Boss2002n marked this conversation as resolved.
Resolved
return load_config_json(f"{cfg_dir}/DEFAULT.json")
0 ...s/triton/configs/gfx1151-MHA-DEFAULT.json → ...gfx1151/triton/attention/mha/DEFAULT.json
View file
Open in desktop
File renamed without changes.
0 ...ton/configs/gfx1250-EXTEND_ATTENTION.json → ...n/attention/extend_attention/DEFAULT.json
View file
Open in desktop
File renamed without changes.
0 ...s/triton/configs/gfx1250-MHA-DEFAULT.json → ...gfx1250/triton/attention/mha/DEFAULT.json
View file
Open in desktop
File renamed without changes.
0 ...figs/gfx1250-MLA_DECODE_ROPE-DEFAULT.json → ...on/attention/mla_decode_rope/DEFAULT.json
View file
Open in desktop
File renamed without changes.
0 ...iton/configs/gfx942-EXTEND_ATTENTION.json → ...n/attention/extend_attention/DEFAULT.json
View file
Open in desktop
File renamed without changes.
0 ...nfigs/hstu_attn/gfx942-HSTU_ATTN_BWD.json → ...iton/attention/hstu_attn_bwd/DEFAULT.json
View file
Open in desktop
File renamed without changes.
0 ...nfigs/hstu_attn/gfx942-HSTU_ATTN_FWD.json → ...iton/attention/hstu_attn_fwd/DEFAULT.json
View file
Open in desktop
File renamed without changes.
0 ...ps/triton/configs/gfx942-MHA-DEFAULT.json → .../gfx942/triton/attention/mha/DEFAULT.json
View file
Open in desktop
File renamed without changes.
0 ...nfigs/gfx942-MLA_DECODE_ROPE-DEFAULT.json → ...on/attention/mla_decode_rope/DEFAULT.json
View file
Open in desktop
File renamed without changes.
0 ...iton/configs/gfx950-EXTEND_ATTENTION.json → ...n/attention/extend_attention/DEFAULT.json
View file
Open in desktop
File renamed without changes.
0 ...nfigs/hstu_attn/gfx950-HSTU_ATTN_BWD.json → ...iton/attention/hstu_attn_bwd/DEFAULT.json
View file
Open in desktop
File renamed without changes.
0 ...nfigs/hstu_attn/gfx950-HSTU_ATTN_FWD.json → ...iton/attention/hstu_attn_fwd/DEFAULT.json
View file
Open in desktop
File renamed without changes.
0 ...ps/triton/configs/gfx950-MHA-DEFAULT.json → .../gfx950/triton/attention/mha/DEFAULT.json
View file
Open in desktop
File renamed without changes.
0 ...nfigs/gfx950-MLA_DECODE_ROPE-DEFAULT.json → ...on/attention/mla_decode_rope/DEFAULT.json
View file
Open in desktop
File renamed without changes.
Loading