[MLA] triton mla ps mi355

[Zzz9990]

Motivation

Technical Details

Test Plan

Test Result

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

[github-actions]

Remaining comments which cannot be posted as a review comment to avoid GitHub Rate Limit

ruff

---

⚠️ [ruff] <F841> _{reported by reviewdog 🐶}
Local variable mask_on is assigned to but never used

aiter/aiter/ops/triton/gluon/mla_decode_fp8.py

Line 208 in 9d3f8b1

mask_on = offs_on < kv_lora_rank

---

⚠️ [ruff] <F401> _{reported by reviewdog 🐶}
typing.Optional imported but unused

aiter/aiter/ops/triton/gluon/mla_decode_mi355.py

Line 25 in 9d3f8b1

from typing import Optional

---

⚠️ [ruff] <F401> _{reported by reviewdog 🐶}
functools imported but unused

aiter/aiter/ops/triton/gluon/mla_decode_mi355.py

Line 26 in 9d3f8b1

import functools

---

⚠️ [ruff] <F401> _{reported by reviewdog 🐶}
json imported but unused

aiter/aiter/ops/triton/gluon/mla_decode_mi355.py

Line 27 in 9d3f8b1

import json

---

⚠️ [ruff] <F401> _{reported by reviewdog 🐶}
triton imported but unused

aiter/aiter/ops/triton/gluon/mla_decode_mi355.py

Line 28 in 9d3f8b1

import triton

---

⚠️ [ruff] <F401> _{reported by reviewdog 🐶}
torch imported but unused

aiter/aiter/ops/triton/gluon/mla_decode_mi355.py

Line 30 in 9d3f8b1

import torch

---

⚠️ [ruff] <F401> _{reported by reviewdog 🐶}
aiter.ops.triton.utils._triton.arch_info imported but unused

aiter/aiter/ops/triton/gluon/mla_decode_mi355.py

Line 31 in 9d3f8b1

import aiter.ops.triton.utils._triton.arch_info as arch_info

---

⚠️ [ruff] <F401> _{reported by reviewdog 🐶}
aiter.ops.triton.utils.core.AITER_TRITON_CONFIGS_PATH imported but unused

aiter/aiter/ops/triton/gluon/mla_decode_mi355.py

Line 32 in 9d3f8b1

from aiter.ops.triton.utils.core import AITER_TRITON_CONFIGS_PATH

---

⚠️ [ruff] <F401> _{reported by reviewdog 🐶}
aiter.dtypes imported but unused

aiter/aiter/ops/triton/gluon/mla_decode_mi355.py

Line 34 in 9d3f8b1

from aiter import dtypes

---

⚠️ [ruff] <F841> _{reported by reviewdog 🐶}
Local variable mask_on is assigned to but never used

aiter/aiter/ops/triton/gluon/mla_decode_mi355.py

Line 207 in 9d3f8b1

mask_on = offs_on < kv_lora_rank

---

⚠️ [ruff] <F841> _{reported by reviewdog 🐶}
Local variable mask_k_c is assigned to but never used

aiter/aiter/ops/triton/gluon/mla_decode_mi355.py

Line 238 in 9d3f8b1

mask_k_c = offs_k_c < kv_lora_rank

---

⚠️ [ruff] <F841> _{reported by reviewdog 🐶}
Local variable mask_qk_r is assigned to but never used

aiter/aiter/ops/triton/gluon/mla_decode_mi355.py

Line 239 in 9d3f8b1

mask_qk_r = offs_qk_r < (kv_lora_rank + qk_rope_head_dim)

---

⚠️ [ruff] <F841> _{reported by reviewdog 🐶}
Local variable mask_k_r is assigned to but never used

aiter/aiter/ops/triton/gluon/mla_decode_mi355.py

Line 240 in 9d3f8b1

mask_k_r = offs_k_r < (kv_lora_rank + qk_rope_head_dim)

---

⚠️ [ruff] <F841> _{reported by reviewdog 🐶}
Local variable shared_v is assigned to but never used

aiter/aiter/ops/triton/gluon/mla_decode_mi355.py

