forked from tinygrad/tinygrad
Compare commits
8
Commits
remove_anchors
...
master
| Author | SHA1 | Date | |
|---|---|---|---|
  sirhcmandGitHub
|
303d1677b3 | ||
|
George HotzandGitHub
|
673c6463f9 | ||
|
wozeparrotandGitHub
|
849074f0db | ||
|
George HotzandGitHub
|
0252cb8fa7 | ||
|
qazalandGitHub
|
cc6d33bde7 | ||
|
qazalandGitHub
|
16c5ff2490 | ||
|
wozeparrotandGitHub
|
1b7f040984 | ||
|
qazalandGitHub
|
39d144546e |
@@ -45,10 +45,6 @@ inputs:
|
||||
description: "Install qemu"
|
||||
required: false
|
||||
default: 'false'
|
||||
docker-qemu:
|
||||
description: "Setup docker to use qemu"
|
||||
required: false
|
||||
default: 'false'
|
||||
runs:
|
||||
using: "composite"
|
||||
steps:
|
||||
@@ -286,9 +282,3 @@ runs:
|
||||
sudo mkdir -p /etc/OpenCL/vendors
|
||||
echo "/usr/lib/libRusticlOpenCL.so" | sudo tee /etc/OpenCL/vendors/rusticl.icd
|
||||
echo "RUSTICL_ENABLE=llvmpipe" >> "$GITHUB_ENV"
|
||||
|
||||
- name: Setup docker to use qemu
|
||||
if: inputs.docker-qemu == 'true'
|
||||
uses: docker/setup-qemu-action@v4
|
||||
with:
|
||||
platforms: arm64
|
||||
|
||||
@@ -94,7 +94,6 @@ jobs:
|
||||
shell: bash -e -o pipefail {0}
|
||||
env:
|
||||
DEV: ${{ matrix.dev }}
|
||||
HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
|
||||
if: github.repository_owner == 'tinygrad'
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
@@ -149,7 +148,6 @@ jobs:
|
||||
shell: bash -e -o pipefail {0}
|
||||
env:
|
||||
DEV: ${{ matrix.dev }}
|
||||
HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
|
||||
if: github.repository_owner == 'tinygrad'
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
@@ -202,7 +200,6 @@ jobs:
|
||||
shell: bash -e -o pipefail {0}
|
||||
env:
|
||||
DEV: ${{ matrix.dev }}
|
||||
HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
|
||||
if: github.repository_owner == 'tinygrad'
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
@@ -252,7 +249,6 @@ jobs:
|
||||
shell: bash -e -o pipefail {0}
|
||||
env:
|
||||
DEV: ${{ matrix.dev }}
|
||||
HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
|
||||
if: github.repository_owner == 'tinygrad'
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
|
||||
+23
-23
@@ -21,7 +21,7 @@ concurrency:
|
||||
jobs:
|
||||
docs:
|
||||
name: Docs
|
||||
runs-on: &linux ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
|
||||
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
|
||||
timeout-minutes: 10
|
||||
env:
|
||||
CHECK_OOB: 0
|
||||
@@ -61,7 +61,7 @@ jobs:
|
||||
|
||||
torchbackend:
|
||||
name: Torch Backend Tests
|
||||
runs-on: *linux
|
||||
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
|
||||
timeout-minutes: 15
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
@@ -88,7 +88,7 @@ jobs:
|
||||
|
||||
bepython:
|
||||
name: Python Backend
|
||||
runs-on: *linux
|
||||
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
|
||||
timeout-minutes: 15
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
@@ -126,7 +126,7 @@ jobs:
|
||||
|
||||
linter:
|
||||
name: Linters
|
||||
runs-on: *linux
|
||||
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
|
||||
timeout-minutes: 10
|
||||
|
||||
steps:
|
||||
@@ -157,7 +157,7 @@ jobs:
|
||||
|
||||
nulltest:
|
||||
name: Null Tests
|
||||
runs-on: *linux
|
||||
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
|
||||
timeout-minutes: 15
|
||||
|
||||
steps:
|
||||
@@ -191,7 +191,7 @@ jobs:
|
||||
|
||||
unittest:
|
||||
name: Unit Tests
|
||||
runs-on: *linux
|
||||
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
|
||||
timeout-minutes: 15
|
||||
|
||||
steps:
|
||||
@@ -228,7 +228,7 @@ jobs:
|
||||
matrix:
|
||||
group: [1, 2]
|
||||
name: SPEC=2 (${{ matrix.group }})
|
||||
runs-on: *linux
|
||||
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
|
||||
timeout-minutes: 15
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
@@ -244,7 +244,7 @@ jobs:
|
||||
|
||||
fuzzing:
|
||||
name: Fuzzing
|
||||
runs-on: *linux
|
||||
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
|
||||
timeout-minutes: 10
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
@@ -260,7 +260,7 @@ jobs:
|
||||
|
||||
testopenclimage:
|
||||
name: CL IMAGE Tests
|
||||
runs-on: *linux
|
||||
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
|
||||
timeout-minutes: 15
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
@@ -280,7 +280,7 @@ jobs:
|
||||
|
||||
testopenpilot:
|
||||
name: openpilot Compile Tests
|
||||
runs-on: *linux
|
||||
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
|
||||
timeout-minutes: 15
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
@@ -309,7 +309,7 @@ jobs:
|
||||
|
||||
testonnxcpu:
|
||||
name: ONNX (CPU) Tests
|
||||
runs-on: *linux
|
||||
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
|
||||
timeout-minutes: 20
|
||||
|
||||
steps:
|
||||
@@ -328,7 +328,7 @@ jobs:
|
||||
|
||||
testoptim:
|
||||
name: Optimization Tests
|
||||
runs-on: *linux
|
||||
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
|
||||
timeout-minutes: 20
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
@@ -360,7 +360,7 @@ jobs:
|
||||
|
||||
testllm:
|
||||
name: Test LLM
|
||||
runs-on: *linux
|
||||
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
|
||||
timeout-minutes: 15
|
||||
env:
|
||||
CHECK_OOB: 0
|
||||
@@ -387,7 +387,7 @@ jobs:
|
||||
|
||||
testmodels:
|
||||
name: Models
|
||||
runs-on: *linux
|
||||
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
|
||||
timeout-minutes: 15
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
@@ -407,7 +407,7 @@ jobs:
|
||||
|
||||
testdsp:
|
||||
name: Linux (DSP)
|
||||
runs-on: *linux
|
||||
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
|
||||
timeout-minutes: 15
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
@@ -435,7 +435,7 @@ jobs:
|
||||
- 'WEBGPU'
|
||||
|
||||
name: Linux (DEV=${{ matrix.dev }})
|
||||
runs-on: *linux
|
||||
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
|
||||
timeout-minutes: 20
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
@@ -461,7 +461,7 @@ jobs:
|
||||
|
||||
testamdasm:
|
||||
name: AMD ASM IDE
|
||||
runs-on: *linux
|
||||
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
|
||||
timeout-minutes: 20
|
||||
env:
|
||||
DEV: MOCKKFD+AMD
|
||||
@@ -507,7 +507,7 @@ jobs:
|
||||
|
||||
hcq2:
|
||||
name: hcq2
|
||||
runs-on: *linux
|
||||
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
|
||||
timeout-minutes: 5
|
||||
steps:
|
||||
- name: Checkout Code
|
||||
@@ -529,7 +529,7 @@ jobs:
|
||||
|
||||
testmockam:
|
||||
name: Linux (am)
|
||||
runs-on: *linux
|
||||
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
|
||||
timeout-minutes: 15
|
||||
env:
|
||||
DEV: MOCKPCI+AMD
|
||||
@@ -565,7 +565,7 @@ jobs:
|
||||
arch: [gfx1100, gfx1201, gfx950]
|
||||
|
||||
name: Linux (${{ matrix.backend }} ${{ matrix.arch }})
|
||||
runs-on: *linux
|
||||
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
|
||||
timeout-minutes: 15
|
||||
env:
|
||||
DEV: MOCKKFD+AMD:${{ matrix.backend == 'amdllvm' && 'LLVM' || '' }}:${{ matrix.arch }}
|
||||
@@ -603,7 +603,7 @@ jobs:
|
||||
backend: [ptx, nv]
|
||||
|
||||
name: Linux (${{ matrix.backend }})
|
||||
runs-on: *linux
|
||||
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
|
||||
timeout-minutes: 20
|
||||
env:
|
||||
FORWARD_ONLY: 1
|
||||
@@ -642,7 +642,7 @@ jobs:
|
||||
- 'NULL:QCOMCL:a630'
|
||||
- 'NULL:NAK:sm_120'
|
||||
name: Compile-only (DEV=${{ matrix.dev }})
|
||||
runs-on: *linux
|
||||
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
|
||||
timeout-minutes: 15
|
||||
env:
|
||||
NULL_ALLOW_COPYOUT: 1
|
||||
@@ -656,7 +656,7 @@ jobs:
|
||||
with:
|
||||
key: compile-${{ matrix.backend }}
|
||||
deps: "testing_unit mesa"
|
||||
docker-qemu: ${{ contains(matrix.dev, 'QCOMCL') }}
|
||||
qemu: ${{ contains(matrix.dev, 'QCOMCL') }}
|
||||
- name: Test IMAGE
|
||||
shell: bash
|
||||
if: contains(matrix.dev, 'a630')
|
||||
|
||||
@@ -182,12 +182,12 @@ class GPTOSS:
|
||||
xq, xk = apply_rotary_emb(xq, xk, freqs_cis)
|
||||
xq, xk, xv = xq.cast(dtypes.bfloat16), xk.cast(dtypes.bfloat16), xv.cast(dtypes.bfloat16) # (B,N,H,D)/(B,N,KV,D)
|
||||
|
||||
if sliding:
|
||||
attn = self._sliding_attention(xq, xk, xv, sinks)
|
||||
elif getenv("HK_FLASH_ATTENTION"):
|
||||
if getenv("HK_FLASH_ATTENTION"):
|
||||
from extra.thunder.amd.fa import flash_attention
|
||||
attn, *_ = flash_attention(xq, xk, xv, is_causal=True, write_flat=True, sinks=sinks)
|
||||
attn, *_ = flash_attention(xq, xk, xv, is_causal=True, write_flat=True, sinks=sinks, window=self.sliding_window if sliding else 0)
|
||||
attn = attn.reshape(bsz, seqlen, self.n_heads * self.head_dim)
|
||||
elif sliding:
|
||||
attn = self._sliding_attention(xq, xk, xv, sinks)
|
||||
else:
|
||||
xqm = xq.reshape(bsz, seqlen, self.n_kv_heads, self.n_rep, self.head_dim).permute(0, 2, 3, 1, 4)
|
||||
xkm, xvm = xk.permute(0, 2, 1, 3).unsqueeze(2), xv.permute(0, 2, 1, 3).unsqueeze(2)
|
||||
|
||||
@@ -122,7 +122,8 @@ def custom_mxfp4_gemm(C:UOp, A:UOp, B:UOp, scale_a:UOp, scale_b:UOp, *extra:UOp,
|
||||
groups_x, groups_y = UOp.special(ceildiv(N, tile_n), "gidx0"), UOp.special(ceildiv(M, tile_m), "gidx1")
|
||||
lds = UOp.placeholder((163840,), dtypes.uint8, 0, AddrSpace.LOCAL)
|
||||
sink = UOp.sink(C.base, A.base, B.base, scale_a.base, scale_b.base, *(x.base for x in extra), lds, threads, groups_x, groups_y,
|
||||
arg=KernelInfo(f"custom_mxfp4_gemm_{M}_{N}_{K}", estimates=Estimates(ops=2*M*N*K)))
|
||||
arg=KernelInfo(f"mxfp4_gemm_{M}_{N}_{K}",
|
||||
estimates=Estimates(ops=2*M*N*K, mem=(M*half_k+N*half_k)*A.dtype.itemsize+M*N*C.dtype.itemsize)))
|
||||
insts = build_kernel(M, N, K, tile_m, tile_n)
|
||||
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple(UOp(Ops.INS, arg=x) for x in insts))))
|
||||
|
||||
|
||||
+55
-9
@@ -110,7 +110,49 @@ def _sharded_empty_like(ref:Tensor, axis:int|None=None) -> Tensor:
|
||||
return _sharded_empty(ref.shape, ref, axis)
|
||||
|
||||
@functools.cache
|
||||
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):
|
||||
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)
|
||||
attn = Tensor(ker.src[1].after(ker), device=ker.src[1].device)
|
||||
@@ -118,6 +160,8 @@ 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
|
||||
@@ -128,8 +172,10 @@ 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))[:3]
|
||||
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]
|
||||
|
||||
if D == 64:
|
||||
dq = dq.reshape(B, H, N//16, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2).permute(0, 1, 2, 8, 9, 10, 11, 3, 4, 6, 7, 5, 12).reshape(B, H, N, D).transpose(1, 2)
|
||||
@@ -149,7 +195,7 @@ def _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, sha
|
||||
return grad
|
||||
|
||||
# TODO: remove write_flat once scheduler can remove reshapes between custom_kernel. TestCustomKernel.test_simple_reshape
|
||||
def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False, write_flat:bool=False, sinks:Tensor|None=None):
|
||||
def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False, write_flat:bool=False, sinks:Tensor|None=None, window:int=0):
|
||||
assert attn_mask is None, "attn_mask not supported"
|
||||
assert is_causal, "only causal attention supported"
|
||||
|
||||
@@ -176,18 +222,18 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
|
||||
attn = _sharded_empty((B, N, H * D), xq, axis=shard_axis) if write_flat else _sharded_empty_like(xq, axis=shard_axis)
|
||||
l_vec = _sharded_empty((B, H, 1, N), xq, dtype=dtypes.float32, axis=shard_axis_t)
|
||||
|
||||
grad = _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)
|
||||
grad = _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=window)
|
||||
|
||||
fwd_inputs = (attn, l_vec, xq, xk, xv) + ((sinks,) if has_sink else ())
|
||||
attn, l_vec = Tensor.custom_kernel(*fwd_inputs, 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, has_sink=has_sink), grad_fxn=grad)[:2]
|
||||
attn, l_vec = Tensor.custom_kernel(*fwd_inputs, 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, has_sink=has_sink, window=window), grad_fxn=grad)[:2]
|
||||
|
||||
return attn, attn, l_vec
|
||||
|
||||
@functools.cache
|
||||
def custom_fa_forward(o:UOp, l_vec:UOp, q:UOp, k:UOp, v:UOp, sinks:UOp|None=None, *, device:str, arch:str, B:int, N:int, H:int, H_KV:int, D:int, has_sink:bool=True):
|
||||
def custom_fa_forward(o:UOp, l_vec:UOp, q:UOp, k:UOp, v:UOp, sinks:UOp|None=None, *, device:str, arch:str, B:int, N:int, H:int, H_KV:int, D:int, has_sink:bool=True, window:int=0):
|
||||
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}", f"-DATTN_D={D}", f"-DATTN_SINK={int(has_sink)}"]
|
||||
f"-DATTN_B={B}", f"-DATTN_N={N}", f"-DATTN_H={H}", f"-DATTN_H_KV={H_KV}", f"-DATTN_D={D}", f"-DATTN_SINK={int(has_sink)}", f"-DWINDOW={window}"]
|
||||
|
||||
Q_BLOCK_SIZE = 32
|
||||
NUM_WARPS = 8
|
||||
@@ -247,10 +293,10 @@ def custom_fa_backward_pre(delta_vec:UOp, dq:UOp, o:UOp, do:UOp, device:str, arc
|
||||
src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=code), UOp(Ops.BINARY, arg=lib)))
|
||||
|
||||
@functools.cache
|
||||
def custom_fa_backward(dq:UOp, dk:UOp, dv:UOp, do:UOp, q:UOp, k:UOp, v:UOp, l_vec:UOp, delta_vec:UOp, device:str, arch:str, B:int, N:int, H:int, H_KV:int, D:int):
|
||||
def custom_fa_backward(dq:UOp, dk:UOp, dv:UOp, do:UOp, q:UOp, k:UOp, v:UOp, l_vec:UOp, delta_vec:UOp, device:str, arch:str, B:int, N:int, H:int, H_KV:int, D:int, window:int=0):
|
||||
code = (pathlib.Path(__file__).parent / "fa_bwd_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}", f"-DATTN_D={D}"]
|
||||
f"-DATTN_B={B}", f"-DATTN_N={N}", f"-DATTN_H={H}", f"-DATTN_H_KV={H_KV}", f"-DATTN_D={D}", f"-DWINDOW={window}"]
|
||||
|
||||
BLOCK_SIZE_KV = 256
|
||||
GROUP_SIZE = H // H_KV
|
||||
|
||||
+17
@@ -0,0 +1,17 @@
|
||||
from tinygrad import Device, Tensor, TinyJit, dtypes
|
||||
from tinygrad.helpers import Timing, Context
|
||||
|
||||
GPUS, DEPTH, SZ = 8, 4, 128 * 2**20
|
||||
WARMUP, ITERS = 3, 5
|
||||
devs = tuple(f"{Device.DEFAULT}:{i}" for i in range(GPUS))
|
||||
bufs = tuple(Tensor.empty(SZ, dtype=dtypes.uint8, device=dev).contiguous().realize() for _ in range(DEPTH) for dev in devs)
|
||||
|
||||
@TinyJit
|
||||
def all_to_all(*srcs:Tensor): return Tensor.realize(*(src.to(dst) for i,src in enumerate(srcs) for j,dst in enumerate(devs) if i % GPUS != j))
|
||||
|
||||
if __name__ == "__main__":
|
||||
with Context(ALL2ALL=1, JIT_BATCH_SIZE=0):
|
||||
for i in range(-WARMUP, ITERS):
|
||||
with Timing("ALL2ALL ", lambda ns: f" {SZ*GPUS*(GPUS-1)*DEPTH/ns:.2f} GB/s", enabled=i>=0):
|
||||
all_to_all(*bufs)
|
||||
for dev in devs: Device[dev].synchronize()
|
||||
@@ -790,10 +790,14 @@ class KFDIface:
|
||||
|
||||
def create_queue(self, queue_type, ring, gart, rptr, wptr, eop_buffer=None, cwsr_buffer=None, ctl_stack_size=0, ctx_save_restore_size=0,
|
||||
xcc_id=0, idx=0):
|
||||
if queue_type == kfd.KFD_IOC_QUEUE_TYPE_SDMA and idx and self.kfd_ver >= (1,17):
|
||||
queue_type = kfd.KFD_IOC_QUEUE_TYPE_SDMA_BY_ENG_ID
|
||||
sdma_engine_id = idx % (self.props['num_sdma_engines'] + self.props.get('num_sdma_xgmi_engines', 0))
|
||||
else: sdma_engine_id = 0
|
||||
queue = kfd.AMDKFD_IOC_CREATE_QUEUE(KFDIface.kfd, ring_base_address=ring.va_addr, ring_size=ring.size, gpu_id=self.gpu_id,
|
||||
queue_type=queue_type, queue_percentage=kfd.KFD_MAX_QUEUE_PERCENTAGE|(xcc_id<<8), queue_priority=getenv("AMD_KFD_QUEUE_PRIORITY", 7),
|
||||
eop_buffer_address=eop_buffer.va_addr if eop_buffer else 0, eop_buffer_size=eop_buffer.size if eop_buffer else 0, ctl_stack_size=ctl_stack_size,
|
||||
ctx_save_restore_address=cwsr_buffer.va_addr if cwsr_buffer else 0, ctx_save_restore_size=ctx_save_restore_size,
|
||||
ctx_save_restore_address=cwsr_buffer.va_addr if cwsr_buffer else 0, ctx_save_restore_size=ctx_save_restore_size, sdma_engine_id=sdma_engine_id,
|
||||
write_pointer_address=gart.va_addr+wptr, read_pointer_address=gart.va_addr+rptr+8*xcc_id)
|
||||
|
||||
if not hasattr(self, 'doorbells'):
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
import ctypes, struct, platform, pathlib, subprocess, sys
|
||||
import ctypes, struct, platform, pathlib, shutil, subprocess, sys, tarfile, tempfile
|
||||
from tinygrad.device import Compiler
|
||||
from tinygrad.helpers import DEBUG, system, fetch, unwrap
|
||||
from tinygrad.runtime.support.compiler_mesa import disas_adreno
|
||||
@@ -11,10 +11,13 @@ class QCOMCompiler(Compiler):
|
||||
def __init__(self, arch:str):
|
||||
assert arch.split(',')[0] == "a630", "only a630 supported"
|
||||
if platform.machine() == "aarch64": self.arch, self.chip_id, self.llvm_inst = arch, 0x6030001, llvm_qcom.cl_compiler_create_llvm_instance()
|
||||
else: self.arch, self.chip_id, self.compiler_process = arch, 0x6030001, subprocess.Popen(
|
||||
(f"docker run --rm -i --platform linux/aarch64 -e PYTHONPATH=/ -e QEMU_CPU=max,pauth=off -v {pathlib.Path(__file__).parents[2]}:/tinygrad "
|
||||
f"-v {fetch('https://github.com/sirhcm/tinydreno/raw/refs/heads/master/libllvm-qcom.so')}:/lib/libllvm-qcom.so python:3.12-slim "
|
||||
f"python /tinygrad/runtime/support/compiler_qcom.py {arch}").split(), stdout=subprocess.PIPE, stdin=subprocess.PIPE, bufsize=0)
|
||||
else:
|
||||
self.arch, self.chip_id, self.fs = arch, 0x6030001, tempfile.TemporaryDirectory()
|
||||
with tarfile.open(fetch('https://git.tinygrad.win/sirhcm/images/releases/download/v2/qcomcl.tar.gz')) as t: t.extractall(fs:=self.fs.name)
|
||||
if (qemu:=shutil.which("qemu-aarch64-static")): argv = f"{qemu} -cpu max,pauth=off -L {fs} {fs}/usr/bin/python3 {__file__} {arch}"
|
||||
else: argv = (f"docker run --rm -i --platform linux/aarch64 -v {fs}/usr:/usr -v {pathlib.Path(__file__).parents[2]}:/tinygrad "
|
||||
f"-e PYTHONPATH=/ -e QEMU_CPU=max,pauth=off gcr.io/distroless/static python3 /tinygrad/runtime/support/compiler_qcom.py {arch}")
|
||||
self.compiler_process = subprocess.Popen(argv.split(), stdout=subprocess.PIPE, stdin=subprocess.PIPE, bufsize=0)
|
||||
super().__init__(f"compile_qcomcl_{arch}")
|
||||
|
||||
def __del__(self): llvm_qcom.cl_compiler_destroy_llvm_instance(self.llvm_inst) if platform.machine() == "aarch64" else self.compiler_process.kill()
|
||||
|
||||
Reference in New Issue
Block a user