[Triton/Gluon] Add 3D neighborhood flash attention (na3d_flash)

[jjuvonen-amd]

Motivation

Video diffusion models such as LTX-2.5 use 3D neighborhood attention in
the VAE diffusion decoder: each token at grid position (t, h, w) attends
only to its (KT x KH x KW) neighbors rather than the full sequence. This
PR adds a BF16 MFMA flash-style kernel for that operation targeting gfx942
and gfx950, plus a reference, tests and benchmarks.

Use case example
LTX-2.5 Diffusion decoder attention processor in xDiT inference framework.

Technical Details

Six new files:

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

_na3d_flash_fwd: flash attention inner kernel.

• Inward border-shift: every window contains exactly KT x KH x KW keys
• Q loaded once into registers; K/V iterated over KT x KH rows
• Online softmax in log2 space (exp2, log2(e) folded into scale factor)
• Flat KT x KH loop with static trip count enables Triton's software pipeliner

aiter/ops/triton/attention/na3d_flash.py

na3d_flash_attn(q, k, v, kernel_size) -
public launcher. Input/output shape: (B, T, H, W, NH, HD) bfloat16,
Q pre-scaled by head_dim**-0.5.
Tuned configs provided for gfx950/gfx942 for ltx-2.5 model default shapes.

aiter/op_tests/triton_tests/attention/test_na3d_flash.py

• _na3d_sdpa_exact(q, k, v, kernel_size, w_offset=0) - Reference method. Exact per-query
  neighborhood SDPA. Groups queries by shared relative-window geometry and
  calls F.scaled_dot_product_attention per group with a cached additive mask.
  With w_offset=0 and full q/k/v this is equivalent to a per-query FP32
  reference but runs on the GPU. With w_offset>0, q is a W sub-slice and
  k/v are the full tensors; neighborhood bounds use global W coordinates.

• test_na3d_flash: 9 shapes (5 LTX-2.5 default + 5 edge cases); reference
  is _na3d_sdpa_exact, atol=5e-2

• test_na3d_sdpa_exact_vs_ref: 6 small shapes; cross-checks _na3d_sdpa_exact
  against the FP32 per-query loop

op_tests/op_benchmarks/triton/bench_na3d_flash.py

• Performance benchmark for 3D neighborhood flash attention (na3d_flash) that compares against _na3d_sdpa_exact reference method from aiter/op_tests/triton_tests/attention/test_na3d_flash.py.

• Measures TFLOPS and TB/s per shape

• aiter/ops/triton/configs/gfx950/triton/attention/na3d_flash/DEFAULT.json: Tuned config for gfx950

• aiter/ops/triton/configs/gfx942/triton/attention/na3d_flash/DEFAULT.json: Tuned config for gfx942

Test Plan

Run pytest on gfx950:

pytest op_tests/triton_tests/attention/test_na3d_flash.py -v

Test Result

All tests pass on gfx950 (MI350X) in ~30 s.

[github-actions]

🏷️ 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:gfx1250-ffm-triton	Run the five-shard gfx1250 FFM Triton test suite
ci:triton-300x	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
multigpu	Aiter multi-GPU tests on the 8-GPU runner
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 5230 --add-label <label>

PR title tags & labels:
Component tags ([Triton/Gluon], [HIP], [CK], [ASM], ...) are added to the PR title and as PR labels automatically from the changed files and re-synced on every push — change-type tags like [fix]/[Perf], op tags like [MLA], and human labels (ci:*) are left untouched. Add the no-auto-title label to opt this PR out.

[Copilot]

🟡 Changes recommended

There are confirmed correctness/robustness issues in the correction/reference tiling logic and kernel-size validation, plus missing kernel repr= wiring needed for consistency with existing tracing conventions.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Pull request overview

Adds a new Triton-based BF16 “flash-style” 3D neighborhood attention implementation (na3d_flash) for video diffusion decoder-style workloads, alongside an exact PyTorch SDPA reference path and a dedicated correctness/benchmark test suite under the attention test category.

Changes:

• Introduces a new autotuned Triton forward kernel for 3D neighborhood attention (_na3d_flash_fwd).
• Adds a public launcher na3d_flash_attn(...) plus an exact grouped-SDPA reference helper (_na3d_sdpa_exact) used for validation and optional rightmost-block correction.
• Adds pytest coverage and a benchmark harness for representative LTX-2.5 and edge-case shapes.

File summaries

