Antigravity Agent
134 lines
5.1 KiB
Python
134 lines
5.1 KiB
Python
"""
|
|
PagedFieldprintAttention: Custom Triton Kernel Implementation
|
|
===========================================================
|
|
|
|
This module implements the fused CUDA/Triton kernel proposed in Paper 02.
|
|
It computes:
|
|
Output = FusedSoftmax( (Q [K, K_anchor]^T) / sqrt(d) ) [V, V_anchor]
|
|
|
|
By computing the anchor contribution explicitly inside the SRAM before
|
|
writing back to HBM, it avoids the catastrophic memory thrashing of unfused dual-attention.
|
|
"""
|
|
|
|
import torch
|
|
import triton
|
|
import triton.language as tl
|
|
|
|
@triton.jit
|
|
def paged_fieldprint_attention_kernel(
|
|
Q, K_ctx, V_ctx, K_anchor, V_anchor, Out,
|
|
stride_qz, stride_qh, stride_qm, stride_qk,
|
|
stride_kz, stride_kh, stride_kn, stride_kk,
|
|
stride_vz, stride_vh, stride_vn, stride_vk,
|
|
stride_ka_z, stride_ka_h, stride_ka_n, stride_ka_k,
|
|
stride_va_z, stride_va_h, stride_va_n, stride_va_k,
|
|
stride_oz, stride_oh, stride_om, stride_ok,
|
|
Z, H, N_CTX, N_ANCHOR,
|
|
BLOCK_M: tl.constexpr, BLOCK_DMODEL: tl.constexpr, BLOCK_N: tl.constexpr,
|
|
):
|
|
start_m = tl.program_id(0)
|
|
off_hz = tl.program_id(1)
|
|
|
|
# Initialize offsets
|
|
q_offset = off_hz * stride_qh
|
|
k_ctx_offset = off_hz * stride_kh
|
|
v_ctx_offset = off_hz * stride_vh
|
|
k_anchor_offset = off_hz * stride_ka_h
|
|
v_anchor_offset = off_hz * stride_va_h
|
|
|
|
# Block pointers
|
|
offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
|
|
offs_n = tl.arange(0, BLOCK_N)
|
|
offs_d = tl.arange(0, BLOCK_DMODEL)
|
|
|
|
q_ptrs = Q + q_offset + offs_m[:, None] * stride_qm + offs_d[None, :] * stride_qk
|
|
o_ptrs = Out + q_offset + offs_m[:, None] * stride_om + offs_d[None, :] * stride_ok
|
|
|
|
# Load Q
|
|
q = tl.load(q_ptrs, mask=offs_m[:, None] < N_CTX, other=0.0)
|
|
|
|
# Initialize accumulators
|
|
m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf")
|
|
l_i = tl.zeros([BLOCK_M], dtype=tl.float32)
|
|
acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32)
|
|
|
|
sm_scale = 1.0 / (BLOCK_DMODEL ** 0.5)
|
|
|
|
# -------------------------------------------------------------
|
|
# PHASE 1: Process the Cryptographic Anchor Tokens
|
|
# -------------------------------------------------------------
|
|
# We process the anchor tokens first. They compete in the softmax.
|
|
for start_n in range(0, N_ANCHOR, BLOCK_N):
|
|
k_ptrs = K_anchor + k_anchor_offset + (start_n + offs_n[None, :]) * stride_ka_n + offs_d[:, None] * stride_ka_k
|
|
v_ptrs = V_anchor + v_anchor_offset + (start_n + offs_n[:, None]) * stride_va_n + offs_d[None, :] * stride_va_k
|
|
|
|
k = tl.load(k_ptrs, mask=(start_n + offs_n[None, :]) < N_ANCHOR, other=0.0)
|
|
qk = tl.dot(q, k) * sm_scale
|
|
|
|
# Softmax logic
|
|
m_ij = tl.maximum(m_i, tl.max(qk, 1))
|
|
p = tl.math.exp(qk - m_ij[:, None])
|
|
l_ij = tl.sum(p, 1)
|
|
|
|
# Scaling previous acc
|
|
alpha = tl.math.exp(m_i - m_ij)
|
|
acc = acc * alpha[:, None]
|
|
|
|
# Update
|
|
v = tl.load(v_ptrs, mask=(start_n + offs_n[:, None]) < N_ANCHOR, other=0.0)
|
|
acc += tl.dot(p.to(tl.float16), v)
|
|
m_i = m_ij
|
|
l_i = l_i * alpha + l_ij
|
|
|
|
# -------------------------------------------------------------
|
|
# PHASE 2: Process the Standard Context Tokens
|
|
# -------------------------------------------------------------
|
|
for start_n in range(0, N_CTX, BLOCK_N):
|
|
k_ptrs = K_ctx + k_ctx_offset + (start_n + offs_n[None, :]) * stride_kn + offs_d[:, None] * stride_kk
|
|
v_ptrs = V_ctx + v_ctx_offset + (start_n + offs_n[:, None]) * stride_vn + offs_d[None, :] * stride_vk
|
|
|
|
k = tl.load(k_ptrs, mask=(start_n + offs_n[None, :]) < N_CTX, other=0.0)
|
|
qk = tl.dot(q, k) * sm_scale
|
|
|
|
# Softmax logic
|
|
m_ij = tl.maximum(m_i, tl.max(qk, 1))
|
|
p = tl.math.exp(qk - m_ij[:, None])
|
|
l_ij = tl.sum(p, 1)
|
|
|
|
alpha = tl.math.exp(m_i - m_ij)
|
|
acc = acc * alpha[:, None]
|
|
|
|
v = tl.load(v_ptrs, mask=(start_n + offs_n[:, None]) < N_CTX, other=0.0)
|
|
acc += tl.dot(p.to(tl.float16), v)
|
|
m_i = m_ij
|
|
l_i = l_i * alpha + l_ij
|
|
|
|
# Normalize and write back
|
|
acc = acc / l_i[:, None]
|
|
tl.store(o_ptrs, acc.to(tl.float16), mask=offs_m[:, None] < N_CTX)
|
|
|
|
def paged_fieldprint_attention(q, k_ctx, v_ctx, k_anchor, v_anchor):
|
|
# Shape expectations: [batch, heads, seq_len, d_model]
|
|
Z, H, N_CTX, D_HEAD = q.shape
|
|
_, _, N_ANCHOR, _ = k_anchor.shape
|
|
|
|
out = torch.empty_like(q)
|
|
|
|
BLOCK_M = 64
|
|
BLOCK_N = 32
|
|
|
|
grid = (triton.cdiv(N_CTX, BLOCK_M), Z * H)
|
|
|
|
paged_fieldprint_attention_kernel[grid](
|
|
q, k_ctx, v_ctx, k_anchor, v_anchor, out,
|
|
q.stride(0), q.stride(1), q.stride(2), q.stride(3),
|
|
k_ctx.stride(0), k_ctx.stride(1), k_ctx.stride(2), k_ctx.stride(3),
|
|
v_ctx.stride(0), v_ctx.stride(1), v_ctx.stride(2), v_ctx.stride(3),
|
|
k_anchor.stride(0), k_anchor.stride(1), k_anchor.stride(2), k_anchor.stride(3),
|
|
v_anchor.stride(0), v_anchor.stride(1), v_anchor.stride(2), v_anchor.stride(3),
|
|
out.stride(0), out.stride(1), out.stride(2), out.stride(3),
|
|
Z, H, N_CTX, N_ANCHOR,
|
|
BLOCK_M=BLOCK_M, BLOCK_DMODEL=D_HEAD, BLOCK_N=BLOCK_N,
|
|
)
|
|
return out
|