Spatial Attention: XCD-aware spatial workgroup mapping for MHA and GQA (SWIZZLE=1)

aiter/ops/triton/_triton_kernels/attention/mha.py

@@ -9,7 +9,7 @@
 
 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
+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
 
@@ -359,6 +359,7 @@ def _attn_fwd(
     VARLEN: tl.constexpr,
     BATCH,
     NUM_XCD: tl.constexpr,
+    SWIZZLE: tl.constexpr,
     USE_INT64_STRIDES: tl.constexpr,
     ENABLE_SINK: tl.constexpr,
     SLIDING_WINDOW: tl.constexpr,
@@ -371,10 +372,22 @@ def _attn_fwd(
     )  # workgroup id ranging: 0,1,2,...., (BATCH * NUM_Q_HEADS * NUM_BLOCKS - 1)
     # num blocks along seqlen
 
-    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
+    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: 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
+        )
 
     # offsets
     offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)

aiter/ops/triton/attention/mha.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: MIT
 # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.
 
+import os
 from typing import Literal, Optional, Tuple, Union
 import torch
 import triton
@@ -47,6 +48,33 @@ def mha_set_use_int64_strides(value: bool):
     _USE_INT64_STRIDES = value
 
 
+_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.
+
+    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:
     return int(window_size[0]) if int(window_size[0]) >= 0 else 0
 
@@ -299,6 +327,7 @@ def _flash_attn_forward(
             VARLEN=is_varlen,
             BATCH=batch,
             NUM_XCD=get_num_xcds(),
+            SWIZZLE=_MHA_SWIZZLE,
             USE_INT64_STRIDES=_USE_INT64_STRIDES,
             ENABLE_SINK=sink is not None,
             SLIDING_WINDOW=sliding_window,

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
@@ -99,3 +98,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

op_tests/test_mha_spatial_swizzle.py

@@ -0,0 +1,114 @@
+# SPDX-License-Identifier: MIT
+# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.
+
+"""
+Correctness tests for remap_workgroup_spatial (AITER_TRITON_MHA_SWIZZLE="spatial").
+
+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
+
+import pytest
+import torch
+
+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: 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 = 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("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"
+
+
+@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)
+    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("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)
+    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("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"
+
+
+# ---------------------------------------------------------------------------
+# 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("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.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("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"