Spatial Attention: XCD-aware spatial workgroup mapping for MHA and GQA (SWIZZLE=1) by mc186 · Pull Request #3936 · ROCm/aiter

mc186 · 2026-06-25T20:00:24Z

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

Motivation

AMD CDNA3/3.5 GPUs (MI300X, MI350X, MI355X) are built from 8 XCDs (chiplets), each with a dedicated 4 MB L2 cache. The hardware assigns workgroup wid to an XCD by wid % NUM_XCDS — round-robin, determ

In standard attention scheduling (SWIZZLE=0), consecutive workgroups cycle through all Q heads before advancing the sequence block. For MHA with many heads, this causes each head's KV data to be spread acros

Spatial scheduling (SWIZZLE=1) re-orders workgroups so that all computation touching a given KV head is pinned to one XCD. Each head's KV tensor stays hot in a single 4 MB slice rather than being replicated

Why GQA is unaffected: models with 8 KV heads (e.g. Llama3, GQA ratio 16) already achieve this locality "accidentally" — round-robin workgroup dispatch across 8 XCDs naturally maps one KV head per XCD. Spati

Changes

`aiter/ops/triton/utils/_triton/pid_preprocessing.py`

Adds remap_workgroup_spatial(wid, NUM_Q_HEADS, NUM_BLOCKS, BATCH, NUM_QUERIES_PER_KV, NUM_XCDS), a unified Triton JIT function with two specialised paths selected at compile time by NUM_QUERIES_PER_KV:

MHA path (NUM_QUERIES_PER_KV == 1): groups ceil(NUM_Q_HEADS / NUM_XCDS) consecutive Q heads onto each XCD. Ordering within each XCD is head-first (all blocks for head i before head i+1).
GQA path (NUM_QUERIES_PER_KV > 1): assigns each KV head exclusively to one XCD so that all Q heads sharing a KV head run on the same XCD. Ordering within each XCD is block-first (all Q heads in the
group process block b together before advancing to b+1), which avoids the causal load-imbalance that head-first ordering creates. Three sub-regimes are handled: HK == NXCD (aligned), HK > NXCD, and HK < NXCD.

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

Adds SWIZZLE: tl.constexpr kernel parameter (string type, values: "default" or "spatial").
Includes tl.static_assert to validate SWIZZLE value at compile time with descriptive error message.
SWIZZLE="default": preserves the existing remap_xcd behaviour (backward-compatible default).
SWIZZLE="spatial": calls remap_workgroup_spatial; NUM_QUERIES_PER_KV is derived from NUM_Q_HEADS // NUM_K_HEADS at compile time, so GQA and MHA specialise independently with zero runtime overhead.

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

Reads AITER_TRITON_MHA_SWIZZLE environment variable at module import (default "default").
Adds mha_set_swizzle(value: Literal["default", "spatial"]) for programmatic control.
Invalid values raise ValueError at import or call time.

Usage

# environment variable (affects all subsequent flash_attn_func calls)
AITER_TRITON_MHA_SWIZZLE=spatial python your_script.py

# programmatic
from aiter.ops.triton.attention.mha import mha_set_swizzle
mha_set_swizzle("spatial")   # enable spatial
mha_set_swizzle("default")   # restore default

Kernel-level benchmark results (AMD MI355X, BF16, batch=1)

MHA — SWIZZLE=0 vs SWIZZLE=1, non-causal

Heads (HQ=HK)	Seq len	SWIZZLE=0 (TFLOPS)	SWIZZLE=1 (TFLOPS)	Delta
128	32K	681	924	+35.7%
128	64K	636	895	+40.7%
128	128K	605	850	+40.5%
64	64K	742	811	+9.2%
16	64K	823	831	+1.0%

End-to-end model results (AMD MI355X, BF16, batch=1, random-init weights)

Whole-model forward-pass latency and socket energy (idle-baseline subtracted) on representative architectures:

Model	Heads	Seq	Speedup	Energy savings
Wan2.2-14B DiT (video diffusion, non-causal)	40	32K	+4.0%	+4.1%
Wan2.2-14B DiT	40	64K	+14.3%	+16.7%
Wan2.2-14B DiT	40	131K	+20.3%	+23.3%
Dense MHA 128h (causal)	128	128K	+10.0%	+10.6%
Dense MHA 128h (causal)	128	512K	+18.0%	+19.8%
Llama3-8B GQA (negative control)	32Q/8KV	32K	+/-0%	+/-0%