Line 1936 in 9d3f8b1

shared_v: gl.constexpr = gl.SwizzledSharedLayout(

---

⚠️ [ruff] <F841> _{reported by reviewdog 🐶}
Local variable mask_on is assigned to but never used

aiter/aiter/ops/triton/gluon/mla_decode_mi355.py

Line 2014 in 9d3f8b1

mask_on = offs_on < kv_lora_rank

---

⚠️ [ruff] <F841> _{reported by reviewdog 🐶}
Local variable mask_k_pe is assigned to but never used

aiter/aiter/ops/triton/gluon/mla_decode_mi355.py

Line 2439 in 9d3f8b1

mask_k_pe = mask_k_id < cur_batch_seq_len

---

⚠️ [ruff] <F401> _{reported by reviewdog 🐶}
typing.Optional imported but unused

aiter/aiter/ops/triton/gluon/mla_decode_ps.py

Line 13 in 9d3f8b1

from typing import Optional

---

⚠️ [ruff] <F401> _{reported by reviewdog 🐶}
functools imported but unused

aiter/aiter/ops/triton/gluon/mla_decode_ps.py

Line 14 in 9d3f8b1

import functools

---

⚠️ [ruff] <F401> _{reported by reviewdog 🐶}
triton.language imported but unused

aiter/aiter/ops/triton/gluon/mla_decode_ps.py

Line 16 in 9d3f8b1

import triton.language as tl

---

⚠️ [ruff] <F401> _{reported by reviewdog 🐶}
torch imported but unused

aiter/aiter/ops/triton/gluon/mla_decode_ps.py

Line 17 in 9d3f8b1

import torch

---

⚠️ [ruff] <F401> _{reported by reviewdog 🐶}
aiter.ops.triton.utils.arch_info imported but unused

aiter/aiter/ops/triton/gluon/mla_decode_ps.py

Line 19 in 9d3f8b1

import aiter.ops.triton.utils.arch_info as arch_info

---

⚠️ [ruff] <F401> _{reported by reviewdog 🐶}
aiter.ops.triton.utils.core.AITER_TRITON_CONFIGS_PATH imported but unused

aiter/aiter/ops/triton/gluon/mla_decode_ps.py

Line 20 in 9d3f8b1

from aiter.ops.triton.utils.core import AITER_TRITON_CONFIGS_PATH

---

⚠️ [ruff] <F401> _{reported by reviewdog 🐶}
aiter.ops.triton.utils.pid_preprocessing.remap_xcd imported but unused

aiter/aiter/ops/triton/gluon/mla_decode_ps.py

Line 21 in 9d3f8b1

from aiter.ops.triton.utils.pid_preprocessing import remap_xcd

---

⚠️ [ruff] <F841> _{reported by reviewdog 🐶}
Local variable token_to_batch_end is assigned to but never used

aiter/aiter/ops/triton/gluon/mla_decode_ps.py

Line 188 in 9d3f8b1

token_to_batch_end = gl.load(work_info_set + work_id * 8 + 6)

---

⚠️ [ruff] <F841> _{reported by reviewdog 🐶}
Local variable num_q_head_blk is assigned to but never used

aiter/aiter/ops/triton/gluon/mla_decode_ps.py

Line 190 in 9d3f8b1

num_q_head_blk = gl.cdiv(q_head_num, BLOCK_H)

---

⚠️ [ruff] <F841> _{reported by reviewdog 🐶}
Local variable mask_k_c is assigned to but never used

aiter/aiter/ops/triton/gluon/mla_decode_ps.py

Line 245 in 9d3f8b1

mask_k_c = offs_k_c < kv_lora_rank

---

⚠️ [ruff] <F841> _{reported by reviewdog 🐶}
Local variable mask_k_r is assigned to but never used

aiter/aiter/ops/triton/gluon/mla_decode_ps.py

Line 247 in 9d3f8b1

mask_k_r = offs_k_r < (kv_lora_rank + qk_rope_head_dim)

---

⚠️ [ruff] <F841> _{reported by reviewdog 🐶}
Local variable offs_n is assigned to but never used

aiter/aiter/ops/triton/gluon/mla_decode_ps.py

Line 284 in 9d3f8b1

offs_n = start_n + gl.arange(0, BLOCK_N,

---

⚠️ [ruff] <F841> _{reported by reviewdog 🐶}
Local variable offs_n_pe is assigned to but never used

aiter/aiter/ops/triton/gluon/mla_decode_ps.py

Line 287 in 9d3f8b1

offs_n_pe = start_n + gl.arange(0, BLOCK_N,

---

⚠️ [ruff] <F841> _{reported by reviewdog 🐶}
Local variable LAST_SPLIT is assigned to but never used

aiter/aiter/ops/triton/mla_decode.py

Line 118 in 9d3f8b1

LAST_SPLIT = split_kv_end == cur_batch_seq_len

---

⚠️ [ruff] <F841> _{reported by reviewdog 🐶}
Local variable k_pe_last_token is assigned to but never used

aiter/aiter/ops/triton/mla_decode.py

Line 119 in 9d3f8b1

k_pe_last_token = tl.zeros([BLOCK_R], dtype=q.dtype)

---

⚠️ [ruff] <E741> _{reported by reviewdog 🐶}
Ambiguous variable name: O

aiter/aiter/ops/triton/mla_decode.py

Line 327 in 9d3f8b1

O,

---

⚠️ [ruff] <F401> _{reported by reviewdog 🐶}
triton.experimental.gluon imported but unused

aiter/aiter/ops/triton/mla_decode_dispatch.py

Line 38 in 9d3f8b1

from triton.experimental import gluon

---

⚠️ [ruff] <E402> _{reported by reviewdog 🐶}
Module level import not at top of file

aiter/aiter/ops/triton/mla_decode_dispatch.py

Lines 72 to 74 in 9d3f8b1

	from aiter.utility.triton.triton_metadata_redirect import (
	AOTMetadataContext,
	)

---

⚠️ [ruff] <F601> _{reported by reviewdog 🐶}
Dictionary key literal (1,) repeated

aiter/aiter/ops/triton/mla_decode_dispatch.py

Line 183 in 9d3f8b1

(1,): [["tt.divisibility", 16]], # kv_indptr

---

⚠️ [ruff] <F601> _{reported by reviewdog 🐶}
Dictionary key literal (2,) repeated

aiter/aiter/ops/triton/mla_decode_dispatch.py

Line 184 in 9d3f8b1

(2,): [["tt.divisibility", 16]], # kv_indices

---

⚠️ [ruff] <F601> _{reported by reviewdog 🐶}
Dictionary key literal (1,) repeated

aiter/aiter/ops/triton/mla_decode_dispatch.py

Line 185 in 9d3f8b1

(1,): [["tt.divisibility", 16]], # Att_Out

---

⚠️ [ruff] <F601> _{reported by reviewdog 🐶}
Dictionary key literal (2,) repeated

aiter/aiter/ops/triton/mla_decode_dispatch.py

Line 186 in 9d3f8b1

(2,): [["tt.divisibility", 16]], # Att_Lse

---

⚠️ [ruff] <E741> _{reported by reviewdog 🐶}
Ambiguous variable name: O

aiter/aiter/ops/triton/mla_decode_dispatch.py

Line 396 in 9d3f8b1

O,

---

⚠️ [ruff] <F401> _{reported by reviewdog 🐶}
typing.Optional imported but unused

aiter/aiter/ops/triton/mla_decode_ps.py

Line 13 in 9d3f8b1

from typing import Optional

---

⚠️ [ruff] <F401> _{reported by reviewdog 🐶}
functools imported but unused

aiter/aiter/ops/triton/mla_decode_ps.py

Line 14 in 9d3f8b1

import functools

---

⚠️ [ruff] <F401> _{reported by reviewdog 🐶}
torch imported but unused

aiter/aiter/ops/triton/mla_decode_ps.py

Line 17 in 9d3f8b1

import torch

---

⚠️ [ruff] <F401> _{reported by reviewdog 🐶}
aiter.ops.triton.utils.arch_info imported but unused

aiter/aiter/ops/triton/mla_decode_ps.py

Line 19 in 9d3f8b1

import aiter.ops.triton.utils.arch_info as arch_info

---

⚠️ [ruff] <F401> _{reported by reviewdog 🐶}
aiter.ops.triton.utils.core.AITER_TRITON_CONFIGS_PATH imported but unused

aiter/aiter/ops/triton/mla_decode_ps.py

Line 20 in 9d3f8b1

from aiter.ops.triton.utils.core import AITER_TRITON_CONFIGS_PATH

---

⚠️ [ruff] <F401> _{reported by reviewdog 🐶}
aiter.ops.triton.utils.pid_preprocessing.remap_xcd imported but unused

aiter/aiter/ops/triton/mla_decode_ps.py

Line 21 in 9d3f8b1

from aiter.ops.triton.utils.pid_preprocessing import remap_xcd

---

⚠️ [ruff] <F841> _{reported by reviewdog 🐶}
Local variable num_q_head_blk is assigned to but never used

aiter/aiter/ops/triton/mla_decode_ps.py

Line 87 in 9d3f8b1

num_q_head_blk = tl.cdiv(q_head_num, BLOCK_H)

---

⚠️ [ruff] <F401> _{reported by reviewdog 🐶}
aiter imported but unused

aiter/op_tests/op_benchmarks/triton/bench_mla_decode_without_rope.py

Line 13 in 9d3f8b1

import aiter

---

⚠️ [ruff] <F841> _{reported by reviewdog 🐶}
Local variable HK is assigned to but never used

aiter/op_tests/op_benchmarks/triton/bench_mla_decode_without_rope.py

Line 113 in 9d3f8b1

HK = 1

---

⚠️ [ruff] <E741> _{reported by reviewdog 🐶}
Ambiguous variable name: l

aiter/op_tests/op_benchmarks/triton/bench_mla_decode_without_rope.py

Line 201 in 9d3f8b1

l = attn_weights_exp.sum(-1)

---

⚠️ [ruff] <F841> _{reported by reviewdog 🐶}
Local variable dtype is assigned to but never used

aiter/op_tests/op_benchmarks/triton/bench_mla_decode_without_rope.py

Line 272 in 9d3f8b1

dtype = args.dtype

---

⚠️ [ruff] <F841> _{reported by reviewdog 🐶}
Local variable max_seqlen_qo is assigned to but never used

aiter/op_tests/op_benchmarks/triton/bench_mla_decode_without_rope.py

Line 439 in 9d3f8b1

max_seqlen_qo = seq_lens_qo.max().item()

---

⚠️ [ruff] <F841> _{reported by reviewdog 🐶}
Local variable new_k_pe_elems is assigned to but never used

aiter/op_tests/op_benchmarks/triton/bench_mla_decode_without_rope.py

Line 496 in 9d3f8b1

new_k_pe_elems = qk_rope_head_dim # to add to kv cache

---

⚠️ [ruff] <F541> _{reported by reviewdog 🐶}
f-string without any placeholders

aiter/op_tests/op_benchmarks/triton/bench_mla_decode_without_rope.py

Line 574 in 9d3f8b1

os.system(f"zip -r mla_aot_kernel mla")

---

⚠️ [ruff] <F821> _{reported by reviewdog 🐶}
Undefined name get_caller_name_no_ext

aiter/op_tests/op_benchmarks/triton/bench_mla_decode_without_rope.py

Line 656 in 9d3f8b1

print_vgpr(fun, get_caller_name_no_ext())

---

⚠️ [ruff] <F841> _{reported by reviewdog 🐶}
Local variable us_aiter is assigned to but never used

aiter/op_tests/test_mla.py

Line 219 in 9d3f8b1

us_aiter = None

---

⚠️ [ruff] <F841> _{reported by reviewdog 🐶}
Local variable S is assigned to but never used

aiter/op_tests/triton_tests/test_mla_decode.py

Line 103 in 9d3f8b1

S = kv_indptr[1].item()

---

⚠️ [ruff] <F841> _{reported by reviewdog 🐶}
Local variable S is assigned to but never used

aiter/op_tests/triton_tests/test_mla_decode.py

Line 181 in 9d3f8b1

S = kv_indptr[1].item()

---

⚠️ [ruff] <F401> _{reported by reviewdog 🐶}
functools imported but unused

aiter/test_ascyn.py

Line 2 in 9d3f8b1

import functools

---

⚠️ [ruff] <F401> _{reported by reviewdog 🐶}
json imported but unused

aiter/test_ascyn.py

Line 3 in 9d3f8b1

import json

---

⚠️ [ruff] <F401> _{reported by reviewdog 🐶}
triton imported but unused

aiter/test_ascyn.py

Line 4 in 9d3f8b1

import triton

---

⚠️ [ruff] <F401> _{reported by reviewdog 🐶}
triton.language imported but unused

aiter/test_ascyn.py

Line 5 in 9d3f8b1

import triton.language as tl

---

⚠️ [ruff] <F401> _{reported by reviewdog 🐶}
aiter.ops.triton.utils._triton.arch_info imported but unused

aiter/test_ascyn.py

Line 7 in 9d3f8b1

import aiter.ops.triton.utils._triton.arch_info as arch_info

---

⚠️ [ruff] <F401> _{reported by reviewdog 🐶}
aiter.ops.triton.utils.core.AITER_TRITON_CONFIGS_PATH imported but unused

aiter/test_ascyn.py

Line 8 in 9d3f8b1

from aiter.ops.triton.utils.core import AITER_TRITON_CONFIGS_PATH

---

⚠️ [ruff] <F401> _{reported by reviewdog 🐶}
aiter.ops.triton.utils._triton.pid_preprocessing.remap_xcd imported but unused

aiter/test_ascyn.py

Line 9 in 9d3f8b1

from aiter.ops.triton.utils._triton.pid_preprocessing import remap_xcd

---

⚠️ [ruff] <F401> _{reported by reviewdog 🐶}
aiter.dtypes imported but unused

aiter/test_ascyn.py

Line 10 in 9d3f8b1

from aiter import dtypes

---

⚠️ [ruff] <F401> _{reported by reviewdog 🐶}
triton.experimental.gluon.language.amd.cdna4.async_copy imported but unused

aiter/test_ascyn.py

Line 15 in 9d3f8b1

from triton.experimental.gluon.language.amd.cdna4 import async_copy as cdna4_async_copy

---

⚠️ [ruff] <F841> _{reported by reviewdog 🐶}
Local variable pid is assigned to but never used

aiter/test_ascyn.py

Line 110 in 9d3f8b1

pid = gl.program_id(0)

---

⚠️ [ruff] <F841> _{reported by reviewdog 🐶}
Local variable blocked_b is assigned to but never used

aiter/test_ascyn.py

Line 118 in 9d3f8b1

blocked_b: gl.constexpr = gl.DistributedLinearLayout(

---

⚠️ [ruff] <F841> _{reported by reviewdog 🐶}
Local variable pid is assigned to but never used

aiter/test_ascyn.py

Line 342 in 9d3f8b1

pid = gl.program_id(0)

---

⚠️ [ruff] <F841> _{reported by reviewdog 🐶}
Local variable blocked_b is assigned to but never used

aiter/test_ascyn.py

Line 351 in 9d3f8b1

blocked_b: gl.constexpr = gl.DistributedLinearLayout(

---

⚠️ [ruff] <F841> _{reported by reviewdog 🐶}
Local variable blocked_a_mk is assigned to but never used

aiter/test_ascyn.py

Line 371 in 9d3f8b1

blocked_a_mk: gl.constexpr = gl.BlockedLayout([4, 8], [2, 32], [4, 1], [1, 0])

---

⚠️ [ruff] <F821> _{reported by reviewdog 🐶}
Undefined name blocked_ld_in

aiter/test_ascyn.py

Line 396 in 9d3f8b1

0, BLOCK_H, layout=gl.SliceLayout(1, blocked_ld_in)

---

⚠️ [ruff] <F821> _{reported by reviewdog 🐶}
Undefined name blocked_ld_in

aiter/test_ascyn.py

Line 400 in 9d3f8b1

0, BLOCK_C, layout=gl.SliceLayout(0, blocked_ld_in)

---

⚠️ [ruff] <F841> _{reported by reviewdog 🐶}
Local variable shared_b is assigned to but never used

aiter/test_ascyn.py

Line 410 in 9d3f8b1

shared_b: gl.constexpr = gl.PaddedSharedLayout(

---

⚠️ [ruff] <F821> _{reported by reviewdog 🐶}
Undefined name shared_in

aiter/test_ascyn.py

Line 418 in 9d3f8b1

K_Buffer.type.element_ty, [BLOCK_H * 4, BLOCK_C], layout=shared_in

---

⚠️ [ruff] <F841> _{reported by reviewdog 🐶}
Local variable H is assigned to but never used

aiter/test_ascyn.py

Line 463 in 9d3f8b1

H = 8

---

⚠️ [ruff] <F841> _{reported by reviewdog 🐶}
Local variable C is assigned to but never used

aiter/test_ascyn.py

Line 473 in 9d3f8b1

C = A @ B.T

[github-actions]

⚠️ [ruff] <F401> _{reported by reviewdog 🐶}
typing.Optional imported but unused

Suggested change

	from typing import Optional
	import functools
	import functools

[github-actions]

⚠️ [ruff] <F401> _{reported by reviewdog 🐶}
functools imported but unused

Suggested change

	import functools
	import json
	import json

[github-actions]

⚠️ [ruff] <F401> _{reported by reviewdog 🐶}
json imported but unused

Suggested change

	import json
	import triton
	import triton

[github-actions]

⚠️ [ruff] <F401> _{reported by reviewdog 🐶}
torch imported but unused

Suggested change

	import torch
	import aiter.ops.triton.utils._triton.arch_info as arch_info
	import aiter.ops.triton.utils._triton.arch_info as arch_info

[github-actions]

⚠️ [ruff] <F401> _{reported by reviewdog 🐶}
aiter.ops.triton.utils._triton.arch_info imported but unused

Suggested change

	import aiter.ops.triton.utils._triton.arch_info as arch_info
	from aiter.ops.triton.utils.core import AITER_TRITON_CONFIGS_PATH
	from aiter.ops.triton.utils.core import AITER_TRITON_CONFIGS_PATH

[github-actions]

⚠️ [ruff] <F401> _{reported by reviewdog 🐶}
triton imported but unused

Suggested change

	import triton
	import triton.language as tl
	import triton.language as tl

[github-actions]

⚠️ [ruff] <F401> _{reported by reviewdog 🐶}
torch imported but unused

Suggested change

	import torch
	from aiter.ops.triton.activation import _tanh
	from aiter.ops.triton.activation import _tanh

[github-actions]

⚠️ [ruff] <F401> _{reported by reviewdog 🐶}
aiter.ops.triton.activation._tanh imported but unused

Suggested change

	from aiter.ops.triton.activation import _tanh
	import aiter.ops.triton.utils.arch_info as arch_info
	import aiter.ops.triton.utils.arch_info as arch_info

[github-actions]

⚠️ [ruff] <F401> _{reported by reviewdog 🐶}
aiter.ops.triton.utils.arch_info imported but unused

Suggested change

	import aiter.ops.triton.utils.arch_info as arch_info
	from aiter.ops.triton.utils.core import AITER_TRITON_CONFIGS_PATH
	from aiter.ops.triton.utils.core import AITER_TRITON_CONFIGS_PATH

[github-actions]

⚠️ [ruff] <F401> _{reported by reviewdog 🐶}
aiter.ops.triton.utils.core.AITER_TRITON_CONFIGS_PATH imported but unused

Suggested change

	from aiter.ops.triton.utils.core import AITER_TRITON_CONFIGS_PATH
	from aiter.ops.triton.utils.pid_preprocessing import remap_xcd
	from aiter.ops.triton.utils.pid_preprocessing import remap_xcd