mirror of
https://github.com/tinygrad/tinygrad.git
synced 2026-08-29 09:36:08 +00:00
Merge branch 'master' into parallel_compile
This commit is contained in:
George Hotz
@@ -4,7 +4,7 @@ inputs:
|
||||
python-version:
|
||||
description: 'Python version to use'
|
||||
required: false
|
||||
default: '' # if you don't set a version, the native python version will be used
|
||||
default: '3.14'
|
||||
key:
|
||||
description: 'Key for the python cache'
|
||||
required: false
|
||||
@@ -59,6 +59,11 @@ runs:
|
||||
echo "OMP_NUM_THREADS=1" >> "$GITHUB_ENV"
|
||||
# no buffers should be over 300MB in CI
|
||||
echo "MAX_BUFFER_SIZE=300000000" >> "$GITHUB_ENV"
|
||||
if [[ "$RUNNER_OS" == "Linux" ]]; then
|
||||
echo "VIRTUAL_ENV=/opt/venv/${{ inputs.python-version }}" >> "$GITHUB_ENV"
|
||||
else
|
||||
echo "VIRTUAL_ENV=${{ github.workspace }}/.venv" >> "$GITHUB_ENV"
|
||||
fi
|
||||
|
||||
- name: Set up uv
|
||||
uses: astral-sh/setup-uv@08807647e7069bb48b6ef5acd8ec9567f424441b
|
||||
@@ -67,7 +72,6 @@ runs:
|
||||
|
||||
- name: Set up Python ${{ inputs.python-version }}
|
||||
uses: actions/setup-python@v6
|
||||
if: inputs.python-version != ''
|
||||
with:
|
||||
python-version: ${{ inputs.python-version }}
|
||||
|
||||
@@ -109,15 +113,15 @@ runs:
|
||||
if: inputs.deps != ''
|
||||
shell: bash
|
||||
run: |
|
||||
uv venv .venv
|
||||
uv venv --allow-existing --python ${{ inputs.python-version }} "$VIRTUAL_ENV"
|
||||
DEPS="${{ inputs.deps }}"
|
||||
uv pip install --python .venv -e ".[${DEPS// /,}]" ${{ inputs.pydeps }} --torch-backend cpu --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/Triton-Nightly/pypi/simple/
|
||||
uv pip install --python "$VIRTUAL_ENV" -e ".[${DEPS// /,}]" ${{ inputs.pydeps }} --torch-backend cpu --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/Triton-Nightly/pypi/simple/
|
||||
- name: Install dependencies in venv (without extra)
|
||||
if: inputs.deps == ''
|
||||
shell: bash
|
||||
run: |
|
||||
uv venv .venv
|
||||
uv pip install --python .venv -e . ${{ inputs.pydeps }}
|
||||
uv venv --allow-existing --python ${{ inputs.python-version }} "$VIRTUAL_ENV"
|
||||
uv pip install --python "$VIRTUAL_ENV" -e . ${{ inputs.pydeps }}
|
||||
- name: Prune uv cache
|
||||
if: github.event_name != 'pull_request'
|
||||
shell: bash
|
||||
@@ -125,11 +129,10 @@ runs:
|
||||
- name: Configure venv
|
||||
shell: bash
|
||||
run: |
|
||||
echo "VIRTUAL_ENV=${{ github.workspace }}/.venv" >> "$GITHUB_ENV"
|
||||
if [[ "$RUNNER_OS" == "Windows" ]]; then
|
||||
echo "${{ github.workspace }}/.venv/Scripts" >> "$GITHUB_PATH"
|
||||
echo "$VIRTUAL_ENV/Scripts" >> "$GITHUB_PATH"
|
||||
else
|
||||
echo "${{ github.workspace }}/.venv/bin" >> "$GITHUB_PATH"
|
||||
echo "$VIRTUAL_ENV/bin" >> "$GITHUB_PATH"
|
||||
fi
|
||||
|
||||
# ******************* apt *******************
|
||||
|
||||
@@ -166,7 +166,7 @@ jobs:
|
||||
uses: ./.github/actions/setup-tinygrad
|
||||
with:
|
||||
key: windows-${{ matrix.dev }}-minimal
|
||||
deps: testing_unit
|
||||
deps: testing_minimal
|
||||
pydeps: ${{ matrix.dev == 'WEBGPU' && 'dawn-python' || '' }}
|
||||
- name: Set env
|
||||
shell: bash
|
||||
|
||||
@@ -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]
|
||||
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -6,6 +6,8 @@ from tinygrad.tensor import _to_np_dtype
|
||||
from tinygrad.runtime.ops_python import from_storage_scalar
|
||||
from tinygrad.renderer.ptx import PTXRenderer
|
||||
from tinygrad.renderer.nir import NIRRenderer
|
||||
from tinygrad.renderer.llvmir import CPULLVMRenderer
|
||||
from tinygrad.renderer.isa.x86 import X86Renderer
|
||||
from tinygrad.uop import Ops
|
||||
import numpy as np
|
||||
import pytest
|
||||
@@ -64,6 +66,8 @@ ht.fp8e5m2fnuz = ht.uint8
|
||||
def universal_test(a, b, dtype, op):
|
||||
if not isinstance(op, tuple): op = (op, op)
|
||||
if op[0] == operator.mod and b == 0: return
|
||||
# TODO: throws floating point exception
|
||||
if isinstance(Device[Device.DEFAULT].renderer, (X86Renderer, CPULLVMRenderer)) and op[0] == operator.mod and a == dtype.min and b == -1: return
|
||||
# lt and max with nan is undefined in tinygrad
|
||||
if op[0] in (operator.lt, Tensor.maximum) and (math.isnan(a) or math.isnan(b)): return
|
||||
ta, tb = Tensor([a], dtype=dtype), Tensor([b], dtype=dtype)
|
||||
|
||||
@@ -106,7 +106,8 @@ class TestHelpers(unittest.TestCase):
|
||||
|
||||
def test_float_to_bf16(self):
|
||||
max_bf16 = torch.finfo(torch.bfloat16).max
|
||||
for a in [1, 1.1, 1234, 23456, -777.777, max_bf16, max_bf16 * 1.00001, -max_bf16, -max_bf16 * 1.00001, math.inf, -math.inf]:
|
||||
for a in [1, 1.1, 1234, 23456, -777.777, max_bf16, max_bf16 * 1.00001, -max_bf16, -max_bf16 * 1.00001,
|
||||
max_bf16 * 2, -max_bf16 * 2, math.inf, -math.inf]:
|
||||
self.assertEqual(float_to_bf16(a), torch.tensor([a], dtype=torch.bfloat16).item())
|
||||
self.assertTrue(math.isnan(float_to_bf16(math.nan)))
|
||||
|
||||
|
||||
@@ -1020,6 +1020,14 @@ class TestSymbolic(unittest.TestCase):
|
||||
self.assertIs(graph_rewrite(cond.where(a, uconst(2)).cast(dtypes.half), sym), cond.where(a.cast(dtypes.half), UOp.const(2, dtypes.half)))
|
||||
self.assertIs(graph_rewrite(cond.where(a, UOp.invalid()).cast(dtypes.half), sym), cond.where(a.cast(dtypes.half), UOp.invalid()))
|
||||
|
||||
def test_where_const_gate_keeps_stated_width(self):
|
||||
a = Variable("a", 0, 3, dtypes.half)
|
||||
self.assertIs(graph_rewrite(UOp.const(True, dtypes.bool).where(uconst(0.0), a), sym), UOp.const(0.0, dtypes.half))
|
||||
self.assertIs(graph_rewrite(UOp.const(True, dtypes.bool).where(uconst(0), Variable("i", 0, 3, dtypes.int)), sym), UOp.const(0, dtypes.int))
|
||||
self.assertIs(graph_rewrite(UOp.const(False, dtypes.bool).where(uconst(0.0), a), sym), a)
|
||||
self.assertIs(graph_rewrite(UOp.const(False, dtypes.bool).where(uconst(0.0), UOp.invalid()), sym), UOp.invalid())
|
||||
self.assertIs(graph_rewrite(UOp.const(True, dtypes.bool).where(uconst(0.0), uconst(1)), sym), uconst(0.0))
|
||||
|
||||
def test_where_merge_branches(self):
|
||||
cond1 = Variable("s", 0, 10) < 6
|
||||
cond2 = Variable("s", 0, 10) > 2
|
||||
|
||||
+1
-1
@@ -221,7 +221,7 @@ def float_to_fp16(x):
|
||||
|
||||
def float_to_bf16(x):
|
||||
if not math.isfinite(x): return x
|
||||
u = struct.unpack('I', struct.pack('f', x))[0]
|
||||
u = struct.unpack('I', struct.pack('f', truncate[dtypes.float](x)))[0]
|
||||
u = (u + 0x7FFF + ((u >> 16) & 1)) & 0xFFFF0000
|
||||
return struct.unpack('f', struct.pack('I', u))[0]
|
||||
|
||||
|
||||
@@ -101,6 +101,11 @@ pm_remove_invalid = PatternMatcher([
|
||||
if any(x.is_invalid for x in s.src) else None),
|
||||
])
|
||||
|
||||
def fold_const_where(gate:UOp, c0:UOp, c1:UOp, w:UOp) -> UOp:
|
||||
# folding a strong dtype WHERE to a weak const branch keeps the strong dtype
|
||||
ret = c0 if gate.val else c1
|
||||
return commit_weak(ret, w.dtype) if ret.op is Ops.CONST and ret.dtype in dtypes.weaks and w.dtype not in dtypes.weaks else ret
|
||||
|
||||
symbolic_simple = pm_data_invalid + PatternMatcher([
|
||||
# ** self folding **
|
||||
(UPat.var("x") + 0, lambda x: x), # x+0 -> x
|
||||
@@ -175,7 +180,7 @@ symbolic_simple = pm_data_invalid + PatternMatcher([
|
||||
# ** simple where folding **
|
||||
# a conditional with the same results either way is a noop, also fold const conditionals
|
||||
(UPat.var().where(UPat.var("val"), UPat.var("val")), lambda val: val),
|
||||
(UPat.cvar("gate").where(UPat.var("c0"), UPat.var("c1")), lambda gate, c0, c1: c0 if gate.val else c1),
|
||||
(UPat.cvar("gate").where(UPat.var("c0"), UPat.var("c1")).named("w"), fold_const_where),
|
||||
# a.where(b.where(c, d), d) -> (a & b).where(c, d)
|
||||
(UPat.var("a").where(UPat.var("b").where(UPat.var("c"), UPat.var("d")), UPat.var("d")), lambda a,b,c,d: (a&b).where(c,d)),
|
||||
# a.where(c, b.where(c, d)) -> (a | b).where(c, d)
|
||||
|
||||
Reference in New Issue
Block a user