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

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

@@ -0,0 +1,183 @@
+# 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).
+- Autotuned over BLOCK_Q (16, 32), num_stages (2, 3, 4), num_warps (4, 8).
+"""
+
+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:
+#   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.
+
+    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]
+
+
+_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(repr=_na3d_flash_fwd_repr)
+def _na3d_flash_fwd(
+    Q_ptr,
+    K_ptr,
+    V_ptr,
+    Out_ptr,
+    stride_bnh,
+    stride_seq,
+    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).
+
+    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
+
+    # 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 * HW + prog_h * W + q_w  # 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, 999999))
+
+    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],
+    )

aiter/ops/triton/attention/na3d_flash.py

@@ -0,0 +1,101 @@
+# 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.ops.triton._triton_kernels.attention.na3d_flash import _na3d_flash_fwd
+from aiter.ops.triton.utils.logger import AiterTritonLogger
+
+_LOGGER = AiterTritonLogger()
+
+
+# ---------------------------------------------------------------------------
+# Public launcher
+# ---------------------------------------------------------------------------
+
+
+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.
+
+    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 (
+        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)."
+
+    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)
+
+    # 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
+        T,
+        H,
+        W,
+        HD=HD,
+        KT=KT,
+        KH=KH,
+        KW=KW,
+    )
+
+    return out_f.reshape(B, NH, T, H, W, HD).permute(0, 2, 3, 4, 1, 5).contiguous()