Skip to content

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

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

Open
Boss2002n wants to merge 1 commit into
base: main
Choose a base branch
from
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
19 changes: 18 additions & 1 deletion aiter/ops/triton/_gluon_kernels/gfx1250/attention/fp8_mqa_logits.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
from aiter.ops.triton._gluon_kernels.gfx950.attention.fp8_mqa_logits import (
_weighted_sum_fma_fold,
)
from aiter.ops.triton.utils._triton.kernel_repr import make_kernel_repr
from aiter.ops.triton.utils.common_utils import strip_annotate

_MAX_PROPAGATE_NAN_ALL = gl.constexpr(PropagateNan.ALL)
Expand Down Expand Up @@ -567,7 +568,23 @@ def mqa_logits_loop_pipelined(
)


@gluon.jit
_gluon_fp8_mqa_logits_kernel_repr = make_kernel_repr(
"_gluon_fp8_mqa_logits_kernel",
[
"NUM_HEADS",
"HEAD_SIZE",
"BLOCK_KV",
"NUM_WARPS",
"NUM_BUFFERS",
"NUM_CHAINS",
"LOOP_VARIANT",
"USE_BUFFER_LOAD",
"USE_BUFFER_STORE",
],
)


@gluon.jit(repr=_gluon_fp8_mqa_logits_kernel_repr)
def _gluon_fp8_mqa_logits_kernel(
Q_ptr, # fp8e4m3 [seq_len, NUM_HEADS, HEAD_SIZE]
KV_ptr, # fp8e4m3 [seq_len_kv, HEAD_SIZE]
Expand Down
21 changes: 20 additions & 1 deletion aiter/ops/triton/_gluon_kernels/gfx950/attention/fp8_mqa_logits.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
from triton.language.core import PropagateNan
from triton.language.core import _aggregate as aggregate

from aiter.ops.triton.utils._triton.kernel_repr import make_kernel_repr
from aiter.ops.triton.utils.common_utils import strip_annotate

_MAX_PROPAGATE_NAN_ALL = gl.constexpr(PropagateNan.ALL)
Expand Down Expand Up @@ -754,7 +755,25 @@ def mqa_logits_loop_double_buf(
)


@gluon.jit
_gluon_fp8_mqa_logits_kernel_repr = make_kernel_repr(
"_gluon_fp8_mqa_logits_kernel",
[
"NUM_HEADS",
"HEAD_SIZE",
"BLOCK_KV",
"NUM_WARPS",
"NUM_BUFFERS",
"NUM_CHAINS",
"USE_BUFFER_LOAD",
"USE_BUFFER_STORE",
"USE_PADDED_SHARED_LAYOUT",
"BLOCK_M",
"MFMA_NONK_DIM",
],
)


@gluon.jit(repr=_gluon_fp8_mqa_logits_kernel_repr)
def _gluon_fp8_mqa_logits_kernel(
Q_ptr, # fp8e4m3 [seq_len, NUM_HEADS, HEAD_SIZE]
KV_ptr, # fp8e4m3 [seq_len_kv, HEAD_SIZE]
Expand Down
11 changes: 10 additions & 1 deletion aiter/ops/triton/_triton_kernels/attention/block_lut.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,8 +10,17 @@
import triton
import triton.language as tl

from aiter.ops.triton.utils._triton.kernel_repr import make_kernel_repr

@triton.jit()
_block_attn_mask_to_lut_kernel_repr = make_kernel_repr(
"_block_attn_mask_to_lut_kernel",
[
"BLOCK_KB",
],
)


@triton.jit(repr=_block_attn_mask_to_lut_kernel_repr)
def _block_attn_mask_to_lut_kernel(
mask_ptr,
lut_start_ptr,
Expand Down
32 changes: 31 additions & 1 deletion aiter/ops/triton/_triton_kernels/attention/fav3_sage_attention.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
from aiter.ops.triton._triton_kernels.flash_attn_triton_amd.common import (
compute_alibi_block,
)
from aiter.ops.triton.utils._triton.kernel_repr import make_kernel_repr


def map_dims(shape, indices):
Expand Down Expand Up @@ -1219,7 +1220,36 @@ def compute_block_masking(
)


@triton.jit
_sage_fwd_repr = make_kernel_repr(
"sage_fwd",
[
"RETURN_LSE",
"HQ",
"HK",
"ACTUAL_BLOCK_DMODEL_QK",
"ACTUAL_BLOCK_DMODEL_V",
"IS_VARLEN",
"IS_CAUSAL",
"USE_SLIDING_WINDOW",
"WINDOW_SIZE_LEFT",
"WINDOW_SIZE_RIGHT",
"BLOCK_M",
"BLOCK_DMODEL_QK",
"BLOCK_DMODEL_V",
"BLOCK_N",
"PRE_LOAD_V",
"USE_BIAS",
"ENABLE_DROPOUT",
"RETURN_SCORES",
"USE_ALIBI",
"USE_EXP2",
"USE_SEQUSED",
"USE_BLOCK_SPARSE",
],
)


@triton.jit(repr=_sage_fwd_repr)
def sage_fwd(
Q,
K,
Expand Down
27 changes: 26 additions & 1 deletion aiter/ops/triton/_triton_kernels/attention/fav3_sage_attention_mxfp4.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
import triton
import triton.language as tl

from aiter.ops.triton.utils._triton.kernel_repr import make_kernel_repr


@triton.jit
def compute_padding_info(seqlen_k, BLOCK_N: tl.constexpr):
Expand Down Expand Up @@ -564,7 +566,30 @@ def _sage_fwd_blocksparse_mask_mxfp4(
return acc, l_i, m_i


@triton.jit
_sage_fwd_mxfp4_repr = make_kernel_repr(
"sage_fwd_mxfp4",
[
"Q_DTYPE_STR",
"K_DTYPE_STR",
"HQ",
"HK",
"ACTUAL_BLOCK_DMODEL_QK",
"ACTUAL_BLOCK_DMODEL_V",
"IS_VARLEN",
"IS_CAUSAL",
"BLOCK_M",
"BLOCK_DMODEL_QK",
"BLOCK_DMODEL_V",
"BLOCK_N",
"PRE_LOAD_V",
"USE_BIAS",
"USE_BLOCK_SPARSE",
"RETURN_LSE",
],
)


@triton.jit(repr=_sage_fwd_mxfp4_repr)
def sage_fwd_mxfp4(
Q,
K,
Expand Down
13 changes: 12 additions & 1 deletion aiter/ops/triton/_triton_kernels/attention/fp8_mqa_logits.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,8 +1,19 @@
import triton
import triton.language as tl

from aiter.ops.triton.utils._triton.kernel_repr import make_kernel_repr

@triton.jit
_fp8_mqa_logits_kernel_repr = make_kernel_repr(
"_fp8_mqa_logits_kernel",
[
"NUM_HEADS",
"HEAD_SIZE",
"BLOCK_KV",
],
)
Comment thread
Copilot marked this conversation as resolved.


@triton.jit(repr=_fp8_mqa_logits_kernel_repr)
def _fp8_mqa_logits_kernel(
Q_ptr, # fp8e4m3 [seq_len, H, D]
KV_ptr, # fp8e4m3 [seq_len_kv, D]
Expand Down
20 changes: 19 additions & 1 deletion aiter/ops/triton/_triton_kernels/attention/lean_atten_paged.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,8 @@
import triton
import triton.language as tl

from aiter.ops.triton.utils._triton.kernel_repr import make_kernel_repr


@triton.jit
def find_group(x):
Expand All @@ -29,7 +31,23 @@ def find_group(x):
return group_id, group_size, total_blocks


@triton.jit
_la_persistent_paged_repr = make_kernel_repr(
"la_persistent_paged",
[
"HEAD_DIM",
"BLOCK_M",
"BLOCK_N",
"batch_size",
"num_m_blocks",
"high_load_wgs",
"max_tiles_per_wg",
"tiles_per_head",
"num_splits",
],
)
Comment thread
Copilot marked this conversation as resolved.


@triton.jit(repr=_la_persistent_paged_repr)
def la_persistent_paged(
Q,
K,
Expand Down
68 changes: 63 additions & 5 deletions aiter/ops/triton/_triton_kernels/attention/pa_mqa_logits.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,13 +4,26 @@
import triton
import triton.language as tl

from aiter.ops.triton.utils._triton.kernel_repr import make_kernel_repr


@triton.jit
def _sum_combine(a, b):
return a + b


@triton.jit
_deepgemm_fp8_paged_mqa_logits_stage1_ragged_k_repr = make_kernel_repr(
"_deepgemm_fp8_paged_mqa_logits_stage1_ragged_k",
[
"ChunkQ",
"ChunkK",
"HiddenDim",
"SplitKV",
],
)


@triton.jit(repr=_deepgemm_fp8_paged_mqa_logits_stage1_ragged_k_repr)
def _deepgemm_fp8_paged_mqa_logits_stage1_ragged_k(
batch_size,
next_n,
Expand Down Expand Up @@ -106,7 +119,18 @@ def _deepgemm_fp8_paged_mqa_logits_stage1_ragged_k(
)


@triton.jit
_deepgemm_fp8_paged_mqa_logits_ragged_k_repr = make_kernel_repr(
"_deepgemm_fp8_paged_mqa_logits_ragged_k",
[
"ChunkQ",
"ChunkK",
"HiddenDim",
"SplitKV",
],
)


@triton.jit(repr=_deepgemm_fp8_paged_mqa_logits_ragged_k_repr)
def _deepgemm_fp8_paged_mqa_logits_ragged_k(
batch_size,
next_n,
Expand Down Expand Up @@ -201,7 +225,19 @@ def _deepgemm_fp8_paged_mqa_logits_ragged_k(
)


@triton.jit
_deepgemm_fp8_paged_mqa_logits_stage1_repr = make_kernel_repr(
"_deepgemm_fp8_paged_mqa_logits_stage1",
[
"ChunkQ",
"ChunkK",
"HiddenDim",
"KVBlockSize",
"SplitKV",
],
)
Comment thread
Copilot marked this conversation as resolved.


@triton.jit(repr=_deepgemm_fp8_paged_mqa_logits_stage1_repr)
def _deepgemm_fp8_paged_mqa_logits_stage1(
batch_size,
next_n,
Expand Down Expand Up @@ -311,7 +347,17 @@ def _deepgemm_fp8_paged_mqa_logits_stage1(
)


@triton.jit
_deepgemm_fp8_paged_mqa_logits_varctx_schedule_repr = make_kernel_repr(
"_deepgemm_fp8_paged_mqa_logits_varctx_schedule",
[
"ChunkK",
"AlignedBatchSize",
"TryCount",
],
)


@triton.jit(repr=_deepgemm_fp8_paged_mqa_logits_varctx_schedule_repr)
def _deepgemm_fp8_paged_mqa_logits_varctx_schedule(
batch_size,
context_len_ptr,
Expand Down Expand Up @@ -357,7 +403,19 @@ def _deepgemm_fp8_paged_mqa_logits_varctx_schedule(
tl.store(safe_chunks_per_cta_ptr, safe_seg_lens)


@triton.jit
_deepgemm_fp8_paged_mqa_logits_repr = make_kernel_repr(
"_deepgemm_fp8_paged_mqa_logits",
[
"ChunkQ",
"ChunkK",
"HiddenDim",
"KVBlockSize",
"SplitKV",
],
)
Comment thread
Copilot marked this conversation as resolved.


@triton.jit(repr=_deepgemm_fp8_paged_mqa_logits_repr)
def _deepgemm_fp8_paged_mqa_logits(
batch_size,
next_n,
Expand Down
15 changes: 14 additions & 1 deletion aiter/ops/triton/_triton_kernels/attention/sparse_attention_dsv4.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,8 @@
import triton
import triton.language as tl

from aiter.ops.triton.utils._triton.kernel_repr import make_kernel_repr

# =====================================================================
# Utility
# =====================================================================
Expand Down Expand Up @@ -261,12 +263,23 @@ def _get_prefill_autotune_configs():
]


_sparse_attn_prefill_kernel_repr = make_kernel_repr(
"_sparse_attn_prefill_kernel",
[
"HAS_ATTN_SINK",
"BLOCK_H",
"BLOCK_D",
"BLOCK_K",
],
)
Comment thread
Copilot marked this conversation as resolved.


@triton.autotune(
configs=_get_prefill_autotune_configs(),
key=["num_heads", "head_dim", "HAS_ATTN_SINK"],
prune_configs_by={"early_config_prune": _prefill_prune_configs},
)
@triton.jit
@triton.jit(repr=_sparse_attn_prefill_kernel_repr)
def _sparse_attn_prefill_kernel(
q_ptr, # [num_queries, num_heads, head_dim]
kv_ptr, # [num_kv, head_dim]
Expand Down
Loading
Loading