Spaces:
Sleeping
Sleeping
File size: 3,523 Bytes
74679d4 | 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 | """Row-invariant decode attention (Triton), usable inside CUDA/HIP graphs.
One program per (query row, head). Row i attends to the logical keys t in [0, pos[i]]
of its layer; `pos` is read from device memory (graph-capturable). Optional tree
addressing: keys t < cl[i] are read at cache index t, keys t >= cl[i] at
anc[i, t - cl[i]] (the slots of the row's own path in a draft tree; its last entry
is the row's own slot). The summation order (blocks of BLOCK_N logical keys) does not
depend on the addresses, so a chain (cl = pos + 1) is bitwise unchanged and a tree
node gets exactly the AR arithmetic of its own path.
"""
from __future__ import annotations
import torch
import triton
import triton.language as tl
@triton.jit
def _row_attention_kernel(Q, K, V, O, pos_ptr, cl_ptr, anc_ptr, sm_scale,
stride_qn, stride_qh, stride_kh, stride_kt,
stride_on, stride_oh,
GROUP: tl.constexpr, HEAD_DIM: tl.constexpr, BLOCK_N: tl.constexpr,
MAXD: tl.constexpr):
i = tl.program_id(0)
h = tl.program_id(1)
kvh = h // GROUP
length = tl.load(pos_ptr + i).to(tl.int32) + 1
cl = tl.load(cl_ptr + i).to(tl.int32)
d = tl.arange(0, HEAD_DIM)
q = tl.load(Q + i * stride_qn + h * stride_qh + d).to(tl.float32) * sm_scale
m_i = tl.full((), -1.0e30, tl.float32)
l_i = tl.full((), 0.0, tl.float32)
acc = tl.zeros((HEAD_DIM,), tl.float32)
k_base = K + kvh * stride_kh
v_base = V + kvh * stride_kh
for s in range(0, length, BLOCK_N):
t = s + tl.arange(0, BLOCK_N)
mask = t < length
e = t - cl
tree = mask & (e >= 0)
a = tl.where(tree, tl.load(anc_ptr + i * MAXD + e, mask=tree, other=0).to(tl.int32), t)
k = tl.load(k_base + a[:, None] * stride_kt + d[None, :], mask=mask[:, None], other=0.0).to(tl.float32)
scores = tl.sum(k * q[None, :], axis=1)
scores = tl.where(mask, scores, -1.0e30)
m_new = tl.maximum(m_i, tl.max(scores, axis=0))
alpha = tl.exp(m_i - m_new)
p = tl.exp(scores - m_new)
p = tl.where(mask, p, 0.0)
v = tl.load(v_base + a[:, None] * stride_kt + d[None, :], mask=mask[:, None], other=0.0).to(tl.float32)
l_i = l_i * alpha + tl.sum(p, axis=0)
acc = acc * alpha + tl.sum(p[:, None] * v, axis=0)
m_i = m_new
out = acc / l_i
tl.store(O + i * stride_on + h * stride_oh + d, out.to(O.dtype.element_ty))
def row_attention(q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor,
pos: torch.Tensor, scale: float, out: torch.Tensor | None = None,
block_n: int = 64, cl: torch.Tensor | None = None,
anc: torch.Tensor | None = None) -> torch.Tensor:
"""q: (n, H, hd); k_cache/v_cache: (Hkv, max_len, hd); pos: int tensor (>= n,) on device.
Returns (n, H, hd): row i attends to positions [0, pos[i]].
"""
n, H, hd = q.shape
Hkv = k_cache.shape[0]
assert k_cache.stride(-1) == 1 and q.stride(-1) == 1
if out is None:
out = torch.empty_like(q)
if cl is None: # plain chains: every key at its own position
cl, anc = pos + 1, pos[:, None]
_row_attention_kernel[(n, H)](
q, k_cache, v_cache, out, pos, cl, anc, scale,
q.stride(0), q.stride(1), k_cache.stride(0), k_cache.stride(1),
out.stride(0), out.stride(1),
GROUP=H // Hkv, HEAD_DIM=hd, BLOCK_N=block_n, MAXD=anc.shape[1],
)
return out
|