Spaces:
Sleeping
Sleeping
| """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 | |
| 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 | |