Correctness

Verified bit-identical output (rel_max = 0.00e+00) vs SWIZZLE=0 across 17 configurations on MI355X (thor-4):

GQA: HK in {2, 4, 8, 16, 32}, causal + non-causal, batch 1 and 4
MHA: HQ=HK in {16, 32, 40, 64, 128}, causal + non-causal, batch 1 and 2

Test script: op_tests/test_mha_spatial_swizzle.py (included in this branch).

github-actions · 2026-06-25T20:01:07Z

🏷️ CI Guide

Runs automatically on every PR:

✅ Pre-checks (submodule verification, code formatting)
✅ Aiter op tests (gfx942 + gfx950)
✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
`ci:triton-300x`	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
`ci:sglang`	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
`ci:atom`	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
`ci:atom_full`	ATOM accuracy suite for PR and main models from ATOM `models_accuracy.json`
`ci:vllm`	vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
`ci:all`	All standard extended tests (excludes `ci:atom_full`)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 3936 --add-label <label>

brunomazzottiamd

Hi @mc186! Thanks for you PR!

Code looks good but I have two blockers:

change in pid preprocessing function that might affect grouped GEMM kernel
global state mutation in unit test

Please check my comments.

I've also left some suggestions regarding the validation of MHA swizzle modes and an attempt to improve code understanding with strings instead of 0 and 1.

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<NXCD. AITER_SWIZZLE=0 (default) preserves the existing remap_xcd behaviour; AITER_SWIZZLE=1 selects the new spatial mapping. The mode can also be set programmatically via mha_set_swizzle(). Correctness verified bit-identical (rel=0.00e+00) across 17 configs covering GQA ratios 2-16, MHA head counts 16-128, causal/non-causal, and batch sizes 1-4 on MI355X (thor-4).

- 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

mc186 · 2026-06-30T15:52:15Z

All feedback addressed in cb2a345:

Reverted tall_xcds ternary to tl.cast() (blocker - grouped GEMM compatibility)
Added autouse fixture to save/restore _MHA_SWIZZLE state (blocker - global state mutation)
Changed SWIZZLE from {0,1} to {"default","spatial"} throughout
Added tl.static_assert with descriptive message in kernel
Renamed env var to AITER_TRITON_MHA_SWIZZLE with hardened parsing
Added validation in mha_set_swizzle() before state mutation
Updated test docstring to reflect new API

All tests pass (22/22), rebased on current main.

brunomazzottiamd · 2026-06-30T19:03:01Z

Hello @mc186. Why did you close this PR? Do you still want to merge it?

mc186 · 2026-06-30T22:10:13Z

I did not mean to close it! I might have hit close instead of comment...

github-actions · 2026-06-30T22:10:33Z

🏷️ CI Guide

Runs automatically on every PR:

✅ Pre-checks (submodule verification, code formatting)
✅ Aiter op tests (gfx942 + gfx950)
✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
`ci:triton-300x`	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
`ci:sglang`	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
`ci:atom`	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
`ci:atom_full`	ATOM accuracy suite for PR and main models from ATOM `models_accuracy.json`
`ci:vllm`	vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
`ci:all`	All standard extended tests (excludes `ci:atom_full`)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 3936 --add-label <label>

brunomazzottiamd

LGTM! Let's get CI green and then merge it.

cagrikymk

LGTM!

azaidy

LGTM based on @cagrikymk's earlier approval

mc186 · 2026-07-23T20:45:17Z

added formatting fixes for failing Black test

cagrikymk · 2026-07-28T18:14:32Z

@mc186
Can you resolve the merge conflicts? Besides that, LGTM!

mc186 · 2026-07-28T19:33:49Z

resolved conflicts

azaidy

LGTM!

mc186 requested a review from a team June 25, 2026 20:00

zufayu requested a review from amd-ruitang3 June 26, 2026 08:53

brunomazzottiamd requested changes Jun 29, 2026

View reviewed changes

mc186 force-pushed the spatial-attention branch from ece927b to cb2a345 Compare June 30, 2026 15:13

