forked from ccf-ai-infra/TileOPs-Metax
65 lines
1.9 KiB
Python
65 lines
1.9 KiB
Python
import torch
|
|
|
|
|
|
def get_tolerances(dtype: torch.dtype) -> dict:
|
|
"""Return atol/rtol dict for GLA correctness tests."""
|
|
if dtype == torch.float32:
|
|
return {"atol": 1e-2, "rtol": 1e-2}
|
|
elif dtype == torch.float16:
|
|
return {"atol": 5e-2, "rtol": 5e-2}
|
|
else: # bfloat16
|
|
return {"atol": 1e-1, "rtol": 1e-1}
|
|
|
|
|
|
def cosine_sim(a: torch.Tensor, b: torch.Tensor) -> float:
|
|
"""Compute cosine similarity between two tensors (flattened)."""
|
|
a_flat = a.float().flatten()
|
|
b_flat = b.float().flatten()
|
|
return (torch.dot(a_flat, b_flat) / (a_flat.norm() * b_flat.norm() + 1e-12)).item()
|
|
|
|
|
|
def gla_fwd_chunked_torch(q, k, v, g, chunk_size, scale=None):
|
|
"""Fully differentiable chunked GLA forward in float32."""
|
|
B, T, H, K = q.shape
|
|
V = v.shape[-1]
|
|
BC = chunk_size
|
|
NC = T // BC
|
|
|
|
if scale is None:
|
|
scale = K ** -0.5
|
|
|
|
q = q.float() * scale
|
|
k = k.float()
|
|
v = v.float()
|
|
g = g.float()
|
|
|
|
g_cum = g.reshape(B, NC, BC, H, K).cumsum(dim=2).reshape(B, T, H, K)
|
|
|
|
h = q.new_zeros(B, H, K, V)
|
|
mask = torch.tril(torch.ones(BC, BC, device=q.device, dtype=torch.float32))
|
|
|
|
o_chunks = []
|
|
for c in range(NC):
|
|
sl = slice(c * BC, (c + 1) * BC)
|
|
qc = q[:, sl, :, :]
|
|
kc = k[:, sl, :, :]
|
|
vc = v[:, sl, :, :]
|
|
gc = g_cum[:, sl, :, :]
|
|
g_last = gc[:, -1:, :, :]
|
|
|
|
q_gated = qc * torch.exp(gc)
|
|
o_inter = torch.einsum("bthk,bhkv->bthv", q_gated, h)
|
|
|
|
k_ungated = kc * torch.exp(-gc)
|
|
A = torch.einsum("bihk,bjhk->bhij", q_gated, k_ungated)
|
|
A = A * mask.unsqueeze(0).unsqueeze(0)
|
|
o_intra = torch.einsum("bhij,bjhv->bihv", A, vc)
|
|
|
|
o_chunks.append(o_inter + o_intra)
|
|
|
|
k_adj = kc * torch.exp(g_last - gc)
|
|
h = h * torch.exp(g_last).permute(0, 2, 3, 1).squeeze(-1).unsqueeze(-1)
|
|
h = h + torch.einsum("bthk,bthv->bhkv", k_adj, vc)
|
|
|
|
return torch.cat(o_chunks, dim=1)
|