File	Description
aiter/ops/triton/_triton_kernels/attention/na3d_flash.py	New autotuned Triton forward kernel for 3D neighborhood flash attention.
aiter/ops/triton/attention/na3d_flash.py	Public Python launcher + exact grouped SDPA reference/correction utilities.
op_tests/triton_tests/attention/test_na3d_flash.py	New correctness tests and benchmarking for na3d_flash and the exact reference.

Review details

• Files reviewed: 3/3 changed files
• Comments generated: 6
• Review effort level: Lite

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Copilot]

🔵 Needs a closer look

The exact-reference tiling logic can underestimate KV footprint for rightmost-block correction (risking OOM/timeouts), and the new kernel/tests violate established wrapper/kernel-boundary + kernel repr conventions.

Review details

Suppressed comments (5)

Previously missed (2) — in code that hasn't changed since the last review.

aiter/ops/triton/attention/na3d_flash.py:65

• _na3d_sdpa_mask does int(en.max()) where en is a CUDA tensor; converting that to int forces a device sync. Since ends is already a Python tuple, you can compute the max end on CPU and avoid the sync.
  aiter/ops/triton/attention/na3d_flash.py:127
• In _na3d_sdpa_exact, the tiling _cost() and split heuristic uses W_sub (and eff_kw = min(KW, W_sub)) to estimate KV footprint. When correct_rightmost_block=True, q_right can have W_sub < KW, but KV still spans a full KW window in global-W coordinates; this can severely underestimate nk, leading to oversized tiles, large intermediate k_s/v_s tensors, and potential OOM/timeouts.

aiter/ops/triton/_triton_kernels/attention/na3d_flash.py:18

• New launchable Triton kernels should set a config-aware repr via make_kernel_repr (see aiter/ops/triton/README.md and the repo instructions). _na3d_flash_fwd currently lacks a repr, which makes profiling/trace names harder to interpret.

import triton
import triton.language as tl

aiter/ops/triton/_triton_kernels/attention/na3d_flash.py:53

• Hook the kernel repr into the JIT decorator so Triton uses it for compiled-kernel naming.

@triton.jit

op_tests/triton_tests/attention/test_na3d_flash.py:279

