tinygrad/tinygrad
2
0
Fork 3
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-08-29 09:36:08 +00:00
Code Issues Packages Projects Releases Wiki Activity

fa: remove swa recompute (#17635)

Browse Source
This commit is contained in:
wozeparrot
2026-08-20 14:26:44 -07:00
committed by GitHub
parent c773891e3f
commit d80254c1d9
2 changed files with 20 additions and 47 deletions
Download Patch File Download Diff File Expand all files Collapse all files
extra/thunder/amd/fa.py
-47
View File
@@ -126,49 +126,6 @@ def fused_qkv_rope(xqkv:Tensor, freqs_cis:Tensor, n_heads:int, n_kv_heads:int, h
def _sharded_empty_like(ref:Tensor, axis:int|None=None) -> Tensor:
return _sharded_empty(ref.shape, ref, axis)
@functools.cache
def _windowed_lse(xq:Tensor, xk:Tensor, sinks, W:int) -> Tensor:
B, N, H, hd = xq.shape
H_KV = xk.shape[2]; R = H // H_KV; nb = N // W; sm = hd ** -0.5
q = xq.reshape(B, N, H_KV, R, hd).permute(0, 2, 3, 1, 4).reshape(B, H_KV, R, nb, W, hd).float()
k = xk.permute(0, 2, 1, 3).reshape(B, H_KV, 1, nb, W, hd).float()
k_prev = k.pad((None, None, None, (1, 0), None, None))[:, :, :, :nb]
sc_d = (q @ k.transpose(-1, -2)) * sm
sc_p = (q @ k_prev.transpose(-1, -2)) * sm
li, lj = Tensor.arange(W).reshape(W, 1), Tensor.arange(W).reshape(1, W)
pv = (Tensor.arange(nb).reshape(nb, 1, 1) >= 1)
sc_d = (lj <= li).where(sc_d, -float("inf"))
sc_p = ((li < lj) & pv).where(sc_p, -float("inf"))
m = sc_d.max(-1, keepdim=True).maximum(sc_p.max(-1, keepdim=True))
if sinks is not None: m = m.maximum(sinks.reshape(1, H_KV, R, 1, 1, 1).float())
denom = (sc_d - m).exp().sum(-1, keepdim=True) + (sc_p - m).exp().sum(-1, keepdim=True)
if sinks is not None: denom = denom + (sinks.reshape(1, H_KV, R, 1, 1, 1).float() - m).exp()
return (m + denom.log()).reshape(B, H, N).unsqueeze(2) # (B, H, 1, N), matches saved l_vec
def _windowed_delta(xq:Tensor, xk:Tensor, xv:Tensor, do:Tensor, sinks, W:int) -> Tensor:
B, N, H, hd = xq.shape
H_KV = xk.shape[2]; R = H // H_KV; nb = N // W; sm = hd ** -0.5
q = xq.reshape(B, N, H_KV, R, hd).permute(0, 2, 3, 1, 4).reshape(B, H_KV, R, nb, W, hd).float()
k = xk.permute(0, 2, 1, 3).reshape(B, H_KV, 1, nb, W, hd).float()
v = xv.permute(0, 2, 1, 3).reshape(B, H_KV, 1, nb, W, hd).float()
dob = do.reshape(B, N, H_KV, R, hd).permute(0, 2, 3, 1, 4).reshape(B, H_KV, R, nb, W, hd).float()
k_prev = k.pad((None, None, None, (1, 0), None, None))[:, :, :, :nb]
v_prev = v.pad((None, None, None, (1, 0), None, None))[:, :, :, :nb]
sc_d = (q @ k.transpose(-1, -2)) * sm
sc_p = (q @ k_prev.transpose(-1, -2)) * sm
li, lj = Tensor.arange(W).reshape(W, 1), Tensor.arange(W).reshape(1, W)
pv = (Tensor.arange(nb).reshape(nb, 1, 1) >= 1)
sc_d = (lj <= li).where(sc_d, -float("inf"))
sc_p = ((li < lj) & pv).where(sc_p, -float("inf"))
m = sc_d.max(-1, keepdim=True).maximum(sc_p.max(-1, keepdim=True))
if sinks is not None: m = m.maximum(sinks.reshape(1, H_KV, R, 1, 1, 1).float())
e_d, e_p = (sc_d - m).exp(), (sc_p - m).exp()
denom = e_d.sum(-1, keepdim=True) + e_p.sum(-1, keepdim=True)
if sinks is not None: denom = denom + (sinks.reshape(1, H_KV, R, 1, 1, 1).float() - m).exp()
o = ((e_d / denom) @ v) + ((e_p / denom) @ v_prev)
delta = (dob * o).sum(-1)
return delta.reshape(B, H, N).unsqueeze(2)
def _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, shard_axis_t, single_device, arch, has_sink, window=0):
def grad(dou:UOp, ker:UOp) -> tuple:
do = Tensor(dou, device=dou.device)
@@ -177,8 +134,6 @@ def _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, sha
xq = Tensor(ker.src[3], device=ker.src[3].device)
xk = Tensor(ker.src[4], device=ker.src[4].device)
xv = Tensor(ker.src[5], device=ker.src[5].device)
if window:
l_vec = _windowed_lse(xq, xk, Tensor(ker.src[6], device=ker.src[6].device) if has_sink else None, window)
dq = _sharded_empty((B, H, N, D), xq, axis=shard_axis_t)
GROUP_SIZE = H_local // H_KV_local
@@ -189,8 +144,6 @@ def _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, sha
# delta_vec = (do * attn).sum(-1, dtype=dtypes.float32).transpose(1, 2).unsqueeze(-2).detach()
delta_vec = _sharded_empty((B, H, 1, N), xq, dtype=dtypes.float32, axis=shard_axis_t)
delta_vec, dq = Tensor.custom_kernel(delta_vec, dq, attn, do, fxn=functools.partial(custom_fa_backward_pre, device=single_device, arch=arch, B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D))[:2]
if window:
delta_vec = _windowed_delta(xq, xk, xv, do, Tensor(ker.src[6], device=ker.src[6].device) if has_sink else None, window)
dq, dk_partial, dv_partial = Tensor.custom_kernel(dq, dk_partial, dv_partial, do, xq, xk, xv, l_vec, delta_vec, fxn=functools.partial(custom_fa_backward, device=single_device, arch=arch, B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D, window=window))[:3]
extra/thunder/amd/fa_fwd_causal.cpp
+20
View File
@@ -269,7 +269,9 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
qo_tile<D, float> q_reg_fl;
load<1, qo_tile<D, float>, _gl_QKVO>(q_reg_fl, g.Qg, {batch_idx, tile_idx, head_idx, 0});
#if !WINDOW
mul(q_reg_fl, q_reg_fl, TEMPERATURE_SCALE); // Use sqrtf for clarity
#endif
copy(q_reg, q_reg_fl);
transpose(q_reg_transposed, q_reg);
@@ -288,6 +290,9 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
zero(att_block[0]);
transpose(k_reg_transposed, k_reg);
mma_AtB(att_block[0], k_reg_transposed, q_reg_transposed, att_block[0]);
#if WINDOW
mul(att_block[0], att_block[0], TEMPERATURE_SCALE);
#endif
__builtin_amdgcn_sched_barrier(0);
if constexpr (causal) {
const int kv_end_pos = (min_tile + 1) * KV_BLOCK_SIZE;
@@ -337,6 +342,9 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
zero(att_block[1]);
transpose(k_reg_transposed, k_reg);
mma_AtB(att_block[1], k_reg_transposed, q_reg_transposed, att_block[1]);
#if WINDOW
mul(att_block[1], att_block[1], TEMPERATURE_SCALE);
#endif
#if WINDOW
// window masks interior tiles that causal skips
mask_kv_tile(att_block[1], tile_idx, j - 2, neg_inf_v, lane);
@@ -401,6 +409,9 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
zero(att_block[0]);
transpose(k_reg_transposed, k_reg);
mma_AtB(att_block[0], k_reg_transposed, q_reg_transposed, att_block[0]);
#if WINDOW
mul(att_block[0], att_block[0], TEMPERATURE_SCALE);
#endif
// Finish softmax for QK1
exp2(att_block[1].tiles[1][0], att_block[1].tiles[1][0]);
mul(norm_vec, norm_vec, scale_vec);
@@ -469,6 +480,9 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
zero(att_block[1]);
transpose(k_reg_transposed, k_reg);
mma_AtB(att_block[1], k_reg_transposed, q_reg_transposed, att_block[1]);
#if WINDOW
mul(att_block[1], att_block[1], TEMPERATURE_SCALE);
#endif
// Finish softmax for QK2
exp2(att_block[0].tiles[1][0], att_block[0].tiles[1][0]);
mul(norm_vec, norm_vec, scale_vec);
@@ -535,6 +549,9 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
zero(att_block[0]);
transpose(k_reg_transposed, k_reg);
mma_AtB(att_block[0], k_reg_transposed, q_reg_transposed, att_block[0]);
#if WINDOW
mul(att_block[0], att_block[0], TEMPERATURE_SCALE);
#endif
// Finish softmax for QK3
exp2(att_block[1].tiles[1][0], att_block[1].tiles[1][0]);
mul(norm_vec, norm_vec, scale_vec);
@@ -597,6 +614,9 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
zero(att_block[1]);
transpose(k_reg_transposed, k_reg);
mma_AtB(att_block[1], k_reg_transposed, q_reg_transposed, att_block[1]);
#if WINDOW
mul(att_block[1], att_block[1], TEMPERATURE_SCALE);
#endif
// Finish softmax for QK4
exp2(att_block[0].tiles[1][0], att_block[0].tiles[1][0]);
mul(norm_vec, norm_vec, scale_vec);