import math, pathlib, functools, struct from tinygrad import Device, Tensor from tinygrad.dtype import DTypeLike, dtypes from tinygrad.helpers import DEBUG from tinygrad.renderer import Estimates from tinygrad.runtime.support.compiler_amd import HIPCCCompiler from tinygrad.runtime.support.elf import elf_loader from tinygrad.uop.ops import UOp, Ops, KernelInfo def _sharded_empty(shape:Tensor, ref:Tensor, axis:int|None, dtype:DTypeLike|None=None) -> Tensor: dtype = dtype or ref.dtype if not isinstance(ref.device, tuple): return Tensor.empty(*shape, dtype=dtype, device=ref.device) shard_axis = ref.uop.axis if axis is None else axis shape = tuple(s // len(ref.device) if i == shard_axis else s for i, s in enumerate(shape)) axis = ref.uop.axis if axis is None else axis return Tensor(Tensor.empty(*shape, dtype=dtype, device=ref.device).uop.multi(axis), dtype=dtype, device=ref.device) def _sharded_empty_like(ref:Tensor, axis:int|None=None) -> Tensor: return _sharded_empty(ref.shape, ref, axis) def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False): assert attn_mask is None, "attn_mask not supported" assert is_causal, "only causal attention supported" xq, xk, xv = xq.transpose(1, 2), xk.transpose(1, 2), xv.transpose(1, 2) B, N, H, D = xq.shape H_KV = xk.shape[2] assert D == 128, "only D=128 supported" num_devices = len(xq.device) if isinstance(xq.device, tuple) else 1 is_dp = xq.uop.axis == 0 is_mp = xq.uop.axis == 2 B_local = B // num_devices if is_dp else B H_local = H // num_devices if is_mp else H H_KV_local = H_KV // num_devices if is_mp else H_KV shard_axis = 0 if is_dp else 2 if is_mp else None shard_axis_t = 0 if is_dp else 1 if is_mp else None if DEBUG >= 2: print(f"Flash Attention {B=} {B_local=} {N=} {H=} {H_local=} {H_KV=} {H_KV_local=} {D=} on {num_devices} devices, {'DP' if is_dp else 'MP' if is_mp else 'no sharding'}") single_device = xq.device[0] if isinstance(xq.device, tuple) else xq.device arch = Device[single_device].renderer.arch attn = _sharded_empty_like(xq, axis=shard_axis) l_vec = _sharded_empty((B, H, 1, N), xq, dtype=dtypes.float32, axis=shard_axis_t) def grad(dou:UOp, _) -> tuple[None, None, UOp, UOp, UOp]: do = Tensor(dou, device=dou.device) dq_in = _sharded_empty((B, H, N, D), xq, axis=shard_axis_t) dq = _sharded_empty_like(xq, axis=shard_axis) dk = _sharded_empty_like(xk, axis=shard_axis) dv = _sharded_empty_like(xv, axis=shard_axis) # 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_in = Tensor.custom_kernel(delta_vec, dq_in, 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] dq_in, dk, dv = Tensor.custom_kernel(dq_in, dk, dv, 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))[:3] # unshuffle dq dq = Tensor.custom_kernel(dq, dq_in, fxn=functools.partial(custom_fa_backward_post, device=single_device, arch=arch, B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D))[0] return None, None, dq.uop, dk.uop, dv.uop attn, l_vec = Tensor.custom_kernel(attn, l_vec, xq, xk, xv, fxn=functools.partial(custom_fa_forward, device=single_device, arch=arch, B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D), grad_fxn=grad)[:2] return attn.transpose(1, 2) @functools.cache def custom_fa_forward(o:UOp, l_vec:UOp, q:UOp, k:UOp, v:UOp, device:str, arch:str, B:int, N:int, H:int, H_KV:int, D:int): code = (pathlib.Path(__file__).parent / "fa_fwd_causal.cpp").read_text() compile_args = [f"-I{(pathlib.Path(__file__).parent / 'include').as_posix()}", "-std=c++20", "-DKITTENS_CDNA4", "-DHIP_ENABLE_WARP_SYNC_BUILTINS", "-ffast-math", f"-DATTN_B={B}", f"-DATTN_N={N}", f"-DATTN_H={H}", f"-DATTN_H_KV={H_KV}"] Q_BLOCK_SIZE = 32 NUM_WARPS = 8 NUM_THREADS = 64 * NUM_WARPS gsz = (H, (math.ceil((N // Q_BLOCK_SIZE) / NUM_WARPS)), B) lsz = (NUM_THREADS, 1, 1) threadIdx_x = UOp.special(lsz[0], "lidx0") blockIdx_x, blockIdx_y, blockIdx_z = UOp.special(gsz[0], "gidx0"), UOp.special(gsz[1], "gidx1"), UOp.special(gsz[2], "gidx2") el = q.dtype.itemsize mem = (2*B*N*H*D + 2*B*N*H_KV*D) * el + B*H*N * l_vec.dtype.itemsize estimates = Estimates(ops=2*B*H*N*N*D, lds=mem, mem=mem) sink = UOp.sink(o.base, l_vec.base, q.base, k.base, v.base, threadIdx_x, blockIdx_x, blockIdx_y, blockIdx_z, arg=KernelInfo(name="custom_fa_forward", estimates=estimates)) lib = HIPCCCompiler(arch, compile_args).compile_cached(code) lib = bytearray(lib) rodata_off = next(sh.header.sh_offset for sh in elf_loader(bytes(lib))[1] if sh.name == ".rodata") struct.pack_into('