• Tests should not import kernel bodies from aiter/ops/triton/_triton_kernels/... (wrapper-only boundary). Here the test imports _na3d_flash_fwd directly just to read best_config; instead, access the already-imported symbol from the public wrapper module.

    from aiter.ops.triton._triton_kernels.attention.na3d_flash import _na3d_flash_fwd
    actual_bq = _na3d_flash_fwd.best_config.kwargs["BLOCK_Q"]
    W_last = ((W - 1) // actual_bq) * actual_bq

• Files reviewed: 3/3 changed files
• Comments generated: 0 new
• Review effort level: Lite

[Copilot]

🟡 Changes recommended

The current kernel program-id to (t,h,w) mapping can cross W-row boundaries and break correctness due to scalar-reduced (t,h) window starts, and the launcher/correction paths also need alignment/guarding for correctness and predictable performance.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Review details

Suppressed comments (4)

Previously missed (2) — in code that hasn't changed since the last review.

aiter/ops/triton/attention/na3d_flash.py:50

• _na3d_sdpa_mask caches GPU-resident additive masks keyed by (rel_bounds, dtype, device). Since the tile heuristic only bounds nq*nk to 2**25, a single cached mask can be ~64MB (bf16), and maxsize=256 can retain many GB of GPU memory across calls/shapes. Consider reducing the cache size (or disabling it) to avoid unexpected OOMs when correct_rightmost_block/_na3d_sdpa_exact are used repeatedly.
  aiter/ops/triton/attention/na3d_flash.py:251
• The launcher grid currently uses cdiv(SEQ, BLOCK_Q), which matches the old 1D q_idx mapping. If the kernel is updated to a (t,h)-row + w-block program-id mapping (to avoid cross-row blocks), the grid must be updated accordingly or pid0 will exceed the valid (THn_wblocks) range and produce incorrect indexing/masking.

This issue also appears on line 273 of the same file.

aiter/ops/triton/_triton_kernels/attention/na3d_flash.py:103

• INF_I is used as a sentinel for masked-off lanes when reducing t_ws/h_ws/w_lo, but the hardcoded value 999999 can be smaller than valid window starts when T/H/W exceed ~1e6, causing wrong window selection. Use a sentinel guaranteed larger than any possible start (e.g., 1<<30).

    INF_I = 999999

aiter/ops/triton/attention/na3d_flash.py:277

• The docstring/tests assume rightmost-block correction only runs when the number of corrected queries is small ("TH(W-W_last) should be <= 100_000"), but the implementation will always recompute the entire last W-block when correct_rightmost_block=True. For large T/H this can trigger an unexpectedly expensive fallback path.

        actual_bq = _na3d_flash_fwd.best_config.kwargs["BLOCK_Q"]
        W_last = ((W - 1) // actual_bq) * actual_bq
        if W_last > 0:
            q_right = q[:, :, :, W_last:, :, :]
            out_right = _na3d_sdpa_exact(q_right, k, v, kernel_size, w_offset=W_last)

• Files reviewed: 3/3 changed files
• Comments generated: 1
• Review effort level: Lite

[Copilot]

🔵 Needs a closer look

The new kernel and correction path depend on mutable Triton autotuner state (best_config) in a way that can be racy and the exact-path mask caching can retain large GPU tensors, both of which warrant careful human review.

Review details

Suppressed comments (1)

Previously missed (1) — in code that hasn't changed since the last review.

aiter/ops/triton/attention/na3d_flash.py:55

• _na3d_sdpa_mask caches full GPU-resident masks (device is part of the cache key). These masks can be large (tiles are chosen up to ~2**25 nq*nk), so lru_cache(maxsize=256) can retain multiple 10s–100s of MB tensors and create GPU memory growth across model runs / long-lived processes.

• Files reviewed: 3/3 changed files
• Comments generated: 2
• Review effort level: Lite

[Copilot]

🟡 Changes recommended

There are correctness and operational issues (KW-dependent config validity/pruning, missing runtime gating/bounds, leaked global test state) plus a mismatch between the PR description and the exposed API/features.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Review details

Suppressed comments (4)

Previously missed (3) — in code that hasn't changed since the last review.

aiter/ops/triton/attention/na3d_flash.py:38

• The PR description documents a public na3d_flash_attn(..., correct_rightmost_block=False) option and an _na3d_sdpa_exact(...) reference inside this module, plus tests for the correction path. The current module only exposes na3d_flash_attn(q,k,v,kernel_size) and contains no correction/reference implementation, so either the description needs updating or the missing API/features need to be added to match it.
  aiter/ops/triton/attention/na3d_flash.py:72
• The launcher does not gate execution to the advertised supported architectures (gfx942/gfx950), and it also doesn’t fail fast when KW exceeds what the available BLOCK_Q/BLOCK_KV configs can cover. Without these guards, callers on other GPUs (or with larger KW) may hit compilation/runtime failures or incorrect results (if an invalid config slips through). Add explicit arch and KW bounds checks alongside the existing asserts.
  aiter/ops/triton/attention/na3d_flash.py:59
• This log call eagerly builds an f-string on every invocation. AiterTritonLogger is designed to support lazy %-style formatting via *args (see utils/logger.py), so avoid f-strings here to reduce overhead when INFO logging is disabled.

op_tests/triton_tests/attention/test_na3d_flash.py:388

• If the module-level torch.set_default_device("cuda") is removed (to avoid leaking global state), main() will start constructing tensors on CPU by default and the benchmark path will fail. Set the default device inside main() so running this file as a script still works.

def main():
    if get_gfx() not in SUPPORTED_GFX:
        aiter.logger.warning("na3d_flash unsupported on %s; skipping", get_gfx())
        return

• Files reviewed: 3/3 changed files
• Comments generated: 2
• Review effort level: Lite

[Copilot]

🔵 Needs a closer look

The new architecture-specific GPU kernel has substantial numerical and performance risk that requires final human review and hardware validation.

Review details

• Files reviewed: 6/6 changed files
• Comments generated: 0 new
• Review effort level: Balanced

[jjuvonen-amd]

Please don't add autotuning in the loop. Check how the other kernels do tuning (they use a config file with the tunings hardcoded).

Thank you for reviewing @vgokhale I now removed the opt-in @triton.autotune entirely and switched to the hardcoded-config pattern.

[Copilot]

Copilot review overview

🔵 Needs a closer look

GPU-specific correctness and performance require head-matched validation on both supported architectures.

Review effort: Balanced
Findings: None

[vgokhale]

@jjuvonen-amd please address the review comments or close this PR if it is no longer needed. Thanks!

[Copilot]

Copilot review overview

🔵 Needs a closer look

The hardware-specific kernel requires final human review and head-matched validation across both supported architectures.

Review effort: Balanced
Findings: None

Previously missed (1)

In code that hasn't changed since last review

Clarify masks are deduplicated, not cached across calls

op_tests/​triton_tests/​attention/​test_na3d_flash.py:149

This says the masks are cached, but _na3d_sdpa_mask is rebuilt for every unique geometry on each call (as its own docstring explains). Describe the masks as deduplicated instead so the reference's memory/performance behavior is accurate.

[Copilot]

Copilot review overview

🟢 Approval recommended

The implementation, configuration, validation, tests, and benchmark form a consistent and complete kernel addition with no unresolved issues found.

Review effort: Balanced
Findings: None

[jjuvonen-amd]

Copilot review overview

🔵 Needs a closer look

GPU-specific correctness and performance require head-matched validation on both supported architectures.

Review effort: Balanced Findings: None

Validated on both supported architectures against this head.

Correctness: gfx950 (MI350X) and gfx942 (MI300X)

Both run the full suite (20/20 passing on each.

Both also pass the CI autotuning gate TRITON_PRINT_AUTOTUNING=1 yields 0 Triton autotuning for function lines, since this PR replaced @triton.autotune with the per-arch JSON config.

Tuning

The two archs genuinely differ. I swept the candidate tiles per KW family on each arch, validating every candidate against the exact reference before timing it:

family	gfx950	gfx942
na3d_flash_small_kw (BLOCK_Q=16, BLOCK_KV=32)	w4 / s2	w4 / s2
na3d_flash_large_kw (BLOCK_Q=32, BLOCK_KV=64)	w4 / s2	w4 / s4

The large-KW divergence is a measured effect: on gfx942 num_stages=4 is 66.5us vs 70.2us for num_stages=2 (5.6%), while on gfx950 the same comparison is within 0.9% and I left it at s2. Each DEFAULT.json now carries values measured on its own architecture.

Performance

bench_na3d_flash.py, bf16, vs the grouped-SDPA path the eager processor falls back to; times include the wrapper's v.contiguous() staging copy, since the benchmark reproduces the decoder's fused-QKV layout:

shape (B×T×H×W, NH=4/8/32, HD=64)	kernel	gfx950 triton / sdpa	gfx942 triton / sdpa
1×18×32×48	3,7,7	1249 / 9946 us (8.0×)	2228 / 10071 us (4.5×)
1×36×64×96	3,5,5	1910 / 22153 us (11.6×)	3448 / 25565 us (7.4×)
1×40×96×96	3,5,5	3319 / 34811 us (10.5×)	5419 / 37428 us (6.9×)
1×16×80×32	11,11,11	1153 / 3345 us (2.9×)	2010 / 4157 us (2.1×)
1×79×192×192	11,11,11	99143 / 178837 us (1.8×)	152070 / 290311 us (1.9×)
1×4×8×128	3,5,27	40 / 363 us (9.1×)	75 / 399 us (5.3×)

---

@vgokhale all review comments have now been addressed. The remaining change request is outdated and has been addressed and resolved with the per-arch JSON configs.

[jjuvonen-amd]

@Boss2002n could you please review if requested changes have been addressed, thanks. This is currently blocking merge.

[Boss2002n]

LGTM! Lets pass the CI before merge

[Copilot]

Copilot review overview

🔵 Needs a closer look

Hardware-specific correctness and performance require completed latest-head validation on both gfx942 and gfx950.

Review effort: Balanced
Findings: None

[Copilot]

Copilot review overview

🔵 Needs a closer look

The new architecture-specific GPU kernel still needs completed gfx942/gfx950 CI and human correctness and performance review.

Review effort: Balanced
Findings: 1

Open (1)

• Update outdated comment on fast subset edge-case coverage · New

[Copilot]

Copilot review overview

🔵 Needs a closer look

The novel GPU kernel requires completed head-matched correctness and performance validation on both supported architectures.

Review effort: Balanced
Findings: None

Resolved since last review (1)

• Update outdated comment on fast subset edge-case coverage

[Copilot]

Copilot review overview

🔵 Needs a closer look

The new architecture-specific GPU kernel requires final hardware-backed correctness and performance review.

Review effort: Balanced
Findings: None

[Copilot]

Copilot review overview

🔵 Needs a closer look

The architecture-specific GPU kernel requires final human review and completed gfx942/gfx950 correctness and performance validation.

Review effort: Balanced
Findings: None

[Copilot]

Copilot review overview

🔵 Needs a closer look

The new GPU kernel and architecture-specific tuning require final human review and reproducible gfx942/gfx950 validation.

Review effort: Balanced
Findings: None

[Copilot]

Copilot review overview

🟡 Changes recommended

The reference tiling heuristic can exceed its memory cap substantially and OOM for valid custom shapes.

Review effort: Balanced
Findings: 1

Open (1)

• Tiling loop stops when max selects an unsplittable axis · New