File size: 5,189 Bytes
b4f5f26 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 | """Local 2D window attention.
Every token attends to all tokens within a square window of Chebyshev radius ``r``
(i.e. ``|dy| <= r and |dx| <= r``) on its 2D grid.
Two numerically equivalent implementations are provided:
* a fused Triton kernel (CUDA only, forward pass only);
* PyTorch ``scaled_dot_product_attention`` with a boolean window mask (CPU, or when
Triton is unavailable).
Both work on a packed batch: every sample is its own token grid of any size.
"""
import torch
import torch.nn.functional as F
try: # Triton is optional; without it the SDPA path is used.
import triton
import triton.language as tl
_HAS_TRITON = True
except Exception: # pragma: no cover
_HAS_TRITON = False
_LOG2E = 1.4426950408889634
# Global switch: "auto" (kernel when possible), "sdpa" (always masked SDPA).
BACKEND = "auto"
if _HAS_TRITON:
@triton.jit
def _window_attn_fwd(Q, K, V, O, DIMS, NH: tl.constexpr, NMAX, WY: tl.constexpr,
WX: tl.constexpr, R: tl.constexpr, D: tl.constexpr,
BR: tl.constexpr, qk_scale):
# One program = one WY x WX tile of queries for one (batch, head).
# Keys are streamed from the tile extended by R on every side ("halo").
t, bh = tl.program_id(0), tl.program_id(1)
b = bh // NH
gyp = tl.load(DIMS + b * 4 + 0)
gxp = tl.load(DIMS + b * 4 + 1)
gyr = tl.load(DIMS + b * 4 + 2)
gxr = tl.load(DIMS + b * 4 + 3)
ntx = gxp // WX
if t >= (gyp // WY) * ntx:
return
ty, tx = t // ntx, t % ntx
base = bh * NMAX * D
SX: tl.constexpr = WX + 2 * R
SY: tl.constexpr = WY + 2 * R
qi = tl.arange(0, WY * WX)
qyr, qxr = qi // WX, qi % WX
qy, qx = ty * WY + qyr, tx * WX + qxr
d = tl.arange(0, D)
qo = base + (qy * gxp + qx)[:, None] * D + d[None, :]
q = tl.load(Q + qo)
okq = (qy < gyr) & (qx < gxr)
kxr = tl.arange(0, BR * SX) % SX
okx = okq[:, None] & (tl.abs(qxr[:, None] - (kxr[None, :] - R)) <= R)
kxg = tx * WX - R + kxr
okg = (kxg >= 0) & (kxg < gxr)
m = tl.full((WY * WX,), -1e30, tl.float32)
l = tl.zeros((WY * WX,), tl.float32)
acc = tl.zeros((WY * WX, D), tl.float32)
for r0 in range(0, SY, BR):
kyr = r0 + tl.arange(0, BR * SX) // SX
kyg = ty * WY - R + kyr
ok = okg & (kyg >= 0) & (kyg < gyr) & (kyr < SY)
off = base + (kyg * gxp + kxg) * D
k = tl.load(K + off[:, None] + d[None, :], mask=ok[:, None], other=0.0)
v = tl.load(V + off[:, None] + d[None, :], mask=ok[:, None], other=0.0)
s = tl.dot(q, tl.trans(k)) * qk_scale
good = okx & (tl.abs(qyr[:, None] - (kyr[None, :] - R)) <= R) & ok[None, :]
s = tl.where(good, s, -1e30)
mn = tl.maximum(m, tl.max(s, 1))
p = tl.where(good, tl.math.exp2(s - mn[:, None]), 0.0)
a = tl.math.exp2(m - mn)
l = l * a + tl.sum(p, 1)
acc = acc * a[:, None] + tl.dot(p.to(v.dtype), v).to(tl.float32)
m = mn
lc = tl.maximum(l, 1e-30)
tl.store(O + qo, (acc / lc[:, None]).to(O.dtype.element_ty))
def _tile_width(r):
# The kernel needs a power-of-two key strip width (WX + 2R).
for wx in (16, 8):
sx = wx + 2 * r
if sx & (sx - 1) == 0:
return wx
raise ValueError(f"window radius {r}: no supported tile width")
@torch._dynamo.disable
def packed_window_attention(q, k, v, r, dims, max_tiles):
"""Windowed attention for a packed batch of differently shaped token grids.
Sample ``b`` is stored row-major in a padded ``gyp x gxp`` grid at the start of the
sequence, with its real ``gyr x gxr`` tokens in the top-left corner; the rest is padding.
dims: (B, 4) int32 ``[gyp, gxp, gyr, gxr]`` (gyp a multiple of 8, gxp of the tile width).
q, k, v: (B, heads, Nmax, head_dim). Padding outputs are zero (kernel) or unused (SDPA).
"""
B, nh, N, D = q.shape
if BACKEND == "auto" and _HAS_TRITON and q.is_cuda and D in (16, 32, 64, 128):
q, k, v = (t.to(torch.bfloat16).contiguous() for t in (q, k, v))
o = torch.zeros_like(q)
_window_attn_fwd[(max_tiles, B * nh)](
q, k, v, o, dims.to(q.device, torch.int32).contiguous(), NH=nh, NMAX=N, WY=8,
WX=_tile_width(r), R=r, D=D, BR=4, qk_scale=D ** -0.5 * _LOG2E, num_warps=8, num_stages=3)
return o
# reference path: per-sample block mask over the packed sequence
n = torch.arange(N, device=q.device)
d = dims.to(q.device).long()
yy, xx = n[None] // d[:, 1:2], n[None] % d[:, 1:2] # (B, N)
valid = (yy < d[:, 2:3]) & (xx < d[:, 3:4])
win = ((yy[:, :, None] - yy[:, None]).abs() <= r) & ((xx[:, :, None] - xx[:, None]).abs() <= r)
m = win & valid[:, None, :]
m = m | (~valid[:, :, None] & torch.eye(N, dtype=torch.bool, device=q.device))
return F.scaled_dot_product_attention(q.contiguous(), k.contiguous(), v.contiguous(),
attn_mask=m[:, None])
|