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

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

@@ -0,0 +1,181 @@
+# 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.
+
+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).
+- Tile (BLOCK_Q, BLOCK_KV, num_warps, num_stages) is read from the per-arch config
+  configs/<arch>/triton/attention/na3d_flash/DEFAULT.json (no runtime autotune).
+"""
+
+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.config_utils import load_config_json, resolve_config_dir
+
+
+@functools.cache
+def _get_config(KW: int) -> dict:
+    """Per-arch tile from configs/<arch>/triton/attention/na3d_flash/DEFAULT.json.
+
+    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.
+    """
+    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(
+    "_na3d_flash_fwd",
+    ["BLOCK_Q", "BLOCK_KV", "KT", "KH", "KW", "HD"],
+)
+
+
+@triton.jit(repr=_na3d_flash_fwd_repr)
+def _na3d_flash_fwd(
+    Q_ptr,
+    K_ptr,
+    V_ptr,
+    Out_ptr,
+    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,
+    W,
+    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. Grid dim[0] = T*H*ceil(W/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)
+    W64 = W.to(tl.int64)
+    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
+    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_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.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)
+
+    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, W))
+
+    hd_offs = tl.arange(0, HD)
+    kv_offs = tl.arange(0, BLOCK_KV)
+
+    # 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
+
+    # 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.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],
+            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],
+    )

aiter/ops/triton/attention/na3d_flash.py

@@ -0,0 +1,175 @@
+# 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 torch
+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 (
+    _get_config,
+    _na3d_flash_fwd,
+)
+from aiter.ops.triton.utils.logger import AiterTritonLogger
+
+_LOGGER = AiterTritonLogger()
+_NA3D_FLASH_ARCHS = ("gfx942", "gfx950")
+
+
+def na3d_flash_attn(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    kernel_size: tuple[int, int, int],
+) -> 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.
+
+    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.
+
+        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
+    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
+    ), 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),
+        KT,
+        KH,
+        KW,
+        SEQ,
+        HD,
+    )
+
+    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 > 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})"
+    assert KW <= 33, (
+        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 (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}."
+    )
+    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; W must be >= the smallest BLOCK_Q (16)."
+
+    # Ensure contiguous layout
+    q = q.contiguous()
+    k = k.contiguous()
+    v = v.contiguous()
+    out = torch.empty_like(q)
+
+    stride_b = SEQ * NH * HD  # elements between batches
+    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 = (T * H * triton.cdiv(W, block_q), B * NH)
+
+    _na3d_flash_fwd[grid](
+        q,
+        k,
+        v,
+        out,
+        stride_b,
+        stride_nh,
+        stride_seq,
+        NH,
+        T,
+        H,
+        W,
+        HD=HD,
+        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

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": 4
+    }
+}

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
+    }
+}