mc186 added 3 commits June 30, 2026 15:15

op_tests: add correctness test for spatial MHA/GQA swizzle

3936719

mc186 closed this Jun 30, 2026

mc186 reopened this Jun 30, 2026

brunomazzottiamd added ci:triton-300x ci:triton-355 labels Jul 1, 2026

brunomazzottiamd previously approved these changes Jul 1, 2026

View reviewed changes

brunomazzottiamd reviewed Jul 1, 2026

View reviewed changes

Comment thread op_tests/triton_tests/attention/test_mha_spatial_swizzle.py

Merge branch 'main' into spatial-attention

7450de8

azaidy requested a review from cagrikymk July 6, 2026 15:47

cagrikymk previously approved these changes Jul 8, 2026

View reviewed changes

mc186 dismissed stale reviews from cagrikymk and brunomazzottiamd via 2ddf0c6 July 8, 2026 19:52

tests: move test_mha_spatial_swizzle to triton_tests/attention

2ddf0c6

mc186 requested a review from cagrikymk July 8, 2026 19:57

zufayu removed the request for review from amd-ruitang3 July 21, 2026 02:40

azaidy previously approved these changes Jul 22, 2026

View reviewed changes

mc186 and others added 2 commits July 23, 2026 13:33

Merge branch 'main' into spatial-attention

77e9bb8

style: apply Black formatting

620fdd0

mc186 dismissed azaidy’s stale review via 620fdd0 July 23, 2026 20:42

mc186 requested a review from azaidy July 23, 2026 20:45

Merge branch 'main' into spatial-attention

b635fed

jhinpan mentioned this pull request Jul 29, 2026

perf(attention): keep a head's Q-blocks on one XCD in the dualwave mapping ROCm/FlyDSL#919

Merged

Merge branch 'main' into spatial-attention

15e2875

cagrikymk previously approved these changes Jul 29, 2026

View reviewed changes

style: fix ruff 0.16 import sorting and style issues

5bc064d

mc186 dismissed cagrikymk’s stale review via 5bc064d July 29, 2026 18:39

mc186 requested a review from cagrikymk July 29, 2026 18:41

cagrikymk approved these changes Jul 29, 2026

View reviewed changes

cagrikymk and others added 2 commits July 29, 2026 22:02

Merge branch 'main' into spatial-attention

973bd2f

Merge branch 'main' into spatial-attention

d3202d4

azaidy approved these changes Aug 5, 2026

View reviewed changes

azaidy merged commit 35458c3 into ROCm:main Aug 5, 2026
102 of 105 checks passed

Ragua1 mentioned this pull request Sep 10, 2026

[Triton/Gluon] [Config] Add a tuned gfx1101 MHA config, split small_head/default #4493

Open

Uh oh!

Conversation

mc186 commented Jun 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

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

Motivation

Changes

aiter/ops/triton/utils/_triton/pid_preprocessing.py

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

aiter/ops/triton/attention/mha.py

Usage

Kernel-level benchmark results (AMD MI355X, BF16, batch=1)

MHA — SWIZZLE=0 vs SWIZZLE=1, non-causal

End-to-end model results (AMD MI355X, BF16, batch=1, random-init weights)

Correctness

Uh oh!

github-actions Bot commented Jun 25, 2026

🏷️ CI Guide

Uh oh!

brunomazzottiamd left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

mc186 commented Jun 30, 2026

Uh oh!

brunomazzottiamd commented Jun 30, 2026

Uh oh!

mc186 commented Jun 30, 2026

Uh oh!

github-actions Bot commented Jun 30, 2026

🏷️ CI Guide

Uh oh!

brunomazzottiamd left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

cagrikymk left a comment

Choose a reason for hiding this comment

Uh oh!

azaidy left a comment

Choose a reason for hiding this comment

Uh oh!

mc186 commented Jul 23, 2026

Uh oh!

cagrikymk commented Jul 28, 2026

Uh oh!

mc186 commented Jul 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

azaidy left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

mc186 commented Jun 25, 2026 •

edited

Loading

`aiter/ops/triton/utils/_triton/pid_preprocessing.py`

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

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

mc186 commented Jul 28, 2026 •

edited

Loading