From 13db5bbc616e616c951017d06830b5779d239914 Mon Sep 17 00:00:00 2001 From: manschou Date: Thu, 25 Jun 2026 18:55:01 +0000 Subject: [PATCH 1/6] attention: add XCD-aware spatial workgroup mapping for MHA and GQA Introduces remap_workgroup_spatial() in pid_preprocessing.py and wires it into the flash-attention forward kernel behind AITER_SWIZZLE=1. MHA path (spatial_mha): groups ceil(NUM_Q_HEADS/NUM_XCDS) consecutive Q heads onto each XCD so each head's full KV tensor is processed by one XCD, keeping it hot in that XCD's 4 MB L2 partition. GQA path (spatial_gqa): assigns each KV head exclusively to one XCD so all Q heads sharing a KQA head run on the same XCD. Ordering within each XCD is block-first, which avoids the causal load-imbalance that head-first ordering creates. Three regimes are handled: HK==NXCD (aligned), HK>NXCD, and HK int: return int(window_size[0]) if int(window_size[0]) >= 0 else 0 @@ -299,6 +309,7 @@ def _flash_attn_forward( VARLEN=is_varlen, BATCH=batch, NUM_XCD=get_num_xcds(), + SWIZZLE=_AITER_SWIZZLE, USE_INT64_STRIDES=_USE_INT64_STRIDES, ENABLE_SINK=sink is not None, SLIDING_WINDOW=sliding_window, diff --git a/aiter/ops/triton/utils/_triton/pid_preprocessing.py b/aiter/ops/triton/utils/_triton/pid_preprocessing.py index 4263e6fa8b5..4a9cc471be2 100644 --- a/aiter/ops/triton/utils/_triton/pid_preprocessing.py +++ b/aiter/ops/triton/utils/_triton/pid_preprocessing.py @@ -1,5 +1,4 @@ # SPDX-License-Identifier: MIT - # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. import triton @@ -34,8 +33,7 @@ def remap_xcd(pid, GRID_MN, NUM_XCDS: tl.constexpr = 8): # We calculate the number of xcds that have pids_per_xcd pids as # tall_xcds tall_xcds = GRID_MN % NUM_XCDS - if tall_xcds == 0: - tall_xcds = tl.cast(NUM_XCDS, tall_xcds.type) + tall_xcds = NUM_XCDS if tall_xcds == 0 else tall_xcds # Compute current XCD and local pid within the XCD xcd = pid % NUM_XCDS local_pid = pid // NUM_XCDS @@ -99,3 +97,116 @@ def pid_grid_3d(pid: int, num_pid_m: int, num_pid_n: int, num_pid_k): pid_k = pid // (num_pid_m * num_pid_n) % num_pid_k return pid_m, pid_n, pid_k + + +@triton.jit +def remap_workgroup_spatial( + wid, + NUM_Q_HEADS, + NUM_BLOCKS, + BATCH, + NUM_QUERIES_PER_KV: tl.constexpr, + NUM_XCDS: tl.constexpr = 8, +): + """ + XCD-aware workgroup mapping for multi-head attention on AMD CDNA3/3.5 GPUs. + + Hardware assigns workgroup `wid` to XCD `wid % NUM_XCDS` (round-robin). + This function remaps (wid) -> (off_q_head, start_m, off_z) so that workgroups + sharing KV data land on the same XCD, keeping that data hot in its 4 MB L2 slice. + + MHA path (NUM_QUERIES_PER_KV == 1): + Groups ceil(NUM_Q_HEADS / NUM_XCDS) consecutive Q heads onto each XCD so each + head's full KV tensor is processed by one XCD only. Within each XCD the + ordering is head-first: all sequence blocks for head i complete before head i+1. + + GQA path (NUM_QUERIES_PER_KV > 1): + Assigns each KV head exclusively to one XCD so all Q heads that share a KV head + run on the same XCD. Within each XCD the ordering is block-first: all Q heads + in the group process sequence block b together before advancing to block b+1. + Block-first avoids the causal load-imbalance that head-first ordering creates. + + Three regimes depending on NUM_KV_HEADS vs NUM_XCDS: + HK == NUM_XCDS : one KV head per XCD (aligned). + HK > NUM_XCDS : ceil(HK / NXCD) KV heads per XCD, each processed fully + before the next, so only one KV head's data is live at a time. + HK < NUM_XCDS : NXCD / HK XCDs share each KV head, splitting the per-group + work evenly across the cluster of co-owners. + + Args: + wid : flat workgroup ID (tl.program_id(0)) + NUM_Q_HEADS : total number of Q heads + NUM_BLOCKS : number of sequence blocks along the Q dimension + BATCH : batch size + NUM_QUERIES_PER_KV: Q heads per KV head (1 for MHA, >1 for GQA) + NUM_XCDS : number of XCDs on the device (8 for MI3xx) + + Returns: + off_q_head : Q head index for this workgroup + start_m : sequence block index for this workgroup + off_z : batch index for this workgroup + """ + xcd = wid % NUM_XCDS + pos = wid // NUM_XCDS + + # ── MHA path (spatial_mha) ──────────────────────────────────────────────── + if NUM_QUERIES_PER_KV == 1: + wgs_per_head = NUM_BLOCKS * BATCH + local_head_idx = pos // wgs_per_head + remainder = pos % wgs_per_head + off_z = remainder // NUM_BLOCKS + start_m = remainder % NUM_BLOCKS + off_q_head = local_head_idx * NUM_XCDS + xcd + # Bounds check for non-divisible NUM_Q_HEADS / NUM_XCDS. + if off_q_head >= NUM_Q_HEADS: + off_q_head = wid % NUM_Q_HEADS + start_m = (wid // NUM_Q_HEADS) % NUM_BLOCKS + off_z = (wid // (NUM_BLOCKS * NUM_Q_HEADS)) % BATCH + return off_q_head, start_m, off_z + + # ── GQA path (spatial_gqa) ──────────────────────────────────────────────── + NUM_KV_HEADS: tl.constexpr = NUM_Q_HEADS // NUM_QUERIES_PER_KV + + if NUM_KV_HEADS >= NUM_XCDS: + NUM_KV_REPLICAS: tl.constexpr = (NUM_KV_HEADS + NUM_XCDS - 1) // NUM_XCDS + wgs_per_block = NUM_QUERIES_PER_KV * BATCH + if NUM_KV_REPLICAS == 1: + # Aligned: one KV head per XCD. + kv_head = xcd % NUM_KV_HEADS + start_m = (pos // wgs_per_block) % NUM_BLOCKS + remainder = pos % wgs_per_block + off_z = remainder // NUM_QUERIES_PER_KV + local_q_in_group = remainder % NUM_QUERIES_PER_KV + off_q_head = kv_head * NUM_QUERIES_PER_KV + local_q_in_group + else: + # HK > NXCD: each XCD owns ceil(HK/NXCD) KV heads processed one at a time. + wgs_per_kv_group = NUM_QUERIES_PER_KV * NUM_BLOCKS * BATCH + kv_slot = pos // wgs_per_kv_group + kv_head = xcd + kv_slot * NUM_XCDS + local = pos % wgs_per_kv_group + start_m = (local // (NUM_QUERIES_PER_KV * BATCH)) % NUM_BLOCKS + remainder = local % (NUM_QUERIES_PER_KV * BATCH) + off_z = remainder // NUM_QUERIES_PER_KV + local_q_in_group = remainder % NUM_QUERIES_PER_KV + off_q_head = kv_head * NUM_QUERIES_PER_KV + local_q_in_group + else: + # HK < NXCD: NXCD/HK XCDs share each KV head, splitting the per-group work. + NXCD_PER_KV: tl.constexpr = NUM_XCDS // NUM_KV_HEADS + xcd_in_cluster = xcd // NUM_KV_HEADS + kv_head = xcd % NUM_KV_HEADS + wgs_per_kv_group = NUM_QUERIES_PER_KV * NUM_BLOCKS * BATCH + wgs_per_xcd_in_cluster = wgs_per_kv_group // NXCD_PER_KV + local = (pos % wgs_per_xcd_in_cluster) + xcd_in_cluster * wgs_per_xcd_in_cluster + start_m = (local // (NUM_QUERIES_PER_KV * BATCH)) % NUM_BLOCKS + remainder = local % (NUM_QUERIES_PER_KV * BATCH) + off_z = remainder // NUM_QUERIES_PER_KV + local_q_in_group = remainder % NUM_QUERIES_PER_KV + off_q_head = kv_head * NUM_QUERIES_PER_KV + local_q_in_group + + # Bounds check for non-divisible edge cases. + if kv_head >= NUM_KV_HEADS or off_q_head >= NUM_Q_HEADS: + off_q_head = wid % NUM_Q_HEADS + start_m = (wid // NUM_Q_HEADS) % NUM_BLOCKS + off_z = (wid // (NUM_BLOCKS * NUM_Q_HEADS)) % BATCH + + return off_q_head, start_m, off_z From 3936719d37239ab3d9df277b73616d7157ad98e1 Mon Sep 17 00:00:00 2001 From: manschou Date: Thu, 25 Jun 2026 19:43:35 +0000 Subject: [PATCH 2/6] op_tests: add correctness test for spatial MHA/GQA swizzle --- op_tests/test_mha_spatial_swizzle.py | 103 +++++++++++++++++++++++++++ 1 file changed, 103 insertions(+) create mode 100644 op_tests/test_mha_spatial_swizzle.py diff --git a/op_tests/test_mha_spatial_swizzle.py b/op_tests/test_mha_spatial_swizzle.py new file mode 100644 index 00000000000..e22b661bb3e --- /dev/null +++ b/op_tests/test_mha_spatial_swizzle.py @@ -0,0 +1,103 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +""" +Correctness tests for remap_workgroup_spatial (AITER_SWIZZLE=1). + +Verifies that SWIZZLE=1 produces bit-identical output to SWIZZLE=0 across +GQA and MHA configurations, causal and non-causal, single and batched. +""" + +import math + +import pytest +import torch + +import aiter.ops.triton.attention.mha as mha_mod + + +def _run(swizzle: int, B: int, Hq: int, Hk: int, S: int, D: int, causal: bool): + mha_mod.mha_set_swizzle(swizzle) + torch.manual_seed(42) + q = torch.randn(B, S, Hq, D, device="cuda", dtype=torch.bfloat16) + k = torch.randn(B, S, Hk, D, device="cuda", dtype=torch.bfloat16) + v = torch.randn(B, S, Hk, D, device="cuda", dtype=torch.bfloat16) + out = mha_mod.flash_attn_func( + q, k, v, softmax_scale=1.0 / math.sqrt(D), causal=causal + ) + torch.cuda.synchronize() + return (out[0] if isinstance(out, tuple) else out).float() + + +# --------------------------------------------------------------------------- +# GQA configs — spatial_gqa path +# --------------------------------------------------------------------------- + +@pytest.mark.parametrize("causal", [True, False]) +@pytest.mark.parametrize("Hq,Hk", [ + (128, 8), # aligned HK == NUM_XCDS + (64, 8), + (32, 8), +]) +def test_spatial_gqa_aligned(Hq, Hk, causal): + """HK == NUM_XCDS (8): one KV head per XCD.""" + out0 = _run(0, 1, Hq, Hk, 8192, 128, causal) + out1 = _run(1, 1, Hq, Hk, 8192, 128, causal) + assert torch.equal(out0, out1), f"HQ={Hq} HK={Hk} causal={causal}: outputs differ" + + +@pytest.mark.parametrize("Hq,Hk", [ + (128, 16), # HK > NUM_XCDS + (128, 32), +]) +def test_spatial_gqa_hk_gt_nxcd(Hq, Hk): + """HK > NUM_XCDS: each XCD owns multiple KV heads.""" + out0 = _run(0, 1, Hq, Hk, 8192, 128, causal=True) + out1 = _run(1, 1, Hq, Hk, 8192, 128, causal=True) + assert torch.equal(out0, out1), f"HQ={Hq} HK={Hk}: outputs differ" + + +@pytest.mark.parametrize("Hq,Hk", [ + (128, 4), # HK < NUM_XCDS + (128, 2), +]) +def test_spatial_gqa_hk_lt_nxcd(Hq, Hk): + """HK < NUM_XCDS: multiple XCDs share each KV head.""" + out0 = _run(0, 1, Hq, Hk, 8192, 128, causal=True) + out1 = _run(1, 1, Hq, Hk, 8192, 128, causal=True) + assert torch.equal(out0, out1), f"HQ={Hq} HK={Hk}: outputs differ" + + +@pytest.mark.parametrize("causal", [True, False]) +def test_spatial_gqa_batched(causal): + """GQA aligned, batch > 1.""" + out0 = _run(0, 4, 128, 8, 8192, 128, causal) + out1 = _run(1, 4, 128, 8, 8192, 128, causal) + assert torch.equal(out0, out1), f"causal={causal}: batched GQA outputs differ" + + +# --------------------------------------------------------------------------- +# MHA configs — spatial_mha path +# --------------------------------------------------------------------------- + +@pytest.mark.parametrize("causal", [True, False]) +@pytest.mark.parametrize("H", [16, 32, 64, 128]) +def test_spatial_mha(H, causal): + """MHA: HQ == HK, various head counts.""" + out0 = _run(0, 1, H, H, 8192, 128, causal) + out1 = _run(1, 1, H, H, 8192, 128, causal) + assert torch.equal(out0, out1), f"H={H} causal={causal}: outputs differ" + + +def test_spatial_mha_nonpow2_heads(): + """MHA with non-power-of-2 head count (e.g. Wan2.1-14B DiT, H=40).""" + out0 = _run(0, 1, 40, 40, 8192, 128, causal=True) + out1 = _run(1, 1, 40, 40, 8192, 128, causal=True) + assert torch.equal(out0, out1), "H=40 (Wan): outputs differ" + + +def test_spatial_mha_batched(): + """MHA, batch > 1.""" + out0 = _run(0, 2, 128, 128, 8192, 128, causal=True) + out1 = _run(1, 2, 128, 128, 8192, 128, causal=True) + assert torch.equal(out0, out1), "batched MHA outputs differ" From cb2a345b0bde3adec44afa4fcba720807d9629e5 Mon Sep 17 00:00:00 2001 From: manschou Date: Tue, 30 Jun 2026 14:58:01 +0000 Subject: [PATCH 3/6] address review feedback - revert tall_xcds ternary to tl.cast() to preserve int64 correctness in grouped GEMM - change SWIZZLE from int {0,1} to string {"default","spatial"} throughout - add tl.static_assert for SWIZZLE value in kernel with descriptive message - rename AITER_SWIZZLE to AITER_TRITON_MHA_SWIZZLE (more specific) - harden env var parsing: validate against allowed set, raise ValueError on bad input - add validation in mha_set_swizzle() before mutating global state - fix test: add autouse fixture to save/restore _MHA_SWIZZLE around each test - update test docstring to reflect new env var name and string values --- .../triton/_triton_kernels/attention/mha.py | 10 +++- aiter/ops/triton/attention/mha.py | 30 ++++++++-- .../triton/utils/_triton/pid_preprocessing.py | 3 +- op_tests/test_mha_spatial_swizzle.py | 59 +++++++++++-------- 4 files changed, 68 insertions(+), 34 deletions(-) diff --git a/aiter/ops/triton/_triton_kernels/attention/mha.py b/aiter/ops/triton/_triton_kernels/attention/mha.py index dafb4a985d0..5675910d845 100644 --- a/aiter/ops/triton/_triton_kernels/attention/mha.py +++ b/aiter/ops/triton/_triton_kernels/attention/mha.py @@ -372,14 +372,18 @@ def _attn_fwd( ) # workgroup id ranging: 0,1,2,...., (BATCH * NUM_Q_HEADS * NUM_BLOCKS - 1) # num blocks along seqlen - if SWIZZLE == 0: - # Baseline: head-first round-robin with XCD-aware head remapping. + tl.static_assert( + SWIZZLE == "default" or SWIZZLE == "spatial", + "SWIZZLE must be 'default' or 'spatial'; set via AITER_TRITON_MHA_SWIZZLE or mha_set_swizzle()", + ) + if SWIZZLE == "default": + # Default: head-first round-robin with XCD-aware head remapping. off_q_head = wid % NUM_Q_HEADS off_q_head = remap_xcd(off_q_head, NUM_Q_HEADS, NUM_XCD) start_m = (wid // NUM_Q_HEADS) % NUM_BLOCKS off_z = (wid // (NUM_BLOCKS * NUM_Q_HEADS)) % BATCH else: - # Spatial (SWIZZLE == 1): XCD-aware KV-head mapping for MHA and GQA. + # Spatial: XCD-aware KV-head mapping for MHA and GQA. NUM_QUERIES_PER_KV: tl.constexpr = NUM_Q_HEADS // NUM_K_HEADS off_q_head, start_m, off_z = remap_workgroup_spatial( wid, NUM_Q_HEADS, NUM_BLOCKS, BATCH, NUM_QUERIES_PER_KV, NUM_XCD diff --git a/aiter/ops/triton/attention/mha.py b/aiter/ops/triton/attention/mha.py index 827bd12a923..8a0e47e00fd 100644 --- a/aiter/ops/triton/attention/mha.py +++ b/aiter/ops/triton/attention/mha.py @@ -48,13 +48,31 @@ def mha_set_use_int64_strides(value: bool): _USE_INT64_STRIDES = value -_AITER_SWIZZLE = int(os.environ.get("AITER_SWIZZLE", "0")) +_MHA_SWIZZLE_VALUES = ("default", "spatial") +_env_swizzle = os.environ.get("AITER_TRITON_MHA_SWIZZLE", "default") +if _env_swizzle not in _MHA_SWIZZLE_VALUES: + raise ValueError( + f"Invalid AITER_TRITON_MHA_SWIZZLE value: {_env_swizzle!r}. " + f"Must be one of {_MHA_SWIZZLE_VALUES}." + ) +_MHA_SWIZZLE: Literal["default", "spatial"] = _env_swizzle +del _env_swizzle + + +def mha_set_swizzle(value: Literal["default", "spatial"]): + """Set MHA workgroup swizzle mode. -def mha_set_swizzle(value: int): - """Set spatial swizzle mode: 0 = baseline (remap_xcd), 1 = spatial MHA/GQA.""" - global _AITER_SWIZZLE - _AITER_SWIZZLE = value + Args: + value: ``"default"`` preserves the existing remap_xcd behaviour. + ``"spatial"`` enables XCD-aware KV-head mapping for MHA and GQA. + """ + if value not in _MHA_SWIZZLE_VALUES: + raise ValueError( + f"Invalid swizzle value: {value!r}. Must be one of {_MHA_SWIZZLE_VALUES}." + ) + global _MHA_SWIZZLE + _MHA_SWIZZLE = value def _get_sliding_window_size(window_size: Tuple[int, int]) -> int: @@ -309,7 +327,7 @@ def _flash_attn_forward( VARLEN=is_varlen, BATCH=batch, NUM_XCD=get_num_xcds(), - SWIZZLE=_AITER_SWIZZLE, + SWIZZLE=_MHA_SWIZZLE, USE_INT64_STRIDES=_USE_INT64_STRIDES, ENABLE_SINK=sink is not None, SLIDING_WINDOW=sliding_window, diff --git a/aiter/ops/triton/utils/_triton/pid_preprocessing.py b/aiter/ops/triton/utils/_triton/pid_preprocessing.py index 4a9cc471be2..d1a2a001d5a 100644 --- a/aiter/ops/triton/utils/_triton/pid_preprocessing.py +++ b/aiter/ops/triton/utils/_triton/pid_preprocessing.py @@ -33,7 +33,8 @@ def remap_xcd(pid, GRID_MN, NUM_XCDS: tl.constexpr = 8): # We calculate the number of xcds that have pids_per_xcd pids as # tall_xcds tall_xcds = GRID_MN % NUM_XCDS - tall_xcds = NUM_XCDS if tall_xcds == 0 else tall_xcds + if tall_xcds == 0: + tall_xcds = tl.cast(NUM_XCDS, tall_xcds.type) # Compute current XCD and local pid within the XCD xcd = pid % NUM_XCDS local_pid = pid // NUM_XCDS diff --git a/op_tests/test_mha_spatial_swizzle.py b/op_tests/test_mha_spatial_swizzle.py index e22b661bb3e..f50c6bbb3d2 100644 --- a/op_tests/test_mha_spatial_swizzle.py +++ b/op_tests/test_mha_spatial_swizzle.py @@ -2,10 +2,14 @@ # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. """ -Correctness tests for remap_workgroup_spatial (AITER_SWIZZLE=1). +Correctness tests for remap_workgroup_spatial (AITER_TRITON_MHA_SWIZZLE="spatial"). -Verifies that SWIZZLE=1 produces bit-identical output to SWIZZLE=0 across -GQA and MHA configurations, causal and non-causal, single and batched. +Verifies that swizzle="spatial" produces bit-identical output to swizzle="default" +across GQA and MHA configurations, causal and non-causal, single and batched. + +Configuration: + Environment variable : AITER_TRITON_MHA_SWIZZLE (values: "default", "spatial") + Programmatic API : mha_set_swizzle("default" | "spatial") """ import math @@ -13,18 +17,25 @@ import pytest import torch -import aiter.ops.triton.attention.mha as mha_mod +from aiter.ops.triton.attention.mha import flash_attn_func, mha_set_swizzle + + +@pytest.fixture(autouse=True) +def restore_swizzle(): + """Save swizzle state before each test and restore it on teardown.""" + from aiter.ops.triton.attention import mha as mha_mod + saved = mha_mod._MHA_SWIZZLE + yield + mha_set_swizzle(saved) -def _run(swizzle: int, B: int, Hq: int, Hk: int, S: int, D: int, causal: bool): - mha_mod.mha_set_swizzle(swizzle) +def _run(swizzle: str, B: int, Hq: int, Hk: int, S: int, D: int, causal: bool): + mha_set_swizzle(swizzle) torch.manual_seed(42) q = torch.randn(B, S, Hq, D, device="cuda", dtype=torch.bfloat16) k = torch.randn(B, S, Hk, D, device="cuda", dtype=torch.bfloat16) v = torch.randn(B, S, Hk, D, device="cuda", dtype=torch.bfloat16) - out = mha_mod.flash_attn_func( - q, k, v, softmax_scale=1.0 / math.sqrt(D), causal=causal - ) + out = flash_attn_func(q, k, v, softmax_scale=1.0 / math.sqrt(D), causal=causal) torch.cuda.synchronize() return (out[0] if isinstance(out, tuple) else out).float() @@ -41,8 +52,8 @@ def _run(swizzle: int, B: int, Hq: int, Hk: int, S: int, D: int, causal: bool): ]) def test_spatial_gqa_aligned(Hq, Hk, causal): """HK == NUM_XCDS (8): one KV head per XCD.""" - out0 = _run(0, 1, Hq, Hk, 8192, 128, causal) - out1 = _run(1, 1, Hq, Hk, 8192, 128, causal) + out0 = _run("default", 1, Hq, Hk, 8192, 128, causal) + out1 = _run("spatial", 1, Hq, Hk, 8192, 128, causal) assert torch.equal(out0, out1), f"HQ={Hq} HK={Hk} causal={causal}: outputs differ" @@ -52,8 +63,8 @@ def test_spatial_gqa_aligned(Hq, Hk, causal): ]) def test_spatial_gqa_hk_gt_nxcd(Hq, Hk): """HK > NUM_XCDS: each XCD owns multiple KV heads.""" - out0 = _run(0, 1, Hq, Hk, 8192, 128, causal=True) - out1 = _run(1, 1, Hq, Hk, 8192, 128, causal=True) + out0 = _run("default", 1, Hq, Hk, 8192, 128, causal=True) + out1 = _run("spatial", 1, Hq, Hk, 8192, 128, causal=True) assert torch.equal(out0, out1), f"HQ={Hq} HK={Hk}: outputs differ" @@ -63,16 +74,16 @@ def test_spatial_gqa_hk_gt_nxcd(Hq, Hk): ]) def test_spatial_gqa_hk_lt_nxcd(Hq, Hk): """HK < NUM_XCDS: multiple XCDs share each KV head.""" - out0 = _run(0, 1, Hq, Hk, 8192, 128, causal=True) - out1 = _run(1, 1, Hq, Hk, 8192, 128, causal=True) + out0 = _run("default", 1, Hq, Hk, 8192, 128, causal=True) + out1 = _run("spatial", 1, Hq, Hk, 8192, 128, causal=True) assert torch.equal(out0, out1), f"HQ={Hq} HK={Hk}: outputs differ" @pytest.mark.parametrize("causal", [True, False]) def test_spatial_gqa_batched(causal): """GQA aligned, batch > 1.""" - out0 = _run(0, 4, 128, 8, 8192, 128, causal) - out1 = _run(1, 4, 128, 8, 8192, 128, causal) + out0 = _run("default", 4, 128, 8, 8192, 128, causal) + out1 = _run("spatial", 4, 128, 8, 8192, 128, causal) assert torch.equal(out0, out1), f"causal={causal}: batched GQA outputs differ" @@ -84,20 +95,20 @@ def test_spatial_gqa_batched(causal): @pytest.mark.parametrize("H", [16, 32, 64, 128]) def test_spatial_mha(H, causal): """MHA: HQ == HK, various head counts.""" - out0 = _run(0, 1, H, H, 8192, 128, causal) - out1 = _run(1, 1, H, H, 8192, 128, causal) + out0 = _run("default", 1, H, H, 8192, 128, causal) + out1 = _run("spatial", 1, H, H, 8192, 128, causal) assert torch.equal(out0, out1), f"H={H} causal={causal}: outputs differ" def test_spatial_mha_nonpow2_heads(): - """MHA with non-power-of-2 head count (e.g. Wan2.1-14B DiT, H=40).""" - out0 = _run(0, 1, 40, 40, 8192, 128, causal=True) - out1 = _run(1, 1, 40, 40, 8192, 128, causal=True) + """MHA with non-power-of-2 head count (e.g. Wan2.2-14B DiT, H=40).""" + out0 = _run("default", 1, 40, 40, 8192, 128, causal=True) + out1 = _run("spatial", 1, 40, 40, 8192, 128, causal=True) assert torch.equal(out0, out1), "H=40 (Wan): outputs differ" def test_spatial_mha_batched(): """MHA, batch > 1.""" - out0 = _run(0, 2, 128, 128, 8192, 128, causal=True) - out1 = _run(1, 2, 128, 128, 8192, 128, causal=True) + out0 = _run("default", 2, 128, 128, 8192, 128, causal=True) + out1 = _run("spatial", 2, 128, 128, 8192, 128, causal=True) assert torch.equal(out0, out1), "batched MHA outputs differ" From 2ddf0c6f5894ba31733843ef3ff6dc52549b509a Mon Sep 17 00:00:00 2001 From: manschou Date: Wed, 8 Jul 2026 19:54:01 +0000 Subject: [PATCH 4/6] tests: move test_mha_spatial_swizzle to triton_tests/attention --- op_tests/{ => triton_tests/attention}/test_mha_spatial_swizzle.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename op_tests/{ => triton_tests/attention}/test_mha_spatial_swizzle.py (100%) diff --git a/op_tests/test_mha_spatial_swizzle.py b/op_tests/triton_tests/attention/test_mha_spatial_swizzle.py similarity index 100% rename from op_tests/test_mha_spatial_swizzle.py rename to op_tests/triton_tests/attention/test_mha_spatial_swizzle.py From 620fdd0ebc7f4ec9848a891c115c86135e48b23a Mon Sep 17 00:00:00 2001 From: manschou Date: Thu, 23 Jul 2026 20:42:37 +0000 Subject: [PATCH 5/6] style: apply Black formatting --- .../triton/_triton_kernels/attention/mha.py | 5 ++- .../attention/test_mha_spatial_swizzle.py | 38 ++++++++++++------- 2 files changed, 29 insertions(+), 14 deletions(-) diff --git a/aiter/ops/triton/_triton_kernels/attention/mha.py b/aiter/ops/triton/_triton_kernels/attention/mha.py index 5675910d845..4df60abb1d4 100644 --- a/aiter/ops/triton/_triton_kernels/attention/mha.py +++ b/aiter/ops/triton/_triton_kernels/attention/mha.py @@ -9,7 +9,10 @@ from aiter.ops.triton.utils._triton import arch_info from aiter.ops.triton.utils.core import AITER_TRITON_CONFIGS_PATH -from aiter.ops.triton.utils._triton.pid_preprocessing import remap_xcd, remap_workgroup_spatial +from aiter.ops.triton.utils._triton.pid_preprocessing import ( + remap_xcd, + remap_workgroup_spatial, +) from aiter.ops.triton.utils._triton.mha_kernel_utils import _compute_fp8_scaling_factors from aiter.ops.triton.utils._triton.kernel_repr import make_kernel_repr diff --git a/op_tests/triton_tests/attention/test_mha_spatial_swizzle.py b/op_tests/triton_tests/attention/test_mha_spatial_swizzle.py index f50c6bbb3d2..3e0cb12e596 100644 --- a/op_tests/triton_tests/attention/test_mha_spatial_swizzle.py +++ b/op_tests/triton_tests/attention/test_mha_spatial_swizzle.py @@ -24,6 +24,7 @@ def restore_swizzle(): """Save swizzle state before each test and restore it on teardown.""" from aiter.ops.triton.attention import mha as mha_mod + saved = mha_mod._MHA_SWIZZLE yield mha_set_swizzle(saved) @@ -44,12 +45,16 @@ def _run(swizzle: str, B: int, Hq: int, Hk: int, S: int, D: int, causal: bool): # GQA configs — spatial_gqa path # --------------------------------------------------------------------------- + @pytest.mark.parametrize("causal", [True, False]) -@pytest.mark.parametrize("Hq,Hk", [ - (128, 8), # aligned HK == NUM_XCDS - (64, 8), - (32, 8), -]) +@pytest.mark.parametrize( + "Hq,Hk", + [ + (128, 8), # aligned HK == NUM_XCDS + (64, 8), + (32, 8), + ], +) def test_spatial_gqa_aligned(Hq, Hk, causal): """HK == NUM_XCDS (8): one KV head per XCD.""" out0 = _run("default", 1, Hq, Hk, 8192, 128, causal) @@ -57,10 +62,13 @@ def test_spatial_gqa_aligned(Hq, Hk, causal): assert torch.equal(out0, out1), f"HQ={Hq} HK={Hk} causal={causal}: outputs differ" -@pytest.mark.parametrize("Hq,Hk", [ - (128, 16), # HK > NUM_XCDS - (128, 32), -]) +@pytest.mark.parametrize( + "Hq,Hk", + [ + (128, 16), # HK > NUM_XCDS + (128, 32), + ], +) def test_spatial_gqa_hk_gt_nxcd(Hq, Hk): """HK > NUM_XCDS: each XCD owns multiple KV heads.""" out0 = _run("default", 1, Hq, Hk, 8192, 128, causal=True) @@ -68,10 +76,13 @@ def test_spatial_gqa_hk_gt_nxcd(Hq, Hk): assert torch.equal(out0, out1), f"HQ={Hq} HK={Hk}: outputs differ" -@pytest.mark.parametrize("Hq,Hk", [ - (128, 4), # HK < NUM_XCDS - (128, 2), -]) +@pytest.mark.parametrize( + "Hq,Hk", + [ + (128, 4), # HK < NUM_XCDS + (128, 2), + ], +) def test_spatial_gqa_hk_lt_nxcd(Hq, Hk): """HK < NUM_XCDS: multiple XCDs share each KV head.""" out0 = _run("default", 1, Hq, Hk, 8192, 128, causal=True) @@ -91,6 +102,7 @@ def test_spatial_gqa_batched(causal): # MHA configs — spatial_mha path # --------------------------------------------------------------------------- + @pytest.mark.parametrize("causal", [True, False]) @pytest.mark.parametrize("H", [16, 32, 64, 128]) def test_spatial_mha(H, causal): From 5bc064db343dae193bc987cb33decad62ba074be Mon Sep 17 00:00:00 2001 From: manschou Date: Wed, 29 Jul 2026 18:38:27 +0000 Subject: [PATCH 6/6] style: fix ruff 0.16 import sorting and style issues --- aiter/ops/triton/_triton_kernels/attention/mha.py | 10 +++++----- aiter/ops/triton/attention/mha.py | 3 ++- 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/aiter/ops/triton/_triton_kernels/attention/mha.py b/aiter/ops/triton/_triton_kernels/attention/mha.py index fc5cb3face3..7b7e3f0481b 100644 --- a/aiter/ops/triton/_triton_kernels/attention/mha.py +++ b/aiter/ops/triton/_triton_kernels/attention/mha.py @@ -9,13 +9,13 @@ import triton.language as tl from aiter.ops.triton.utils._triton import arch_info -from aiter.ops.triton.utils.core import AITER_TRITON_CONFIGS_PATH +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, remap_workgroup_spatial, + remap_xcd, ) -from aiter.ops.triton.utils._triton.mha_kernel_utils import _compute_fp8_scaling_factors -from aiter.ops.triton.utils._triton.kernel_repr import make_kernel_repr +from aiter.ops.triton.utils.core import AITER_TRITON_CONFIGS_PATH @triton.jit @@ -887,7 +887,7 @@ def _attn_fwd( end_m_idx = (start_m + 1) * BLOCK_M start_m_idx = start_m * BLOCK_M causal_start_idx = seqlen_q - seqlen_k - if IS_CAUSAL and (causal_start_idx > start_m_idx and causal_start_idx < end_m_idx): + if IS_CAUSAL and causal_start_idx > start_m_idx and causal_start_idx < end_m_idx: out_mask_boundary = tl.full( (BLOCK_DMODEL_POW2,), causal_start_idx, dtype=tl.int32 ) diff --git a/aiter/ops/triton/attention/mha.py b/aiter/ops/triton/attention/mha.py index 3b58b994860..b5db75ad890 100644 --- a/aiter/ops/triton/attention/mha.py +++ b/aiter/ops/triton/attention/mha.py @@ -2,7 +2,8 @@ # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. import os -from typing import Literal, Optional, Tuple, Union +from typing import Literal + import torch import triton import triton.language as tl