From 9b301cbc77954aa011bd4354b6b14fa1f5a02f94 Mon Sep 17 00:00:00 2001 From: jjuvonen-amd Date: Thu, 3 Sep 2026 06:01:20 +0000 Subject: [PATCH 01/48] add 3d neighborhood flash attention Signed-off-by: jjuvonen-amd --- .../_triton_kernels/attention/na3d_flash.py | 170 +++++++++ aiter/ops/triton/attention/na3d_flash.py | 265 +++++++++++++ .../triton_tests/attention/test_na3d_flash.py | 357 ++++++++++++++++++ 3 files changed, 792 insertions(+) create mode 100644 aiter/ops/triton/_triton_kernels/attention/na3d_flash.py create mode 100644 aiter/ops/triton/attention/na3d_flash.py create mode 100644 op_tests/triton_tests/attention/test_na3d_flash.py diff --git a/aiter/ops/triton/_triton_kernels/attention/na3d_flash.py b/aiter/ops/triton/_triton_kernels/attention/na3d_flash.py new file mode 100644 index 00000000000..e8a6614a0e1 --- /dev/null +++ b/aiter/ops/triton/_triton_kernels/attention/na3d_flash.py @@ -0,0 +1,170 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. +"""Triton flash-style 3D neighborhood attention. + +Each query token attends to the (KT x KH x KW) neighborhood centered on its +(t, h, w) grid position, with inward border-shift so every window has the full +KT x KH x KW keys. The kernel processes BLOCK_Q consecutive queries per program; +all queries in one block share the same (t, h) row (requires W >= BLOCK_Q). + +Algorithm: +- Q is loaded once into registers and held for all KT x KH inner iterations. +- Each iteration loads one (t_kv, h_kv) row of K and V at the shared W window. +- Running online softmax in log2 space (exp2, log2(e) folded into scale). +- Autotuned over BLOCK_Q (16, 32), num_stages (2, 3, 4), num_warps (4, 8). +""" + +import triton +import triton.language as tl + +# Autotune configs. +# BLOCK_KV = next_pow2(BLOCK_Q + KW - 1) covers the union of all BLOCK_Q +# queries' W windows in one chunk: +# BLOCK_Q=16 -> BLOCK_KV=32 +# BLOCK_Q=32 -> BLOCK_KV=64 +# The W >= BLOCK_Q constraint is enforced by the pruner below. +_CONFIGS = [ + triton.Config({"BLOCK_Q": 16, "BLOCK_KV": 32, "num_stages": 2}, num_warps=4), + triton.Config({"BLOCK_Q": 16, "BLOCK_KV": 32, "num_stages": 2}, num_warps=8), + triton.Config({"BLOCK_Q": 16, "BLOCK_KV": 32, "num_stages": 3}, num_warps=4), + triton.Config({"BLOCK_Q": 16, "BLOCK_KV": 32, "num_stages": 3}, num_warps=8), + triton.Config({"BLOCK_Q": 16, "BLOCK_KV": 32, "num_stages": 4}, num_warps=4), + triton.Config({"BLOCK_Q": 16, "BLOCK_KV": 32, "num_stages": 4}, num_warps=8), + # BLOCK_Q=32 halves program count but doubles BLOCK_KV (more masked compute). + triton.Config({"BLOCK_Q": 32, "BLOCK_KV": 64, "num_stages": 2}, num_warps=4), + triton.Config({"BLOCK_Q": 32, "BLOCK_KV": 64, "num_stages": 2}, num_warps=8), +] + + +def _prune_configs(configs, named_args, **kwargs): + """Remove configs where BLOCK_Q > W. + + Ensures all queries in one block share the same (t, h) row (W >= BLOCK_Q required). + """ + W = named_args["W"] + return [c for c in configs if c.kwargs["BLOCK_Q"] <= W] + + +@triton.autotune( + configs=_CONFIGS, + key=["KT", "KH", "KW", "HD", "W"], + prune_configs_by={"early_config_prune": _prune_configs}, +) +@triton.jit +def _na3d_flash_fwd( + Q_ptr, + K_ptr, + V_ptr, + Out_ptr, + stride_bnh, + stride_seq, + T, + H, + W, + SEQ, + HD: tl.constexpr, # head dimension: any power-of-2 value + KT: tl.constexpr, # neighborhood depth: constexpr -> static loop bound + KH: tl.constexpr, # neighborhood height: constexpr -> static loop bound + KW: tl.constexpr, # neighborhood width: used in W-direction masking + BLOCK_Q: tl.constexpr, + BLOCK_KV: tl.constexpr, # power-of-2, >= BLOCK_Q + KW - 1 +): + """Flash-attention inner kernel. Requires W >= BLOCK_Q (same (t,h) row per block). + + See _prune_configs for the W >= BLOCK_Q constraint. + """ + pid_q = tl.program_id(0) + pid_bnh = tl.program_id(1) + HW = H * W + + q_start = pid_q * BLOCK_Q + q_offs = tl.arange(0, BLOCK_Q) + q_idx = q_start + q_offs + q_mask = q_idx < SEQ + + q_t = q_idx // HW + q_h = (q_idx % HW) // W + q_w = q_idx % W + + # Inward-shifted centered neighborhood window starts. + q_t_ws = tl.minimum(tl.maximum(q_t - KT // 2, 0), T - KT) + q_h_ws = tl.minimum(tl.maximum(q_h - KH // 2, 0), H - KH) + q_w_ws = tl.minimum(tl.maximum(q_w - KW // 2, 0), W - KW) + + # Scalar window starts (shared across all BLOCK_Q queries in this block). + INF_I = 999999 + t_ws = tl.min(tl.where(q_mask, q_t_ws, INF_I)) + h_ws = tl.min(tl.where(q_mask, q_h_ws, INF_I)) + w_lo = tl.min(tl.where(q_mask, q_w_ws, INF_I)) + + hd_offs = tl.arange(0, HD) + kv_offs = tl.arange(0, BLOCK_KV) + + # Cast to int64: pid_bnh (int32) x stride_bnh can exceed int32 range. + base = pid_bnh.to(tl.int64) * stride_bnh + kv_w = w_lo + kv_offs + kv_ok = kv_w < W # W-boundary guard, constant across the KT x KH loop + + # Q loaded once into registers for all KT x KH iterations. + Q_tile = tl.load( + Q_ptr + base + q_idx[:, None] * stride_seq + hd_offs[None, :], + mask=q_mask[:, None], + other=0.0, + ).to(tl.bfloat16) + + m_i = tl.full((BLOCK_Q,), -3e38, dtype=tl.float32) + l_i = tl.zeros((BLOCK_Q,), dtype=tl.float32) + acc = tl.zeros((BLOCK_Q, HD), dtype=tl.float32) + + # Flat KT x KH loop with static trip count - enables Triton's software pipeliner. + # valid_w, kv_ok, and q_mask are loop-invariant; the compiler's LICM handles them + # without explicit pre-computation (pre-computing would raise VGPR pressure and + # reduce occupancy on CDNA). + for kv_idx in range(KT * KH): + dt = kv_idx // KH + dh = kv_idx % KH + t_kv = t_ws + dt + h_kv = h_ws + dh + + row_base = t_kv * HW + h_kv * W + kv_flat = row_base + kv_w + + K_T = tl.load( + K_ptr + base + kv_flat[None, :] * stride_seq + hd_offs[:, None], + mask=kv_ok[None, :], + other=0.0, + ).to(tl.bfloat16) + + V_tile = tl.load( + V_ptr + base + kv_flat[:, None] * stride_seq + hd_offs[None, :], + mask=kv_ok[:, None], + other=0.0, + ).to(tl.bfloat16) + + # QK in log2 space: fold log2(e) into the scale to use exp2 throughout. + scores = tl.dot(Q_tile, K_T, out_dtype=tl.float32) * 1.4426950408889634 + + valid_w = (kv_w[None, :] >= q_w_ws[:, None]) & ( + kv_w[None, :] < q_w_ws[:, None] + KW + ) + scores = tl.where( + valid_w & kv_ok[None, :] & q_mask[:, None], scores, float("-inf") + ) + + m_new = tl.maximum(m_i, tl.max(scores, axis=1)) + exp_s = tl.exp2(scores - m_new[:, None]) + exp_diff = tl.exp2(m_i - m_new) + l_i = l_i * exp_diff + tl.sum(exp_s, axis=1) + acc = acc * exp_diff[:, None] + tl.dot( + exp_s.to(tl.bfloat16), V_tile, out_dtype=tl.float32 + ) + m_i = m_new + + safe_l = tl.where(l_i > 0, l_i, 1.0) + acc = acc / safe_l[:, None] + + tl.store( + Out_ptr + base + q_idx[:, None] * stride_seq + hd_offs[None, :], + acc.to(tl.bfloat16), + mask=q_mask[:, None], + ) diff --git a/aiter/ops/triton/attention/na3d_flash.py b/aiter/ops/triton/attention/na3d_flash.py new file mode 100644 index 00000000000..e7405ed3402 --- /dev/null +++ b/aiter/ops/triton/attention/na3d_flash.py @@ -0,0 +1,265 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. +"""3D neighborhood flash attention launcher. + +Each query token at grid position (t, h, w) attends to the KT x KH x KW neighborhood +centered on it, with inward border-shift so every window has exactly KT x KH x KW keys. + +Usage:: + + from aiter.ops.triton.attention.na3d_flash import na3d_flash_attn + + out = na3d_flash_attn(q, k, v, kernel_size=(11, 11, 11)) + # q, k, v: (B, T, H, W, NH, HD) bfloat16, Q already scaled by head_dim**-0.5 + # out : (B, T, H, W, NH, HD) bfloat16 +""" + +from __future__ import annotations + +import functools +import math + +import torch +import torch.nn.functional as F +import triton + +from aiter.ops.triton._triton_kernels.attention.na3d_flash import _na3d_flash_fwd +from aiter.ops.triton.utils.logger import AiterTritonLogger + +_LOGGER = AiterTritonLogger() + + +# --------------------------------------------------------------------------- +# Exact per-query neighborhood SDPA (generic; used for rightmost-block correction) +# --------------------------------------------------------------------------- + +def _window_bounds_1d(length: int, kernel: int) -> tuple[list[int], list[int]]: + """Inward-shifted window (start, end) for each position along one axis.""" + lo = length - kernel + half = kernel // 2 + starts: list[int] = [] + ends: list[int] = [] + for i in range(length): + s = min(max(i - half, 0), lo) + starts.append(s) + ends.append(s + kernel) + return starts, ends + + +@functools.lru_cache(maxsize=256) +def _na3d_sdpa_mask( + rel_bounds: tuple[tuple[tuple[int, ...], tuple[int, ...]], ...], + dtype: torch.dtype, + device: torch.device, +) -> torch.Tensor: + """Cached additive ``[1, 1, Nq, Nk]`` mask for one tile-geometry group. + + Masks are reused across repeated calls with the same geometry (e.g. all + 8 diff_block layers at W=192 share the same rel_bounds). + """ + bools = [] + for starts, ends in rel_bounds: + st = torch.tensor(starts, device=device) + en = torch.tensor(ends, device=device) + kj = torch.arange(int(en.max()), device=device) + bools.append((kj[None, :] >= st[:, None]) & (kj[None, :] < en[:, None])) + visible = ( + bools[0][:, None, None, :, None, None] + & bools[1][None, :, None, None, :, None] + & bools[2][None, None, :, None, None, :] + ) + nq = math.prod(visible.shape[:3]) + nk = math.prod(visible.shape[3:]) + mask = torch.zeros((nq, nk), dtype=dtype, device=device) + mask.masked_fill_(~visible.reshape(nq, nk), torch.finfo(dtype).min) + return mask.reshape(1, 1, nq, nk) + + +def _na3d_sdpa_exact( + q_right: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + kernel_size: tuple[int, int, int], + w_offset: int = 0, +) -> torch.Tensor: + """Exact per-query 3D neighborhood attention via grouped tiled SDPA. + + When called with the full q tensor and w_offset=0, computes exact neighborhood + attention for all W positions. When called with a W sub-slice and non-zero + w_offset, K/V are drawn from the full k/v tensors using global W coordinates - + used by correct_rightmost_block. + + Args: + q_right : ``(B, T, H, W_sub, NH, HD)`` bfloat16, Q pre-scaled. + k, v : ``(B, T, H, W_full, NH, HD)`` bfloat16. + kernel_size: ``(KT, KH, KW)`` neighborhood window. + w_offset : Global W start of q_right; 0 for a full-tensor call. + + Returns: + ``(B, T, H, W_sub, NH, HD)`` bfloat16. + """ + B, T, H, W_sub, NH, HD = q_right.shape + W_full = k.shape[3] + KT, KH, KW = kernel_size + device = q_right.device + + bt = _window_bounds_1d(T, min(KT, T)) + bh = _window_bounds_1d(H, min(KH, H)) + bw = _window_bounds_1d(W_full, min(KW, W_full)) # global W bounds + + # Pick tile sizes relative to q_right's actual dims (W_sub is small). + eff_kt, eff_kh, eff_kw = min(KT, T), min(KH, H), min(KW, W_sub) + + def _cost(ts: list[int]) -> int: + nq = math.prod(ts) + nk = math.prod( + min(d, t + kk - 1) + for t, kk, d in zip(ts, [eff_kt, eff_kh, eff_kw], [T, H, W_sub]) + ) + return nq * nk + + tiles = [T, H, W_sub] + while _cost(tiles) > 2**25 and max(tiles) > 1: + i = max(range(3), key=lambda a: tiles[a] / [eff_kt, eff_kh, eff_kw][a]) + if tiles[i] <= 1: + break + tiles[i] = max(1, (tiles[i] + 1) // 2) + tile_t, tile_h, tile_w = tiles + + groups: dict = {} + for t0 in range(0, T, tile_t): + t1 = min(t0 + tile_t, T) + rt0, rt1 = bt[0][t0], bt[1][t1 - 1] + rel_t = ( + tuple(s - rt0 for s in bt[0][t0:t1]), + tuple(e - rt0 for e in bt[1][t0:t1]), + ) + for h0 in range(0, H, tile_h): + h1 = min(h0 + tile_h, H) + rh0, rh1 = bh[0][h0], bh[1][h1 - 1] + rel_h = ( + tuple(s - rh0 for s in bh[0][h0:h1]), + tuple(e - rh0 for e in bh[1][h0:h1]), + ) + for w0 in range(0, W_sub, tile_w): + w1 = min(w0 + tile_w, W_sub) + w0g, w1g = w_offset + w0, w_offset + w1 - 1 + rw0, rw1 = bw[0][w0g], bw[1][w1g] + rel_w = ( + tuple(s - rw0 for s in bw[0][w0g : w1g + 1]), + tuple(e - rw0 for e in bw[1][w0g : w1g + 1]), + ) + groups.setdefault((rel_t, rel_h, rel_w), []).append( + ( + (slice(t0, t1), slice(h0, h1), slice(w0, w1)), + (slice(rt0, rt1), slice(rh0, rh1), slice(rw0, rw1)), + ) + ) + + out = torch.empty((B, T, H, W_sub, NH, HD), device=device, dtype=v.dtype) + kv_budget = 2**28 + for rel, tiles_list in groups.items(): + mask = _na3d_sdpa_mask(rel, q_right.dtype, device) + nq, nk = mask.shape[2], mask.shape[3] + g_max = max(1, kv_budget // max(1, B * NH * nk * HD * 2)) + qs0, _ = tiles_list[0] + tq = qs0[0].stop - qs0[0].start + th_sz = qs0[1].stop - qs0[1].start + tw_sz = qs0[2].stop - qs0[2].start + for c0 in range(0, len(tiles_list), g_max): + chunk = tiles_list[c0 : c0 + g_max] + g = len(chunk) + q_s = torch.stack([q_right[:, qs[0], qs[1], qs[2]] for qs, _ in chunk]) + k_s = torch.stack([k[:, rs[0], rs[1], rs[2]] for _, rs in chunk]) + v_s = torch.stack([v[:, rs[0], rs[1], rs[2]] for _, rs in chunk]) + q_s = q_s.permute(0, 1, 5, 2, 3, 4, 6).reshape(g * B, NH, nq, HD) + k_s = k_s.permute(0, 1, 5, 2, 3, 4, 6).reshape(g * B, NH, nk, HD) + v_s = v_s.permute(0, 1, 5, 2, 3, 4, 6).reshape(g * B, NH, nk, HD) + o = F.scaled_dot_product_attention(q_s, k_s, v_s, attn_mask=mask, scale=1.0) + o = o.view(g, B, NH, tq, th_sz, tw_sz, HD).permute(0, 1, 3, 4, 5, 2, 6) + for i, (qs, _) in enumerate(chunk): + out[:, qs[0], qs[1], qs[2]] = o[i] + return out + + +# --------------------------------------------------------------------------- +# Public launcher +# --------------------------------------------------------------------------- + +def na3d_flash_attn( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + kernel_size: tuple[int, int, int], + correct_rightmost_block: bool = False, +) -> torch.Tensor: + """Flash-style 3D neighborhood attention. + + Args: + q, k, v : ``(B, T, H, W, NH, HD)`` bfloat16, channels-last. + Q is expected to be pre-scaled by ``head_dim ** -0.5``. + kernel_size : ``(KT, KH, KW)`` neighborhood window. + correct_rightmost_block : When True, overwrite the rightmost BLOCK_Q W + positions with per-query exact SDPA (via _na3d_sdpa_exact). + The flash kernel loads ``BF16(0)`` for out-of-bounds K/V positions in the + last W-block (``kv_ok=False``), which causes a ~0.002 rounding divergence + from AOTriton at those positions. Off by default. + + Returns: + Output tensor ``(B, T, H, W, NH, HD)`` bfloat16. + + Notes: + ``W >= 16`` is required so that all queries in a BLOCK_Q=16 program share + the same (t, h) grid row. The autotune pruner enforces this. + """ + B, T, H, W, NH, HD = q.shape + KT, KH, KW = kernel_size + SEQ = T * H * W + + _LOGGER.info( + f"NA3D_FLASH_FWD: q={tuple(q.shape)} kernel=({KT},{KH},{KW}) SEQ={SEQ} HD={HD}" + ) + + assert q.dtype == torch.bfloat16, "na3d_flash_attn: inputs must be bfloat16" + assert HD & (HD - 1) == 0, f"head_dim {HD} must be a power of 2" + assert W >= 16, f"W={W} is too small; kernel requires W >= BLOCK_Q (default 16)." + + def _flat(t: torch.Tensor) -> torch.Tensor: + """(B, T, H, W, NH, HD) -> (B*NH, SEQ, HD) contiguous.""" + return t.permute(0, 4, 1, 2, 3, 5).reshape(B * NH, SEQ, HD).contiguous() + + q_f, k_f, v_f = _flat(q), _flat(k), _flat(v) + out_f = torch.empty_like(q_f) + + # Lambda grid: Triton passes the autotuned BLOCK_Q via meta. + grid = lambda meta: (triton.cdiv(SEQ, meta["BLOCK_Q"]), B * NH) + + _na3d_flash_fwd[grid]( + q_f, + k_f, + v_f, + out_f, + SEQ * HD, # stride_bnh + HD, # stride_seq + T, + H, + W, + SEQ, + HD=HD, + KT=KT, + KH=KH, + KW=KW, + ) + + out = out_f.reshape(B, NH, T, H, W, HD).permute(0, 2, 3, 4, 1, 5).contiguous() + + if correct_rightmost_block: + # Use the BLOCK_Q actually selected by autotune for this (KT,KH,KW,W) key. + 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) + out[:, :, :, W_last:, :, :] = out_right + + return out diff --git a/op_tests/triton_tests/attention/test_na3d_flash.py b/op_tests/triton_tests/attention/test_na3d_flash.py new file mode 100644 index 00000000000..3f951577948 --- /dev/null +++ b/op_tests/triton_tests/attention/test_na3d_flash.py @@ -0,0 +1,357 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. +"""Test and benchmark for 3D neighborhood flash attention (na3d_flash). + +Covers representative LTX-2.5 VAE decode shapes plus edge cases that exercise +boundary masking and non-aligned sequence lengths. + +Each shape is compared against a pure-PyTorch tiled SDPA reference that implements +the same inward-shifted centered window. TFLOPS and TB/s are reported per shape. +""" + +import argparse + +import pandas as pd +import pytest +import torch +import torch.nn.functional as F + +import aiter +from aiter import dtypes +from aiter.jit.utils.chip_info import get_gfx +from aiter.ops.triton.attention.na3d_flash import _na3d_sdpa_exact, na3d_flash_attn +from aiter.test_common import benchmark, checkAllclose, run_perftest + +torch.set_default_device("cuda") + +SUPPORTED_GFX = ["gfx942", "gfx950"] + + +# --------------------------------------------------------------------------- +# Pure-PyTorch reference (not timed, not in the table) +# --------------------------------------------------------------------------- +def _window_bounds(length: int, kernel: int) -> tuple[list[int], list[int]]: + """Per-index (start, end) of the inward-shifted centered neighborhood window.""" + lo = length - kernel + half = kernel // 2 + starts, ends = [], [] + for i in range(length): + start = min(max(i - half, 0), lo) + starts.append(start) + ends.append(start + kernel) + return starts, ends + + +def na3d_sdpa_ref( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + kernel_size: tuple[int, int, int], +) -> torch.Tensor: + """FP32 per-query neighborhood attention - ground-truth cross-check only. + + Simple per-query loop used exclusively in ``test_na3d_sdpa_exact_vs_ref`` + to verify that ``_na3d_sdpa_exact`` agrees with this FP32 baseline. + For routine correctness tests use ``_na3d_sdpa_exact`` (GPU-grouped, fast). + + Args: + q, k, v : ``(B, T, H, W, NH, HD)`` bfloat16 - Q pre-scaled. + kernel_size : ``(KT, KH, KW)``. + Returns: + Output ``(B, T, H, W, NH, HD)`` bfloat16. + """ + B, T, H, W, NH, HD = q.shape + KT, KH, KW = kernel_size + + bounds_t = _window_bounds(T, KT) + bounds_h = _window_bounds(H, KH) + bounds_w = _window_bounds(W, KW) + + out = torch.empty_like(q) + for t in range(T): + for h in range(H): + for w in range(W): + ts, te = bounds_t[0][t], bounds_t[1][t] + hs, he = bounds_h[0][h], bounds_h[1][h] + ws, we = bounds_w[0][w], bounds_w[1][w] + # q_tok: (B, NH, 1, HD); kv_k/kv_v: (B, NH, K, HD) + q_tok = q[:, t, h, w, :, :].unsqueeze(2).float() + kv_k = ( + k[:, ts:te, hs:he, ws:we, :, :] + .reshape(B, -1, NH, HD) + .permute(0, 2, 1, 3) + .float() + ) + kv_v = ( + v[:, ts:te, hs:he, ws:we, :, :] + .reshape(B, -1, NH, HD) + .permute(0, 2, 1, 3) + .float() + ) + o = F.scaled_dot_product_attention(q_tok, kv_k, kv_v, scale=1.0) + out[:, t, h, w, :, :] = o.squeeze(2).to(q.dtype) + return out + + +# --------------------------------------------------------------------------- +# Default LTX-2.5 and edge-case shapes +# --------------------------------------------------------------------------- +# Each tuple: (B, T, H, W, NH, HD, KT, KH, KW) +_DEFAULT_SHAPES = [ + # --- LTX-2.5 default shapes (B=1) --- + (1, 18, 32, 48, 32, 64, 3, 7, 7), # det-0 full vol + (1, 36, 64, 96, 8, 64, 3, 5, 5), # det-2 full vol + (1, 40, 96, 96, 8, 64, 3, 5, 5), # det-3 tile (T=40, H=W=96) + (1, 16, 80, 32, 4, 64, 11, 11, 11), # diff-5 small tile + (1, 79, 192, 192, 4, 64, 11, 11, 11), # diff-5 dominant tile + # --- Edge cases for correctness --- + # SEQ % 16 != 0: exercises q_mask for the partial last program block. + # SEQ = 3*5*17 = 255, 255 % 16 = 15. + (1, 3, 5, 17, 4, 64, 3, 5, 5), + # T == KT: every temporal token is at a border, so the T-clamping + # (tl.minimum(..., T - KT)) is the active constraint for all queries. + (1, 11, 16, 16, 4, 64, 11, 11, 11), + # W == 16 == BLOCK_Q minimum: upper half of every BLOCK_KV=32 chunk is + # masked by kv_ok, exercising W-boundary masking on every inner iteration. + (1, 16, 16, 16, 4, 64, 11, 11, 11), + # B > 1: tests pid_bnh addressing across batch elements. + (2, 16, 16, 32, 4, 64, 3, 5, 5), +] + + +# --------------------------------------------------------------------------- +# Fast subset for cross-validation against na3d_sdpa_ref (FP32 per-query loop). +# Only shapes with SEQ <= 50K to keep na3d_sdpa_ref within a few seconds each. +# Covers all 4 edge-case categories plus 2 LTX-2.5 default shapes. +# Large shapes (36,64,96), (40,96,96), (79,192,192) are excluded here; they +# are already exercised in test_na3d_flash via the faster _na3d_sdpa_exact ref. +# --------------------------------------------------------------------------- +_SDPA_FAST_SHAPES = [ + (1, 18, 32, 48, 32, 64, 3, 7, 7), # det-0 (SEQ = 27 648) + (1, 16, 80, 32, 4, 64, 11, 11, 11), # diff-5 small tile (SEQ = 40 960) + (1, 3, 5, 17, 4, 64, 3, 5, 5), # edge: SEQ % 16 != 0 (SEQ = 255) + (1, 11, 16, 16, 4, 64, 11, 11, 11), # edge: T == KT (SEQ = 2 816) + (1, 16, 16, 16, 4, 64, 11, 11, 11), # edge: W == 16 (SEQ = 4 096) + (2, 16, 16, 32, 4, 64, 3, 5, 5), # edge: B > 1 (SEQ = 8 192) +] + + +# --------------------------------------------------------------------------- +# Shapes to test correct_rightmost_block +# --------------------------------------------------------------------------- +_CORRECTION_SHAPES = [ + (1, 16, 80, 32, 4, 64, 11, 11, 11), # diff-5 small tile + (1, 16, 80, 192, 4, 64, 11, 11, 11), # diff-5 wide tile + (1, 18, 32, 48, 32, 64, 3, 7, 7), # det-0 + (1, 36, 64, 96, 8, 64, 3, 5, 5), # det-2 +] + + +# --------------------------------------------------------------------------- +# Pytest correctness tests +# --------------------------------------------------------------------------- +@pytest.mark.parametrize("B,T,H,W,NH,HD,KT,KH,KW", _DEFAULT_SHAPES) +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +def test_na3d_flash(B, T, H, W, NH, HD, KT, KH, KW, dtype): + if get_gfx() not in SUPPORTED_GFX: + pytest.skip(f"na3d_flash unsupported on {get_gfx()}") + + scale = HD**-0.5 + q = torch.randn(B, T, H, W, NH, HD, dtype=dtype) * scale + k = torch.randn(B, T, H, W, NH, HD, dtype=dtype) + v = torch.randn(B, T, H, W, NH, HD, dtype=dtype) + + ref = _na3d_sdpa_exact(q, k, v, kernel_size=(KT, KH, KW)) + out = na3d_flash_attn(q, k, v, kernel_size=(KT, KH, KW)) + + checkAllclose( + ref.float(), + out.float(), + rtol=1e-2, + atol=5e-2, + msg=f"na3d_flash (T={T},H={H},W={W},k=({KT},{KH},{KW}))", + ) + + +# --------------------------------------------------------------------------- +# Benchmark function (one row per shape) +# --------------------------------------------------------------------------- +@benchmark() +def bench_na3d_flash(B, T, H, W, NH, HD, KT, KH, KW, dtype): + SEQ = T * H * W + C = NH * HD + K = KT * KH * KW # neighborhood size + + # Build inputs matching the real decoder call: pre-scaled Q, BF16, channels-last. + scale = HD**-0.5 + q = torch.randn(B, T, H, W, NH, HD, dtype=dtype) * scale + k = torch.randn(B, T, H, W, NH, HD, dtype=dtype) + v = torch.randn(B, T, H, W, NH, HD, dtype=dtype) + + # Reference (not timed): GPU-grouped exact SDPA, faster than na3d_sdpa_ref loop. + ref = _na3d_sdpa_exact(q, k, v, kernel_size=(KT, KH, KW)) + + # FLOPs: QK and AV dot products over K neighbors per query. + # QK: 2 * B * SEQ * K * C + # AV: 2 * B * SEQ * K * C + flops = 4 * B * SEQ * K * C + # Bytes: Q loaded once, K and V reloaded per (t_kv, h_kv) row. + elem = q.element_size() + nbytes = (B * SEQ * C + 2 * B * SEQ * K * C) * elem + + candidates = { + "triton": lambda: na3d_flash_attn(q, k, v, kernel_size=(KT, KH, KW)), + } + + ret = {"gfx": get_gfx()} + for name, fn in candidates.items(): + out, us = run_perftest(fn) + err = checkAllclose( + ref.float(), + out.float(), + rtol=1e-2, + atol=5e-2, + msg=f"{name}: na3d_flash (T={T},H={H},W={W},k=({KT},{KH},{KW}))", + ) + ret[f"{name} us"] = us + ret[f"{name} TFLOPS"] = flops / us / 1e6 + ret[f"{name} TB/s"] = nbytes / us / 1e6 + ret[f"{name} err"] = err + return ret + + +# --------------------------------------------------------------------------- +# _na3d_sdpa_exact tests +# --------------------------------------------------------------------------- + +@pytest.mark.parametrize("B,T,H,W,NH,HD,KT,KH,KW", _SDPA_FAST_SHAPES) +def test_na3d_sdpa_exact_vs_ref(B, T, H, W, NH, HD, KT, KH, KW): + """``_na3d_sdpa_exact`` agrees with ``na3d_sdpa_ref`` (FP32 per-query loop). + + Both compute the same inward-shifted neighborhood attention; the difference is + numerical precision (BF16 vs FP32) and implementation (grouped tiled SDPA vs + per-query loop). The tolerance is the same as ``test_na3d_flash``. + """ + if get_gfx() not in SUPPORTED_GFX: + pytest.skip(f"na3d_flash unsupported on {get_gfx()}") + + scale = HD**-0.5 + q = torch.randn(B, T, H, W, NH, HD, dtype=torch.bfloat16) * scale + k = torch.randn(B, T, H, W, NH, HD, dtype=torch.bfloat16) + v = torch.randn(B, T, H, W, NH, HD, dtype=torch.bfloat16) + + ref = na3d_sdpa_ref(q, k, v, kernel_size=(KT, KH, KW)) + exact = _na3d_sdpa_exact(q, k, v, kernel_size=(KT, KH, KW)) # w_offset=0 default + + checkAllclose( + ref.float(), + exact.float(), + rtol=1e-2, + atol=5e-2, + msg=f"_na3d_sdpa_exact vs na3d_sdpa_ref (T={T},H={H},W={W},k=({KT},{KH},{KW}))", + ) + + +@pytest.mark.parametrize("B,T,H,W,NH,HD,KT,KH,KW", _CORRECTION_SHAPES) +def test_correct_rightmost_block(B, T, H, W, NH, HD, KT, KH, KW): + """Verify the ``correct_rightmost_block`` docstring claim. + + Asserts: + 1. Interior W positions are bit-identical before and after correction. + 2. Rightmost W positions differ (correction fires and changes output). + 3. Corrected rightmost positions match ``_na3d_sdpa_exact`` applied to + just those queries, confirming the kv_ok=False BF16(0) MFMA divergence + is eliminated. + """ + if get_gfx() not in SUPPORTED_GFX: + pytest.skip(f"na3d_flash unsupported on {get_gfx()}") + + scale = HD**-0.5 + q = torch.randn(B, T, H, W, NH, HD, dtype=torch.bfloat16) * scale + k = torch.randn(B, T, H, W, NH, HD, dtype=torch.bfloat16) + v = torch.randn(B, T, H, W, NH, HD, dtype=torch.bfloat16) + + out_base = na3d_flash_attn(q, k, v, (KT, KH, KW), correct_rightmost_block=False) + out_corr = na3d_flash_attn(q, k, v, (KT, KH, KW), correct_rightmost_block=True) + + 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 + + diff = (out_corr.float() - out_base.float()).abs() + + if W_last > 0: + # Claim 1: interior positions unchanged. + interior_max = diff[:, :, :, :W_last, :, :].max().item() + assert interior_max == 0.0, ( + f"Interior W positions changed (max_diff={interior_max:.6f}) - " + f"correction must not touch positions 0..{W_last - 1}" + ) + + # Claim 2: rightmost positions differ (correction fired). + rightmost_max = diff[:, :, :, W_last:, :, :].max().item() + assert rightmost_max > 0.0, ( + f"Rightmost W positions unchanged - " + f"T*H*(W-W_last)={T*H*(W-W_last)} should be <= 100_000 so correction fires" + ) + + # Claim 3: corrected positions match _na3d_sdpa_exact. + q_right = q[:, :, :, W_last:, :, :] + exact_right = _na3d_sdpa_exact(q_right, k, v, (KT, KH, KW), w_offset=W_last) + checkAllclose( + exact_right.float(), + out_corr[:, :, :, W_last:, :, :].float(), + rtol=1e-2, + atol=5e-2, + msg=( + f"corrected rightmost vs _na3d_sdpa_exact " + f"(T={T},H={H},W={W},W_last={W_last},k=({KT},{KH},{KW}))" + ), + ) + + +def main(): + if get_gfx() not in SUPPORTED_GFX: + aiter.logger.warning("na3d_flash unsupported on %s; skipping", get_gfx()) + return + + parser = argparse.ArgumentParser( + formatter_class=argparse.RawTextHelpFormatter, + description="Benchmark 3D neighborhood flash attention", + ) + parser.add_argument( + "-d", + "--dtype", + type=dtypes.str2Dtype, + nargs="*", + default=[dtypes.d_dtypes["bf16"]], + help="Input dtype (default: bf16)", + ) + parser.add_argument( + "-s", + "--shapes", + type=dtypes.str2tuple, + nargs="*", + default=_DEFAULT_SHAPES, + help=( + "Shapes as (B,T,H,W,NH,HD,KT,KH,KW) tuples.\n" + "e.g.: -s '(1,79,192,192,4,64,11,11,11)'" + ), + ) + args = parser.parse_args() + + for dtype in args.dtype: + rows = [] + for shape in args.shapes: + B, T, H, W, NH, HD, KT, KH, KW = shape + rows.append(bench_na3d_flash(B, T, H, W, NH, HD, KT, KH, KW, dtype)) + df = pd.DataFrame(rows) + aiter.logger.info( + "na3d_flash summary (%s):\n%s", + str(dtype).replace("torch.", ""), + df.to_markdown(index=False), + ) + + +if __name__ == "__main__": + main() From 19f8b88826a476511de51fd0b61730890a0986a0 Mon Sep 17 00:00:00 2001 From: Joni Juvonen Date: Thu, 3 Sep 2026 10:14:57 +0300 Subject: [PATCH 02/48] na3d_flash kernel: input type and bounds validation Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- aiter/ops/triton/attention/na3d_flash.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/aiter/ops/triton/attention/na3d_flash.py b/aiter/ops/triton/attention/na3d_flash.py index e7405ed3402..dc80a5c80ba 100644 --- a/aiter/ops/triton/attention/na3d_flash.py +++ b/aiter/ops/triton/attention/na3d_flash.py @@ -221,6 +221,13 @@ def na3d_flash_attn( ) assert q.dtype == torch.bfloat16, "na3d_flash_attn: inputs must be bfloat16" + assert k.dtype == q.dtype and v.dtype == q.dtype, "na3d_flash_attn: q/k/v must share dtype" + assert k.shape == q.shape and v.shape == q.shape, ( + "na3d_flash_attn: q/k/v must have shape (B, T, H, W, NH, HD)" + ) + assert KT <= T and KH <= H and KW <= W, ( + f"na3d_flash_attn: kernel_size=({KT},{KH},{KW}) must be <= (T,H,W)=({T},{H},{W})" + ) assert HD & (HD - 1) == 0, f"head_dim {HD} must be a power of 2" assert W >= 16, f"W={W} is too small; kernel requires W >= BLOCK_Q (default 16)." From 85671cc5bd33c57edc91b579d154dc83fbff18c7 Mon Sep 17 00:00:00 2001 From: jjuvonen-amd Date: Thu, 3 Sep 2026 07:25:52 +0000 Subject: [PATCH 03/48] use repr= in jit to include autotuned configs in kernel names Signed-off-by: jjuvonen-amd --- aiter/ops/triton/_triton_kernels/attention/na3d_flash.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/aiter/ops/triton/_triton_kernels/attention/na3d_flash.py b/aiter/ops/triton/_triton_kernels/attention/na3d_flash.py index e8a6614a0e1..97c2a2f88f4 100644 --- a/aiter/ops/triton/_triton_kernels/attention/na3d_flash.py +++ b/aiter/ops/triton/_triton_kernels/attention/na3d_flash.py @@ -17,6 +17,7 @@ import triton import triton.language as tl +from aiter.ops.triton.utils._triton.kernel_repr import make_kernel_repr # Autotune configs. # BLOCK_KV = next_pow2(BLOCK_Q + KW - 1) covers the union of all BLOCK_Q # queries' W windows in one chunk: @@ -44,13 +45,17 @@ def _prune_configs(configs, named_args, **kwargs): W = named_args["W"] return [c for c in configs if c.kwargs["BLOCK_Q"] <= W] +_na3d_flash_fwd_repr = make_kernel_repr( + "_na3d_flash_fwd", + ["BLOCK_Q", "BLOCK_KV", "KT", "KH", "KW", "HD"], +) @triton.autotune( configs=_CONFIGS, key=["KT", "KH", "KW", "HD", "W"], prune_configs_by={"early_config_prune": _prune_configs}, ) -@triton.jit +@triton.jit(repr=_na3d_flash_fwd_repr) def _na3d_flash_fwd( Q_ptr, K_ptr, From 90a4b4eb60269bdf72f9f0af3308c87f37d03c4c Mon Sep 17 00:00:00 2001 From: jjuvonen-amd Date: Thu, 3 Sep 2026 07:40:47 +0000 Subject: [PATCH 04/48] fix cost heuristics in _na3d_sdpa_exact, correct docstring for rightmost block Signed-off-by: jjuvonen-amd --- aiter/ops/triton/attention/na3d_flash.py | 32 +++++++++++++----------- 1 file changed, 18 insertions(+), 14 deletions(-) diff --git a/aiter/ops/triton/attention/na3d_flash.py b/aiter/ops/triton/attention/na3d_flash.py index dc80a5c80ba..fa63ec99314 100644 --- a/aiter/ops/triton/attention/na3d_flash.py +++ b/aiter/ops/triton/attention/na3d_flash.py @@ -61,7 +61,7 @@ def _na3d_sdpa_mask( for starts, ends in rel_bounds: st = torch.tensor(starts, device=device) en = torch.tensor(ends, device=device) - kj = torch.arange(int(en.max()), device=device) + kj = torch.arange(max(ends), device=device) bools.append((kj[None, :] >= st[:, None]) & (kj[None, :] < en[:, None])) visible = ( bools[0][:, None, None, :, None, None] @@ -107,16 +107,18 @@ def _na3d_sdpa_exact( bh = _window_bounds_1d(H, min(KH, H)) bw = _window_bounds_1d(W_full, min(KW, W_full)) # global W bounds - # Pick tile sizes relative to q_right's actual dims (W_sub is small). - eff_kt, eff_kh, eff_kw = min(KT, T), min(KH, H), min(KW, W_sub) + # Pick tile sizes based on q_right, + # but estimate K/V union sizes using the full W. + eff_kt, eff_kh = min(KT, T), min(KH, H) + eff_kw = min(KW, W_full) def _cost(ts: list[int]) -> int: - nq = math.prod(ts) - nk = math.prod( - min(d, t + kk - 1) - for t, kk, d in zip(ts, [eff_kt, eff_kh, eff_kw], [T, H, W_sub]) - ) - return nq * nk + tile_t, tile_h, tile_w = ts + nq = tile_t * tile_h * tile_w + nk_t = min(T, tile_t + eff_kt - 1) + nk_h = min(H, tile_h + eff_kh - 1) + nk_w = min(W_full, tile_w + eff_kw - 1) + return nq * nk_t * nk_h * nk_w tiles = [T, H, W_sub] while _cost(tiles) > 2**25 and max(tiles) > 1: @@ -199,11 +201,13 @@ def na3d_flash_attn( q, k, v : ``(B, T, H, W, NH, HD)`` bfloat16, channels-last. Q is expected to be pre-scaled by ``head_dim ** -0.5``. kernel_size : ``(KT, KH, KW)`` neighborhood window. - correct_rightmost_block : When True, overwrite the rightmost BLOCK_Q W - positions with per-query exact SDPA (via _na3d_sdpa_exact). - The flash kernel loads ``BF16(0)`` for out-of-bounds K/V positions in the - last W-block (``kv_ok=False``), which causes a ~0.002 rounding divergence - from AOTriton at those positions. Off by default. + correct_rightmost_block : When True, overwrite the last W-block + (``W_last:``) with per-query exact SDPA (via _na3d_sdpa_exact), + where ``W_last = floor((W - 1) / BLOCK_Q) * BLOCK_Q`` for + the autotuned BLOCK_Q. The flash kernel loads ``BF16(0)`` for + out-of-bounds K/V positions in the last W-block (``kv_ok=False``), + which causes a ~0.002 rounding divergence from AOTriton + at those positions. Off by default. Returns: Output tensor ``(B, T, H, W, NH, HD)`` bfloat16. From 508338dee8c8359aad9defeb3cc940b69a8f2c32 Mon Sep 17 00:00:00 2001 From: jjuvonen-amd Date: Thu, 3 Sep 2026 07:44:15 +0000 Subject: [PATCH 05/48] avoid importing from kernels module in test Signed-off-by: jjuvonen-amd --- op_tests/triton_tests/attention/test_na3d_flash.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/op_tests/triton_tests/attention/test_na3d_flash.py b/op_tests/triton_tests/attention/test_na3d_flash.py index 3f951577948..697e42abb17 100644 --- a/op_tests/triton_tests/attention/test_na3d_flash.py +++ b/op_tests/triton_tests/attention/test_na3d_flash.py @@ -274,8 +274,8 @@ def test_correct_rightmost_block(B, T, H, W, NH, HD, KT, KH, KW): out_base = na3d_flash_attn(q, k, v, (KT, KH, KW), correct_rightmost_block=False) out_corr = na3d_flash_attn(q, k, v, (KT, KH, KW), correct_rightmost_block=True) - from aiter.ops.triton._triton_kernels.attention.na3d_flash import _na3d_flash_fwd - actual_bq = _na3d_flash_fwd.best_config.kwargs["BLOCK_Q"] + import aiter.ops.triton.attention.na3d_flash as na3d_flash_mod + actual_bq = na3d_flash_mod._na3d_flash_fwd.best_config.kwargs["BLOCK_Q"] W_last = ((W - 1) // actual_bq) * actual_bq diff = (out_corr.float() - out_base.float()).abs() From c52655112fcf85762fbee448b55b65370d9e9036 Mon Sep 17 00:00:00 2001 From: jjuvonen-amd Date: Thu, 3 Sep 2026 07:57:40 +0000 Subject: [PATCH 06/48] black format Signed-off-by: jjuvonen-amd --- .../_triton_kernels/attention/na3d_flash.py | 3 ++ aiter/ops/triton/attention/na3d_flash.py | 26 +++++++++------- .../triton_tests/attention/test_na3d_flash.py | 30 ++++++++++--------- 3 files changed, 34 insertions(+), 25 deletions(-) diff --git a/aiter/ops/triton/_triton_kernels/attention/na3d_flash.py b/aiter/ops/triton/_triton_kernels/attention/na3d_flash.py index 97c2a2f88f4..758cf336b98 100644 --- a/aiter/ops/triton/_triton_kernels/attention/na3d_flash.py +++ b/aiter/ops/triton/_triton_kernels/attention/na3d_flash.py @@ -18,6 +18,7 @@ import triton.language as tl from aiter.ops.triton.utils._triton.kernel_repr import make_kernel_repr + # Autotune configs. # BLOCK_KV = next_pow2(BLOCK_Q + KW - 1) covers the union of all BLOCK_Q # queries' W windows in one chunk: @@ -45,11 +46,13 @@ def _prune_configs(configs, named_args, **kwargs): W = named_args["W"] return [c for c in configs if c.kwargs["BLOCK_Q"] <= W] + _na3d_flash_fwd_repr = make_kernel_repr( "_na3d_flash_fwd", ["BLOCK_Q", "BLOCK_KV", "KT", "KH", "KW", "HD"], ) + @triton.autotune( configs=_CONFIGS, key=["KT", "KH", "KW", "HD", "W"], diff --git a/aiter/ops/triton/attention/na3d_flash.py b/aiter/ops/triton/attention/na3d_flash.py index fa63ec99314..74b76913268 100644 --- a/aiter/ops/triton/attention/na3d_flash.py +++ b/aiter/ops/triton/attention/na3d_flash.py @@ -33,6 +33,7 @@ # Exact per-query neighborhood SDPA (generic; used for rightmost-block correction) # --------------------------------------------------------------------------- + def _window_bounds_1d(length: int, kernel: int) -> tuple[list[int], list[int]]: """Inward-shifted window (start, end) for each position along one axis.""" lo = length - kernel @@ -107,7 +108,7 @@ def _na3d_sdpa_exact( bh = _window_bounds_1d(H, min(KH, H)) bw = _window_bounds_1d(W_full, min(KW, W_full)) # global W bounds - # Pick tile sizes based on q_right, + # Pick tile sizes based on q_right, # but estimate K/V union sizes using the full W. eff_kt, eff_kh = min(KT, T), min(KH, H) eff_kw = min(KW, W_full) @@ -188,6 +189,7 @@ def _cost(ts: list[int]) -> int: # Public launcher # --------------------------------------------------------------------------- + def na3d_flash_attn( q: torch.Tensor, k: torch.Tensor, @@ -203,10 +205,10 @@ def na3d_flash_attn( kernel_size : ``(KT, KH, KW)`` neighborhood window. correct_rightmost_block : When True, overwrite the last W-block (``W_last:``) with per-query exact SDPA (via _na3d_sdpa_exact), - where ``W_last = floor((W - 1) / BLOCK_Q) * BLOCK_Q`` for - the autotuned BLOCK_Q. The flash kernel loads ``BF16(0)`` for + where ``W_last = floor((W - 1) / BLOCK_Q) * BLOCK_Q`` for + the autotuned BLOCK_Q. The flash kernel loads ``BF16(0)`` for out-of-bounds K/V positions in the last W-block (``kv_ok=False``), - which causes a ~0.002 rounding divergence from AOTriton + which causes a ~0.002 rounding divergence from AOTriton at those positions. Off by default. Returns: @@ -225,13 +227,15 @@ def na3d_flash_attn( ) assert q.dtype == torch.bfloat16, "na3d_flash_attn: inputs must be bfloat16" - assert k.dtype == q.dtype and v.dtype == q.dtype, "na3d_flash_attn: q/k/v must share dtype" - assert k.shape == q.shape and v.shape == q.shape, ( - "na3d_flash_attn: q/k/v must have shape (B, T, H, W, NH, HD)" - ) - assert KT <= T and KH <= H and KW <= W, ( - f"na3d_flash_attn: kernel_size=({KT},{KH},{KW}) must be <= (T,H,W)=({T},{H},{W})" - ) + assert ( + k.dtype == q.dtype and v.dtype == q.dtype + ), "na3d_flash_attn: q/k/v must share dtype" + assert ( + k.shape == q.shape and v.shape == q.shape + ), "na3d_flash_attn: q/k/v must have shape (B, T, H, W, NH, HD)" + assert ( + KT <= T and KH <= H and KW <= W + ), f"na3d_flash_attn: kernel_size=({KT},{KH},{KW}) must be <= (T,H,W)=({T},{H},{W})" assert HD & (HD - 1) == 0, f"head_dim {HD} must be a power of 2" assert W >= 16, f"W={W} is too small; kernel requires W >= BLOCK_Q (default 16)." diff --git a/op_tests/triton_tests/attention/test_na3d_flash.py b/op_tests/triton_tests/attention/test_na3d_flash.py index 697e42abb17..a3da9c0cbd6 100644 --- a/op_tests/triton_tests/attention/test_na3d_flash.py +++ b/op_tests/triton_tests/attention/test_na3d_flash.py @@ -99,10 +99,10 @@ def na3d_sdpa_ref( # Each tuple: (B, T, H, W, NH, HD, KT, KH, KW) _DEFAULT_SHAPES = [ # --- LTX-2.5 default shapes (B=1) --- - (1, 18, 32, 48, 32, 64, 3, 7, 7), # det-0 full vol - (1, 36, 64, 96, 8, 64, 3, 5, 5), # det-2 full vol - (1, 40, 96, 96, 8, 64, 3, 5, 5), # det-3 tile (T=40, H=W=96) - (1, 16, 80, 32, 4, 64, 11, 11, 11), # diff-5 small tile + (1, 18, 32, 48, 32, 64, 3, 7, 7), # det-0 full vol + (1, 36, 64, 96, 8, 64, 3, 5, 5), # det-2 full vol + (1, 40, 96, 96, 8, 64, 3, 5, 5), # det-3 tile (T=40, H=W=96) + (1, 16, 80, 32, 4, 64, 11, 11, 11), # diff-5 small tile (1, 79, 192, 192, 4, 64, 11, 11, 11), # diff-5 dominant tile # --- Edge cases for correctness --- # SEQ % 16 != 0: exercises q_mask for the partial last program block. @@ -127,12 +127,12 @@ def na3d_sdpa_ref( # are already exercised in test_na3d_flash via the faster _na3d_sdpa_exact ref. # --------------------------------------------------------------------------- _SDPA_FAST_SHAPES = [ - (1, 18, 32, 48, 32, 64, 3, 7, 7), # det-0 (SEQ = 27 648) - (1, 16, 80, 32, 4, 64, 11, 11, 11), # diff-5 small tile (SEQ = 40 960) - (1, 3, 5, 17, 4, 64, 3, 5, 5), # edge: SEQ % 16 != 0 (SEQ = 255) - (1, 11, 16, 16, 4, 64, 11, 11, 11), # edge: T == KT (SEQ = 2 816) - (1, 16, 16, 16, 4, 64, 11, 11, 11), # edge: W == 16 (SEQ = 4 096) - (2, 16, 16, 32, 4, 64, 3, 5, 5), # edge: B > 1 (SEQ = 8 192) + (1, 18, 32, 48, 32, 64, 3, 7, 7), # det-0 (SEQ = 27 648) + (1, 16, 80, 32, 4, 64, 11, 11, 11), # diff-5 small tile (SEQ = 40 960) + (1, 3, 5, 17, 4, 64, 3, 5, 5), # edge: SEQ % 16 != 0 (SEQ = 255) + (1, 11, 16, 16, 4, 64, 11, 11, 11), # edge: T == KT (SEQ = 2 816) + (1, 16, 16, 16, 4, 64, 11, 11, 11), # edge: W == 16 (SEQ = 4 096) + (2, 16, 16, 32, 4, 64, 3, 5, 5), # edge: B > 1 (SEQ = 8 192) ] @@ -140,10 +140,10 @@ def na3d_sdpa_ref( # Shapes to test correct_rightmost_block # --------------------------------------------------------------------------- _CORRECTION_SHAPES = [ - (1, 16, 80, 32, 4, 64, 11, 11, 11), # diff-5 small tile - (1, 16, 80, 192, 4, 64, 11, 11, 11), # diff-5 wide tile - (1, 18, 32, 48, 32, 64, 3, 7, 7), # det-0 - (1, 36, 64, 96, 8, 64, 3, 5, 5), # det-2 + (1, 16, 80, 32, 4, 64, 11, 11, 11), # diff-5 small tile + (1, 16, 80, 192, 4, 64, 11, 11, 11), # diff-5 wide tile + (1, 18, 32, 48, 32, 64, 3, 7, 7), # det-0 + (1, 36, 64, 96, 8, 64, 3, 5, 5), # det-2 ] @@ -224,6 +224,7 @@ def bench_na3d_flash(B, T, H, W, NH, HD, KT, KH, KW, dtype): # _na3d_sdpa_exact tests # --------------------------------------------------------------------------- + @pytest.mark.parametrize("B,T,H,W,NH,HD,KT,KH,KW", _SDPA_FAST_SHAPES) def test_na3d_sdpa_exact_vs_ref(B, T, H, W, NH, HD, KT, KH, KW): """``_na3d_sdpa_exact`` agrees with ``na3d_sdpa_ref`` (FP32 per-query loop). @@ -275,6 +276,7 @@ def test_correct_rightmost_block(B, T, H, W, NH, HD, KT, KH, KW): out_corr = na3d_flash_attn(q, k, v, (KT, KH, KW), correct_rightmost_block=True) import aiter.ops.triton.attention.na3d_flash as na3d_flash_mod + actual_bq = na3d_flash_mod._na3d_flash_fwd.best_config.kwargs["BLOCK_Q"] W_last = ((W - 1) // actual_bq) * actual_bq From b448a1b592c7843532a15ac9f14c574e4fa5eb8a Mon Sep 17 00:00:00 2001 From: jjuvonen-amd Date: Thu, 3 Sep 2026 10:05:31 +0000 Subject: [PATCH 07/48] fix cross-row bug when W % BLOCK_Q != 0 and add test shape to validate fix Signed-off-by: jjuvonen-amd --- .../_triton_kernels/attention/na3d_flash.py | 43 +++++++++++-------- aiter/ops/triton/attention/na3d_flash.py | 6 +-- .../triton_tests/attention/test_na3d_flash.py | 2 + 3 files changed, 29 insertions(+), 22 deletions(-) diff --git a/aiter/ops/triton/_triton_kernels/attention/na3d_flash.py b/aiter/ops/triton/_triton_kernels/attention/na3d_flash.py index 758cf336b98..5254b4c426f 100644 --- a/aiter/ops/triton/_triton_kernels/attention/na3d_flash.py +++ b/aiter/ops/triton/_triton_kernels/attention/na3d_flash.py @@ -4,10 +4,10 @@ Each query token attends to the (KT x KH x KW) neighborhood centered on its (t, h, w) grid position, with inward border-shift so every window has the full -KT x KH x KW keys. The kernel processes BLOCK_Q consecutive queries per program; -all queries in one block share the same (t, h) row (requires W >= BLOCK_Q). +KT x KH x KW keys. Algorithm: +- Grid dim[0] = T × H × ⌈W / BLOCK_Q⌉; each program covers one (t, h) row. - Q is loaded once into registers and held for all KT x KH inner iterations. - Each iteration loads one (t_kv, h_kv) row of K and V at the shared W window. - Running online softmax in log2 space (exp2, log2(e) folded into scale). @@ -41,7 +41,9 @@ def _prune_configs(configs, named_args, **kwargs): """Remove configs where BLOCK_Q > W. - Ensures all queries in one block share the same (t, h) row (W >= BLOCK_Q required). + Performance constraint: BLOCK_Q > W would leave most of every block masked, + wasting compute. Correctness is guaranteed by the row-decomposed grid + regardless of whether W is divisible by BLOCK_Q. """ W = named_args["W"] return [c for c in configs if c.kwargs["BLOCK_Q"] <= W] @@ -69,7 +71,6 @@ def _na3d_flash_fwd( T, H, W, - SEQ, HD: tl.constexpr, # head dimension: any power-of-2 value KT: tl.constexpr, # neighborhood depth: constexpr -> static loop bound KH: tl.constexpr, # neighborhood height: constexpr -> static loop bound @@ -77,33 +78,37 @@ def _na3d_flash_fwd( BLOCK_Q: tl.constexpr, BLOCK_KV: tl.constexpr, # power-of-2, >= BLOCK_Q + KW - 1 ): - """Flash-attention inner kernel. Requires W >= BLOCK_Q (same (t,h) row per block). + """Flash-attention inner kernel. Grid dim[0] = T*H*ceil(W/BLOCK_Q). - See _prune_configs for the W >= BLOCK_Q constraint. + Each program is assigned to exactly one (t, h) row via decomposition of + pid_q into (row_idx, w_block_idx). This guarantees a shared (t, h) window + start regardless of whether W is divisible by BLOCK_Q. """ pid_q = tl.program_id(0) pid_bnh = tl.program_id(1) HW = H * W - q_start = pid_q * BLOCK_Q + # Decompose pid_q into (t,h) row and W-block within that row. + W_blocks = (W + BLOCK_Q - 1) // BLOCK_Q # programs per (t,h) row + row_idx = pid_q // W_blocks # which (t,h) row (t*H + h) + w_bid = pid_q % W_blocks # which W-block in that row + + prog_t = row_idx // H # scalar: same t for every query in this program + prog_h = row_idx % H # scalar: same h for every query in this program + q_offs = tl.arange(0, BLOCK_Q) - q_idx = q_start + q_offs - q_mask = q_idx < SEQ + q_w = w_bid * BLOCK_Q + q_offs # W positions for this block + q_mask = q_w < W # last block in a row may be partial - q_t = q_idx // HW - q_h = (q_idx % HW) // W - q_w = q_idx % W + q_idx = prog_t * HW + prog_h * W + q_w # for Q load and Out store # Inward-shifted centered neighborhood window starts. - q_t_ws = tl.minimum(tl.maximum(q_t - KT // 2, 0), T - KT) - q_h_ws = tl.minimum(tl.maximum(q_h - KH // 2, 0), H - KH) + # prog_t / prog_h are scalars; t_ws / h_ws need no tl.min reduction. q_w_ws = tl.minimum(tl.maximum(q_w - KW // 2, 0), W - KW) - # Scalar window starts (shared across all BLOCK_Q queries in this block). - INF_I = 999999 - t_ws = tl.min(tl.where(q_mask, q_t_ws, INF_I)) - h_ws = tl.min(tl.where(q_mask, q_h_ws, INF_I)) - w_lo = tl.min(tl.where(q_mask, q_w_ws, INF_I)) + t_ws = tl.minimum(tl.maximum(prog_t - KT // 2, 0), T - KT) # scalar + h_ws = tl.minimum(tl.maximum(prog_h - KH // 2, 0), H - KH) # scalar + w_lo = tl.min(tl.where(q_mask, q_w_ws, 999999)) hd_offs = tl.arange(0, HD) kv_offs = tl.arange(0, BLOCK_KV) diff --git a/aiter/ops/triton/attention/na3d_flash.py b/aiter/ops/triton/attention/na3d_flash.py index 74b76913268..b2567defbf6 100644 --- a/aiter/ops/triton/attention/na3d_flash.py +++ b/aiter/ops/triton/attention/na3d_flash.py @@ -246,8 +246,9 @@ def _flat(t: torch.Tensor) -> torch.Tensor: q_f, k_f, v_f = _flat(q), _flat(k), _flat(v) out_f = torch.empty_like(q_f) - # Lambda grid: Triton passes the autotuned BLOCK_Q via meta. - grid = lambda meta: (triton.cdiv(SEQ, meta["BLOCK_Q"]), B * NH) + # Grid: one program per (t, h) row per W-block. This guarantees each + # program covers queries from exactly one (t, h) row regardless of W % BLOCK_Q. + grid = lambda meta: (T * H * triton.cdiv(W, meta["BLOCK_Q"]), B * NH) _na3d_flash_fwd[grid]( q_f, @@ -259,7 +260,6 @@ def _flat(t: torch.Tensor) -> torch.Tensor: T, H, W, - SEQ, HD=HD, KT=KT, KH=KH, diff --git a/op_tests/triton_tests/attention/test_na3d_flash.py b/op_tests/triton_tests/attention/test_na3d_flash.py index a3da9c0cbd6..66a832c2bae 100644 --- a/op_tests/triton_tests/attention/test_na3d_flash.py +++ b/op_tests/triton_tests/attention/test_na3d_flash.py @@ -116,6 +116,8 @@ def na3d_sdpa_ref( (1, 16, 16, 16, 4, 64, 11, 11, 11), # B > 1: tests pid_bnh addressing across batch elements. (2, 16, 16, 32, 4, 64, 3, 5, 5), + # Cross-(t,h)-row: W=33 is not a multiple of BLOCK_Q=16. + (1, 4, 8, 33, 4, 64, 3, 5, 5), ] From 6dda9d38d2a6d1baf81d0e6800e3add717ad72ab Mon Sep 17 00:00:00 2001 From: jjuvonen-amd Date: Thu, 3 Sep 2026 10:56:52 +0000 Subject: [PATCH 08/48] remove unnecessary correct_rightmost_block parameter and related tests Signed-off-by: jjuvonen-amd --- aiter/ops/triton/attention/na3d_flash.py | 181 +-------------- .../triton_tests/attention/test_na3d_flash.py | 209 ++++++++++++------ 2 files changed, 139 insertions(+), 251 deletions(-) diff --git a/aiter/ops/triton/attention/na3d_flash.py b/aiter/ops/triton/attention/na3d_flash.py index b2567defbf6..a55621cf6aa 100644 --- a/aiter/ops/triton/attention/na3d_flash.py +++ b/aiter/ops/triton/attention/na3d_flash.py @@ -16,11 +16,7 @@ from __future__ import annotations -import functools -import math - import torch -import torch.nn.functional as F import triton from aiter.ops.triton._triton_kernels.attention.na3d_flash import _na3d_flash_fwd @@ -29,162 +25,6 @@ _LOGGER = AiterTritonLogger() -# --------------------------------------------------------------------------- -# Exact per-query neighborhood SDPA (generic; used for rightmost-block correction) -# --------------------------------------------------------------------------- - - -def _window_bounds_1d(length: int, kernel: int) -> tuple[list[int], list[int]]: - """Inward-shifted window (start, end) for each position along one axis.""" - lo = length - kernel - half = kernel // 2 - starts: list[int] = [] - ends: list[int] = [] - for i in range(length): - s = min(max(i - half, 0), lo) - starts.append(s) - ends.append(s + kernel) - return starts, ends - - -@functools.lru_cache(maxsize=256) -def _na3d_sdpa_mask( - rel_bounds: tuple[tuple[tuple[int, ...], tuple[int, ...]], ...], - dtype: torch.dtype, - device: torch.device, -) -> torch.Tensor: - """Cached additive ``[1, 1, Nq, Nk]`` mask for one tile-geometry group. - - Masks are reused across repeated calls with the same geometry (e.g. all - 8 diff_block layers at W=192 share the same rel_bounds). - """ - bools = [] - for starts, ends in rel_bounds: - st = torch.tensor(starts, device=device) - en = torch.tensor(ends, device=device) - kj = torch.arange(max(ends), device=device) - bools.append((kj[None, :] >= st[:, None]) & (kj[None, :] < en[:, None])) - visible = ( - bools[0][:, None, None, :, None, None] - & bools[1][None, :, None, None, :, None] - & bools[2][None, None, :, None, None, :] - ) - nq = math.prod(visible.shape[:3]) - nk = math.prod(visible.shape[3:]) - mask = torch.zeros((nq, nk), dtype=dtype, device=device) - mask.masked_fill_(~visible.reshape(nq, nk), torch.finfo(dtype).min) - return mask.reshape(1, 1, nq, nk) - - -def _na3d_sdpa_exact( - q_right: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - kernel_size: tuple[int, int, int], - w_offset: int = 0, -) -> torch.Tensor: - """Exact per-query 3D neighborhood attention via grouped tiled SDPA. - - When called with the full q tensor and w_offset=0, computes exact neighborhood - attention for all W positions. When called with a W sub-slice and non-zero - w_offset, K/V are drawn from the full k/v tensors using global W coordinates - - used by correct_rightmost_block. - - Args: - q_right : ``(B, T, H, W_sub, NH, HD)`` bfloat16, Q pre-scaled. - k, v : ``(B, T, H, W_full, NH, HD)`` bfloat16. - kernel_size: ``(KT, KH, KW)`` neighborhood window. - w_offset : Global W start of q_right; 0 for a full-tensor call. - - Returns: - ``(B, T, H, W_sub, NH, HD)`` bfloat16. - """ - B, T, H, W_sub, NH, HD = q_right.shape - W_full = k.shape[3] - KT, KH, KW = kernel_size - device = q_right.device - - bt = _window_bounds_1d(T, min(KT, T)) - bh = _window_bounds_1d(H, min(KH, H)) - bw = _window_bounds_1d(W_full, min(KW, W_full)) # global W bounds - - # Pick tile sizes based on q_right, - # but estimate K/V union sizes using the full W. - eff_kt, eff_kh = min(KT, T), min(KH, H) - eff_kw = min(KW, W_full) - - def _cost(ts: list[int]) -> int: - tile_t, tile_h, tile_w = ts - nq = tile_t * tile_h * tile_w - nk_t = min(T, tile_t + eff_kt - 1) - nk_h = min(H, tile_h + eff_kh - 1) - nk_w = min(W_full, tile_w + eff_kw - 1) - return nq * nk_t * nk_h * nk_w - - tiles = [T, H, W_sub] - while _cost(tiles) > 2**25 and max(tiles) > 1: - i = max(range(3), key=lambda a: tiles[a] / [eff_kt, eff_kh, eff_kw][a]) - if tiles[i] <= 1: - break - tiles[i] = max(1, (tiles[i] + 1) // 2) - tile_t, tile_h, tile_w = tiles - - groups: dict = {} - for t0 in range(0, T, tile_t): - t1 = min(t0 + tile_t, T) - rt0, rt1 = bt[0][t0], bt[1][t1 - 1] - rel_t = ( - tuple(s - rt0 for s in bt[0][t0:t1]), - tuple(e - rt0 for e in bt[1][t0:t1]), - ) - for h0 in range(0, H, tile_h): - h1 = min(h0 + tile_h, H) - rh0, rh1 = bh[0][h0], bh[1][h1 - 1] - rel_h = ( - tuple(s - rh0 for s in bh[0][h0:h1]), - tuple(e - rh0 for e in bh[1][h0:h1]), - ) - for w0 in range(0, W_sub, tile_w): - w1 = min(w0 + tile_w, W_sub) - w0g, w1g = w_offset + w0, w_offset + w1 - 1 - rw0, rw1 = bw[0][w0g], bw[1][w1g] - rel_w = ( - tuple(s - rw0 for s in bw[0][w0g : w1g + 1]), - tuple(e - rw0 for e in bw[1][w0g : w1g + 1]), - ) - groups.setdefault((rel_t, rel_h, rel_w), []).append( - ( - (slice(t0, t1), slice(h0, h1), slice(w0, w1)), - (slice(rt0, rt1), slice(rh0, rh1), slice(rw0, rw1)), - ) - ) - - out = torch.empty((B, T, H, W_sub, NH, HD), device=device, dtype=v.dtype) - kv_budget = 2**28 - for rel, tiles_list in groups.items(): - mask = _na3d_sdpa_mask(rel, q_right.dtype, device) - nq, nk = mask.shape[2], mask.shape[3] - g_max = max(1, kv_budget // max(1, B * NH * nk * HD * 2)) - qs0, _ = tiles_list[0] - tq = qs0[0].stop - qs0[0].start - th_sz = qs0[1].stop - qs0[1].start - tw_sz = qs0[2].stop - qs0[2].start - for c0 in range(0, len(tiles_list), g_max): - chunk = tiles_list[c0 : c0 + g_max] - g = len(chunk) - q_s = torch.stack([q_right[:, qs[0], qs[1], qs[2]] for qs, _ in chunk]) - k_s = torch.stack([k[:, rs[0], rs[1], rs[2]] for _, rs in chunk]) - v_s = torch.stack([v[:, rs[0], rs[1], rs[2]] for _, rs in chunk]) - q_s = q_s.permute(0, 1, 5, 2, 3, 4, 6).reshape(g * B, NH, nq, HD) - k_s = k_s.permute(0, 1, 5, 2, 3, 4, 6).reshape(g * B, NH, nk, HD) - v_s = v_s.permute(0, 1, 5, 2, 3, 4, 6).reshape(g * B, NH, nk, HD) - o = F.scaled_dot_product_attention(q_s, k_s, v_s, attn_mask=mask, scale=1.0) - o = o.view(g, B, NH, tq, th_sz, tw_sz, HD).permute(0, 1, 3, 4, 5, 2, 6) - for i, (qs, _) in enumerate(chunk): - out[:, qs[0], qs[1], qs[2]] = o[i] - return out - - # --------------------------------------------------------------------------- # Public launcher # --------------------------------------------------------------------------- @@ -195,7 +35,6 @@ def na3d_flash_attn( k: torch.Tensor, v: torch.Tensor, kernel_size: tuple[int, int, int], - correct_rightmost_block: bool = False, ) -> torch.Tensor: """Flash-style 3D neighborhood attention. @@ -203,13 +42,6 @@ def na3d_flash_attn( q, k, v : ``(B, T, H, W, NH, HD)`` bfloat16, channels-last. Q is expected to be pre-scaled by ``head_dim ** -0.5``. kernel_size : ``(KT, KH, KW)`` neighborhood window. - correct_rightmost_block : When True, overwrite the last W-block - (``W_last:``) with per-query exact SDPA (via _na3d_sdpa_exact), - where ``W_last = floor((W - 1) / BLOCK_Q) * BLOCK_Q`` for - the autotuned BLOCK_Q. The flash kernel loads ``BF16(0)`` for - out-of-bounds K/V positions in the last W-block (``kv_ok=False``), - which causes a ~0.002 rounding divergence from AOTriton - at those positions. Off by default. Returns: Output tensor ``(B, T, H, W, NH, HD)`` bfloat16. @@ -266,15 +98,4 @@ def _flat(t: torch.Tensor) -> torch.Tensor: KW=KW, ) - out = out_f.reshape(B, NH, T, H, W, HD).permute(0, 2, 3, 4, 1, 5).contiguous() - - if correct_rightmost_block: - # Use the BLOCK_Q actually selected by autotune for this (KT,KH,KW,W) key. - 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) - out[:, :, :, W_last:, :, :] = out_right - - return out + return out_f.reshape(B, NH, T, H, W, HD).permute(0, 2, 3, 4, 1, 5).contiguous() diff --git a/op_tests/triton_tests/attention/test_na3d_flash.py b/op_tests/triton_tests/attention/test_na3d_flash.py index 66a832c2bae..831803376ba 100644 --- a/op_tests/triton_tests/attention/test_na3d_flash.py +++ b/op_tests/triton_tests/attention/test_na3d_flash.py @@ -10,6 +10,8 @@ """ import argparse +import functools +import math import pandas as pd import pytest @@ -19,7 +21,7 @@ import aiter from aiter import dtypes from aiter.jit.utils.chip_info import get_gfx -from aiter.ops.triton.attention.na3d_flash import _na3d_sdpa_exact, na3d_flash_attn +from aiter.ops.triton.attention.na3d_flash import na3d_flash_attn from aiter.test_common import benchmark, checkAllclose, run_perftest torch.set_default_device("cuda") @@ -93,6 +95,141 @@ def na3d_sdpa_ref( return out +# --------------------------------------------------------------------------- +# GPU-grouped exact reference (fast alternative to na3d_sdpa_ref) +# --------------------------------------------------------------------------- + + +@functools.lru_cache(maxsize=256) +def _na3d_sdpa_mask( + rel_bounds: tuple[tuple[tuple[int, ...], tuple[int, ...]], ...], + dtype: torch.dtype, + device: torch.device, +) -> torch.Tensor: + """Cached additive ``[1, 1, Nq, Nk]`` mask for one tile-geometry group.""" + bools = [] + for starts, ends in rel_bounds: + st = torch.tensor(starts, device=device) + en = torch.tensor(ends, device=device) + kj = torch.arange(max(ends), device=device) + bools.append((kj[None, :] >= st[:, None]) & (kj[None, :] < en[:, None])) + visible = ( + bools[0][:, None, None, :, None, None] + & bools[1][None, :, None, None, :, None] + & bools[2][None, None, :, None, None, :] + ) + nq = math.prod(visible.shape[:3]) + nk = math.prod(visible.shape[3:]) + mask = torch.zeros((nq, nk), dtype=dtype, device=device) + mask.masked_fill_(~visible.reshape(nq, nk), torch.finfo(dtype).min) + return mask.reshape(1, 1, nq, nk) + + +def _na3d_sdpa_exact( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + kernel_size: tuple[int, int, int], + w_offset: int = 0, +) -> torch.Tensor: + """Exact per-query 3D neighborhood attention via grouped tiled SDPA. + + Generic GPU-grouped reference function. With ``w_offset=0`` and full + ``q/k/v`` this computes exact neighborhood attention for all W positions + and is equivalent to ``na3d_sdpa_ref`` but GPU-accelerated with cached masks. + + Args: + q : ``(B, T, H, W_sub, NH, HD)`` bfloat16, Q pre-scaled. + k, v : ``(B, T, H, W_full, NH, HD)`` bfloat16. + kernel_size: ``(KT, KH, KW)`` neighborhood window. + w_offset : Global W start of q; 0 for a full-tensor call. + + Returns: + ``(B, T, H, W_sub, NH, HD)`` bfloat16. + """ + B, T, H, W_sub, NH, HD = q.shape + W_full = k.shape[3] + KT, KH, KW = kernel_size + device = q.device + + bt = _window_bounds(T, min(KT, T)) + bh = _window_bounds(H, min(KH, H)) + bw = _window_bounds(W_full, min(KW, W_full)) + + eff_kt, eff_kh, eff_kw = min(KT, T), min(KH, H), min(KW, W_full) + + def _cost(ts: list[int]) -> int: + tile_t, tile_h, tile_w = ts + nq = tile_t * tile_h * tile_w + nk_t = min(T, tile_t + eff_kt - 1) + nk_h = min(H, tile_h + eff_kh - 1) + nk_w = min(W_full, tile_w + eff_kw - 1) + return nq * nk_t * nk_h * nk_w + + tiles = [T, H, W_sub] + while _cost(tiles) > 2**25 and max(tiles) > 1: + i = max(range(3), key=lambda a: tiles[a] / [eff_kt, eff_kh, eff_kw][a]) + if tiles[i] <= 1: + break + tiles[i] = max(1, (tiles[i] + 1) // 2) + tile_t, tile_h, tile_w = tiles + + groups: dict = {} + for t0 in range(0, T, tile_t): + t1 = min(t0 + tile_t, T) + rt0, rt1 = bt[0][t0], bt[1][t1 - 1] + rel_t = ( + tuple(s - rt0 for s in bt[0][t0:t1]), + tuple(e - rt0 for e in bt[1][t0:t1]), + ) + for h0 in range(0, H, tile_h): + h1 = min(h0 + tile_h, H) + rh0, rh1 = bh[0][h0], bh[1][h1 - 1] + rel_h = ( + tuple(s - rh0 for s in bh[0][h0:h1]), + tuple(e - rh0 for e in bh[1][h0:h1]), + ) + for w0 in range(0, W_sub, tile_w): + w1 = min(w0 + tile_w, W_sub) + w0g, w1g = w_offset + w0, w_offset + w1 - 1 + rw0, rw1 = bw[0][w0g], bw[1][w1g] + rel_w = ( + tuple(s - rw0 for s in bw[0][w0g : w1g + 1]), + tuple(e - rw0 for e in bw[1][w0g : w1g + 1]), + ) + groups.setdefault((rel_t, rel_h, rel_w), []).append( + ( + (slice(t0, t1), slice(h0, h1), slice(w0, w1)), + (slice(rt0, rt1), slice(rh0, rh1), slice(rw0, rw1)), + ) + ) + + out = torch.empty((B, T, H, W_sub, NH, HD), device=device, dtype=v.dtype) + kv_budget = 2**28 + for rel, tiles_list in groups.items(): + mask = _na3d_sdpa_mask(rel, q.dtype, device) + nq, nk = mask.shape[2], mask.shape[3] + g_max = max(1, kv_budget // max(1, B * NH * nk * HD * 2)) + qs0, _ = tiles_list[0] + tq = qs0[0].stop - qs0[0].start + th_sz = qs0[1].stop - qs0[1].start + tw_sz = qs0[2].stop - qs0[2].start + for c0 in range(0, len(tiles_list), g_max): + chunk = tiles_list[c0 : c0 + g_max] + g = len(chunk) + q_s = torch.stack([q[:, qs[0], qs[1], qs[2]] for qs, _ in chunk]) + k_s = torch.stack([k[:, rs[0], rs[1], rs[2]] for _, rs in chunk]) + v_s = torch.stack([v[:, rs[0], rs[1], rs[2]] for _, rs in chunk]) + q_s = q_s.permute(0, 1, 5, 2, 3, 4, 6).reshape(g * B, NH, nq, HD) + k_s = k_s.permute(0, 1, 5, 2, 3, 4, 6).reshape(g * B, NH, nk, HD) + v_s = v_s.permute(0, 1, 5, 2, 3, 4, 6).reshape(g * B, NH, nk, HD) + o = F.scaled_dot_product_attention(q_s, k_s, v_s, attn_mask=mask, scale=1.0) + o = o.view(g, B, NH, tq, th_sz, tw_sz, HD).permute(0, 1, 3, 4, 5, 2, 6) + for i, (qs, _) in enumerate(chunk): + out[:, qs[0], qs[1], qs[2]] = o[i] + return out + + # --------------------------------------------------------------------------- # Default LTX-2.5 and edge-case shapes # --------------------------------------------------------------------------- @@ -138,17 +275,6 @@ def na3d_sdpa_ref( ] -# --------------------------------------------------------------------------- -# Shapes to test correct_rightmost_block -# --------------------------------------------------------------------------- -_CORRECTION_SHAPES = [ - (1, 16, 80, 32, 4, 64, 11, 11, 11), # diff-5 small tile - (1, 16, 80, 192, 4, 64, 11, 11, 11), # diff-5 wide tile - (1, 18, 32, 48, 32, 64, 3, 7, 7), # det-0 - (1, 36, 64, 96, 8, 64, 3, 5, 5), # det-2 -] - - # --------------------------------------------------------------------------- # Pytest correctness tests # --------------------------------------------------------------------------- @@ -255,65 +381,6 @@ def test_na3d_sdpa_exact_vs_ref(B, T, H, W, NH, HD, KT, KH, KW): ) -@pytest.mark.parametrize("B,T,H,W,NH,HD,KT,KH,KW", _CORRECTION_SHAPES) -def test_correct_rightmost_block(B, T, H, W, NH, HD, KT, KH, KW): - """Verify the ``correct_rightmost_block`` docstring claim. - - Asserts: - 1. Interior W positions are bit-identical before and after correction. - 2. Rightmost W positions differ (correction fires and changes output). - 3. Corrected rightmost positions match ``_na3d_sdpa_exact`` applied to - just those queries, confirming the kv_ok=False BF16(0) MFMA divergence - is eliminated. - """ - if get_gfx() not in SUPPORTED_GFX: - pytest.skip(f"na3d_flash unsupported on {get_gfx()}") - - scale = HD**-0.5 - q = torch.randn(B, T, H, W, NH, HD, dtype=torch.bfloat16) * scale - k = torch.randn(B, T, H, W, NH, HD, dtype=torch.bfloat16) - v = torch.randn(B, T, H, W, NH, HD, dtype=torch.bfloat16) - - out_base = na3d_flash_attn(q, k, v, (KT, KH, KW), correct_rightmost_block=False) - out_corr = na3d_flash_attn(q, k, v, (KT, KH, KW), correct_rightmost_block=True) - - import aiter.ops.triton.attention.na3d_flash as na3d_flash_mod - - actual_bq = na3d_flash_mod._na3d_flash_fwd.best_config.kwargs["BLOCK_Q"] - W_last = ((W - 1) // actual_bq) * actual_bq - - diff = (out_corr.float() - out_base.float()).abs() - - if W_last > 0: - # Claim 1: interior positions unchanged. - interior_max = diff[:, :, :, :W_last, :, :].max().item() - assert interior_max == 0.0, ( - f"Interior W positions changed (max_diff={interior_max:.6f}) - " - f"correction must not touch positions 0..{W_last - 1}" - ) - - # Claim 2: rightmost positions differ (correction fired). - rightmost_max = diff[:, :, :, W_last:, :, :].max().item() - assert rightmost_max > 0.0, ( - f"Rightmost W positions unchanged - " - f"T*H*(W-W_last)={T*H*(W-W_last)} should be <= 100_000 so correction fires" - ) - - # Claim 3: corrected positions match _na3d_sdpa_exact. - q_right = q[:, :, :, W_last:, :, :] - exact_right = _na3d_sdpa_exact(q_right, k, v, (KT, KH, KW), w_offset=W_last) - checkAllclose( - exact_right.float(), - out_corr[:, :, :, W_last:, :, :].float(), - rtol=1e-2, - atol=5e-2, - msg=( - f"corrected rightmost vs _na3d_sdpa_exact " - f"(T={T},H={H},W={W},W_last={W_last},k=({KT},{KH},{KW}))" - ), - ) - - def main(): if get_gfx() not in SUPPORTED_GFX: aiter.logger.warning("na3d_flash unsupported on %s; skipping", get_gfx()) From 2617cc8493a84c225b7dc3201eb79b793e5649b6 Mon Sep 17 00:00:00 2001 From: Joni Juvonen Date: Thu, 3 Sep 2026 14:21:55 +0300 Subject: [PATCH 09/48] Prune autotune configs where BLOCK_KV is too small for KW Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- .../triton/_triton_kernels/attention/na3d_flash.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/aiter/ops/triton/_triton_kernels/attention/na3d_flash.py b/aiter/ops/triton/_triton_kernels/attention/na3d_flash.py index 5254b4c426f..98a53adeba3 100644 --- a/aiter/ops/triton/_triton_kernels/attention/na3d_flash.py +++ b/aiter/ops/triton/_triton_kernels/attention/na3d_flash.py @@ -39,14 +39,19 @@ def _prune_configs(configs, named_args, **kwargs): - """Remove configs where BLOCK_Q > W. + """Remove configs where BLOCK_Q > W or BLOCK_KV is too small for KW. Performance constraint: BLOCK_Q > W would leave most of every block masked, - wasting compute. Correctness is guaranteed by the row-decomposed grid - regardless of whether W is divisible by BLOCK_Q. + wasting compute. Correctness also requires BLOCK_KV >= (BLOCK_Q + KW - 1) + so the KV tile covers the union of all BLOCK_Q query windows. """ W = named_args["W"] - return [c for c in configs if c.kwargs["BLOCK_Q"] <= W] + KW = named_args["KW"] + return [ + c + for c in configs + if c.kwargs["BLOCK_Q"] <= W and c.kwargs["BLOCK_KV"] >= c.kwargs["BLOCK_Q"] + KW - 1 + ] _na3d_flash_fwd_repr = make_kernel_repr( From 3cf7288177751bf6c39ceb5ace853bf4bec11af6 Mon Sep 17 00:00:00 2001 From: jjuvonen-amd Date: Thu, 3 Sep 2026 11:36:23 +0000 Subject: [PATCH 10/48] fix global state leak in pytest, fix config prune conditions Signed-off-by: jjuvonen-amd --- .../triton/_triton_kernels/attention/na3d_flash.py | 5 +++-- aiter/ops/triton/attention/na3d_flash.py | 5 ----- op_tests/triton_tests/attention/test_na3d_flash.py | 13 ++++++++++++- 3 files changed, 15 insertions(+), 8 deletions(-) diff --git a/aiter/ops/triton/_triton_kernels/attention/na3d_flash.py b/aiter/ops/triton/_triton_kernels/attention/na3d_flash.py index 98a53adeba3..f233d6d3d34 100644 --- a/aiter/ops/triton/_triton_kernels/attention/na3d_flash.py +++ b/aiter/ops/triton/_triton_kernels/attention/na3d_flash.py @@ -46,11 +46,12 @@ def _prune_configs(configs, named_args, **kwargs): so the KV tile covers the union of all BLOCK_Q query windows. """ W = named_args["W"] - KW = named_args["KW"] + KW = kwargs.get("KW") # constexpr in kwargs return [ c for c in configs - if c.kwargs["BLOCK_Q"] <= W and c.kwargs["BLOCK_KV"] >= c.kwargs["BLOCK_Q"] + KW - 1 + if c.kwargs["BLOCK_Q"] <= W + and (KW is None or c.kwargs["BLOCK_KV"] >= c.kwargs["BLOCK_Q"] + KW - 1) ] diff --git a/aiter/ops/triton/attention/na3d_flash.py b/aiter/ops/triton/attention/na3d_flash.py index a55621cf6aa..70b70eda2aa 100644 --- a/aiter/ops/triton/attention/na3d_flash.py +++ b/aiter/ops/triton/attention/na3d_flash.py @@ -25,11 +25,6 @@ _LOGGER = AiterTritonLogger() -# --------------------------------------------------------------------------- -# Public launcher -# --------------------------------------------------------------------------- - - def na3d_flash_attn( q: torch.Tensor, k: torch.Tensor, diff --git a/op_tests/triton_tests/attention/test_na3d_flash.py b/op_tests/triton_tests/attention/test_na3d_flash.py index 831803376ba..f892502cbec 100644 --- a/op_tests/triton_tests/attention/test_na3d_flash.py +++ b/op_tests/triton_tests/attention/test_na3d_flash.py @@ -24,7 +24,18 @@ from aiter.ops.triton.attention.na3d_flash import na3d_flash_attn from aiter.test_common import benchmark, checkAllclose, run_perftest -torch.set_default_device("cuda") + +@pytest.fixture(autouse=True) +def _na3d_default_device(): + if not torch.cuda.is_available(): + pytest.skip("CUDA device required for na3d_flash tests") + prev_device = torch.get_default_device() + try: + torch.set_default_device("cuda") + yield + finally: + torch.set_default_device(prev_device) + SUPPORTED_GFX = ["gfx942", "gfx950"] From 510e65a10e1e0a0fc6077bd23b914aa231eca537 Mon Sep 17 00:00:00 2001 From: jjuvonen-amd Date: Thu, 3 Sep 2026 11:39:19 +0000 Subject: [PATCH 11/48] use actually highest value W in min(), instead of arbitary high value Signed-off-by: jjuvonen-amd --- aiter/ops/triton/_triton_kernels/attention/na3d_flash.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/aiter/ops/triton/_triton_kernels/attention/na3d_flash.py b/aiter/ops/triton/_triton_kernels/attention/na3d_flash.py index f233d6d3d34..3802f5877f2 100644 --- a/aiter/ops/triton/_triton_kernels/attention/na3d_flash.py +++ b/aiter/ops/triton/_triton_kernels/attention/na3d_flash.py @@ -114,7 +114,7 @@ def _na3d_flash_fwd( t_ws = tl.minimum(tl.maximum(prog_t - KT // 2, 0), T - KT) # scalar h_ws = tl.minimum(tl.maximum(prog_h - KH // 2, 0), H - KH) # scalar - w_lo = tl.min(tl.where(q_mask, q_w_ws, 999999)) + w_lo = tl.min(tl.where(q_mask, q_w_ws, W)) hd_offs = tl.arange(0, HD) kv_offs = tl.arange(0, BLOCK_KV) From e00b5d3202348c68bf9d9199bc269a987af46d43 Mon Sep 17 00:00:00 2001 From: jjuvonen-amd Date: Thu, 3 Sep 2026 11:56:00 +0000 Subject: [PATCH 12/48] Add a fast-fail validation to attn call Signed-off-by: jjuvonen-amd --- aiter/ops/triton/attention/na3d_flash.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/aiter/ops/triton/attention/na3d_flash.py b/aiter/ops/triton/attention/na3d_flash.py index 70b70eda2aa..58aac8f5a2e 100644 --- a/aiter/ops/triton/attention/na3d_flash.py +++ b/aiter/ops/triton/attention/na3d_flash.py @@ -63,6 +63,13 @@ def na3d_flash_attn( assert ( KT <= T and KH <= H and KW <= W ), f"na3d_flash_attn: kernel_size=({KT},{KH},{KW}) must be <= (T,H,W)=({T},{H},{W})" + assert ( + KT <= T and KH <= H and KW <= W + ), f"na3d_flash_attn: kernel_size=({KT},{KH},{KW}) must be <= (T,H,W)=({T},{H},{W})" + assert KW <= 33, ( + f"na3d_flash_attn: KW={KW} is too large for the current autotune configs " + f"(max supported KW is 33 with BLOCK_Q=32/BLOCK_KV=64)." + ) assert HD & (HD - 1) == 0, f"head_dim {HD} must be a power of 2" assert W >= 16, f"W={W} is too small; kernel requires W >= BLOCK_Q (default 16)." From 0846c2369047c478f7e665b4a82d6a1ff8663ab2 Mon Sep 17 00:00:00 2001 From: Joni Juvonen Date: Thu, 3 Sep 2026 15:05:35 +0300 Subject: [PATCH 13/48] remove duplicate check Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- aiter/ops/triton/attention/na3d_flash.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/aiter/ops/triton/attention/na3d_flash.py b/aiter/ops/triton/attention/na3d_flash.py index 58aac8f5a2e..535009a1619 100644 --- a/aiter/ops/triton/attention/na3d_flash.py +++ b/aiter/ops/triton/attention/na3d_flash.py @@ -63,9 +63,6 @@ def na3d_flash_attn( assert ( KT <= T and KH <= H and KW <= W ), f"na3d_flash_attn: kernel_size=({KT},{KH},{KW}) must be <= (T,H,W)=({T},{H},{W})" - assert ( - KT <= T and KH <= H and KW <= W - ), f"na3d_flash_attn: kernel_size=({KT},{KH},{KW}) must be <= (T,H,W)=({T},{H},{W})" assert KW <= 33, ( f"na3d_flash_attn: KW={KW} is too large for the current autotune configs " f"(max supported KW is 33 with BLOCK_Q=32/BLOCK_KV=64)." From af1eab58c8a05907fe13ff0c4e6455a4a2ef909d Mon Sep 17 00:00:00 2001 From: Joni Juvonen Date: Thu, 3 Sep 2026 15:18:32 +0300 Subject: [PATCH 14/48] cast indices to int64 to avoid potential OOB reads/writes Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- aiter/ops/triton/_triton_kernels/attention/na3d_flash.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/aiter/ops/triton/_triton_kernels/attention/na3d_flash.py b/aiter/ops/triton/_triton_kernels/attention/na3d_flash.py index 3802f5877f2..4f03a2101f5 100644 --- a/aiter/ops/triton/_triton_kernels/attention/na3d_flash.py +++ b/aiter/ops/triton/_triton_kernels/attention/na3d_flash.py @@ -92,7 +92,9 @@ def _na3d_flash_fwd( """ pid_q = tl.program_id(0) pid_bnh = tl.program_id(1) - HW = H * W + H64 = H.to(tl.int64) + W64 = W.to(tl.int64) + HW = H64 * W64 # Decompose pid_q into (t,h) row and W-block within that row. W_blocks = (W + BLOCK_Q - 1) // BLOCK_Q # programs per (t,h) row @@ -106,8 +108,9 @@ def _na3d_flash_fwd( q_w = w_bid * BLOCK_Q + q_offs # W positions for this block q_mask = q_w < W # last block in a row may be partial - q_idx = prog_t * HW + prog_h * W + q_w # for Q load and Out store - + q_idx = ( + prog_t.to(tl.int64) * HW + prog_h.to(tl.int64) * W64 + q_w.to(tl.int64) + ) # for Q load and Out store # Inward-shifted centered neighborhood window starts. # prog_t / prog_h are scalars; t_ws / h_ws need no tl.min reduction. q_w_ws = tl.minimum(tl.maximum(q_w - KW // 2, 0), W - KW) From 028f9cad8721696bdeab2b3311b43f8ad0816e70 Mon Sep 17 00:00:00 2001 From: Joni Juvonen Date: Thu, 3 Sep 2026 15:31:57 +0300 Subject: [PATCH 15/48] cast int multiplications to int64 to avoid potential OOB reads/writes Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- aiter/ops/triton/_triton_kernels/attention/na3d_flash.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/aiter/ops/triton/_triton_kernels/attention/na3d_flash.py b/aiter/ops/triton/_triton_kernels/attention/na3d_flash.py index 4f03a2101f5..9ddfde518ef 100644 --- a/aiter/ops/triton/_triton_kernels/attention/na3d_flash.py +++ b/aiter/ops/triton/_triton_kernels/attention/na3d_flash.py @@ -148,8 +148,8 @@ def _na3d_flash_fwd( t_kv = t_ws + dt h_kv = h_ws + dh - row_base = t_kv * HW + h_kv * W - kv_flat = row_base + kv_w + row_base = t_kv * HW + h_kv.to(tl.int64) * W64 + kv_flat = row_base + kv_w.to(tl.int64) K_T = tl.load( K_ptr + base + kv_flat[None, :] * stride_seq + hd_offs[:, None], From bdd1f759a919c3620488d12ebced8ec1ead48286 Mon Sep 17 00:00:00 2001 From: jjuvonen-amd Date: Thu, 3 Sep 2026 12:44:09 +0000 Subject: [PATCH 16/48] cast stride_bnh to int64 to avoid exceeding int32 range Signed-off-by: jjuvonen-amd --- aiter/ops/triton/_triton_kernels/attention/na3d_flash.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/aiter/ops/triton/_triton_kernels/attention/na3d_flash.py b/aiter/ops/triton/_triton_kernels/attention/na3d_flash.py index 9ddfde518ef..4aa8c687a20 100644 --- a/aiter/ops/triton/_triton_kernels/attention/na3d_flash.py +++ b/aiter/ops/triton/_triton_kernels/attention/na3d_flash.py @@ -122,8 +122,8 @@ def _na3d_flash_fwd( hd_offs = tl.arange(0, HD) kv_offs = tl.arange(0, BLOCK_KV) - # Cast to int64: pid_bnh (int32) x stride_bnh can exceed int32 range. - base = pid_bnh.to(tl.int64) * stride_bnh + # Cast to int64 to avoid exceeding int32 range + base = pid_bnh.to(tl.int64) * stride_bnh.to(tl.int64) kv_w = w_lo + kv_offs kv_ok = kv_w < W # W-boundary guard, constant across the KT x KH loop From 82cfbbe91b26823fa017a483bb44476302c0b3bc Mon Sep 17 00:00:00 2001 From: jjuvonen-amd Date: Fri, 4 Sep 2026 04:49:04 +0000 Subject: [PATCH 17/48] drop lru_cache on _na3d_sdpa_mask in test file Signed-off-by: jjuvonen-amd --- op_tests/triton_tests/attention/test_na3d_flash.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/op_tests/triton_tests/attention/test_na3d_flash.py b/op_tests/triton_tests/attention/test_na3d_flash.py index f892502cbec..250f57c50f6 100644 --- a/op_tests/triton_tests/attention/test_na3d_flash.py +++ b/op_tests/triton_tests/attention/test_na3d_flash.py @@ -10,7 +10,6 @@ """ import argparse -import functools import math import pandas as pd @@ -111,7 +110,6 @@ def na3d_sdpa_ref( # --------------------------------------------------------------------------- -@functools.lru_cache(maxsize=256) def _na3d_sdpa_mask( rel_bounds: tuple[tuple[tuple[int, ...], tuple[int, ...]], ...], dtype: torch.dtype, From 400e9658b0a4fd651017d6c14be6009d22a138b1 Mon Sep 17 00:00:00 2001 From: jjuvonen-amd Date: Fri, 4 Sep 2026 04:58:18 +0000 Subject: [PATCH 18/48] replace f-string logging with lazy %-style logging Signed-off-by: jjuvonen-amd --- aiter/ops/triton/attention/na3d_flash.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/aiter/ops/triton/attention/na3d_flash.py b/aiter/ops/triton/attention/na3d_flash.py index 535009a1619..08302a549b5 100644 --- a/aiter/ops/triton/attention/na3d_flash.py +++ b/aiter/ops/triton/attention/na3d_flash.py @@ -50,7 +50,13 @@ def na3d_flash_attn( SEQ = T * H * W _LOGGER.info( - f"NA3D_FLASH_FWD: q={tuple(q.shape)} kernel=({KT},{KH},{KW}) SEQ={SEQ} HD={HD}" + "NA3D_FLASH_FWD: q=%s kernel=(%d,%d,%d) SEQ=%d HD=%d", + tuple(q.shape), + KT, + KH, + KW, + SEQ, + HD, ) assert q.dtype == torch.bfloat16, "na3d_flash_attn: inputs must be bfloat16" From 1a3857ade7a965fcdbada145346be1688a31e4e2 Mon Sep 17 00:00:00 2001 From: jjuvonen-amd Date: Fri, 4 Sep 2026 06:12:21 +0000 Subject: [PATCH 19/48] add arch guard to public na3d_flash_attn launcher Signed-off-by: jjuvonen-amd --- aiter/ops/triton/attention/na3d_flash.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/aiter/ops/triton/attention/na3d_flash.py b/aiter/ops/triton/attention/na3d_flash.py index 08302a549b5..e58321f52d0 100644 --- a/aiter/ops/triton/attention/na3d_flash.py +++ b/aiter/ops/triton/attention/na3d_flash.py @@ -19,10 +19,12 @@ import torch import triton +from aiter.jit.utils.chip_info import get_gfx from aiter.ops.triton._triton_kernels.attention.na3d_flash import _na3d_flash_fwd from aiter.ops.triton.utils.logger import AiterTritonLogger _LOGGER = AiterTritonLogger() +_NA3D_FLASH_ARCHS = ("gfx942", "gfx950") def na3d_flash_attn( @@ -49,6 +51,11 @@ def na3d_flash_attn( KT, KH, KW = kernel_size SEQ = T * H * W + gfx = get_gfx() + assert ( + gfx in _NA3D_FLASH_ARCHS + ), f"na3d_flash_attn is only supported on {_NA3D_FLASH_ARCHS}; got {gfx}." + _LOGGER.info( "NA3D_FLASH_FWD: q=%s kernel=(%d,%d,%d) SEQ=%d HD=%d", tuple(q.shape), From 44e8fe4b3b3f42bed6f6068d210f8c7eaebfa566 Mon Sep 17 00:00:00 2001 From: jjuvonen-amd Date: Fri, 4 Sep 2026 06:37:16 +0000 Subject: [PATCH 20/48] split test_na3d_flash.py into separate bench and test files Signed-off-by: jjuvonen-amd --- .../op_benchmarks/triton/bench_na3d_flash.py | 127 ++++++++++++++++++ .../triton_tests/attention/test_na3d_flash.py | 113 ++-------------- 2 files changed, 137 insertions(+), 103 deletions(-) create mode 100644 op_tests/op_benchmarks/triton/bench_na3d_flash.py diff --git a/op_tests/op_benchmarks/triton/bench_na3d_flash.py b/op_tests/op_benchmarks/triton/bench_na3d_flash.py new file mode 100644 index 00000000000..a4a72d4fba8 --- /dev/null +++ b/op_tests/op_benchmarks/triton/bench_na3d_flash.py @@ -0,0 +1,127 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. +"""Performance benchmark for 3D neighborhood flash attention (na3d_flash). + +Usage examples +-------------- +# Sweep default LTX-2.5 shapes, BF16 +python bench_na3d_flash.py + +# Custom shapes +python bench_na3d_flash.py -s '(1,79,192,192,4,64,11,11,11)' + +# FP16 (if supported) +python bench_na3d_flash.py -d fp16 +""" + +import argparse + +import pandas as pd +import torch + +import aiter +from aiter import dtypes +from aiter.jit.utils.chip_info import get_gfx +from aiter.ops.triton.attention.na3d_flash import na3d_flash_attn +from aiter.test_common import benchmark, checkAllclose, run_perftest +from op_tests.triton_tests.attention.test_na3d_flash import ( + _DEFAULT_SHAPES, + _na3d_sdpa_exact, +) + +SUPPORTED_GFX = ["gfx942", "gfx950"] + + +@benchmark() +def bench_na3d_flash(B, T, H, W, NH, HD, KT, KH, KW, dtype): + SEQ = T * H * W + C = NH * HD + K = KT * KH * KW # neighborhood size + + # Build inputs matching the real decoder call: pre-scaled Q, BF16, channels-last. + scale = HD**-0.5 + q = torch.randn(B, T, H, W, NH, HD, dtype=dtype, device="cuda") * scale + k = torch.randn(B, T, H, W, NH, HD, dtype=dtype, device="cuda") + v = torch.randn(B, T, H, W, NH, HD, dtype=dtype, device="cuda") + + # Reference: grouped-SDPA exact path (same as the eager processor fallback). + # Also used as a timed candidate so the table shows the real kernel speedup. + ref = _na3d_sdpa_exact(q, k, v, kernel_size=(KT, KH, KW)) + + # FLOPs/bytes model: logical neighborhood cost. + # sdpa_exact does extra masked work, so its effective TFLOPS is lower, + # which makes the comparison meaningful. + # QK: 2 * B * SEQ * K * C + # AV: 2 * B * SEQ * K * C + flops = 4 * B * SEQ * K * C + # Bytes: Q loaded once, K and V reloaded per (t_kv, h_kv) row. + elem = q.element_size() + nbytes = (B * SEQ * C + 2 * B * SEQ * K * C) * elem + + candidates = { + "triton": lambda: na3d_flash_attn(q, k, v, kernel_size=(KT, KH, KW)), + "sdpa_exact": lambda: _na3d_sdpa_exact(q, k, v, kernel_size=(KT, KH, KW)), + } + + ret = {"gfx": get_gfx()} + for name, fn in candidates.items(): + out, us = run_perftest(fn) + err = checkAllclose( + ref.float(), + out.float(), + rtol=1e-2, + atol=5e-2, + msg=f"{name}: na3d_flash (T={T},H={H},W={W},k=({KT},{KH},{KW}))", + ) + ret[f"{name} us"] = us + ret[f"{name} TFLOPS"] = flops / us / 1e6 + ret[f"{name} TB/s"] = nbytes / us / 1e6 + ret[f"{name} err"] = err + return ret + + +def main(): + if get_gfx() not in SUPPORTED_GFX: + aiter.logger.warning("na3d_flash unsupported on %s; skipping", get_gfx()) + return + + parser = argparse.ArgumentParser( + formatter_class=argparse.RawTextHelpFormatter, + description="Benchmark 3D neighborhood flash attention", + ) + parser.add_argument( + "-d", + "--dtype", + type=dtypes.str2Dtype, + nargs="*", + default=[dtypes.d_dtypes["bf16"]], + help="Input dtype (default: bf16)", + ) + parser.add_argument( + "-s", + "--shapes", + type=dtypes.str2tuple, + nargs="*", + default=_DEFAULT_SHAPES, + help=( + "Shapes as (B,T,H,W,NH,HD,KT,KH,KW) tuples.\n" + "e.g.: -s '(1,79,192,192,4,64,11,11,11)'" + ), + ) + args = parser.parse_args() + + for dtype in args.dtype: + rows = [] + for shape in args.shapes: + B, T, H, W, NH, HD, KT, KH, KW = shape + rows.append(bench_na3d_flash(B, T, H, W, NH, HD, KT, KH, KW, dtype)) + df = pd.DataFrame(rows) + aiter.logger.info( + "na3d_flash summary (%s):\n%s", + str(dtype).replace("torch.", ""), + df.to_markdown(index=False), + ) + + +if __name__ == "__main__": + main() diff --git a/op_tests/triton_tests/attention/test_na3d_flash.py b/op_tests/triton_tests/attention/test_na3d_flash.py index 250f57c50f6..2ca6fd39236 100644 --- a/op_tests/triton_tests/attention/test_na3d_flash.py +++ b/op_tests/triton_tests/attention/test_na3d_flash.py @@ -1,27 +1,23 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. -"""Test and benchmark for 3D neighborhood flash attention (na3d_flash). +"""Correctness tests for 3D neighborhood flash attention (na3d_flash). Covers representative LTX-2.5 VAE decode shapes plus edge cases that exercise boundary masking and non-aligned sequence lengths. -Each shape is compared against a pure-PyTorch tiled SDPA reference that implements -the same inward-shifted centered window. TFLOPS and TB/s are reported per shape. +Each shape is compared against a pure-PyTorch tiled SDPA reference. +Performance benchmarks live in op_tests/op_benchmarks/triton/bench_na3d_flash.py. """ -import argparse import math -import pandas as pd import pytest import torch import torch.nn.functional as F -import aiter -from aiter import dtypes from aiter.jit.utils.chip_info import get_gfx from aiter.ops.triton.attention.na3d_flash import na3d_flash_attn -from aiter.test_common import benchmark, checkAllclose, run_perftest +from aiter.test_common import checkAllclose @pytest.fixture(autouse=True) @@ -115,7 +111,12 @@ def _na3d_sdpa_mask( dtype: torch.dtype, device: torch.device, ) -> torch.Tensor: - """Cached additive ``[1, 1, Nq, Nk]`` mask for one tile-geometry group.""" + """Build an additive ``[1, 1, Nq, Nk]`` mask for one tile-geometry group. + + Not cached: ``_na3d_sdpa_exact`` deduplicates geometries in the ``groups`` + dict before calling this, so each unique geometry is computed exactly once + per call. Caching would only retain GPU tensors across tests with no benefit. + """ bools = [] for starts, ends in rel_bounds: st = torch.tensor(starts, device=device) @@ -310,53 +311,6 @@ def test_na3d_flash(B, T, H, W, NH, HD, KT, KH, KW, dtype): ) -# --------------------------------------------------------------------------- -# Benchmark function (one row per shape) -# --------------------------------------------------------------------------- -@benchmark() -def bench_na3d_flash(B, T, H, W, NH, HD, KT, KH, KW, dtype): - SEQ = T * H * W - C = NH * HD - K = KT * KH * KW # neighborhood size - - # Build inputs matching the real decoder call: pre-scaled Q, BF16, channels-last. - scale = HD**-0.5 - q = torch.randn(B, T, H, W, NH, HD, dtype=dtype) * scale - k = torch.randn(B, T, H, W, NH, HD, dtype=dtype) - v = torch.randn(B, T, H, W, NH, HD, dtype=dtype) - - # Reference (not timed): GPU-grouped exact SDPA, faster than na3d_sdpa_ref loop. - ref = _na3d_sdpa_exact(q, k, v, kernel_size=(KT, KH, KW)) - - # FLOPs: QK and AV dot products over K neighbors per query. - # QK: 2 * B * SEQ * K * C - # AV: 2 * B * SEQ * K * C - flops = 4 * B * SEQ * K * C - # Bytes: Q loaded once, K and V reloaded per (t_kv, h_kv) row. - elem = q.element_size() - nbytes = (B * SEQ * C + 2 * B * SEQ * K * C) * elem - - candidates = { - "triton": lambda: na3d_flash_attn(q, k, v, kernel_size=(KT, KH, KW)), - } - - ret = {"gfx": get_gfx()} - for name, fn in candidates.items(): - out, us = run_perftest(fn) - err = checkAllclose( - ref.float(), - out.float(), - rtol=1e-2, - atol=5e-2, - msg=f"{name}: na3d_flash (T={T},H={H},W={W},k=({KT},{KH},{KW}))", - ) - ret[f"{name} us"] = us - ret[f"{name} TFLOPS"] = flops / us / 1e6 - ret[f"{name} TB/s"] = nbytes / us / 1e6 - ret[f"{name} err"] = err - return ret - - # --------------------------------------------------------------------------- # _na3d_sdpa_exact tests # --------------------------------------------------------------------------- @@ -388,50 +342,3 @@ def test_na3d_sdpa_exact_vs_ref(B, T, H, W, NH, HD, KT, KH, KW): atol=5e-2, msg=f"_na3d_sdpa_exact vs na3d_sdpa_ref (T={T},H={H},W={W},k=({KT},{KH},{KW}))", ) - - -def main(): - if get_gfx() not in SUPPORTED_GFX: - aiter.logger.warning("na3d_flash unsupported on %s; skipping", get_gfx()) - return - - parser = argparse.ArgumentParser( - formatter_class=argparse.RawTextHelpFormatter, - description="Benchmark 3D neighborhood flash attention", - ) - parser.add_argument( - "-d", - "--dtype", - type=dtypes.str2Dtype, - nargs="*", - default=[dtypes.d_dtypes["bf16"]], - help="Input dtype (default: bf16)", - ) - parser.add_argument( - "-s", - "--shapes", - type=dtypes.str2tuple, - nargs="*", - default=_DEFAULT_SHAPES, - help=( - "Shapes as (B,T,H,W,NH,HD,KT,KH,KW) tuples.\n" - "e.g.: -s '(1,79,192,192,4,64,11,11,11)'" - ), - ) - args = parser.parse_args() - - for dtype in args.dtype: - rows = [] - for shape in args.shapes: - B, T, H, W, NH, HD, KT, KH, KW = shape - rows.append(bench_na3d_flash(B, T, H, W, NH, HD, KT, KH, KW, dtype)) - df = pd.DataFrame(rows) - aiter.logger.info( - "na3d_flash summary (%s):\n%s", - str(dtype).replace("torch.", ""), - df.to_markdown(index=False), - ) - - -if __name__ == "__main__": - main() From fa74f050be3a10d0fa4ff55369037a75582919b3 Mon Sep 17 00:00:00 2001 From: jjuvonen-amd Date: Sun, 6 Sep 2026 07:22:24 +0000 Subject: [PATCH 21/48] perf opt: na3d_flash to use default (B,T,H,W,NH,HD) layout to avoid permutes Signed-off-by: jjuvonen-amd --- .../_triton_kernels/attention/na3d_flash.py | 25 +++++++++---- aiter/ops/triton/attention/na3d_flash.py | 35 +++++++++++-------- 2 files changed, 39 insertions(+), 21 deletions(-) diff --git a/aiter/ops/triton/_triton_kernels/attention/na3d_flash.py b/aiter/ops/triton/_triton_kernels/attention/na3d_flash.py index 4aa8c687a20..e7fef877014 100644 --- a/aiter/ops/triton/_triton_kernels/attention/na3d_flash.py +++ b/aiter/ops/triton/_triton_kernels/attention/na3d_flash.py @@ -72,8 +72,10 @@ def _na3d_flash_fwd( K_ptr, V_ptr, Out_ptr, - stride_bnh, - stride_seq, + stride_b, # B stride = SEQ * NH * HD (elements between batches) + stride_nh, # NH stride = HD (elements between heads) + stride_seq, # seq stride = NH * HD (elements between tokens) + NH, # number of attention heads (runtime) T, H, W, @@ -86,9 +88,13 @@ def _na3d_flash_fwd( ): """Flash-attention inner kernel. Grid dim[0] = T*H*ceil(W/BLOCK_Q). - Each program is assigned to exactly one (t, h) row via decomposition of - pid_q into (row_idx, w_block_idx). This guarantees a shared (t, h) window - start regardless of whether W is divisible by BLOCK_Q. + Operates on (B, T, H, W, NH, HD) layout. + Each program is assigned to exactly one (t, h) row via + decomposition of pid_q into (row_idx, w_block_idx). + + Addressing: for batch b and head nh, + element (b, t, h, w, nh, hd) = ptr + b*stride_b + seq*stride_seq + nh*stride_nh + hd + where seq = t*H*W + h*W + w. """ pid_q = tl.program_id(0) pid_bnh = tl.program_id(1) @@ -122,8 +128,13 @@ def _na3d_flash_fwd( hd_offs = tl.arange(0, HD) kv_offs = tl.arange(0, BLOCK_KV) - # Cast to int64 to avoid exceeding int32 range - base = pid_bnh.to(tl.int64) * stride_bnh.to(tl.int64) + # Decompose pid_bnh into (b, nh) for base address. + # int64: stride_b = SEQ * NH * HD can exceed int32 range for large volumes. + b_idx = pid_bnh // NH + nh_idx = pid_bnh % NH + base = b_idx.to(tl.int64) * stride_b.to(tl.int64) + nh_idx.to( + tl.int64 + ) * stride_nh.to(tl.int64) kv_w = w_lo + kv_offs kv_ok = kv_w < W # W-boundary guard, constant across the KT x KH loop diff --git a/aiter/ops/triton/attention/na3d_flash.py b/aiter/ops/triton/attention/na3d_flash.py index e58321f52d0..9680b52ec5f 100644 --- a/aiter/ops/triton/attention/na3d_flash.py +++ b/aiter/ops/triton/attention/na3d_flash.py @@ -44,8 +44,10 @@ def na3d_flash_attn( Output tensor ``(B, T, H, W, NH, HD)`` bfloat16. Notes: - ``W >= 16`` is required so that all queries in a BLOCK_Q=16 program share - the same (t, h) grid row. The autotune pruner enforces this. + Input tensors are consumed in their native ``(B, T, H, W, NH, HD)`` layout + without staging copies. Only a ``contiguous()`` call is made if a tensor is + not already contiguous in that order. ``W >= 16`` is required so that all + queries in a BLOCK_Q=16 program share the same (t, h) grid row. """ B, T, H, W, NH, HD = q.shape KT, KH, KW = kernel_size @@ -83,24 +85,29 @@ def na3d_flash_attn( assert HD & (HD - 1) == 0, f"head_dim {HD} must be a power of 2" assert W >= 16, f"W={W} is too small; kernel requires W >= BLOCK_Q (default 16)." - def _flat(t: torch.Tensor) -> torch.Tensor: - """(B, T, H, W, NH, HD) -> (B*NH, SEQ, HD) contiguous.""" - return t.permute(0, 4, 1, 2, 3, 5).reshape(B * NH, SEQ, HD).contiguous() + # Ensure contiguous layout + q = q.contiguous() + k = k.contiguous() + v = v.contiguous() + out = torch.empty_like(q) - q_f, k_f, v_f = _flat(q), _flat(k), _flat(v) - out_f = torch.empty_like(q_f) + stride_b = SEQ * NH * HD # elements between batches + stride_nh = HD # elements between heads + stride_seq = NH * HD # elements between tokens # Grid: one program per (t, h) row per W-block. This guarantees each # program covers queries from exactly one (t, h) row regardless of W % BLOCK_Q. grid = lambda meta: (T * H * triton.cdiv(W, meta["BLOCK_Q"]), B * NH) _na3d_flash_fwd[grid]( - q_f, - k_f, - v_f, - out_f, - SEQ * HD, # stride_bnh - HD, # stride_seq + q, + k, + v, + out, + stride_b, + stride_nh, + stride_seq, + NH, T, H, W, @@ -110,4 +117,4 @@ def _flat(t: torch.Tensor) -> torch.Tensor: KW=KW, ) - return out_f.reshape(B, NH, T, H, W, HD).permute(0, 2, 3, 4, 1, 5).contiguous() + return out From 4497db9eabf5d60e4477104bd54a72f02c519157 Mon Sep 17 00:00:00 2001 From: jjuvonen-amd Date: Sun, 6 Sep 2026 07:43:35 +0000 Subject: [PATCH 22/48] assert gap in autotune configs with descriptive message Signed-off-by: jjuvonen-amd --- aiter/ops/triton/attention/na3d_flash.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/aiter/ops/triton/attention/na3d_flash.py b/aiter/ops/triton/attention/na3d_flash.py index 9680b52ec5f..94c0df593d9 100644 --- a/aiter/ops/triton/attention/na3d_flash.py +++ b/aiter/ops/triton/attention/na3d_flash.py @@ -82,6 +82,12 @@ def na3d_flash_attn( f"na3d_flash_attn: KW={KW} is too large for the current autotune configs " f"(max supported KW is 33 with BLOCK_Q=32/BLOCK_KV=64)." ) + # Config coverage: BLOCK_Q=16/BLOCK_KV=32 supports KW <= 17; larger KW requires + # BLOCK_Q=32/BLOCK_KV=64 which the pruner keeps only when W >= 32. + assert KW <= 17 or W >= 32, ( + f"na3d_flash_attn: KW={KW} > 17 requires W >= 32 " + f"(only the BLOCK_Q=32/BLOCK_KV=64 config covers KW > 17); got W={W}." + ) assert HD & (HD - 1) == 0, f"head_dim {HD} must be a power of 2" assert W >= 16, f"W={W} is too small; kernel requires W >= BLOCK_Q (default 16)." From 1b9077896a9d54b8e58451a1ff5c4af423ec6b43 Mon Sep 17 00:00:00 2001 From: jjuvonen-amd Date: Sun, 6 Sep 2026 07:44:52 +0000 Subject: [PATCH 23/48] benchmark only valid dtypes (currently bf16) and skip others with warning Signed-off-by: jjuvonen-amd --- op_tests/op_benchmarks/triton/bench_na3d_flash.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/op_tests/op_benchmarks/triton/bench_na3d_flash.py b/op_tests/op_benchmarks/triton/bench_na3d_flash.py index a4a72d4fba8..43e9b602b05 100644 --- a/op_tests/op_benchmarks/triton/bench_na3d_flash.py +++ b/op_tests/op_benchmarks/triton/bench_na3d_flash.py @@ -30,6 +30,7 @@ ) SUPPORTED_GFX = ["gfx942", "gfx950"] +_SUPPORTED_DTYPES = (torch.bfloat16,) @benchmark() @@ -111,6 +112,13 @@ def main(): args = parser.parse_args() for dtype in args.dtype: + if dtype not in _SUPPORTED_DTYPES: + aiter.logger.warning( + "na3d_flash only supports %s; skipping dtype=%s", + _SUPPORTED_DTYPES, + str(dtype).replace("torch.", ""), + ) + continue rows = [] for shape in args.shapes: B, T, H, W, NH, HD, KT, KH, KW = shape From 9ac34867c9d179f6c96dc85674dca6c98c8b4ea2 Mon Sep 17 00:00:00 2001 From: jjuvonen-amd Date: Sun, 6 Sep 2026 07:47:47 +0000 Subject: [PATCH 24/48] remove unsupported fp16 dtype example from bench docstring Signed-off-by: jjuvonen-amd --- op_tests/op_benchmarks/triton/bench_na3d_flash.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/op_tests/op_benchmarks/triton/bench_na3d_flash.py b/op_tests/op_benchmarks/triton/bench_na3d_flash.py index 43e9b602b05..1fcaef1a9bf 100644 --- a/op_tests/op_benchmarks/triton/bench_na3d_flash.py +++ b/op_tests/op_benchmarks/triton/bench_na3d_flash.py @@ -9,9 +9,6 @@ # Custom shapes python bench_na3d_flash.py -s '(1,79,192,192,4,64,11,11,11)' - -# FP16 (if supported) -python bench_na3d_flash.py -d fp16 """ import argparse From 3cedf1d5c549e0d51657271820b837e922c9d88e Mon Sep 17 00:00:00 2001 From: jjuvonen-amd Date: Sun, 6 Sep 2026 14:31:18 +0000 Subject: [PATCH 25/48] annotate strides as tl.int64 and refactor bench file gfx check Signed-off-by: jjuvonen-amd --- aiter/ops/triton/_triton_kernels/attention/na3d_flash.py | 6 +++--- op_tests/op_benchmarks/triton/bench_na3d_flash.py | 9 +++++++-- 2 files changed, 10 insertions(+), 5 deletions(-) diff --git a/aiter/ops/triton/_triton_kernels/attention/na3d_flash.py b/aiter/ops/triton/_triton_kernels/attention/na3d_flash.py index e7fef877014..5de0ae8f7c2 100644 --- a/aiter/ops/triton/_triton_kernels/attention/na3d_flash.py +++ b/aiter/ops/triton/_triton_kernels/attention/na3d_flash.py @@ -72,9 +72,9 @@ def _na3d_flash_fwd( K_ptr, V_ptr, Out_ptr, - stride_b, # B stride = SEQ * NH * HD (elements between batches) - stride_nh, # NH stride = HD (elements between heads) - stride_seq, # seq stride = NH * HD (elements between tokens) + stride_b: tl.int64, # B stride = SEQ * NH * HD (elements between batches) + stride_nh: tl.int64, # NH stride = HD (elements between heads) + stride_seq: tl.int64, # seq stride = NH * HD (elements between tokens) NH, # number of attention heads (runtime) T, H, diff --git a/op_tests/op_benchmarks/triton/bench_na3d_flash.py b/op_tests/op_benchmarks/triton/bench_na3d_flash.py index 1fcaef1a9bf..c1dc1192935 100644 --- a/op_tests/op_benchmarks/triton/bench_na3d_flash.py +++ b/op_tests/op_benchmarks/triton/bench_na3d_flash.py @@ -79,8 +79,13 @@ def bench_na3d_flash(B, T, H, W, NH, HD, KT, KH, KW, dtype): def main(): - if get_gfx() not in SUPPORTED_GFX: - aiter.logger.warning("na3d_flash unsupported on %s; skipping", get_gfx()) + try: + gfx = get_gfx() + except RuntimeError as e: + aiter.logger.warning("na3d_flash: could not detect GPU arch, skipping (%s)", e) + return + if gfx not in SUPPORTED_GFX: + aiter.logger.warning("na3d_flash unsupported on %s; skipping", gfx) return parser = argparse.ArgumentParser( From aa8e3fdac572f9040821e632a4e857719a15628f Mon Sep 17 00:00:00 2001 From: jjuvonen-amd Date: Mon, 7 Sep 2026 06:17:13 +0000 Subject: [PATCH 26/48] renamed autotune configs for clarity Signed-off-by: jjuvonen-amd --- aiter/ops/triton/_triton_kernels/attention/na3d_flash.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/aiter/ops/triton/_triton_kernels/attention/na3d_flash.py b/aiter/ops/triton/_triton_kernels/attention/na3d_flash.py index 5de0ae8f7c2..a93908737df 100644 --- a/aiter/ops/triton/_triton_kernels/attention/na3d_flash.py +++ b/aiter/ops/triton/_triton_kernels/attention/na3d_flash.py @@ -25,7 +25,7 @@ # BLOCK_Q=16 -> BLOCK_KV=32 # BLOCK_Q=32 -> BLOCK_KV=64 # The W >= BLOCK_Q constraint is enforced by the pruner below. -_CONFIGS = [ +_NA3D_AUTOTUNE_CONFIGS = [ triton.Config({"BLOCK_Q": 16, "BLOCK_KV": 32, "num_stages": 2}, num_warps=4), triton.Config({"BLOCK_Q": 16, "BLOCK_KV": 32, "num_stages": 2}, num_warps=8), triton.Config({"BLOCK_Q": 16, "BLOCK_KV": 32, "num_stages": 3}, num_warps=4), @@ -62,7 +62,7 @@ def _prune_configs(configs, named_args, **kwargs): @triton.autotune( - configs=_CONFIGS, + configs=_NA3D_AUTOTUNE_CONFIGS, key=["KT", "KH", "KW", "HD", "W"], prune_configs_by={"early_config_prune": _prune_configs}, ) From 98607a74856068a1884f46c12dfe551713b582c4 Mon Sep 17 00:00:00 2001 From: jjuvonen-amd Date: Mon, 7 Sep 2026 06:24:31 +0000 Subject: [PATCH 27/48] assert lower bounds of KT, KH and KW Signed-off-by: jjuvonen-amd --- aiter/ops/triton/attention/na3d_flash.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/aiter/ops/triton/attention/na3d_flash.py b/aiter/ops/triton/attention/na3d_flash.py index 94c0df593d9..26d34d9e42c 100644 --- a/aiter/ops/triton/attention/na3d_flash.py +++ b/aiter/ops/triton/attention/na3d_flash.py @@ -75,6 +75,9 @@ def na3d_flash_attn( assert ( k.shape == q.shape and v.shape == q.shape ), "na3d_flash_attn: q/k/v must have shape (B, T, H, W, NH, HD)" + assert ( + KT > 0 and KH > 0 and KW > 0 + ), f"na3d_flash_attn: kernel_size dimensions must be positive; got ({KT},{KH},{KW})" assert ( KT <= T and KH <= H and KW <= W ), f"na3d_flash_attn: kernel_size=({KT},{KH},{KW}) must be <= (T,H,W)=({T},{H},{W})" From 17c39a791eab158d522a36dcbbd87bb21fb833b2 Mon Sep 17 00:00:00 2001 From: jjuvonen-amd Date: Mon, 7 Sep 2026 08:18:17 +0000 Subject: [PATCH 28/48] assert HD >= 16 in na3d_flash Signed-off-by: jjuvonen-amd --- aiter/ops/triton/attention/na3d_flash.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/aiter/ops/triton/attention/na3d_flash.py b/aiter/ops/triton/attention/na3d_flash.py index 26d34d9e42c..02fdea9e2f3 100644 --- a/aiter/ops/triton/attention/na3d_flash.py +++ b/aiter/ops/triton/attention/na3d_flash.py @@ -91,7 +91,10 @@ def na3d_flash_attn( f"na3d_flash_attn: KW={KW} > 17 requires W >= 32 " f"(only the BLOCK_Q=32/BLOCK_KV=64 config covers KW > 17); got W={W}." ) - assert HD & (HD - 1) == 0, f"head_dim {HD} must be a power of 2" + assert HD >= 16 and HD & (HD - 1) == 0, ( + f"head_dim {HD} must be a power of 2 and >= 16 " + f"(tl.dot requires matrix dimensions of at least 16)." + ) assert W >= 16, f"W={W} is too small; kernel requires W >= BLOCK_Q (default 16)." # Ensure contiguous layout From 9035e3b27f02bf219435f5c50a5c914ca7b7b681 Mon Sep 17 00:00:00 2001 From: jjuvonen-amd Date: Mon, 7 Sep 2026 08:22:01 +0000 Subject: [PATCH 29/48] verify that q/k/v are CUDA tensors Signed-off-by: jjuvonen-amd --- aiter/ops/triton/attention/na3d_flash.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/aiter/ops/triton/attention/na3d_flash.py b/aiter/ops/triton/attention/na3d_flash.py index 02fdea9e2f3..b21e738746d 100644 --- a/aiter/ops/triton/attention/na3d_flash.py +++ b/aiter/ops/triton/attention/na3d_flash.py @@ -53,6 +53,19 @@ def na3d_flash_attn( KT, KH, KW = kernel_size SEQ = T * H * W + assert ( + q.is_cuda and k.is_cuda and v.is_cuda + ), "na3d_flash_attn requires CUDA/HIP tensors" + assert q.device == k.device == v.device, ( + f"na3d_flash_attn: q/k/v must be on the same device; got " + f"q={q.device}, k={k.device}, v={v.device}" + ) + current_device = torch.cuda.current_device() + assert q.device.index == current_device, ( + f"na3d_flash_attn: kernel will launch on cuda:{current_device}, " + f"but inputs are on {q.device}" + ) + gfx = get_gfx() assert ( gfx in _NA3D_FLASH_ARCHS From ce22c6932b9f3930b13a12c39fc16337dc2610c8 Mon Sep 17 00:00:00 2001 From: jjuvonen-amd Date: Mon, 7 Sep 2026 08:29:17 +0000 Subject: [PATCH 30/48] fix num_stages forwarding bug in autotune configs Signed-off-by: jjuvonen-amd --- .../_triton_kernels/attention/na3d_flash.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/aiter/ops/triton/_triton_kernels/attention/na3d_flash.py b/aiter/ops/triton/_triton_kernels/attention/na3d_flash.py index a93908737df..928ec777f63 100644 --- a/aiter/ops/triton/_triton_kernels/attention/na3d_flash.py +++ b/aiter/ops/triton/_triton_kernels/attention/na3d_flash.py @@ -26,15 +26,15 @@ # BLOCK_Q=32 -> BLOCK_KV=64 # The W >= BLOCK_Q constraint is enforced by the pruner below. _NA3D_AUTOTUNE_CONFIGS = [ - triton.Config({"BLOCK_Q": 16, "BLOCK_KV": 32, "num_stages": 2}, num_warps=4), - triton.Config({"BLOCK_Q": 16, "BLOCK_KV": 32, "num_stages": 2}, num_warps=8), - triton.Config({"BLOCK_Q": 16, "BLOCK_KV": 32, "num_stages": 3}, num_warps=4), - triton.Config({"BLOCK_Q": 16, "BLOCK_KV": 32, "num_stages": 3}, num_warps=8), - triton.Config({"BLOCK_Q": 16, "BLOCK_KV": 32, "num_stages": 4}, num_warps=4), - triton.Config({"BLOCK_Q": 16, "BLOCK_KV": 32, "num_stages": 4}, num_warps=8), + triton.Config({"BLOCK_Q": 16, "BLOCK_KV": 32}, num_warps=4, num_stages=2), + triton.Config({"BLOCK_Q": 16, "BLOCK_KV": 32}, num_warps=8, num_stages=2), + triton.Config({"BLOCK_Q": 16, "BLOCK_KV": 32}, num_warps=4, num_stages=3), + triton.Config({"BLOCK_Q": 16, "BLOCK_KV": 32}, num_warps=8, num_stages=3), + triton.Config({"BLOCK_Q": 16, "BLOCK_KV": 32}, num_warps=4, num_stages=4), + triton.Config({"BLOCK_Q": 16, "BLOCK_KV": 32}, num_warps=8, num_stages=4), # BLOCK_Q=32 halves program count but doubles BLOCK_KV (more masked compute). - triton.Config({"BLOCK_Q": 32, "BLOCK_KV": 64, "num_stages": 2}, num_warps=4), - triton.Config({"BLOCK_Q": 32, "BLOCK_KV": 64, "num_stages": 2}, num_warps=8), + triton.Config({"BLOCK_Q": 32, "BLOCK_KV": 64}, num_warps=4, num_stages=2), + triton.Config({"BLOCK_Q": 32, "BLOCK_KV": 64}, num_warps=8, num_stages=2), ] From d9b9afd63802a0ea965274164bc8e0c80d91512e Mon Sep 17 00:00:00 2001 From: Joni Juvonen Date: Mon, 7 Sep 2026 11:38:51 +0300 Subject: [PATCH 31/48] use runtime variant of get_gfx Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- aiter/ops/triton/attention/na3d_flash.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/aiter/ops/triton/attention/na3d_flash.py b/aiter/ops/triton/attention/na3d_flash.py index b21e738746d..f431297169c 100644 --- a/aiter/ops/triton/attention/na3d_flash.py +++ b/aiter/ops/triton/attention/na3d_flash.py @@ -19,7 +19,7 @@ import torch import triton -from aiter.jit.utils.chip_info import get_gfx +from aiter.jit.utils.chip_info import get_gfx_runtime as get_gfx from aiter.ops.triton._triton_kernels.attention.na3d_flash import _na3d_flash_fwd from aiter.ops.triton.utils.logger import AiterTritonLogger From 2c18894384e120d17e0f1644c46b44a9a1854296 Mon Sep 17 00:00:00 2001 From: jjuvonen-amd Date: Mon, 7 Sep 2026 09:42:56 +0000 Subject: [PATCH 32/48] fix checkAllClose asserts in tests Signed-off-by: jjuvonen-amd --- .../triton_tests/attention/test_na3d_flash.py | 30 +++++++++++-------- 1 file changed, 18 insertions(+), 12 deletions(-) diff --git a/op_tests/triton_tests/attention/test_na3d_flash.py b/op_tests/triton_tests/attention/test_na3d_flash.py index 2ca6fd39236..5fe370fb6d5 100644 --- a/op_tests/triton_tests/attention/test_na3d_flash.py +++ b/op_tests/triton_tests/attention/test_na3d_flash.py @@ -302,12 +302,15 @@ def test_na3d_flash(B, T, H, W, NH, HD, KT, KH, KW, dtype): ref = _na3d_sdpa_exact(q, k, v, kernel_size=(KT, KH, KW)) out = na3d_flash_attn(q, k, v, kernel_size=(KT, KH, KW)) - checkAllclose( - ref.float(), - out.float(), - rtol=1e-2, - atol=5e-2, - msg=f"na3d_flash (T={T},H={H},W={W},k=({KT},{KH},{KW}))", + assert ( + checkAllclose( + ref.float(), + out.float(), + rtol=1e-2, + atol=5e-2, + msg=f"na3d_flash (T={T},H={H},W={W},k=({KT},{KH},{KW}))", + ) + <= 0.05 ) @@ -335,10 +338,13 @@ def test_na3d_sdpa_exact_vs_ref(B, T, H, W, NH, HD, KT, KH, KW): ref = na3d_sdpa_ref(q, k, v, kernel_size=(KT, KH, KW)) exact = _na3d_sdpa_exact(q, k, v, kernel_size=(KT, KH, KW)) # w_offset=0 default - checkAllclose( - ref.float(), - exact.float(), - rtol=1e-2, - atol=5e-2, - msg=f"_na3d_sdpa_exact vs na3d_sdpa_ref (T={T},H={H},W={W},k=({KT},{KH},{KW}))", + assert ( + checkAllclose( + ref.float(), + exact.float(), + rtol=1e-2, + atol=5e-2, + msg=f"_na3d_sdpa_exact vs na3d_sdpa_ref (T={T},H={H},W={W},k=({KT},{KH},{KW}))", + ) + <= 0.05 ) From 085f542d30115fa8c748e2895c24f5b5b994956c Mon Sep 17 00:00:00 2001 From: jjuvonen-amd Date: Mon, 7 Sep 2026 09:53:38 +0000 Subject: [PATCH 33/48] tighten test boundary check from 5% mismatch to 1% Signed-off-by: jjuvonen-amd --- op_tests/triton_tests/attention/test_na3d_flash.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/op_tests/triton_tests/attention/test_na3d_flash.py b/op_tests/triton_tests/attention/test_na3d_flash.py index 5fe370fb6d5..767284e222a 100644 --- a/op_tests/triton_tests/attention/test_na3d_flash.py +++ b/op_tests/triton_tests/attention/test_na3d_flash.py @@ -310,7 +310,7 @@ def test_na3d_flash(B, T, H, W, NH, HD, KT, KH, KW, dtype): atol=5e-2, msg=f"na3d_flash (T={T},H={H},W={W},k=({KT},{KH},{KW}))", ) - <= 0.05 + <= 0.01 ) From c81c79a69767f17ba359ba06bbc4574a89b913e0 Mon Sep 17 00:00:00 2001 From: jjuvonen-amd Date: Tue, 8 Sep 2026 05:07:08 +0000 Subject: [PATCH 34/48] correct HBM formula in bench to report effective rate Signed-off-by: jjuvonen-amd --- .../op_benchmarks/triton/bench_na3d_flash.py | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) diff --git a/op_tests/op_benchmarks/triton/bench_na3d_flash.py b/op_tests/op_benchmarks/triton/bench_na3d_flash.py index c1dc1192935..58c9849a276 100644 --- a/op_tests/op_benchmarks/triton/bench_na3d_flash.py +++ b/op_tests/op_benchmarks/triton/bench_na3d_flash.py @@ -46,15 +46,20 @@ def bench_na3d_flash(B, T, H, W, NH, HD, KT, KH, KW, dtype): # Also used as a timed candidate so the table shows the real kernel speedup. ref = _na3d_sdpa_exact(q, k, v, kernel_size=(KT, KH, KW)) - # FLOPs/bytes model: logical neighborhood cost. - # sdpa_exact does extra masked work, so its effective TFLOPS is lower, - # which makes the comparison meaningful. + # Logical FLOPs: QK and AV dot products over K neighbors per query. # QK: 2 * B * SEQ * K * C # AV: 2 * B * SEQ * K * C flops = 4 * B * SEQ * K * C - # Bytes: Q loaded once, K and V reloaded per (t_kv, h_kv) row. + # Logical bytes: minimum data the operation must touch if every element is + # read/written exactly once: Q, K, V each read once per query, output written once. + # This is an effective/logical rate, not actual HBM traffic (which varies by + # implementation and tile size). The same formula is used for both candidates so + # the comparison is fair; sdpa_exact's lower effective TB/s reflects its higher + # real bandwidth cost. + # reads : Q (SEQ*C) + K (SEQ*K*C) + V (SEQ*K*C) + # writes: Out (SEQ*C) elem = q.element_size() - nbytes = (B * SEQ * C + 2 * B * SEQ * K * C) * elem + nbytes = (2 * B * SEQ * C + 2 * B * SEQ * K * C) * elem candidates = { "triton": lambda: na3d_flash_attn(q, k, v, kernel_size=(KT, KH, KW)), @@ -73,7 +78,7 @@ def bench_na3d_flash(B, T, H, W, NH, HD, KT, KH, KW, dtype): ) ret[f"{name} us"] = us ret[f"{name} TFLOPS"] = flops / us / 1e6 - ret[f"{name} TB/s"] = nbytes / us / 1e6 + ret[f"{name} eff TB/s"] = nbytes / us / 1e6 ret[f"{name} err"] = err return ret From c6e7afefbf057e5ae7fb6886c670f7ef76b394c7 Mon Sep 17 00:00:00 2001 From: jjuvonen-amd Date: Tue, 8 Sep 2026 05:12:07 +0000 Subject: [PATCH 35/48] add test shapes for KW > 17 edge case Signed-off-by: jjuvonen-amd --- op_tests/triton_tests/attention/test_na3d_flash.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/op_tests/triton_tests/attention/test_na3d_flash.py b/op_tests/triton_tests/attention/test_na3d_flash.py index 767284e222a..1ecb12be3a4 100644 --- a/op_tests/triton_tests/attention/test_na3d_flash.py +++ b/op_tests/triton_tests/attention/test_na3d_flash.py @@ -265,6 +265,10 @@ def _cost(ts: list[int]) -> int: (2, 16, 16, 32, 4, 64, 3, 5, 5), # Cross-(t,h)-row: W=33 is not a multiple of BLOCK_Q=16. (1, 4, 8, 33, 4, 64, 3, 5, 5), + # KW > 17 forces BLOCK_Q=32/BLOCK_KV=64: BLOCK_Q=16/BLOCK_KV=32 is pruned + # because 32 < BLOCK_Q + KW - 1 = 16 + 27 - 1 = 42. Only the 32/64 configs + # survive; exercises large-KW masking and the W-boundary of BLOCK_KV=64. + (1, 4, 8, 32, 4, 64, 3, 5, 27), ] @@ -282,6 +286,7 @@ def _cost(ts: list[int]) -> int: (1, 11, 16, 16, 4, 64, 11, 11, 11), # edge: T == KT (SEQ = 2 816) (1, 16, 16, 16, 4, 64, 11, 11, 11), # edge: W == 16 (SEQ = 4 096) (2, 16, 16, 32, 4, 64, 3, 5, 5), # edge: B > 1 (SEQ = 8 192) + (1, 4, 8, 32, 4, 64, 3, 5, 27), # edge: KW=27>17, BLOCK_Q=32 only (SEQ = 1 024) ] From a97d6c62a87ee74f2666201a0c3f4b378504cbcd Mon Sep 17 00:00:00 2001 From: jjuvonen-amd Date: Tue, 8 Sep 2026 05:14:33 +0000 Subject: [PATCH 36/48] document na3d_flash_attn is forward-only Signed-off-by: jjuvonen-amd --- aiter/ops/triton/attention/na3d_flash.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/aiter/ops/triton/attention/na3d_flash.py b/aiter/ops/triton/attention/na3d_flash.py index f431297169c..185f561eaf7 100644 --- a/aiter/ops/triton/attention/na3d_flash.py +++ b/aiter/ops/triton/attention/na3d_flash.py @@ -44,6 +44,10 @@ def na3d_flash_attn( Output tensor ``(B, T, H, W, NH, HD)`` bfloat16. Notes: + **Inference / forward-pass only.** No autograd backward is implemented. + Inputs with ``requires_grad=True`` will produce an output that is detached + from the autograd graph, silently dropping gradients. + Input tensors are consumed in their native ``(B, T, H, W, NH, HD)`` layout without staging copies. Only a ``contiguous()`` call is made if a tensor is not already contiguous in that order. ``W >= 16`` is required so that all From 24f3120a6a27ee0f14aa77ad88c191b85d71fb54 Mon Sep 17 00:00:00 2001 From: jjuvonen-amd Date: Tue, 8 Sep 2026 06:24:54 +0000 Subject: [PATCH 37/48] use runtime gfx check in bench and tests Signed-off-by: jjuvonen-amd --- op_tests/op_benchmarks/triton/bench_na3d_flash.py | 4 ++-- op_tests/triton_tests/attention/test_na3d_flash.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/op_tests/op_benchmarks/triton/bench_na3d_flash.py b/op_tests/op_benchmarks/triton/bench_na3d_flash.py index 58c9849a276..310796350c8 100644 --- a/op_tests/op_benchmarks/triton/bench_na3d_flash.py +++ b/op_tests/op_benchmarks/triton/bench_na3d_flash.py @@ -18,7 +18,7 @@ import aiter from aiter import dtypes -from aiter.jit.utils.chip_info import get_gfx +from aiter.jit.utils.chip_info import get_gfx_runtime as get_gfx from aiter.ops.triton.attention.na3d_flash import na3d_flash_attn from aiter.test_common import benchmark, checkAllclose, run_perftest from op_tests.triton_tests.attention.test_na3d_flash import ( @@ -86,7 +86,7 @@ def bench_na3d_flash(B, T, H, W, NH, HD, KT, KH, KW, dtype): def main(): try: gfx = get_gfx() - except RuntimeError as e: + except (RuntimeError, KeyError) as e: aiter.logger.warning("na3d_flash: could not detect GPU arch, skipping (%s)", e) return if gfx not in SUPPORTED_GFX: diff --git a/op_tests/triton_tests/attention/test_na3d_flash.py b/op_tests/triton_tests/attention/test_na3d_flash.py index 1ecb12be3a4..5604db9225d 100644 --- a/op_tests/triton_tests/attention/test_na3d_flash.py +++ b/op_tests/triton_tests/attention/test_na3d_flash.py @@ -15,7 +15,7 @@ import torch import torch.nn.functional as F -from aiter.jit.utils.chip_info import get_gfx +from aiter.jit.utils.chip_info import get_gfx_runtime as get_gfx from aiter.ops.triton.attention.na3d_flash import na3d_flash_attn from aiter.test_common import checkAllclose From 44c90fd983b4234ee03846ad711379a55429446e Mon Sep 17 00:00:00 2001 From: jjuvonen-amd Date: Wed, 9 Sep 2026 14:01:24 +0000 Subject: [PATCH 38/48] fix W.to() crash when W=1, add test shape with W=1 Signed-off-by: jjuvonen-amd --- aiter/ops/triton/_triton_kernels/attention/na3d_flash.py | 3 +-- op_tests/triton_tests/attention/test_na3d_flash.py | 4 ++++ 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/aiter/ops/triton/_triton_kernels/attention/na3d_flash.py b/aiter/ops/triton/_triton_kernels/attention/na3d_flash.py index 928ec777f63..76eca0c27f0 100644 --- a/aiter/ops/triton/_triton_kernels/attention/na3d_flash.py +++ b/aiter/ops/triton/_triton_kernels/attention/na3d_flash.py @@ -98,9 +98,8 @@ def _na3d_flash_fwd( """ pid_q = tl.program_id(0) pid_bnh = tl.program_id(1) - H64 = H.to(tl.int64) W64 = W.to(tl.int64) - HW = H64 * W64 + HW = H * W64 # promotes to int64 # Decompose pid_q into (t,h) row and W-block within that row. W_blocks = (W + BLOCK_Q - 1) // BLOCK_Q # programs per (t,h) row diff --git a/op_tests/triton_tests/attention/test_na3d_flash.py b/op_tests/triton_tests/attention/test_na3d_flash.py index 5604db9225d..a63da28af8b 100644 --- a/op_tests/triton_tests/attention/test_na3d_flash.py +++ b/op_tests/triton_tests/attention/test_na3d_flash.py @@ -269,6 +269,9 @@ def _cost(ts: list[int]) -> int: # because 32 < BLOCK_Q + KW - 1 = 16 + 27 - 1 = 42. Only the 32/64 configs # survive; exercises large-KW masking and the W-boundary of BLOCK_KV=64. (1, 4, 8, 32, 4, 64, 3, 5, 27), + # H=1 degenerate case: Triton specialises runtime args equal to 1 to a Python + # constexpr (no .to() method), crashing H.to(tl.int64). KH=1 is required (KH<=H). + (1, 4, 1, 32, 4, 64, 3, 1, 5), ] @@ -287,6 +290,7 @@ def _cost(ts: list[int]) -> int: (1, 16, 16, 16, 4, 64, 11, 11, 11), # edge: W == 16 (SEQ = 4 096) (2, 16, 16, 32, 4, 64, 3, 5, 5), # edge: B > 1 (SEQ = 8 192) (1, 4, 8, 32, 4, 64, 3, 5, 27), # edge: KW=27>17, BLOCK_Q=32 only (SEQ = 1 024) + (1, 4, 1, 32, 4, 64, 3, 1, 5), # edge: H=1, Triton constexpr (SEQ = 128) ] From e9eff6885a8f159234b703f3d92c23269ab9937c Mon Sep 17 00:00:00 2001 From: jjuvonen-amd Date: Mon, 14 Sep 2026 13:46:48 +0000 Subject: [PATCH 39/48] set default configs and keep opt-in autotune Signed-off-by: jjuvonen-amd --- .../_triton_kernels/attention/na3d_flash.py | 93 +++++++++++++++---- .../triton/attention/na3d_flash/DEFAULT.json | 14 +++ .../triton/attention/na3d_flash/DEFAULT.json | 14 +++ 3 files changed, 102 insertions(+), 19 deletions(-) create mode 100644 aiter/ops/triton/configs/gfx942/triton/attention/na3d_flash/DEFAULT.json create mode 100644 aiter/ops/triton/configs/gfx950/triton/attention/na3d_flash/DEFAULT.json diff --git a/aiter/ops/triton/_triton_kernels/attention/na3d_flash.py b/aiter/ops/triton/_triton_kernels/attention/na3d_flash.py index 76eca0c27f0..f0c1b1df33b 100644 --- a/aiter/ops/triton/_triton_kernels/attention/na3d_flash.py +++ b/aiter/ops/triton/_triton_kernels/attention/na3d_flash.py @@ -11,31 +11,86 @@ - Q is loaded once into registers and held for all KT x KH inner iterations. - Each iteration loads one (t_kv, h_kv) row of K and V at the shared W window. - Running online softmax in log2 space (exp2, log2(e) folded into scale). -- Autotuned over BLOCK_Q (16, 32), num_stages (2, 3, 4), num_warps (4, 8). +- Optionally autotuned (NA3D_FLASH_TRITON_AUTOTUNE=1) over BLOCK_Q (16, 32), + num_stages (2, 3, 4), num_warps (4, 8). """ +import os + import triton import triton.language as tl from aiter.ops.triton.utils._triton.kernel_repr import make_kernel_repr +from aiter.ops.triton.utils.tuned_config_utils import get_tuned_kernel_config + + +def _na3d_autotune_configs(): + """Full autotune search space (opt-in via NA3D_FLASH_TRITON_AUTOTUNE=1). + + BLOCK_KV = next_pow2(BLOCK_Q + KW - 1) covers the union of all BLOCK_Q queries' + W windows in one chunk (BLOCK_Q=16 -> 32, BLOCK_Q=32 -> 64). The BLOCK_Q=32 + family halves program count but doubles BLOCK_KV (more masked compute), so it + is only swept at num_stages=2. W >= BLOCK_Q is enforced by the pruner below. + """ + configs = [] + for block_q, block_kv, stages in ((16, 32, (2, 3, 4)), (32, 64, (2,))): + for ns in stages: + for nw in (4, 8): + configs.append( + triton.Config( + {"BLOCK_Q": block_q, "BLOCK_KV": block_kv}, + num_warps=nw, + num_stages=ns, + ) + ) + return configs + + +_NA3D_AUTOTUNE_CONFIGS = _na3d_autotune_configs() + +# Live-search opt-in. Default (unset) uses the per-arch tiles published in +# configs//triton/attention/na3d_flash/DEFAULT.json; set to 1 to re-tune +# over the full _na3d_autotune_configs() search space. +_NA3D_FLASH_TRITON_AUTOTUNE = os.getenv("NA3D_FLASH_TRITON_AUTOTUNE", "0").lower() in ( + "1", + "true", + "yes", + "on", +) + +# Always-launchable fallbacks used when an arch has no published DEFAULT.json entry. +# One per KW-family: the pruner keeps BLOCK_Q=16 for KW <= 17 and BLOCK_Q=32 for +# KW > 17 (BLOCK_KV must cover BLOCK_Q + KW - 1). +_FALLBACK_SMALL_KW = triton.Config( + {"BLOCK_Q": 16, "BLOCK_KV": 32}, num_warps=4, num_stages=2 +) +_FALLBACK_LARGE_KW = triton.Config( + {"BLOCK_Q": 32, "BLOCK_KV": 64}, num_warps=4, num_stages=2 +) -# Autotune configs. -# BLOCK_KV = next_pow2(BLOCK_Q + KW - 1) covers the union of all BLOCK_Q -# queries' W windows in one chunk: -# BLOCK_Q=16 -> BLOCK_KV=32 -# BLOCK_Q=32 -> BLOCK_KV=64 -# The W >= BLOCK_Q constraint is enforced by the pruner below. -_NA3D_AUTOTUNE_CONFIGS = [ - triton.Config({"BLOCK_Q": 16, "BLOCK_KV": 32}, num_warps=4, num_stages=2), - triton.Config({"BLOCK_Q": 16, "BLOCK_KV": 32}, num_warps=8, num_stages=2), - triton.Config({"BLOCK_Q": 16, "BLOCK_KV": 32}, num_warps=4, num_stages=3), - triton.Config({"BLOCK_Q": 16, "BLOCK_KV": 32}, num_warps=8, num_stages=3), - triton.Config({"BLOCK_Q": 16, "BLOCK_KV": 32}, num_warps=4, num_stages=4), - triton.Config({"BLOCK_Q": 16, "BLOCK_KV": 32}, num_warps=8, num_stages=4), - # BLOCK_Q=32 halves program count but doubles BLOCK_KV (more masked compute). - triton.Config({"BLOCK_Q": 32, "BLOCK_KV": 64}, num_warps=4, num_stages=2), - triton.Config({"BLOCK_Q": 32, "BLOCK_KV": 64}, num_warps=8, num_stages=2), -] + +def _na3d_default_configs(): + """Per-arch default tiles from JSON, one per KW-family. + + Both are returned; the pruner keeps only the family valid for the runtime KW + (BLOCK_Q=16 for KW <= 17, BLOCK_Q=32 for KW > 17), so this is at most a 2-way + select for KW <= 17 and a single config for KW > 17 -- not a full search. + A single pinned tile cannot serve both branches (a KW > 17 shape with only the + BLOCK_Q=16 tile would leave the pruner with an empty config set). + """ + return [ + get_tuned_kernel_config( + "attention", "NA3D_FLASH", "na3d_flash_small_kw", _FALLBACK_SMALL_KW + ), + get_tuned_kernel_config( + "attention", "NA3D_FLASH", "na3d_flash_large_kw", _FALLBACK_LARGE_KW + ), + ] + + +_NA3D_ACTIVE_CONFIGS = ( + _NA3D_AUTOTUNE_CONFIGS if _NA3D_FLASH_TRITON_AUTOTUNE else _na3d_default_configs() +) def _prune_configs(configs, named_args, **kwargs): @@ -62,7 +117,7 @@ def _prune_configs(configs, named_args, **kwargs): @triton.autotune( - configs=_NA3D_AUTOTUNE_CONFIGS, + configs=_NA3D_ACTIVE_CONFIGS, key=["KT", "KH", "KW", "HD", "W"], prune_configs_by={"early_config_prune": _prune_configs}, ) diff --git a/aiter/ops/triton/configs/gfx942/triton/attention/na3d_flash/DEFAULT.json b/aiter/ops/triton/configs/gfx942/triton/attention/na3d_flash/DEFAULT.json new file mode 100644 index 00000000000..7876b638b05 --- /dev/null +++ b/aiter/ops/triton/configs/gfx942/triton/attention/na3d_flash/DEFAULT.json @@ -0,0 +1,14 @@ +{ + "na3d_flash_small_kw": { + "BLOCK_Q": 16, + "BLOCK_KV": 32, + "num_warps": 4, + "num_stages": 2 + }, + "na3d_flash_large_kw": { + "BLOCK_Q": 32, + "BLOCK_KV": 64, + "num_warps": 4, + "num_stages": 2 + } +} diff --git a/aiter/ops/triton/configs/gfx950/triton/attention/na3d_flash/DEFAULT.json b/aiter/ops/triton/configs/gfx950/triton/attention/na3d_flash/DEFAULT.json new file mode 100644 index 00000000000..7876b638b05 --- /dev/null +++ b/aiter/ops/triton/configs/gfx950/triton/attention/na3d_flash/DEFAULT.json @@ -0,0 +1,14 @@ +{ + "na3d_flash_small_kw": { + "BLOCK_Q": 16, + "BLOCK_KV": 32, + "num_warps": 4, + "num_stages": 2 + }, + "na3d_flash_large_kw": { + "BLOCK_Q": 32, + "BLOCK_KV": 64, + "num_warps": 4, + "num_stages": 2 + } +} From c307a74b4292412547570b540929fb67eb1ada04 Mon Sep 17 00:00:00 2001 From: jjuvonen-amd Date: Mon, 14 Sep 2026 14:12:12 +0000 Subject: [PATCH 40/48] pruner to read named_args and kwargs; derive triton fallback configs from generator Signed-off-by: jjuvonen-amd --- .../_triton_kernels/attention/na3d_flash.py | 25 ++++++++++++------- 1 file changed, 16 insertions(+), 9 deletions(-) diff --git a/aiter/ops/triton/_triton_kernels/attention/na3d_flash.py b/aiter/ops/triton/_triton_kernels/attention/na3d_flash.py index f0c1b1df33b..65e20d08c3d 100644 --- a/aiter/ops/triton/_triton_kernels/attention/na3d_flash.py +++ b/aiter/ops/triton/_triton_kernels/attention/na3d_flash.py @@ -58,14 +58,18 @@ def _na3d_autotune_configs(): "on", ) -# Always-launchable fallbacks used when an arch has no published DEFAULT.json entry. -# One per KW-family: the pruner keeps BLOCK_Q=16 for KW <= 17 and BLOCK_Q=32 for -# KW > 17 (BLOCK_KV must cover BLOCK_Q + KW - 1). -_FALLBACK_SMALL_KW = triton.Config( - {"BLOCK_Q": 16, "BLOCK_KV": 32}, num_warps=4, num_stages=2 +# Always-launchable per-KW-family fallbacks, used only when an arch has no published +# DEFAULT.json entry (no accelerator at import, unpublished arch, or unreadable file). +# get_tuned_kernel_config() requires a Python triton.Config fallback by design: this is +# the mandated not-JSON path -- a conservative correctness floor, not a copy of the tuned +# defaults. Derived from the single search-space source so there are no parallel literals +# to drift: the minimal tile (fewest num_stages) of each BLOCK_Q family. The pruner then +# keeps BLOCK_Q=16 for KW <= 17 and BLOCK_Q=32 for KW > 17. +_FALLBACK_SMALL_KW = next( + c for c in _NA3D_AUTOTUNE_CONFIGS if c.kwargs["BLOCK_Q"] == 16 ) -_FALLBACK_LARGE_KW = triton.Config( - {"BLOCK_Q": 32, "BLOCK_KV": 64}, num_warps=4, num_stages=2 +_FALLBACK_LARGE_KW = next( + c for c in _NA3D_AUTOTUNE_CONFIGS if c.kwargs["BLOCK_Q"] == 32 ) @@ -100,8 +104,11 @@ def _prune_configs(configs, named_args, **kwargs): wasting compute. Correctness also requires BLOCK_KV >= (BLOCK_Q + KW - 1) so the KV tile covers the union of all BLOCK_Q query windows. """ - W = named_args["W"] - KW = kwargs.get("KW") # constexpr in kwargs + # Triton routes positional launch args into named_args and keyword launch args + # into kwargs. This launcher passes W positionally and KW as a keyword, but read + # from both so the size check stays active regardless of how either is passed. + W = named_args.get("W", kwargs.get("W")) + KW = named_args.get("KW", kwargs.get("KW")) return [ c for c in configs From bf3765343216f406e2c597ce2c371e94e92b63f8 Mon Sep 17 00:00:00 2001 From: jjuvonen-amd Date: Thu, 17 Sep 2026 17:02:43 +0000 Subject: [PATCH 41/48] removed opt-in @triton.autotune Signed-off-by: jjuvonen-amd --- .../_triton_kernels/attention/na3d_flash.py | 113 +++--------------- aiter/ops/triton/attention/na3d_flash.py | 23 +++- 2 files changed, 33 insertions(+), 103 deletions(-) diff --git a/aiter/ops/triton/_triton_kernels/attention/na3d_flash.py b/aiter/ops/triton/_triton_kernels/attention/na3d_flash.py index 65e20d08c3d..512c624a097 100644 --- a/aiter/ops/triton/_triton_kernels/attention/na3d_flash.py +++ b/aiter/ops/triton/_triton_kernels/attention/na3d_flash.py @@ -11,110 +11,32 @@ - Q is loaded once into registers and held for all KT x KH inner iterations. - Each iteration loads one (t_kv, h_kv) row of K and V at the shared W window. - Running online softmax in log2 space (exp2, log2(e) folded into scale). -- Optionally autotuned (NA3D_FLASH_TRITON_AUTOTUNE=1) over BLOCK_Q (16, 32), - num_stages (2, 3, 4), num_warps (4, 8). +- Tile (BLOCK_Q, BLOCK_KV, num_warps, num_stages) is read from the per-arch config + configs//triton/attention/na3d_flash/DEFAULT.json (no runtime autotune). """ -import os +import functools import triton import triton.language as tl from aiter.ops.triton.utils._triton.kernel_repr import make_kernel_repr -from aiter.ops.triton.utils.tuned_config_utils import get_tuned_kernel_config +from aiter.ops.triton.utils.config_utils import load_config_json, resolve_config_dir -def _na3d_autotune_configs(): - """Full autotune search space (opt-in via NA3D_FLASH_TRITON_AUTOTUNE=1). +@functools.cache +def _get_config(KW: int) -> dict: + """Per-arch tile from configs//triton/attention/na3d_flash/DEFAULT.json. - BLOCK_KV = next_pow2(BLOCK_Q + KW - 1) covers the union of all BLOCK_Q queries' - W windows in one chunk (BLOCK_Q=16 -> 32, BLOCK_Q=32 -> 64). The BLOCK_Q=32 - family halves program count but doubles BLOCK_KV (more masked compute), so it - is only swept at num_stages=2. W >= BLOCK_Q is enforced by the pruner below. + Two published entries cover the two KW families: + - na3d_flash_small_kw (BLOCK_Q=16, BLOCK_KV=32) for KW <= 17 + - na3d_flash_large_kw (BLOCK_Q=32, BLOCK_KV=64) for KW in [18, 33] + BLOCK_KV must satisfy BLOCK_KV >= BLOCK_Q + KW - 1 so the KV tile covers the + union of all BLOCK_Q query windows; both entries meet that at their KW bound. """ - configs = [] - for block_q, block_kv, stages in ((16, 32, (2, 3, 4)), (32, 64, (2,))): - for ns in stages: - for nw in (4, 8): - configs.append( - triton.Config( - {"BLOCK_Q": block_q, "BLOCK_KV": block_kv}, - num_warps=nw, - num_stages=ns, - ) - ) - return configs - - -_NA3D_AUTOTUNE_CONFIGS = _na3d_autotune_configs() - -# Live-search opt-in. Default (unset) uses the per-arch tiles published in -# configs//triton/attention/na3d_flash/DEFAULT.json; set to 1 to re-tune -# over the full _na3d_autotune_configs() search space. -_NA3D_FLASH_TRITON_AUTOTUNE = os.getenv("NA3D_FLASH_TRITON_AUTOTUNE", "0").lower() in ( - "1", - "true", - "yes", - "on", -) - -# Always-launchable per-KW-family fallbacks, used only when an arch has no published -# DEFAULT.json entry (no accelerator at import, unpublished arch, or unreadable file). -# get_tuned_kernel_config() requires a Python triton.Config fallback by design: this is -# the mandated not-JSON path -- a conservative correctness floor, not a copy of the tuned -# defaults. Derived from the single search-space source so there are no parallel literals -# to drift: the minimal tile (fewest num_stages) of each BLOCK_Q family. The pruner then -# keeps BLOCK_Q=16 for KW <= 17 and BLOCK_Q=32 for KW > 17. -_FALLBACK_SMALL_KW = next( - c for c in _NA3D_AUTOTUNE_CONFIGS if c.kwargs["BLOCK_Q"] == 16 -) -_FALLBACK_LARGE_KW = next( - c for c in _NA3D_AUTOTUNE_CONFIGS if c.kwargs["BLOCK_Q"] == 32 -) - - -def _na3d_default_configs(): - """Per-arch default tiles from JSON, one per KW-family. - - Both are returned; the pruner keeps only the family valid for the runtime KW - (BLOCK_Q=16 for KW <= 17, BLOCK_Q=32 for KW > 17), so this is at most a 2-way - select for KW <= 17 and a single config for KW > 17 -- not a full search. - A single pinned tile cannot serve both branches (a KW > 17 shape with only the - BLOCK_Q=16 tile would leave the pruner with an empty config set). - """ - return [ - get_tuned_kernel_config( - "attention", "NA3D_FLASH", "na3d_flash_small_kw", _FALLBACK_SMALL_KW - ), - get_tuned_kernel_config( - "attention", "NA3D_FLASH", "na3d_flash_large_kw", _FALLBACK_LARGE_KW - ), - ] - - -_NA3D_ACTIVE_CONFIGS = ( - _NA3D_AUTOTUNE_CONFIGS if _NA3D_FLASH_TRITON_AUTOTUNE else _na3d_default_configs() -) - - -def _prune_configs(configs, named_args, **kwargs): - """Remove configs where BLOCK_Q > W or BLOCK_KV is too small for KW. - - Performance constraint: BLOCK_Q > W would leave most of every block masked, - wasting compute. Correctness also requires BLOCK_KV >= (BLOCK_Q + KW - 1) - so the KV tile covers the union of all BLOCK_Q query windows. - """ - # Triton routes positional launch args into named_args and keyword launch args - # into kwargs. This launcher passes W positionally and KW as a keyword, but read - # from both so the size check stays active regardless of how either is passed. - W = named_args.get("W", kwargs.get("W")) - KW = named_args.get("KW", kwargs.get("KW")) - return [ - c - for c in configs - if c.kwargs["BLOCK_Q"] <= W - and (KW is None or c.kwargs["BLOCK_KV"] >= c.kwargs["BLOCK_Q"] + KW - 1) - ] + cfg_dir = resolve_config_dir("attention", "NA3D_FLASH", backend="triton") + config = load_config_json(f"{cfg_dir}/DEFAULT.json") + return config["na3d_flash_small_kw" if KW <= 17 else "na3d_flash_large_kw"] _na3d_flash_fwd_repr = make_kernel_repr( @@ -123,11 +45,6 @@ def _prune_configs(configs, named_args, **kwargs): ) -@triton.autotune( - configs=_NA3D_ACTIVE_CONFIGS, - key=["KT", "KH", "KW", "HD", "W"], - prune_configs_by={"early_config_prune": _prune_configs}, -) @triton.jit(repr=_na3d_flash_fwd_repr) def _na3d_flash_fwd( Q_ptr, diff --git a/aiter/ops/triton/attention/na3d_flash.py b/aiter/ops/triton/attention/na3d_flash.py index 185f561eaf7..6cc45989299 100644 --- a/aiter/ops/triton/attention/na3d_flash.py +++ b/aiter/ops/triton/attention/na3d_flash.py @@ -20,7 +20,10 @@ import triton from aiter.jit.utils.chip_info import get_gfx_runtime as get_gfx -from aiter.ops.triton._triton_kernels.attention.na3d_flash import _na3d_flash_fwd +from aiter.ops.triton._triton_kernels.attention.na3d_flash import ( + _get_config, + _na3d_flash_fwd, +) from aiter.ops.triton.utils.logger import AiterTritonLogger _LOGGER = AiterTritonLogger() @@ -99,11 +102,11 @@ def na3d_flash_attn( KT <= T and KH <= H and KW <= W ), f"na3d_flash_attn: kernel_size=({KT},{KH},{KW}) must be <= (T,H,W)=({T},{H},{W})" assert KW <= 33, ( - f"na3d_flash_attn: KW={KW} is too large for the current autotune configs " + f"na3d_flash_attn: KW={KW} is too large for the published configs " f"(max supported KW is 33 with BLOCK_Q=32/BLOCK_KV=64)." ) - # Config coverage: BLOCK_Q=16/BLOCK_KV=32 supports KW <= 17; larger KW requires - # BLOCK_Q=32/BLOCK_KV=64 which the pruner keeps only when W >= 32. + # Config coverage: BLOCK_Q=16/BLOCK_KV=32 (na3d_flash_small_kw) supports KW <= 17; + # larger KW uses BLOCK_Q=32/BLOCK_KV=64 (na3d_flash_large_kw), which needs W >= 32. assert KW <= 17 or W >= 32, ( f"na3d_flash_attn: KW={KW} > 17 requires W >= 32 " f"(only the BLOCK_Q=32/BLOCK_KV=64 config covers KW > 17); got W={W}." @@ -124,9 +127,15 @@ def na3d_flash_attn( stride_nh = HD # elements between heads stride_seq = NH * HD # elements between tokens + # Per-arch tile from the config file (no runtime autotune). The KW-based + # selection guarantees BLOCK_Q <= W and BLOCK_KV >= BLOCK_Q + KW - 1 given the + # asserts above (W >= 16; KW <= 17 or W >= 32; KW <= 33). + cfg = _get_config(KW) + block_q = cfg["BLOCK_Q"] + # Grid: one program per (t, h) row per W-block. This guarantees each # program covers queries from exactly one (t, h) row regardless of W % BLOCK_Q. - grid = lambda meta: (T * H * triton.cdiv(W, meta["BLOCK_Q"]), B * NH) + grid = (T * H * triton.cdiv(W, block_q), B * NH) _na3d_flash_fwd[grid]( q, @@ -144,6 +153,10 @@ def na3d_flash_attn( KT=KT, KH=KH, KW=KW, + BLOCK_Q=block_q, + BLOCK_KV=cfg["BLOCK_KV"], + num_warps=cfg["num_warps"], + num_stages=cfg["num_stages"], ) return out From a80699cf237bf946d9316aba0c357208e2e2344b Mon Sep 17 00:00:00 2001 From: jjuvonen-amd Date: Thu, 17 Sep 2026 17:04:15 +0000 Subject: [PATCH 42/48] test(na3d_flash): widen large-KW shape to W=128 to catch wrong config Signed-off-by: jjuvonen-amd --- op_tests/triton_tests/attention/test_na3d_flash.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/op_tests/triton_tests/attention/test_na3d_flash.py b/op_tests/triton_tests/attention/test_na3d_flash.py index a63da28af8b..3f06488f95e 100644 --- a/op_tests/triton_tests/attention/test_na3d_flash.py +++ b/op_tests/triton_tests/attention/test_na3d_flash.py @@ -265,10 +265,12 @@ def _cost(ts: list[int]) -> int: (2, 16, 16, 32, 4, 64, 3, 5, 5), # Cross-(t,h)-row: W=33 is not a multiple of BLOCK_Q=16. (1, 4, 8, 33, 4, 64, 3, 5, 5), - # KW > 17 forces BLOCK_Q=32/BLOCK_KV=64: BLOCK_Q=16/BLOCK_KV=32 is pruned - # because 32 < BLOCK_Q + KW - 1 = 16 + 27 - 1 = 42. Only the 32/64 configs - # survive; exercises large-KW masking and the W-boundary of BLOCK_KV=64. - (1, 4, 8, 32, 4, 64, 3, 5, 27), + # KW > 17 must select the BLOCK_Q=32/BLOCK_KV=64 tile. W is chosen large enough + # that each BLOCK_Q=16 block's KV union exceeds BLOCK_KV=32, so a wrongly-selected + # narrow tile corrupts ~9% of outputs (> the 5% test tolerance) and fails -- the + # test detects incorrect large-window config selection. Small W (e.g. 32) keeps + # every union within 32 columns via border-clamping and would hide the error. + (1, 4, 8, 128, 4, 64, 3, 5, 27), # H=1 degenerate case: Triton specialises runtime args equal to 1 to a Python # constexpr (no .to() method), crashing H.to(tl.int64). KH=1 is required (KH<=H). (1, 4, 1, 32, 4, 64, 3, 1, 5), @@ -289,7 +291,7 @@ def _cost(ts: list[int]) -> int: (1, 11, 16, 16, 4, 64, 11, 11, 11), # edge: T == KT (SEQ = 2 816) (1, 16, 16, 16, 4, 64, 11, 11, 11), # edge: W == 16 (SEQ = 4 096) (2, 16, 16, 32, 4, 64, 3, 5, 5), # edge: B > 1 (SEQ = 8 192) - (1, 4, 8, 32, 4, 64, 3, 5, 27), # edge: KW=27>17, BLOCK_Q=32 only (SEQ = 1 024) + (1, 4, 8, 48, 4, 64, 3, 5, 27), # ref cross-check at large KW=27 (SEQ = 1 536) (1, 4, 1, 32, 4, 64, 3, 1, 5), # edge: H=1, Triton constexpr (SEQ = 128) ] From fd4c13a1c9babb602d75862113575064e14641f9 Mon Sep 17 00:00:00 2001 From: jjuvonen-amd Date: Thu, 17 Sep 2026 17:26:17 +0000 Subject: [PATCH 43/48] test(na3d_flash): use assertAllclose so failures keep diagnostics Signed-off-by: jjuvonen-amd --- .../triton_tests/attention/test_na3d_flash.py | 37 +++++++++---------- 1 file changed, 18 insertions(+), 19 deletions(-) diff --git a/op_tests/triton_tests/attention/test_na3d_flash.py b/op_tests/triton_tests/attention/test_na3d_flash.py index 3f06488f95e..528b3c387b8 100644 --- a/op_tests/triton_tests/attention/test_na3d_flash.py +++ b/op_tests/triton_tests/attention/test_na3d_flash.py @@ -17,7 +17,7 @@ from aiter.jit.utils.chip_info import get_gfx_runtime as get_gfx from aiter.ops.triton.attention.na3d_flash import na3d_flash_attn -from aiter.test_common import checkAllclose +from aiter.test_common import assertAllclose @pytest.fixture(autouse=True) @@ -313,15 +313,16 @@ def test_na3d_flash(B, T, H, W, NH, HD, KT, KH, KW, dtype): ref = _na3d_sdpa_exact(q, k, v, kernel_size=(KT, KH, KW)) out = na3d_flash_attn(q, k, v, kernel_size=(KT, KH, KW)) - assert ( - checkAllclose( - ref.float(), - out.float(), - rtol=1e-2, - atol=5e-2, - msg=f"na3d_flash (T={T},H={H},W={W},k=({KT},{KH},{KW}))", - ) - <= 0.01 + # assertAllclose threads one tol_err_ratio into both the assertion and + # checkAllclose's diagnostics, so a 1-5% regression fails AND logs the mismatch + # report at ERROR (survives pytest's WARNING-level logger). + assertAllclose( + ref.float(), + out.float(), + rtol=1e-2, + atol=5e-2, + tol_err_ratio=0.01, + msg=f"na3d_flash (T={T},H={H},W={W},k=({KT},{KH},{KW}))", ) @@ -349,13 +350,11 @@ def test_na3d_sdpa_exact_vs_ref(B, T, H, W, NH, HD, KT, KH, KW): ref = na3d_sdpa_ref(q, k, v, kernel_size=(KT, KH, KW)) exact = _na3d_sdpa_exact(q, k, v, kernel_size=(KT, KH, KW)) # w_offset=0 default - assert ( - checkAllclose( - ref.float(), - exact.float(), - rtol=1e-2, - atol=5e-2, - msg=f"_na3d_sdpa_exact vs na3d_sdpa_ref (T={T},H={H},W={W},k=({KT},{KH},{KW}))", - ) - <= 0.05 + assertAllclose( + ref.float(), + exact.float(), + rtol=1e-2, + atol=5e-2, + tol_err_ratio=0.05, + msg=f"_na3d_sdpa_exact vs na3d_sdpa_ref (T={T},H={H},W={W},k=({KT},{KH},{KW}))", ) From 4fad40404b91235f95cc3a19845b1bc8c6a8b504 Mon Sep 17 00:00:00 2001 From: jjuvonen-amd Date: Thu, 17 Sep 2026 17:40:41 +0000 Subject: [PATCH 44/48] bench(na3d_flash): match decoder v layout so timing includes staging copy Signed-off-by: jjuvonen-amd --- op_tests/op_benchmarks/triton/bench_na3d_flash.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/op_tests/op_benchmarks/triton/bench_na3d_flash.py b/op_tests/op_benchmarks/triton/bench_na3d_flash.py index 310796350c8..cd770d660ac 100644 --- a/op_tests/op_benchmarks/triton/bench_na3d_flash.py +++ b/op_tests/op_benchmarks/triton/bench_na3d_flash.py @@ -36,11 +36,18 @@ def bench_na3d_flash(B, T, H, W, NH, HD, KT, KH, KW, dtype): C = NH * HD K = KT * KH * KW # neighborhood size - # Build inputs matching the real decoder call: pre-scaled Q, BF16, channels-last. + # Match the LTX-2.5 decoder layout so the timed call includes the wrapper + # overhead users actually pay. In the fused-QKV processor q/k pass through + # norm + rope (fresh contiguous tensors), but v is returned directly as + # qkv.chunk(3, dim=-1)[2].view(...), i.e. NON-contiguous (W-stride = 3*NH*HD). + # na3d_flash_attn's v.contiguous() is therefore a real full-volume staging copy + # here, not the no-op it would be for a freshly allocated contiguous v. scale = HD**-0.5 q = torch.randn(B, T, H, W, NH, HD, dtype=dtype, device="cuda") * scale k = torch.randn(B, T, H, W, NH, HD, dtype=dtype, device="cuda") - v = torch.randn(B, T, H, W, NH, HD, dtype=dtype, device="cuda") + qkv = torch.randn(B, T, H, W, 3 * NH * HD, dtype=dtype, device="cuda") + v = qkv.chunk(3, dim=-1)[2].view(B, T, H, W, NH, HD) + assert not v.is_contiguous() # consumer layout: v.contiguous() copies in-call # Reference: grouped-SDPA exact path (same as the eager processor fallback). # Also used as a timed candidate so the table shows the real kernel speedup. From 35a5f84bfe9ff87bb8e8ddb43bc8b7b4f4a7a7d9 Mon Sep 17 00:00:00 2001 From: jjuvonen-amd Date: Thu, 17 Sep 2026 17:41:10 +0000 Subject: [PATCH 45/48] docs(na3d_flash): document full supported domain in the launcher docstring Signed-off-by: jjuvonen-amd --- aiter/ops/triton/attention/na3d_flash.py | 23 ++++++++++++++++++----- 1 file changed, 18 insertions(+), 5 deletions(-) diff --git a/aiter/ops/triton/attention/na3d_flash.py b/aiter/ops/triton/attention/na3d_flash.py index 6cc45989299..b82537707b4 100644 --- a/aiter/ops/triton/attention/na3d_flash.py +++ b/aiter/ops/triton/attention/na3d_flash.py @@ -46,15 +46,28 @@ def na3d_flash_attn( Returns: Output tensor ``(B, T, H, W, NH, HD)`` bfloat16. + Supported domain (all asserted): + - Device : running arch in ``{gfx942, gfx950}``; q/k/v on the current + CUDA/HIP device. + - dtype : ``bfloat16`` for q, k, and v (shared). + - Shape : k and v match q's ``(B, T, H, W, NH, HD)``. + - head_dim ``HD`` : power of two and ``>= 16`` (``tl.dot`` needs matrix + dims of at least 16). + - Window : ``0 < KT <= T``, ``0 < KH <= H``, ``0 < KW <= W``. + - Width : ``W >= 16`` (the smallest configured ``BLOCK_Q``). ``KW <= 33``, + and ``KW > 17`` additionally requires ``W >= 32`` -- large KW is served + only by the ``BLOCK_Q=32/BLOCK_KV=64`` config, whose ``BLOCK_KV`` must + cover ``BLOCK_Q + KW - 1``. + Notes: **Inference / forward-pass only.** No autograd backward is implemented. Inputs with ``requires_grad=True`` will produce an output that is detached from the autograd graph, silently dropping gradients. - Input tensors are consumed in their native ``(B, T, H, W, NH, HD)`` layout - without staging copies. Only a ``contiguous()`` call is made if a tensor is - not already contiguous in that order. ``W >= 16`` is required so that all - queries in a BLOCK_Q=16 program share the same (t, h) grid row. + Tensors are consumed in their native ``(B, T, H, W, NH, HD)`` layout; a + ``contiguous()`` copy is made only for an input not already contiguous in + that order. The tile (BLOCK_Q, BLOCK_KV, num_warps, num_stages) is read + from the per-arch config file (no runtime autotune). """ B, T, H, W, NH, HD = q.shape KT, KH, KW = kernel_size @@ -115,7 +128,7 @@ def na3d_flash_attn( f"head_dim {HD} must be a power of 2 and >= 16 " f"(tl.dot requires matrix dimensions of at least 16)." ) - assert W >= 16, f"W={W} is too small; kernel requires W >= BLOCK_Q (default 16)." + assert W >= 16, f"W={W} is too small; W must be >= the smallest BLOCK_Q (16)." # Ensure contiguous layout q = q.contiguous() From 23b9b4301a7fd9822cf4fb49709cf1cad01612ed Mon Sep 17 00:00:00 2001 From: jjuvonen-amd Date: Sat, 26 Sep 2026 15:11:33 +0000 Subject: [PATCH 46/48] tune(na3d_flash): use measured gfx942 tile for the large-KW config (num_stages=4 is 5.6% faster) Signed-off-by: jjuvonen-amd --- .../configs/gfx942/triton/attention/na3d_flash/DEFAULT.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/aiter/ops/triton/configs/gfx942/triton/attention/na3d_flash/DEFAULT.json b/aiter/ops/triton/configs/gfx942/triton/attention/na3d_flash/DEFAULT.json index 7876b638b05..579e0f94006 100644 --- a/aiter/ops/triton/configs/gfx942/triton/attention/na3d_flash/DEFAULT.json +++ b/aiter/ops/triton/configs/gfx942/triton/attention/na3d_flash/DEFAULT.json @@ -9,6 +9,6 @@ "BLOCK_Q": 32, "BLOCK_KV": 64, "num_warps": 4, - "num_stages": 2 + "num_stages": 4 } } From a597430da4915b75507479e33f76a111b33a0f9b Mon Sep 17 00:00:00 2001 From: jjuvonen-amd Date: Sat, 26 Sep 2026 15:38:33 +0000 Subject: [PATCH 47/48] docs(na3d_flash): describe reference masks as deduplicated, not cached Signed-off-by: jjuvonen-amd --- op_tests/triton_tests/attention/test_na3d_flash.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/op_tests/triton_tests/attention/test_na3d_flash.py b/op_tests/triton_tests/attention/test_na3d_flash.py index 528b3c387b8..f4027e12661 100644 --- a/op_tests/triton_tests/attention/test_na3d_flash.py +++ b/op_tests/triton_tests/attention/test_na3d_flash.py @@ -146,7 +146,10 @@ def _na3d_sdpa_exact( Generic GPU-grouped reference function. With ``w_offset=0`` and full ``q/k/v`` this computes exact neighborhood attention for all W positions - and is equivalent to ``na3d_sdpa_ref`` but GPU-accelerated with cached masks. + and is equivalent to ``na3d_sdpa_ref``, but GPU-accelerated: queries are + grouped by tile geometry so each distinct geometry builds its additive mask + once per call (deduplicated, not cached across calls -- see + ``_na3d_sdpa_mask``), and the masks are freed when this function returns. Args: q : ``(B, T, H, W_sub, NH, HD)`` bfloat16, Q pre-scaled. From 26481fb8aae81078737a957d555cd93e13c9b2c6 Mon Sep 17 00:00:00 2001 From: Joni Juvonen Date: Wed, 30 Sep 2026 16:37:00 +0300 Subject: [PATCH 48/48] Update comments to reflect edge-case coverage changes Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- op_tests/triton_tests/attention/test_na3d_flash.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/op_tests/triton_tests/attention/test_na3d_flash.py b/op_tests/triton_tests/attention/test_na3d_flash.py index f4027e12661..63d0fe49f56 100644 --- a/op_tests/triton_tests/attention/test_na3d_flash.py +++ b/op_tests/triton_tests/attention/test_na3d_flash.py @@ -283,7 +283,8 @@ def _cost(ts: list[int]) -> int: # --------------------------------------------------------------------------- # Fast subset for cross-validation against na3d_sdpa_ref (FP32 per-query loop). # Only shapes with SEQ <= 50K to keep na3d_sdpa_ref within a few seconds each. -# Covers all 4 edge-case categories plus 2 LTX-2.5 default shapes. +# Covers 6 edge-case categories plus 2 LTX-2.5 default shapes; the W=33 +# cross-row case is covered only by test_na3d_flash. # Large shapes (36,64,96), (40,96,96), (79,192,192) are excluded here; they # are already exercised in test_na3d_flash via the faster _na3d_sdpa_exact ref. # ---------------------------------------------------------------------------