J3ff/IQ.Pilot
1
0
Fork 0
forked from IQ.Lvbs/IQ.Pilot
Code Pull Requests Activity

IQ.Pilot Prebuilt Release @ ab07000

Browse Source
This commit is contained in:
IQ.Lvbs history cleanup
2026-08-22 23:42:42 -05:00
commit 9f9c9a70cc
3729 changed files with 778697 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

0
tinygrad_repo/test/testextra/__init__.py Normal file
View File

111
tinygrad_repo/test/testextra/test_bench_log.py Normal file
View File

@@ -0,0 +1,111 @@
import unittest, time
from unittest.case import skipIf
from extra.bench_log import BenchEvent, InstantBenchEvent, WallTimeEvent, KernelTimeEvent, log_event_instant, _events, clear_events
from tinygrad.helpers import Context, DEV
from tinygrad.tensor import Tensor
from tinygrad.device import Device
# WEBGPU kernel timing not supported, ocelot CUDA is inaccurate
_SKIP_KERNEL_TIMING = Device.DEFAULT == "WEBGPU" or (Device.DEFAULT == "CUDA" and DEV.interface.startswith("MOCK"))
class TestBenchLog(unittest.TestCase):
def setUp(self):
clear_events()
def test_log_single_wall_time(self):
for event in BenchEvent:
with WallTimeEvent(event):
time.sleep(0.1)
# check event list
for event in BenchEvent:
self.assertEqual(len(_events[event]["wall"]), 1)
self.assertGreater(_events[event]["wall"][0], 0)
def test_log_double_wall_time(self):
for event in BenchEvent:
with WallTimeEvent(event):
time.sleep(0.1)
for event in reversed(BenchEvent):
with WallTimeEvent(event):
time.sleep(0.2)
# check event list
for event in BenchEvent:
self.assertEqual(len(_events[event]["wall"]), 2)
self.assertGreater(_events[event]["wall"][0], 0)
self.assertGreater(_events[event]["wall"][1], 0)
@skipIf(_SKIP_KERNEL_TIMING, "ci timing is not accurate")
def test_log_single_kernel_time(self):
wall_times = []
with Context(DEBUG=2):
for event in BenchEvent:
with KernelTimeEvent(event):
st = time.perf_counter()
Tensor.rand(32, 32).sum().realize().item()
wall_times.append(time.perf_counter() - st)
# check event list
for event in BenchEvent:
self.assertEqual(len(_events[event]["kernel"]), 1)
self.assertLess(_events[event]["kernel"][0], wall_times[0])
self.assertGreater(_events[event]["kernel"][0], 0)
@skipIf(_SKIP_KERNEL_TIMING, "ci cuda timing is not accurate")
def test_interleaved_wall_kernel_time(self):
wall_times = []
with Context(DEBUG=2):
for event in BenchEvent:
with KernelTimeEvent(event):
st = time.perf_counter()
Tensor.rand(32, 32).sum().realize().item()
wall_times.append(time.perf_counter() - st)
with WallTimeEvent(event):
st = time.perf_counter()
Tensor.rand(32, 32).sum().realize().item()
wall_times.append(time.perf_counter() - st)
# check event list
for event in BenchEvent:
self.assertEqual(len(_events[event]["wall"]), 1)
self.assertEqual(len(_events[event]["kernel"]), 1)
self.assertLess(_events[event]["kernel"][0], wall_times[0])
self.assertGreater(_events[event]["kernel"][0], 0)
@skipIf(_SKIP_KERNEL_TIMING, "ci cuda timing is not accurate")
def test_stacked_wall_kernel_time(self):
with Context(DEBUG=2):
for event in BenchEvent:
with KernelTimeEvent(event):
with WallTimeEvent(event):
Tensor.rand(32, 32).sum().realize().item()
for event in BenchEvent:
with WallTimeEvent(event):
with KernelTimeEvent(event):
Tensor.rand(32, 32).sum().realize().item()
for event in BenchEvent:
self.assertEqual(len(_events[event]["wall"]), 2)
self.assertEqual(len(_events[event]["kernel"]), 2)
self.assertLess(_events[event]["kernel"][0], _events[event]["wall"][0])
self.assertGreater(_events[event]["kernel"][0], 0)
self.assertLess(_events[event]["kernel"][1], _events[event]["wall"][1])
self.assertGreater(_events[event]["kernel"][1], 0)
def test_log_instant_event(self):
for event in InstantBenchEvent:
log_event_instant(event, 1000)
# check event list
for event in InstantBenchEvent:
self.assertEqual(len(_events[event]), 1)
self.assertEqual(_events[event][0], 1000)
if __name__ == '__main__':
unittest.main()

82
tinygrad_repo/test/testextra/test_export_model.py Normal file
View File

@@ -0,0 +1,82 @@
import unittest
from extra.export_model import export_model, EXPORT_SUPPORTED_DEVICE
from tinygrad.tensor import Tensor
from tinygrad.device import Device
from tinygrad.nn import Linear
from tinygrad.nn.state import get_state_dict
from tinygrad import dtypes
import json
class MockMultiInputModel:
def forward(self, x1, x2, x3):
return x1 + x2 + x3
class MockMultiOutputModel:
def __call__(self, x1):
return x1 + 2.0, x1.pad(((0, 0), (0, 1))) + 1.0
# TODO: move compile_efficientnet tests here
@unittest.skipUnless(Device.DEFAULT in EXPORT_SUPPORTED_DEVICE, f"Model export is not supported on {Device.DEFAULT}")
class TextModelExport(unittest.TestCase):
def test_multi_input_model_export(self):
model = MockMultiInputModel()
inputs = [Tensor.rand(2,2), Tensor.rand(2,2), Tensor.rand(2,2)]
prg, inp_sizes, _, _ = export_model(model, "", *inputs)
prg = json.loads(prg)
assert len(inputs) == len(prg["inputs"]) == len(inp_sizes), f"Model and exported inputs don't match: mdl={len(inputs)}, prg={len(prg['inputs'])}, inp_sizes={len(inp_sizes)}" # noqa: E501
for i in range(len(inputs)):
assert f"input{i}" in inp_sizes, f"input{i} not captured in inp_sizes"
assert f"input{i}" in prg["buffers"], f"input{i} not captured in exported buffers"
for i, exported_input in enumerate(prg["inputs"]):
assert inputs[i].dtype.name == exported_input["dtype"], f"Model and exported input dtype don't match: mdl={inputs[i].dtype.name}, prg={exported_input['dtype']}" # noqa: E501
def test_multi_output_model_export(self):
model = MockMultiOutputModel()
input_tensor = Tensor.rand(2,2)
outputs = model(input_tensor)
prg, _, out_sizes, _ = export_model(model, "", input_tensor)
prg = json.loads(prg)
assert len(outputs) == len(prg["outputs"]) == len(out_sizes), f"Model and exported outputs don't match: mdl={len(outputs)}, prg={len(prg['outputs'])}, inp_sizes={len(out_sizes)}" # noqa: E501
for i in range(len(outputs)):
assert f"output{i}" in out_sizes, f"output{i} not captured in out_sizes"
assert f"output{i}" in prg["buffers"], f"output{i} not captured in exported buffers"
for i, exported_output in enumerate(prg["outputs"]):
assert outputs[i].dtype.name == exported_output["dtype"], f"Model and exported output dtype don't match: mdl={outputs[i].dtype.name}, prg={exported_output['dtype']}" # noqa: E501
@unittest.skipUnless(Device.DEFAULT == "WEBGPU", "Testing WebGPU specific model export behavior")
class TextModelExportWebGPU(unittest.TestCase):
def test_exported_input_output_dtypes(self):
class MyModel:
def forward(self, *inputs): return tuple([(inp+2).cast(inp.dtype) for inp in inputs])
model = MyModel()
# [:-1] because "ulong" and "long" is not supported
inputs = [Tensor.randn(2, dtype=dt) for dt in dtypes.uints[:-1] + dtypes.sints[:-1] + (dtypes.bool, dtypes.float)]
prg, _, _, _ = export_model(model, "webgpu", *inputs)
expected_buffer_types = ["Uint"]*len(dtypes.uints[:-1]) + ["Int"]*len(dtypes.sints[:-1]) + ["Int", "Float"]
for i, expected_buffer_type in enumerate(expected_buffer_types):
dt = inputs[i].dtype
expected_arr_prefix = f"{expected_buffer_type}{dt.itemsize*8}"
# test input buffers
self.assertIn(f"new {expected_arr_prefix}Array(gpuWriteBuffer{i}.getMappedRange()).set(_input{i});", prg)
# test output buffers
self.assertIn(f"const resultBuffer{i} = new {expected_arr_prefix}Array(gpuReadBuffer{i}.size/{dt.itemsize});", prg)
self.assertIn(f"resultBuffer{i}.set(new {expected_arr_prefix}Array(gpuReadBuffer{i}.getMappedRange()));", prg)
def test_weights_bound_to_safetensor(self):
# regression test: every weight ended up as createEmptyBuf (zero-init) instead of createWeightBuf
class MyModel:
def __init__(self): self.fc1, self.fc2 = Linear(4, 8), Linear(8, 2)
def forward(self, x): return self.fc2(self.fc1(x).relu())
model = MyModel()
for t in get_state_dict(model).values(): t.realize()
prg, _, _, _ = export_model(model, "webgpu", Tensor.randn(1, 4))
self.assertEqual(prg.count("createWeightBuf("), len(get_state_dict(model)))
if __name__ == '__main__':
unittest.main()

14
tinygrad_repo/test/testextra/test_f16_decompress.py Normal file
View File

@@ -0,0 +1,14 @@
import unittest
from extra.f16_decompress import u32_to_f16
from tinygrad.tensor import Tensor
from tinygrad import dtypes
import numpy as np
class TestF16Decompression(unittest.TestCase):
def test_u32_to_f16(self):
a = Tensor.randn(50, dtype=dtypes.float16)
f16_as_u32 = a.bitcast(dtypes.uint32)
f16 = u32_to_f16(f16_as_u32)
ref = a.numpy()
out = f16.numpy().astype(np.float16)
np.testing.assert_allclose(out, ref)

90
tinygrad_repo/test/testextra/test_fp8_linear.py Normal file
View File

@@ -0,0 +1,90 @@
#!/usr/bin/env python
import unittest
import numpy as np
from tinygrad import Tensor, dtypes, Device
from tinygrad.nn import Linear
from extra.fp8.fp8_linear import FP8Linear, convert_to_float8_training
from test.helpers import not_support_multi_device, needs_second_gpu
BS, T, in_dim, out_dim = 16, 4, 128, 128
@unittest.skipUnless(dtypes.fp8e4m3 in Device[Device.DEFAULT].renderer.supported_dtypes(), f"no fp8e4m3 on {Device.DEFAULT}")
class TestFP8Linear(unittest.TestCase):
def setUp(self):
Tensor.manual_seed(42)
def _test_forward(self, shape, in_features, out_features):
fp8_layer = FP8Linear(in_features, out_features)
normal_layer = Linear(in_features, out_features)
weight = Tensor.randn(out_features, in_features, dtype=dtypes.float32) * 0.2
bias = Tensor.randn(out_features, dtype=dtypes.float32) * 0.2
fp8_layer.weight.assign(weight)
normal_layer.weight.assign(weight)
fp8_layer.bias.assign(bias)
normal_layer.bias.assign(bias)
x = Tensor.randn(*shape, dtype=dtypes.float32) * 0.2
y_fp8, y_normal = fp8_layer(x), normal_layer(x)
np.testing.assert_allclose(y_fp8.numpy(), y_normal.numpy(), rtol=0.1, atol=0.1)
def _test_backward(self, shape, in_features, out_features):
fp8_layer = FP8Linear(in_features, out_features)
normal_layer = Linear(in_features, out_features)
weight = Tensor.randn(out_features, in_features, dtype=dtypes.float32) * 0.2
bias = Tensor.randn(out_features, dtype=dtypes.float32) * 0.2
fp8_layer.weight, normal_layer.weight = weight.detach(), weight.detach()
fp8_layer.bias, normal_layer.bias = bias.detach(), bias.detach()
x_fp8 = Tensor.randn(*shape, dtype=dtypes.float32) * 0.2
x_normal = x_fp8.detach()
fp8_layer(x_fp8).sum().backward()
normal_layer(x_normal).sum().backward()
np.testing.assert_allclose(x_fp8.grad.numpy(), x_normal.grad.numpy(), rtol=1.0, atol=0.1)
np.testing.assert_allclose(fp8_layer.weight.grad.numpy(), normal_layer.weight.grad.numpy(), rtol=1.0, atol=0.1)
def test_forward_2d(self): self._test_forward((BS, in_dim), in_dim, out_dim)
def test_forward_3d(self): self._test_forward((BS, T, in_dim), in_dim, out_dim)
def test_backward_2d(self): self._test_backward((BS, in_dim), in_dim, out_dim)
def test_backward_3d(self): self._test_backward((BS, T, in_dim), in_dim, out_dim)
def test_filter(self):
class Model:
def __init__(self):
self.fc1 = Linear(32, 16)
self.fc2 = Linear(16, 8)
def __call__(self, x):
return self.fc2(self.fc1(x).relu())
model = Model()
x = Tensor.randn(16, 32)
y_before = model(x).numpy()
convert_to_float8_training(model, module_filter_fn=lambda _, fqn: "fc1" in fqn)
self.assertIsInstance(model.fc1, FP8Linear)
self.assertNotIsInstance(model.fc2, FP8Linear)
y_after = model(x).numpy()
np.testing.assert_allclose(y_after, y_before, rtol=0.1, atol=0.1)
@needs_second_gpu
@unittest.skipIf(not_support_multi_device(), "no multi")
def test_multi_gpu(self):
GPUS = tuple(f"{Device.DEFAULT}:{i}" for i in range(2))
fp8_layer = FP8Linear(in_dim, out_dim)
normal_layer = Linear(in_dim, out_dim)
weight = Tensor.randn(out_dim, in_dim, dtype=dtypes.float32) * 0.2
bias = Tensor.randn(out_dim, dtype=dtypes.float32) * 0.2
fp8_layer.weight.assign(weight)
fp8_layer.bias.assign(bias)
normal_layer.weight.assign(weight)
normal_layer.bias.assign(bias)
fp8_layer.weight.to_(GPUS)
fp8_layer.bias.to_(GPUS)
normal_layer.weight.to_(GPUS)
normal_layer.bias.to_(GPUS)
x = Tensor.randn(BS*2, in_dim, dtype=dtypes.float32) * 0.2
x_sharded = x.detach()
x = x.shard_(GPUS, axis=0)
y_normal = normal_layer(x).realize()
x_sharded.shard_(GPUS, axis=0)
y_fp8 = fp8_layer(x_sharded).realize()
np.testing.assert_allclose(y_fp8.numpy(), y_normal.numpy(), rtol=0.1, atol=0.1)
if __name__ == '__main__':
unittest.main()

87
tinygrad_repo/test/testextra/test_hevc.py Normal file
View File

@@ -0,0 +1,87 @@
import unittest
from tinygrad import Tensor, Device, dtypes
from tinygrad.helpers import fetch, round_up
from extra.hevc.hevc import parse_hevc_file_headers, nv_gpu
from extra.hevc.decode import hevc_decode
class TestHevc(unittest.TestCase):
def test_hevc_parser(self):
url = "https://github.com/haraschax/filedump/raw/09a497959f7fa6fd8dba501a25f2cdb3a41ecb12/comma_video.hevc"
dat = fetch(url, headers={"Range": f"bytes=0-{512<<10}"}).read_bytes()
opaque, frame_info, w, h, luma_w, luma_h, chroma_off = parse_hevc_file_headers(dat, device=Device.DEFAULT)
def _test_common(frame, bts):
self.assertEqual(frame0.pic_width_in_luma_samples, 1952)
self.assertEqual(frame0.pic_height_in_luma_samples, 1216)
self.assertEqual(frame0.chroma_format_idc, 1)
self.assertEqual(frame0.bit_depth_luma, 8)
self.assertEqual(frame0.bit_depth_chroma, 8)
self.assertEqual(frame0.log2_min_luma_coding_block_size, 3)
self.assertEqual(frame0.log2_max_luma_coding_block_size, 5)
self.assertEqual(frame0.log2_min_transform_block_size, 2)
self.assertEqual(frame0.log2_max_transform_block_size, 5)
self.assertEqual(frame0.num_tile_columns, 3)
self.assertEqual(frame0.num_tile_rows, 1)
self.assertEqual(frame0.colMvBuffersize, 589)
self.assertEqual(frame0.HevcSaoBufferOffset, 2888)
self.assertEqual(frame0.HevcBsdCtrlOffset, 25992)
self.assertEqual(frame0.v1.hevc_main10_444_ext.HevcFltAboveOffset, 26714)
self.assertEqual(frame0.v1.hevc_main10_444_ext.HevcSaoAboveOffset, 36214)
# tiles
self.assertEqual(bytes(bts[0x200:0x210]), b'\x18\x00&\x00\x18\x00&\x00\r\x00&\x00\x00\x00\x00\x00')
frame0 = nv_gpu.nvdec_hevc_pic_s.from_buffer(opaque[0].data())
_test_common(frame0, opaque[0].data())
self.assertEqual(frame0.stream_len, 148063)
self.assertEqual(frame0.IDR_picture_flag, 1)
self.assertEqual(frame0.RAP_picture_flag, 1)
self.assertEqual(frame0.sw_hdr_skip_length, 0)
self.assertEqual(frame0.num_ref_frames, 0)
frame1 = nv_gpu.nvdec_hevc_pic_s.from_buffer(opaque[1].data())
_test_common(frame1, opaque[1].data())
self.assertEqual(frame1.stream_len, 57110)
self.assertEqual(frame1.IDR_picture_flag, 0)
self.assertEqual(frame1.RAP_picture_flag, 0)
self.assertEqual(frame1.sw_hdr_skip_length, 9)
self.assertEqual(frame1.num_ref_frames, 1)
self.assertEqual(list(frame1.initreflistidxl0), [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0])
self.assertEqual(list(frame1.initreflistidxl1), [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0])
self.assertEqual(list(frame1.RefDiffPicOrderCnts), [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0])
frame3 = nv_gpu.nvdec_hevc_pic_s.from_buffer(opaque[3].data())
_test_common(frame3, opaque[3].data())
self.assertEqual(frame3.stream_len, 47036)
self.assertEqual(frame3.IDR_picture_flag, 0)
self.assertEqual(frame3.RAP_picture_flag, 0)
self.assertEqual(frame3.sw_hdr_skip_length, 9)
self.assertEqual(frame3.num_ref_frames, 1)
self.assertEqual(list(frame3.initreflistidxl0), [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0])
self.assertEqual(list(frame3.initreflistidxl1), [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0])
self.assertEqual(list(frame3.RefDiffPicOrderCnts), [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0])
@unittest.skipUnless(Device.DEFAULT == "NV", "NV only")
def test_hevc_decode(self):
url = "https://github.com/haraschax/filedump/raw/09a497959f7fa6fd8dba501a25f2cdb3a41ecb12/comma_video.hevc"
dat = fetch(url, headers={"Range": f"bytes=0-{512<<10}"}).read_bytes()
opaque, frame_info, w, h, luma_w, luma_h, chroma_off = parse_hevc_file_headers(dat)
frame_info = frame_info[:4]
out_image_size = luma_h + (luma_h + 1) // 2, round_up(luma_w, 64)
hevc_tensor = Tensor(dat, device="NV")
opaque_nv = opaque.to("NV").contiguous().realize()
frames = list(hevc_decode(hevc_tensor, opaque_nv, frame_info, luma_h, luma_w))
Device.default.synchronize()
self.assertEqual(len(frames), 4)
for f in frames:
self.assertEqual(f.shape, out_image_size)
self.assertEqual(f.dtype, dtypes.uint8)
self.assertEqual(f.device, "NV")
if __name__ == "__main__":
unittest.main()

226
tinygrad_repo/test/testextra/test_hk_fa.py Normal file
View File

@@ -0,0 +1,226 @@
import unittest, time
from tinygrad import Tensor, Device, dtypes, Context
from tinygrad.engine.jit import TinyJit
import numpy as np
from extra.thunder.amd.fa import flash_attention
def assert_allclose(cmp:Tensor, ref:Tensor, **kwargs) -> None:
if Device.DEFAULT == "NULL": Tensor.realize(cmp, ref)
else: np.testing.assert_allclose(cmp.numpy(), ref.numpy(), **kwargs)
class TestFA(unittest.TestCase):
def setUp(self):
arch = Device[Device.DEFAULT].renderer.target.arch
if not arch.startswith("gfx9"):
self.skipTest(f"arch {arch} not supported")
def test_fast_fa_causal(self):
B, N, H, H_KV, D = 1, 8192, 32, 8, 128
with Context(DEBUG=0):
q = Tensor.randn(B, N, H, D, dtype=dtypes.bfloat16).contiguous()
k = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16).contiguous()
v = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16).contiguous()
Tensor.realize(q, k, v)
q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
fa_jitted = TinyJit(flash_attention)
for _ in range(10):
st = time.perf_counter()
out = fa_jitted(q, k, v, is_causal=True)
et = time.perf_counter() - st
attn_flops = 2 * B * H * N * N * D + \
4 * B * H * N * N + \
2 * B * H * N * N * D
print(f"{attn_flops/(et*1e9):2f} GFLOPS")
out = out.float().transpose(1, 2)
ref = q.scaled_dot_product_attention(k, v, is_causal=True, enable_gqa=True).float().transpose(1, 2)
assert_allclose(out, ref, atol=2e-2, rtol=2e-2)
def test_fast_fa_bwd_causal(self):
Tensor.manual_seed(42)
B, N, H, H_KV, D = 1, 8192, 32, 8, 128
with Context(DEBUG=0):
q = Tensor.randn(B, N, H, D, dtype=dtypes.bfloat16).contiguous()
k = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16).contiguous()
v = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16).contiguous()
Tensor.realize(q, k, v)
do = Tensor.ones(B, N, H, D, dtype=dtypes.float32).contiguous()
Tensor.realize(do)
q_, k_, v_ = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
out = flash_attention(q_, k_, v_, is_causal=True)
out = out.float().transpose(1, 2)
out.backward(do)
Tensor.realize(q.grad, k.grad, v.grad)
with Context(DEBUG=0):
q_ref = q.detach().clone()
k_ref = k.detach().clone()
v_ref = v.detach().clone()
Tensor.realize(q_ref, k_ref, v_ref)
q_ref_, k_ref_, v_ref_ = q_ref.transpose(1, 2), k_ref.transpose(1, 2), v_ref.transpose(1, 2)
ref = q_ref_.scaled_dot_product_attention(k_ref_, v_ref_, is_causal=True, enable_gqa=True)
ref = ref.float().transpose(1, 2)
ref.backward(do)
Tensor.realize(q_ref.grad, k_ref.grad, v_ref.grad)
assert_allclose(q.grad, q_ref.grad, atol=2e-2, rtol=2e-2)
assert_allclose(v.grad, v_ref.grad, atol=2e-2, rtol=2e-2)
assert_allclose(k.grad, k_ref.grad, atol=6e-2, rtol=2e-2)
def test_fast_fa_bwd_causal_jitted(self):
Tensor.manual_seed(42)
B, N, H, H_KV, D = 1, 8192, 32, 8, 128
with Context(DEBUG=0):
q = Tensor.randn(B, N, H, D, dtype=dtypes.bfloat16).contiguous()
k = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16).contiguous()
v = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16).contiguous()
Tensor.realize(q, k, v)
do = Tensor.ones(B, N, H, D, dtype=dtypes.float32).contiguous()
Tensor.realize(do)
def fn(q, k, v, do):
q_, k_, v_ = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
out = flash_attention(q_, k_, v_, is_causal=True)
out = out.float().transpose(1, 2)
out.backward(do)
Tensor.realize(out, q.grad, k.grad, v.grad)
return q.grad, k.grad, v.grad
fn_jitted = TinyJit(fn)
for _ in range(10):
q = Tensor.randn(B, N, H, D, dtype=dtypes.bfloat16).contiguous()
k = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16).contiguous()
v = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16).contiguous()
Tensor.realize(q, k, v)
do = Tensor.ones(B, N, H, D, dtype=dtypes.float32).contiguous()
Tensor.realize(do)
q.grad, k.grad, v.grad = fn_jitted(q, k, v, do)
with Context(DEBUG=0):
q_ref = q.detach().clone()
k_ref = k.detach().clone()
v_ref = v.detach().clone()
Tensor.realize(q_ref, k_ref, v_ref)
q_ref_, k_ref_, v_ref_ = q_ref.transpose(1, 2), k_ref.transpose(1, 2), v_ref.transpose(1, 2)
ref = flash_attention(q_ref_, k_ref_, v_ref_, is_causal=True)
ref = ref.float().transpose(1, 2)
ref.backward(do)
Tensor.realize(q_ref.grad, k_ref.grad, v_ref.grad)
assert_allclose(q.grad, q_ref.grad, atol=3e-3, rtol=3e-3)
assert_allclose(k.grad, k_ref.grad, atol=1e-5, rtol=1e-5)
assert_allclose(v.grad, v_ref.grad, atol=1e-5, rtol=1e-5)
def test_fast_fa_bwd_dp(self):
Tensor.manual_seed(42)
B, N, H, H_KV, D = 2, 1024, 32, 8, 128
GPUS = tuple(f"AMD:{i}" for i in range(B))
with Context(DEBUG=0):
base_q = Tensor.randn(B, N, H, D, dtype=dtypes.bfloat16).contiguous()
base_k = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16).contiguous()
base_v = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16).contiguous()
base_do = Tensor.ones(B, N, H, D, dtype=dtypes.float32).contiguous()
with Context(DEBUG=0):
q = base_q.clone().shard(GPUS, axis=0)
k = base_k.clone().shard(GPUS, axis=0)
v = base_v.clone().shard(GPUS, axis=0)
Tensor.realize(q, k, v)
do = base_do.clone().shard(GPUS, axis=0)
Tensor.realize(do)
q_, k_, v_ = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
out = flash_attention(q_, k_, v_, is_causal=True)
out = out.float().transpose(1, 2)
out.backward(do)
Tensor.realize(q.grad, k.grad, v.grad)
with Context(DEBUG=0):
q_ref = base_q.clone()
k_ref = base_k.clone()
v_ref = base_v.clone()
Tensor.realize(q_ref, k_ref, v_ref)
do_ref = base_do.clone()
Tensor.realize(do_ref)
q_ref_, k_ref_, v_ref_ = q_ref.transpose(1, 2), k_ref.transpose(1, 2), v_ref.transpose(1, 2)
ref = flash_attention(q_ref_, k_ref_, v_ref_, is_causal=True)
ref = ref.float().transpose(1, 2)
ref.backward(do_ref)
Tensor.realize(q_ref.grad, k_ref.grad, v_ref.grad)
assert_allclose(q.grad, q_ref.grad, atol=1e-5, rtol=1e-5)
assert_allclose(v.grad, v_ref.grad, atol=1e-5, rtol=1e-5)
assert_allclose(k.grad, k_ref.grad, atol=1e-5, rtol=1e-5)
def test_fast_fa_bwd_mp(self):
Tensor.manual_seed(42)
B, N, H, H_KV, D = 2, 1024, 32, 8, 128
GPUS = tuple(f"AMD:{i}" for i in range(B))
with Context(DEBUG=0):
base_q = Tensor.randn(B, N, H, D, dtype=dtypes.bfloat16).contiguous()
base_k = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16).contiguous()
base_v = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16).contiguous()
base_do = Tensor.ones(B, N, H, D, dtype=dtypes.float32).contiguous()
with Context(DEBUG=0):
q = base_q.clone().shard(GPUS, axis=2)
k = base_k.clone().shard(GPUS, axis=2)
v = base_v.clone().shard(GPUS, axis=2)
Tensor.realize(q, k, v)
do = base_do.clone().shard(GPUS, axis=2)
Tensor.realize(do)
q_, k_, v_ = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
out = flash_attention(q_, k_, v_, is_causal=True)
out = out.float().transpose(1, 2)
out.backward(do)
Tensor.realize(q.grad, k.grad, v.grad)
with Context(DEBUG=0):
q_ref = base_q.clone()
k_ref = base_k.clone()
v_ref = base_v.clone()
Tensor.realize(q_ref, k_ref, v_ref)
do_ref = base_do.clone()
Tensor.realize(do_ref)
q_ref_, k_ref_, v_ref_ = q_ref.transpose(1, 2), k_ref.transpose(1, 2), v_ref.transpose(1, 2)
ref = flash_attention(q_ref_, k_ref_, v_ref_, is_causal=True)
ref = ref.float().transpose(1, 2)
ref.backward(do_ref)
Tensor.realize(q_ref.grad, k_ref.grad, v_ref.grad)
assert_allclose(q.grad, q_ref.grad, atol=1e-5, rtol=1e-5)
assert_allclose(v.grad, v_ref.grad, atol=1e-5, rtol=1e-5)
assert_allclose(k.grad, k_ref.grad, atol=1e-5, rtol=1e-5)
if __name__ == "__main__":
unittest.main()

125
tinygrad_repo/test/testextra/test_lr_scheduler.py Normal file
View File

@@ -0,0 +1,125 @@
import numpy as np
import torch
import unittest
from tinygrad.tensor import Tensor
from tinygrad.nn.state import get_parameters
from tinygrad.nn.optim import Adam, SGD
from tinygrad.helpers import DEBUG
from extra.lr_scheduler import MultiStepLR, ReduceLROnPlateau, CosineAnnealingLR, OneCycleLR
from extra.training import train, evaluate
from extra.datasets import fetch_mnist
np.random.seed(1337)
Tensor.manual_seed(1337)
X_train, Y_train, X_test, Y_test = fetch_mnist()
class TinyBobNet:
def __init__(self):
self.l1 = Tensor.scaled_uniform(784, 128)
self.l2 = Tensor.scaled_uniform(128, 10)
def parameters(self):
return get_parameters(self)
def forward(self, x):
return x.dot(self.l1).relu().dot(self.l2).log_softmax()
def lr_scheduler_training(sched_fn=None, args=None):
model = TinyBobNet()
optim = Adam(model.parameters(), lr=0.01)
if sched_fn is not None: sched = sched_fn(optim, **args)
for _ in range(25):
train(model, X_train, Y_train, optim, 100)
if sched_fn is not None:
if isinstance(sched, ReduceLROnPlateau):
sched.step(evaluate(model, X_test, Y_test))
else:
sched.step()
return evaluate(model, X_test, Y_test)
def current_lr(optim): return optim.param_groups[0]['lr'] if hasattr(optim, 'param_groups') else optim.lr
def get_lrs(optim, sched, epochs, steps=1, accs=None):
lr = current_lr(optim)
if not isinstance(lr, float): lr = lr.numpy()[0]
lrs = [lr]
for e in range(epochs):
for _ in range(steps):
optim.step()
sched.step() if accs is None else sched.step(accs[e])
lr = current_lr(optim)
if not isinstance(lr, float): lr = lr.numpy()[0]
lrs.append(lr)
return lrs
class TestLrScheduler(unittest.TestCase):
def setUp(self):
self.old_training = Tensor.training
Tensor.training = True
def tearDown(self):
Tensor.training = self.old_training
def _test_lr_scheduler(self, tinygrad_sched, torch_sched, epochs, opts, atol, rtol, adam=True):
accs = opts.pop('accs', None)
test_tensor = Tensor([0.]) # NOTE: optimizers are broken on 0-dim tensors because it broadcasts to [lr]
test_tensor.mean().backward()
if adam:
tinygrad_optim, torch_optim = Adam([test_tensor], lr=0.01), torch.optim.Adam([torch.tensor([0.], requires_grad=True)], lr=0.01)
else:
tinygrad_optim, torch_optim = SGD([test_tensor], lr=0.01), torch.optim.SGD([torch.tensor([0.], requires_grad=True)], lr=0.01)
tinygrad_sched, torch_sched = tinygrad_sched(tinygrad_optim, **opts), torch_sched(torch_optim, **opts)
tinygrad_lrs = get_lrs(tinygrad_optim, tinygrad_sched, epochs, accs=accs)
torch_lrs = get_lrs(torch_optim, torch_sched, epochs, accs=accs)
np.testing.assert_allclose(tinygrad_lrs, torch_lrs, atol=atol, rtol=rtol)
def _test_multisteplr(self, epochs, opts, atol, rtol, adam=True):
self._test_lr_scheduler(MultiStepLR, torch.optim.lr_scheduler.MultiStepLR, epochs, opts, atol, rtol, adam=adam)
def _test_reducelronplateau(self, epochs, opts, atol, rtol):
opts['accs'] = np.random.randn(epochs)
self._test_lr_scheduler(ReduceLROnPlateau, torch.optim.lr_scheduler.ReduceLROnPlateau, epochs, opts, atol, rtol)
def _test_cosineannealinglr(self, epochs, opts, atol, rtol):
opts['T_max'] = epochs
self._test_lr_scheduler(CosineAnnealingLR, torch.optim.lr_scheduler.CosineAnnealingLR, epochs, opts, atol, rtol)
def _test_onecyclelr(self, epochs, opts, atol, rtol):
opts['total_steps'] = epochs
self._test_lr_scheduler(OneCycleLR, torch.optim.lr_scheduler.OneCycleLR, epochs, opts, atol, rtol)
def test_multisteplr(self): self._test_multisteplr(10, {'milestones': [1, 2, 7]}, 1e-6, 1e-6)
def test_multisteplr_gamma(self): self._test_multisteplr(10, {'milestones': [1, 2, 7], 'gamma': 0.1337}, 1e-6, 1e-6)
def test_reducelronplateau(self): self._test_reducelronplateau(100, {}, 1e-6, 1e-6)
def test_reducelronplateau_max(self): self._test_reducelronplateau(100, {'mode': 'max'}, 1e-6, 1e-6)
def test_reducelronplateau_factor(self): self._test_reducelronplateau(100, {'factor': 0.1337}, 1e-6, 1e-6)
def test_reducelronplateau_patience(self): self._test_reducelronplateau(100, {'patience': 3}, 1e-6, 1e-6)
def test_reducelronplateau_threshold(self): self._test_reducelronplateau(100, {'threshold': 1e-6}, 1e-6, 1e-6)
def test_reducelronplateau_threshold_mode(self): self._test_reducelronplateau(100, {'threshold_mode': 'abs'}, 1e-6, 1e-6)
def test_cosineannealinglr(self): self._test_cosineannealinglr(100, {}, 1e-6, 1e-6)
def test_cosineannealinglr_eta_min(self): self._test_cosineannealinglr(100, {'eta_min': 0.001}, 1e-6, 1e-6)
def test_multistep_2step(self):
# was making this fail with LRU=1, some issue with epoch_counter
if DEBUG>=2: print("first")
self._test_multisteplr(1, {'milestones': [1]}, 1e-6, 1e-6, adam=False)
if DEBUG>=2: print("second")
self._test_multisteplr(1, {'milestones': [1], 'gamma': 0.133}, 1e-6, 1e-6, adam=False)
if DEBUG>=2: print("third")
def test_onecyclelr(self): self._test_onecyclelr(100, {'pct_start': 0.3, 'anneal_strategy': 'linear',
'cycle_momentum': False, 'div_factor': 25.0,
'final_div_factor': 10000.0, 'max_lr':1e-5}, 1e-6, 1e-6)
@unittest.skip("slow")
def test_training(self):
without = lr_scheduler_training()
sched_fns = [MultiStepLR, ReduceLROnPlateau, CosineAnnealingLR, OneCycleLR]
argss = [{'milestones': [5, 7, 10, 15], 'gamma': 0.5}, {'factor': 0.5, 'patience': 2}, {'T_max': 25, 'eta_min': 0.001},
{'pct_start': 0.3, 'anneal_strategy': 'linear', 'cycle_momentum': False, 'div_factor': 25.0, 'final_div_factor': 10000.0,
'max_lr':1e-5, 'total_steps': 25}]
for sched_fn, args in zip(sched_fns, argss):
with_sched = lr_scheduler_training(sched_fn, args)
assert with_sched > without
if __name__ == '__main__':
unittest.main()

13
tinygrad_repo/test/testextra/test_mockgpu.py Normal file
View File

@@ -0,0 +1,13 @@
from tinygrad.helpers import DEV
import unittest, importlib
@unittest.skipUnless(DEV.interface.startswith("MOCK"), 'Testing mockgpu')
class TestMockGPU(unittest.TestCase):
# https://github.com/tinygrad/tinygrad/pull/7627
def test_import_typing_extensions(self):
import test.mockgpu.mockgpu # noqa: F401 # pylint: disable=unused-import
import typing_extensions
importlib.reload(typing_extensions) # pytest imports typing_extension before mockgpu
if __name__ == '__main__':
unittest.main()

988
tinygrad_repo/test/testextra/test_tk.py Normal file
View File

@@ -0,0 +1,988 @@
import unittest, math, time
from tinygrad import Tensor, Device, dtypes, Context, GlobalCounters
from tinygrad.uop.ops import UOp, Ops
from tinygrad.engine.realize import run_linear
from tinygrad.engine.jit import TinyJit
import numpy as np
from extra.thunder.tiny.tk import WARP_THREADS
from extra.thunder.tiny.tk.kernel import Kernel
from extra.thunder.tiny.tk.tiles import ST_16X32, RT_16X32, RT_16X16, TileLayout
def assert_allclose(cmp:Tensor, ref:Tensor, **kwargs) -> None:
if Device.DEFAULT == "NULL": Tensor.realize(cmp, ref)
else: np.testing.assert_allclose(cmp.numpy(), ref.numpy(), **kwargs)
@unittest.skip("TODO: broken after ranges on store instead of after")
class TestTK(unittest.TestCase):
def setUp(self):
arch = Device[Device.DEFAULT].renderer.target.arch
if not arch.startswith("gfx9"):
self.skipTest(f"arch {arch} not supported")
def test_simple_matmul(self):
N = 8192
BLOCK_SIZE = 64
with Kernel("simple_matmul", (N // BLOCK_SIZE, N // BLOCK_SIZE, 1), WARP_THREADS) as ker:
warp = ker.warp
c = ker.gl((1, 1, N, N), dtypes.float32)
a = ker.gl((1, 1, N, N), dtypes.bfloat16)
b = ker.gl((1, 1, N, N), dtypes.bfloat16)
a_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE), dtypes.bfloat16)
b_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE), dtypes.bfloat16)
c_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
a_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.bfloat16)
b_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.bfloat16, TileLayout.COL)
c_reg_col = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32, TileLayout.COL)
c_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
col, row = ker.blockIdx_x, ker.blockIdx_y
c_reg_col = warp.zero(c_reg_col)
for tile in ker.range(N // BLOCK_SIZE):
a_smem = warp.load(a_smem, a, (), (0, 0, row, tile), axis=2)
b_smem = warp.load(b_smem, b, (), (0, 0, tile, col), axis=2)
a_reg = warp.load(a_reg, a_smem)
b_reg = warp.load(b_reg, b_smem)
c_reg_col = warp.mma_AB(c_reg_col, a_reg, b_reg)
c_reg_col = ker.endrange()
c_smem = warp.store(c_smem, c_reg_col)
c_reg = warp.load(c_reg, c_smem)
c = warp.store(c, c_reg, (0, 0, row, col), (), axis=2)
sink = ker.finish()
with Context(DEBUG=0):
a = Tensor.rand(1, 1, N, N, dtype="bfloat16").contiguous()
b = Tensor.rand(1, 1, N, N, dtype="bfloat16").contiguous()
c = Tensor.empty(1, 1, N, N, dtype="float32")
Tensor.realize(a, b, c)
linear = UOp(Ops.LINEAR, src=(sink.call(*[t.uop.buf_uop for t in (c, a, b)]),))
for _ in range(5): run_linear(linear, update_stats=False)
c = c.float()
ref = a.matmul(b, dtype=dtypes.float32).float()
assert_allclose(c, ref)
def test_simple_matmul_transposed(self):
N = 8192
BLOCK_N, BLOCK_M, BLOCK_K = 64, 64, 128
with Kernel("simple_matmul_transposed", (N // BLOCK_N, N // BLOCK_M, 1), WARP_THREADS) as ker:
warp = ker.warp
c = ker.gl((1, 1, N, N), dtypes.float32)
a = ker.gl((1, 1, N, N), dtypes.bfloat16)
b = ker.gl((1, 1, N, N), dtypes.bfloat16)
a_smem = ker.st((BLOCK_N, BLOCK_K), dtypes.bfloat16, base_shape=ST_16X32)
b_smem = ker.st((BLOCK_M, BLOCK_K), dtypes.bfloat16, base_shape=ST_16X32)
a_reg = ker.rt((BLOCK_N, BLOCK_K), dtypes.bfloat16, base_shape=RT_16X32)
b_reg = ker.rt((BLOCK_M, BLOCK_K), dtypes.bfloat16, base_shape=RT_16X32)
c_reg = ker.rt((BLOCK_N, BLOCK_M), dtypes.float32, TileLayout.COL, base_shape=RT_16X16)
col, row = ker.blockIdx_x, ker.blockIdx_y
c_reg = warp.zero(c_reg)
for tile in ker.range(N // BLOCK_K):
a_smem = warp.load(a_smem, a, (), (0, 0, row, tile), axis=2)
b_smem = warp.load(b_smem, b, (), (0, 0, col, tile), axis=2)
a_reg = warp.load(a_reg, a_smem)
b_reg = warp.load(b_reg, b_smem)
c_reg = warp.mma_ABt(c_reg, a_reg, b_reg)
c_reg = ker.endrange()
c = warp.store(c, c_reg, (0, 0, row, col), (), axis=2)
sink = ker.finish()
with Context(DEBUG=0):
a = Tensor.rand(1, 1, N, N, dtype="bfloat16").contiguous()
b = Tensor.rand(1, 1, N, N, dtype="bfloat16").contiguous()
c = Tensor.empty(1, 1, N, N, dtype="float32")
Tensor.realize(a, b, c)
linear = UOp(Ops.LINEAR, src=(sink.call(*[t.uop.buf_uop for t in (c, a, b)]),))
for _ in range(5): run_linear(linear, update_stats=False)
c = c.float()
ref = a.matmul(b.transpose(2, 3), dtype=dtypes.float32).float()
assert_allclose(c, ref)
def test_load_store(self):
N = 64
BLOCK_SIZE = 32
with Kernel("load_store", (N // BLOCK_SIZE, N // BLOCK_SIZE, 1), WARP_THREADS) as ker:
warp = ker.warp
b = ker.gl((1, 1, N, N), dtypes.float32)
a = ker.gl((1, 1, N, N), dtypes.float32)
a_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
a_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
b_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
col, row = ker.blockIdx_x, ker.blockIdx_y
a_smem = warp.load(a_smem, a, (), (0, 0, row, col), axis=2)
a_reg = warp.load(a_reg, a_smem)
b_reg = warp.copy(b_reg, a_reg)
b = warp.store(b, b_reg, (0, 0, row, col), (), axis=2)
sink = ker.finish()
with Context(DEBUG=0):
a = Tensor.rand(1, 1, N, N, dtype="float32").contiguous()
b = Tensor.empty(1, 1, N, N, dtype="float32")
Tensor.realize(a, b)
linear = UOp(Ops.LINEAR, src=(sink.call(*[t.uop.buf_uop for t in (b, a)]),))
for _ in range(5): run_linear(linear, update_stats=False)
b = b.float()
ref = a.float()
assert_allclose(b, ref)
def test_load_store_local_hop(self):
N = 64
BLOCK_SIZE = 32
with Kernel("load_store_local_hop", (N // BLOCK_SIZE, N // BLOCK_SIZE, 1), WARP_THREADS) as ker:
warp = ker.warp
b = ker.gl((1, 1, N, N), dtypes.float32)
a = ker.gl((1, 1, N, N), dtypes.float32)
a_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
b_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
a_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
b_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
col, row = ker.blockIdx_x, ker.blockIdx_y
a_smem = warp.load(a_smem, a, (), (0, 0, row, col), axis=2)
a_reg = warp.load(a_reg, a_smem)
b_reg = warp.copy(b_reg, a_reg)
b_smem = warp.store(b_smem, b_reg)
b_reg = warp.load(b_reg, b_smem)
b = warp.store(b, b_reg, (0, 0, row, col), (), axis=2)
sink = ker.finish()
with Context(DEBUG=0):
a = Tensor.rand(1, 1, N, N, dtype="float32").contiguous()
b = Tensor.empty(1, 1, N, N, dtype="float32")
Tensor.realize(a, b)
linear = UOp(Ops.LINEAR, src=(sink.call(*[t.uop.buf_uop for t in (b, a)]),))
for _ in range(5): run_linear(linear, update_stats=False)
b = b.float()
ref = a.float()
assert_allclose(b, ref)
def test_load_store_multioutput(self):
N = 64
BLOCK_SIZE = 32
with Kernel("load_store_multioutput", (N // BLOCK_SIZE, N // BLOCK_SIZE, 1), WARP_THREADS) as ker:
warp = ker.warp
b = ker.gl((1, 1, N, N), dtypes.float32)
c = ker.gl((1, 1, N, N), dtypes.float32)
a = ker.gl((1, 1, N, N), dtypes.float32)
a_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
b_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
a_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
b_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
col, row = ker.blockIdx_x, ker.blockIdx_y
a_smem = warp.load(a_smem, a, (), (0, 0, row, col), axis=2)
a_reg = warp.load(a_reg, a_smem)
b_reg = warp.copy(b_reg, a_reg)
b_smem = warp.store(b_smem, b_reg)
b_reg = warp.load(b_reg, b_smem)
b = warp.store(b, b_reg, (0, 0, row, col), (), axis=2)
c = warp.store(c, b_reg, (0, 0, row, col), (), axis=2)
sink = ker.finish(2)
with Context(DEBUG=0):
a = Tensor.rand(1, 1, N, N, dtype="float32").contiguous()
b = Tensor.empty(1, 1, N, N, dtype="float32")
c = Tensor.empty(1, 1, N, N, dtype="float32")
Tensor.realize(a, b, c)
linear = UOp(Ops.LINEAR, src=(sink.call(*[t.uop.buf_uop for t in (b, c, a)]),))
for _ in range(5): run_linear(linear, update_stats=False)
b = b.float()
c = c.float()
ref = a.float()
assert_allclose(b, ref)
assert_allclose(c, ref)
def test_load_store_group(self):
N = 1024
BLOCK_SIZE = 64
NUM_WORKERS = 4
with Kernel("load_store_group", (N // (BLOCK_SIZE * NUM_WORKERS), N // BLOCK_SIZE, 1), WARP_THREADS * NUM_WORKERS) as ker:
warp = ker.warp
group = ker.group(NUM_WORKERS)
b = ker.gl((1, 1, N, N), dtypes.float32)
a = ker.gl((1, 1, N, N), dtypes.float32)
a_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE * NUM_WORKERS), dtypes.float32)
a_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
b_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
col, row = ker.blockIdx_x, ker.blockIdx_y
a_smem = group.load(a_smem, a, (), (0, 0, row, col), axis=2)
a_reg = warp.load(a_reg, a_smem, (), (0, ker.warpid,))
b_reg = warp.copy(b_reg, a_reg)
b = warp.store(b, b_reg, (0, 0, row, col * NUM_WORKERS + ker.warpid), (), axis=2)
sink = ker.finish()
with Context(DEBUG=0):
a = Tensor.rand(1, 1, N, N, dtype="float32").contiguous()
b = Tensor.empty(1, 1, N, N, dtype="float32")
Tensor.realize(a, b)
linear = UOp(Ops.LINEAR, src=(sink.call(*[t.uop.buf_uop for t in (b, a)]),))
for _ in range(5): run_linear(linear, update_stats=False)
b = b.float()
ref = a.float()
assert_allclose(b, ref)
def test_add(self):
N = 64
BLOCK_SIZE = 32
with Kernel("add", (1, 1, 1), WARP_THREADS) as ker:
warp = ker.warp
b = ker.gl((1, 1, N, N), dtypes.float32)
a = ker.gl((1, 1, N, N), dtypes.float32)
a_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
a_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
for tile_row in ker.range(N // BLOCK_SIZE):
for tile_col in ker.range(N // BLOCK_SIZE):
a_smem = warp.load(a_smem, a, (), (0, 0, tile_row, tile_col), axis=2)
a_reg = warp.load(a_reg, a_smem)
a_reg += 1
b = warp.store(b, a_reg, (0, 0, tile_row, tile_col), (), axis=2)
sink = ker.finish()
with Context(DEBUG=0):
a = Tensor.rand(1, 1, N, N, dtype="float32").contiguous()
b = Tensor.empty(1, 1, N, N, dtype="float32")
Tensor.realize(a, b)
linear = UOp(Ops.LINEAR, src=(sink.call(*[t.uop.buf_uop for t in (b, a)]),))
for _ in range(5): run_linear(linear, update_stats=False)
b = b.float()
ref = a.float() + 1
assert_allclose(b, ref)
def test_max(self):
N = 64
BLOCK_SIZE = 32
with Kernel("max", (1, 1, 1), WARP_THREADS) as ker:
warp = ker.warp
b = ker.gl((1, 1, N, N), dtypes.float32)
a = ker.gl((1, 1, N, N), dtypes.float32)
a_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
a_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32, TileLayout.COL)
b_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32, TileLayout.COL)
max_reg = ker.rv(BLOCK_SIZE, dtypes.float32)
for tile_col in ker.range(N // BLOCK_SIZE):
max_reg = warp.neg_inf(max_reg.after(tile_col))
for tile_row in ker.range(N // BLOCK_SIZE):
a_smem = warp.load(a_smem, a, (), (0, 0, tile_row, tile_col), axis=2)
a_reg = warp.load(a_reg, a_smem)
max_reg = warp.col_reduce(max_reg, a_reg, lambda a, b: a.maximum(b), init_value=-math.inf)
max_reg = ker.endrange()
b_reg = warp.map(b_reg, lambda _, idx: max_reg[idx[1], 0])
for tile_row in ker.range(N // BLOCK_SIZE):
b = warp.store(b, b_reg, (0, 0, tile_row, tile_col), (), axis=2)
sink = ker.finish()
with Context(DEBUG=0):
a = Tensor.rand(1, 1, N, N, dtype="float32").contiguous()
b = Tensor.empty(1, 1, N, N, dtype="float32")
Tensor.realize(a, b)
linear = UOp(Ops.LINEAR, src=(sink.call(*[t.uop.buf_uop for t in (b, a)]),))
for _ in range(5): run_linear(linear, update_stats=False)
b = b.float()
ref = a.float().max(axis=2, keepdim=True).expand(a.shape)
assert_allclose(b, ref)
def test_max_nonsquare(self):
N, M = 32, 128
BLOCK_N, BLOCK_M = 16, 64
with Kernel("max_nonsquare", (1, 1, 1), WARP_THREADS) as ker:
warp = ker.warp
b = ker.gl((1, 1, N, M), dtypes.float32)
a = ker.gl((1, 1, N, M), dtypes.float32)
a_smem = ker.st((BLOCK_N, BLOCK_M), dtypes.float32)
a_reg = ker.rt((BLOCK_N, BLOCK_M), dtypes.float32, TileLayout.COL)
b_reg = ker.rt((BLOCK_N, BLOCK_M), dtypes.float32, TileLayout.COL)
max_reg = ker.rv(BLOCK_M, dtypes.float32)
for tile_col in ker.range(M // BLOCK_M):
max_reg = warp.neg_inf(max_reg.after(tile_col))
for tile_row in ker.range(N // BLOCK_N):
a_smem = warp.load(a_smem, a, (), (0, 0, tile_row, tile_col), axis=2)
a_reg = warp.load(a_reg, a_smem)
max_reg = warp.col_reduce(max_reg, a_reg, lambda a, b: a.maximum(b), init_value=-math.inf)
max_reg = ker.endrange()
b_reg = warp.map(b_reg, lambda _, idx: max_reg[idx[1], 0])
for tile_row in ker.range(N // BLOCK_N):
b = warp.store(b, b_reg, (0, 0, tile_row, tile_col), (), axis=2)
sink = ker.finish()
with Context(DEBUG=0):
a = Tensor.rand(1, 1, N, M, dtype="float32").contiguous()
b = Tensor.empty(1, 1, N, M, dtype="float32")
Tensor.realize(a, b)
linear = UOp(Ops.LINEAR, src=(sink.call(*[t.uop.buf_uop for t in (b, a)]),))
for _ in range(5): run_linear(linear, update_stats=False)
b = b.float()
ref = a.float().max(axis=2, keepdim=True).expand(a.shape)
assert_allclose(b, ref)
def test_sum(self):
N = 64
BLOCK_SIZE = 32
with Kernel("sum", (1, 1, 1), WARP_THREADS) as ker:
warp = ker.warp
b = ker.gl((1, 1, N, N), dtypes.float32)
a = ker.gl((1, 1, N, N), dtypes.float32)
a_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
a_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32, TileLayout.COL)
b_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32, TileLayout.COL)
sum_reg = ker.rv(BLOCK_SIZE, dtypes.float32)
for tile_col in ker.range(N // BLOCK_SIZE):
sum_reg = warp.zero(sum_reg.after(tile_col))
for tile_row in ker.range(N // BLOCK_SIZE):
a_smem = warp.load(a_smem, a, (), (0, 0, tile_row, tile_col), axis=2)
a_reg = warp.load(a_reg, a_smem)
sum_reg = warp.col_reduce(sum_reg, a_reg, lambda a, b: a + b)
sum_reg = ker.endrange()
b_reg = warp.map(b_reg, lambda _, idx: sum_reg[idx[1], 0])
for tile_row in ker.range(N // BLOCK_SIZE):
b = warp.store(b, b_reg, (0, 0, tile_row, tile_col), (), axis=2)
sink = ker.finish()
with Context(DEBUG=0):
a = Tensor.rand(1, 1, N, N, dtype="float32").contiguous()
b = Tensor.empty(1, 1, N, N, dtype="float32")
Tensor.realize(a, b)
linear = UOp(Ops.LINEAR, src=(sink.call(*[t.uop.buf_uop for t in (b, a)]),))
for _ in range(5): run_linear(linear, update_stats=False)
b = b.float()
ref = a.float().sum(axis=2, keepdim=True).expand(a.shape)
assert_allclose(b, ref, atol=1e-5, rtol=1e-5)
def test_sum_nonsquare(self):
N, M = 32, 128
BLOCK_N, BLOCK_M = 16, 64
with Kernel("sum_nonsquare", (1, 1, 1), WARP_THREADS) as ker:
warp = ker.warp
b = ker.gl((1, 1, N, M), dtypes.float32)
a = ker.gl((1, 1, N, M), dtypes.float32)
a_smem = ker.st((BLOCK_N, BLOCK_M), dtypes.float32)
a_reg = ker.rt((BLOCK_N, BLOCK_M), dtypes.float32, TileLayout.COL)
b_reg = ker.rt((BLOCK_N, BLOCK_M), dtypes.float32, TileLayout.COL)
sum_reg = ker.rv(BLOCK_M, dtypes.float32)
for tile_col in ker.range(M // BLOCK_M):
sum_reg = warp.zero(sum_reg.after(tile_col))
for tile_row in ker.range(N // BLOCK_N):
a_smem = warp.load(a_smem, a, (), (0, 0, tile_row, tile_col), axis=2)
a_reg = warp.load(a_reg, a_smem)
sum_reg = warp.col_reduce(sum_reg, a_reg, lambda a, b: a + b)
sum_reg = ker.endrange()
b_reg = warp.map(b_reg, lambda _, idx: sum_reg[idx[1], 0])
for tile_row in ker.range(N // BLOCK_N):
b = warp.store(b, b_reg, (0, 0, tile_row, tile_col), (), axis=2)
sink = ker.finish()
with Context(DEBUG=0):
a = Tensor.rand(1, 1, N, M, dtype="float32").contiguous()
b = Tensor.empty(1, 1, N, M, dtype="float32")
Tensor.realize(a, b)
linear = UOp(Ops.LINEAR, src=(sink.call(*[t.uop.buf_uop for t in (b, a)]),))
for _ in range(5): run_linear(linear, update_stats=False)
b = b.float()
ref = a.float().sum(axis=2, keepdim=True).expand(a.shape)
assert_allclose(b, ref, atol=1e-5, rtol=1e-5)
def test_softmax(self):
N = 64
BLOCK_SIZE = 32
with Kernel("softmax", (1, 1, 1), WARP_THREADS) as ker:
warp = ker.warp
b = ker.gl((1, 1, BLOCK_SIZE, N), dtypes.float32)
a = ker.gl((1, 1, BLOCK_SIZE, N), dtypes.float32)
a_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
a_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
max_vec_last = ker.rv(BLOCK_SIZE, dtypes.float32)
max_vec = ker.rv(BLOCK_SIZE, dtypes.float32)
norm_vec = ker.rv(BLOCK_SIZE, dtypes.float32)
max_vec = warp.neg_inf(max_vec)
norm_vec = warp.zero(norm_vec)
for tile_col in ker.range(N // BLOCK_SIZE):
a_smem_ = warp.load(a_smem, a, (), (0, 0, 0, tile_col), axis=2)
a_reg_ = warp.load(a_reg, a_smem_)
a_reg_ *= 1.0 / math.log(2)
max_vec_last = warp.copy(max_vec_last.after(tile_col), max_vec)
max_vec = warp.row_reduce(max_vec.after(max_vec_last), a_reg_, lambda a, b: a.maximum(b), init_value=-math.inf)
a_reg_ = (a_reg_ - max_vec).exp2()
max_vec_last = (max_vec_last - max_vec).exp2()
norm_vec *= max_vec_last
norm_vec = warp.row_reduce(norm_vec, a_reg_, lambda a, b: a + b)
norm_vec = ker.endrange()
max_vec = max_vec.after(norm_vec)
for tile_col in ker.range(N // BLOCK_SIZE):
a_smem_ = warp.load(a_smem, a, (), (0, 0, 0, tile_col), axis=2)
a_reg_ = warp.load(a_reg, a_smem_)
a_reg_ *= 1.0 / math.log(2)
a_reg_ = (a_reg_ - max_vec).exp2()
a_reg_ /= norm_vec
b = warp.store(b, a_reg_, (0, 0, 0, tile_col), (), axis=2)
sink = ker.finish()
with Context(DEBUG=0):
a = Tensor.rand(1, 1, BLOCK_SIZE, N, dtype="float32")
b = Tensor.empty(1, 1, BLOCK_SIZE, N, dtype="float32")
Tensor.realize(a, b)
linear = UOp(Ops.LINEAR, src=(sink.call(*[t.uop.buf_uop for t in (b, a)]),))
for _ in range(5): run_linear(linear, update_stats=False)
b = b.float()
ref = a.float().softmax(axis=3)
assert_allclose(b, ref, atol=1e-5, rtol=1e-5)
def test_softmax_col(self):
N = 64
BLOCK_SIZE = 32
with Kernel("softmax_col", (1, 1, 1), WARP_THREADS) as ker:
warp = ker.warp
b = ker.gl((1, 1, N, BLOCK_SIZE), dtypes.float32)
a = ker.gl((1, 1, N, BLOCK_SIZE), dtypes.float32)
a_smem = ker.st((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32)
a_reg = ker.rt((BLOCK_SIZE, BLOCK_SIZE), dtypes.float32, TileLayout.COL)
max_vec_last = ker.rv(BLOCK_SIZE, dtypes.float32)
max_vec = ker.rv(BLOCK_SIZE, dtypes.float32)
norm_vec = ker.rv(BLOCK_SIZE, dtypes.float32)
max_vec = warp.neg_inf(max_vec)
norm_vec = warp.zero(norm_vec)
for tile_row in ker.range(N // BLOCK_SIZE):
a_smem_ = warp.load(a_smem, a, (), (0, 0, tile_row, 0), axis=2)
a_reg_ = warp.load(a_reg, a_smem_)
a_reg_ *= 1.0 / math.log(2)
max_vec_last = warp.copy(max_vec_last.after(tile_row), max_vec)
max_vec = warp.col_reduce(max_vec.after(max_vec_last), a_reg_, lambda a, b: a.maximum(b), init_value=-math.inf)
a_reg_ = (a_reg_ - max_vec).exp2()
max_vec_last = (max_vec_last - max_vec).exp2()
norm_vec *= max_vec_last
norm_vec = warp.col_reduce(norm_vec, a_reg_, lambda a, b: a + b)
norm_vec = ker.endrange()
max_vec = max_vec.after(norm_vec)
for tile_row in ker.range(N // BLOCK_SIZE):
a_smem_ = warp.load(a_smem, a, (), (0, 0, tile_row, 0), axis=2)
a_reg_ = warp.load(a_reg.after(norm_vec), a_smem_)
a_reg_ *= 1.0 / math.log(2)
a_reg_ = (a_reg_ - max_vec).exp2()
a_reg_ /= norm_vec
b = warp.store(b, a_reg_, (0, 0, tile_row, 0), (), axis=2)
sink = ker.finish()
with Context(DEBUG=0):
a = Tensor.rand(1, 1, N, BLOCK_SIZE, dtype="float32")
b = Tensor.empty(1, 1, N, BLOCK_SIZE, dtype="float32")
Tensor.realize(a, b)
linear = UOp(Ops.LINEAR, src=(sink.call(*[t.uop.buf_uop for t in (b, a)]),))
for _ in range(5): run_linear(linear, update_stats=False)
b = b.float()
ref = a.float().softmax(axis=2)
assert_allclose(b, ref, atol=1e-5, rtol=1e-5)
def test_fa(self):
NUM_WORKERS = 1
B, N, H, H_KV, D = 2, 8192, 32, 8, 128
Q_BLOCK_SIZE = 16
KV_BLOCK_SIZE = 16
GROUP_SIZE = H // H_KV
with Kernel("fa", (H, N // (Q_BLOCK_SIZE*NUM_WORKERS), B), NUM_WORKERS * WARP_THREADS) as ker:
warp = ker.warp
# kernel
o = ker.gl((B, N, H, D), dtypes.bfloat16)
q = ker.gl((B, N, H, D), dtypes.bfloat16)
k = ker.gl((B, N, H_KV, D), dtypes.bfloat16)
v = ker.gl((B, N, H_KV, D), dtypes.bfloat16)
head = ker.blockIdx_x
head_kv = head // GROUP_SIZE
batch = ker.blockIdx_z
q_seq = ker.blockIdx_y * NUM_WORKERS + ker.warpid
k_smem = ker.st((KV_BLOCK_SIZE, D), dtypes.bfloat16)
v_smem = ker.st((KV_BLOCK_SIZE, D), dtypes.bfloat16)
q_reg_fl = ker.rt((Q_BLOCK_SIZE, D), dtypes.float32)
q_reg = ker.rt((Q_BLOCK_SIZE, D), dtypes.bfloat16)
q_reg_transposed = ker.rt((D, Q_BLOCK_SIZE), dtypes.bfloat16, TileLayout.COL)
k_reg = ker.rt((KV_BLOCK_SIZE, D), dtypes.bfloat16)
k_reg_transposed = ker.rt((D, KV_BLOCK_SIZE), dtypes.bfloat16, TileLayout.COL)
v_reg = ker.rt((KV_BLOCK_SIZE, D), dtypes.bfloat16, TileLayout.COL)
o_reg = ker.rt((D, Q_BLOCK_SIZE), dtypes.float32, TileLayout.COL)
o_reg_transposed = ker.rt((Q_BLOCK_SIZE, D), dtypes.float32)
att_block = ker.rt((KV_BLOCK_SIZE, Q_BLOCK_SIZE), dtypes.float32, TileLayout.COL)
att_block_mma = ker.rt((KV_BLOCK_SIZE, Q_BLOCK_SIZE), dtypes.bfloat16, TileLayout.COL)
max_vec_last = ker.rv(KV_BLOCK_SIZE, dtypes.float32)
max_vec = ker.rv(KV_BLOCK_SIZE, dtypes.float32)
norm_vec = ker.rv(KV_BLOCK_SIZE, dtypes.float32)
scale_vec = ker.rv(KV_BLOCK_SIZE, dtypes.float32)
max_vec = warp.neg_inf(max_vec)
norm_vec = warp.zero(norm_vec)
o_reg = warp.zero(o_reg)
scale_vec = warp.ones(scale_vec)
# load q tile
q_reg_fl = warp.load(q_reg_fl, q, (), (batch, q_seq, head, 0), axis=1)
q_reg_fl *= (1.0 / math.sqrt(D)) * (1.0 / math.log(2))
q_reg = warp.copy(q_reg, q_reg_fl)
q_reg_transposed = warp.transpose(q_reg_transposed, q_reg)
for kv_idx in ker.range(N // KV_BLOCK_SIZE):
k_smem = warp.load(k_smem, k, (), (batch, kv_idx, head_kv, 0), axis=1)
v_smem = warp.load(v_smem, v, (), (batch, kv_idx, head_kv, 0), axis=1)
k_reg = warp.load(k_reg, k_smem)
v_reg = warp.load(v_reg, v_smem)
# mma qk^t
att_block = warp.zero(att_block.after(kv_idx))
k_reg_transposed = warp.transpose(k_reg_transposed, k_reg)
att_block = warp.mma_AtB(att_block, k_reg_transposed, q_reg_transposed)
# mask for causal
q_base = q_seq * Q_BLOCK_SIZE + (warp.laneid % 16)
kv_base = kv_idx * KV_BLOCK_SIZE + (warp.laneid // 16) * 4
att_block = warp.map(att_block,
lambda x, idx: ((kv_base + idx[0]*16 + idx[2]) > (q_base + idx[1]*16)).alu(Ops.WHERE, UOp.ufix(x._uop, -math.inf), x))
# softmax
max_vec_last = warp.copy(max_vec_last.after(kv_idx), max_vec)
max_vec = warp.row_reduce(max_vec.after(max_vec_last), att_block, lambda a, b: a.maximum(b), init_value=-math.inf)
scale_vec = warp.map(scale_vec.after(max_vec_last, max_vec), lambda _, idx: max_vec_last[*idx] - max_vec[*idx])
scale_vec = scale_vec.exp2()
o_reg *= scale_vec
norm_vec *= scale_vec
att_block -= max_vec
att_block = att_block.exp2()
norm_vec = warp.row_reduce(norm_vec.after(scale_vec), att_block, lambda a, b: a + b)
# mma av
att_block_mma = warp.copy(att_block_mma.after(kv_idx, norm_vec), att_block)
o_reg = warp.mma_AtB(o_reg, v_reg, att_block_mma)
o_reg = ker.endrange()
norm_vec = norm_vec.after(o_reg)
o_reg /= norm_vec
o_reg_transposed = warp.transpose(o_reg_transposed, o_reg)
o = warp.store(o, o_reg_transposed, (batch, q_seq, head, 0), (), axis=1)
sink = ker.finish()
with Context(DEBUG=0):
q = Tensor.randn(B, N, H, D, dtype=dtypes.bfloat16).contiguous()
k = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16).contiguous()
v = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16).contiguous()
out = Tensor.empty(B, N, H, D, dtype=dtypes.bfloat16)
Tensor.realize(q, k, v, out)
linear = UOp(Ops.LINEAR, src=(sink.call(*[t.uop.buf_uop for t in (out, q, k, v)]),))
for _ in range(5):
GlobalCounters.reset()
with Context(DEBUG=2): run_linear(linear)
et = GlobalCounters.time_sum_s
attn_flops = 2 * B * H * N * N * D + \
4 * B * H * N * N + \
2 * B * H * N * N * D
print(f"{attn_flops/(et*1e9):2f} GFLOPS")
out = out.float()
q_permuted = q.permute(0, 2, 1, 3)
k_permuted = k.permute(0, 2, 1, 3)
v_permuted = v.permute(0, 2, 1, 3)
ref = q_permuted.scaled_dot_product_attention(k_permuted, v_permuted, is_causal=True, enable_gqa=True).float()
ref = ref.permute(0, 2, 1, 3)
assert_allclose(out, ref, atol=2e-2, rtol=2e-2)
def test_fast_fa(self):
from extra.thunder.tiny.fa import flash_attention
B, N, H, H_KV, D = 2, 8192, 32, 8, 128
with Context(DEBUG=0):
q = Tensor.randn(B, N, H, D, dtype=dtypes.bfloat16).contiguous()
k = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16).contiguous()
v = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16).contiguous()
Tensor.realize(q, k, v)
q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
fa_jitted = TinyJit(flash_attention)
for _ in range(10):
st = time.perf_counter()
out = fa_jitted(q, k, v, is_causal=False)
et = time.perf_counter() - st
attn_flops = 2 * B * H * N * N * D + \
4 * B * H * N * N + \
2 * B * H * N * N * D
print(f"{attn_flops/(et*1e9):2f} GFLOPS")
out = out.float().transpose(1, 2)
ref = q.scaled_dot_product_attention(k, v, is_causal=False, enable_gqa=True).float().transpose(1, 2)
assert_allclose(out, ref, atol=2e-2, rtol=2e-2)
def test_fast_fa_causal(self):
from extra.thunder.tiny.fa import flash_attention
B, N, H, H_KV, D = 2, 8192, 32, 8, 128
with Context(DEBUG=0):
q = Tensor.randn(B, N, H, D, dtype=dtypes.bfloat16).contiguous()
k = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16).contiguous()
v = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16).contiguous()
Tensor.realize(q, k, v)
q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
fa_jitted = TinyJit(flash_attention)
for _ in range(10):
st = time.perf_counter()
out = fa_jitted(q, k, v, is_causal=True)
et = time.perf_counter() - st
attn_flops = 2 * B * H * N * N * D + \
4 * B * H * N * N + \
2 * B * H * N * N * D
print(f"{attn_flops/(et*1e9):2f} GFLOPS")
out = out.float().transpose(1, 2)
ref = q.scaled_dot_product_attention(k, v, is_causal=True, enable_gqa=True).float().transpose(1, 2)
assert_allclose(out, ref, atol=2e-2, rtol=2e-2)
def test_fast_fa_bwd(self):
from extra.thunder.tiny.fa import flash_attention
Tensor.manual_seed(42)
B, N, H, H_KV, D = 1, 32, 2, 1, 32
with Context(DEBUG=0):
q = Tensor.randn(B, N, H, D, dtype=dtypes.bfloat16).contiguous()
k = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16).contiguous()
v = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16).contiguous()
Tensor.realize(q, k, v)
do = Tensor.ones(B, N, H, D, dtype=dtypes.float32).contiguous()
Tensor.realize(do)
q_, k_, v_ = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
out = flash_attention(q_, k_, v_)
out = out.float().transpose(1, 2)
out.backward(do)
Tensor.realize(q.grad, k.grad, v.grad)
with Context(DEBUG=0):
q_ref = q.detach().clone()
k_ref = k.detach().clone()
v_ref = v.detach().clone()
Tensor.realize(q_ref, k_ref, v_ref)
q_ref_, k_ref_, v_ref_ = q_ref.transpose(1, 2), k_ref.transpose(1, 2), v_ref.transpose(1, 2)
ref = q_ref_.scaled_dot_product_attention(k_ref_, v_ref_)
ref = ref.float().transpose(1, 2)
ref.backward(do)
Tensor.realize(q_ref.grad, k_ref.grad, v_ref.grad)
assert_allclose(q.grad, q_ref.grad, atol=2e-2, rtol=2e-2)
assert_allclose(v.grad, v_ref.grad, atol=2e-2, rtol=2e-2)
assert_allclose(k.grad, k_ref.grad, atol=5e-2, rtol=2e-2)
def test_fast_fa_bwd_causal(self):
from extra.thunder.tiny.fa import flash_attention
Tensor.manual_seed(42)
B, N, H, H_KV, D = 1, 8192, 32, 32, 128
with Context(DEBUG=0):
q = Tensor.randn(B, N, H, D, dtype=dtypes.bfloat16).contiguous()
k = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16).contiguous()
v = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16).contiguous()
Tensor.realize(q, k, v)
do = Tensor.ones(B, N, H, D, dtype=dtypes.float32).contiguous()
Tensor.realize(do)
q_, k_, v_ = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
out = flash_attention(q_, k_, v_, is_causal=True)
out = out.float().transpose(1, 2)
out.backward(do)
Tensor.realize(q.grad, k.grad, v.grad)
with Context(DEBUG=0):
q_ref = q.detach().clone()
k_ref = k.detach().clone()
v_ref = v.detach().clone()
Tensor.realize(q_ref, k_ref, v_ref)
q_ref_, k_ref_, v_ref_ = q_ref.transpose(1, 2), k_ref.transpose(1, 2), v_ref.transpose(1, 2)
ref = q_ref_.scaled_dot_product_attention(k_ref_, v_ref_, is_causal=True, enable_gqa=True)
ref = ref.float().transpose(1, 2)
ref.backward(do)
Tensor.realize(q_ref.grad, k_ref.grad, v_ref.grad)
assert_allclose(q.grad, q_ref.grad, atol=2e-2, rtol=2e-2)
assert_allclose(v.grad, v_ref.grad, atol=2e-2, rtol=2e-2)
assert_allclose(k.grad, k_ref.grad, atol=6e-2, rtol=2e-2)
def test_fast_fa_bwd_causal_jitted(self):
from extra.thunder.tiny.fa import flash_attention
Tensor.manual_seed(42)
B, N, H, H_KV, D = 1, 8192, 32, 32, 128
with Context(DEBUG=0):
q = Tensor.randn(B, N, H, D, dtype=dtypes.bfloat16).contiguous()
k = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16).contiguous()
v = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16).contiguous()
Tensor.realize(q, k, v)
do = Tensor.ones(B, N, H, D, dtype=dtypes.float32).contiguous()
Tensor.realize(do)
def fn(q, k, v, do):
q_, k_, v_ = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
out = flash_attention(q_, k_, v_, is_causal=True)
out = out.float().transpose(1, 2)
out.backward(do)
Tensor.realize(out, q.grad, k.grad, v.grad)
return q.grad, k.grad, v.grad
fn_jitted = TinyJit(fn)
for _ in range(10):
q = Tensor.randn(B, N, H, D, dtype=dtypes.bfloat16).contiguous()
k = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16).contiguous()
v = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16).contiguous()
Tensor.realize(q, k, v)
do = Tensor.ones(B, N, H, D, dtype=dtypes.float32).contiguous()
Tensor.realize(do)
q.grad, k.grad, v.grad = fn_jitted(q, k, v, do)
with Context(DEBUG=0):
q_ref = q.detach().clone()
k_ref = k.detach().clone()
v_ref = v.detach().clone()
Tensor.realize(q_ref, k_ref, v_ref)
q_ref_, k_ref_, v_ref_ = q_ref.transpose(1, 2), k_ref.transpose(1, 2), v_ref.transpose(1, 2)
ref = flash_attention(q_ref_, k_ref_, v_ref_, is_causal=True)
ref = ref.float().transpose(1, 2)
ref.backward(do)
Tensor.realize(q_ref.grad, k_ref.grad, v_ref.grad)
assert_allclose(q.grad, q_ref.grad, atol=1e-5, rtol=1e-5)
assert_allclose(k.grad, k_ref.grad, atol=1e-5, rtol=1e-5)
assert_allclose(v.grad, v_ref.grad, atol=1e-5, rtol=1e-5)
def test_fast_fa_bwd_multidevice(self):
from extra.thunder.tiny.fa import flash_attention
Tensor.manual_seed(42)
B, N, H, H_KV, D = 2, 1024, 32, 32, 128
GPUS = tuple(f"{Device.DEFAULT}:{i}" for i in range(B))
with Context(DEBUG=0):
base_q = Tensor.randn(B, N, H, D, dtype=dtypes.bfloat16).contiguous()
base_k = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16).contiguous()
base_v = Tensor.randn(B, N, H_KV, D, dtype=dtypes.bfloat16).contiguous()
base_do = Tensor.ones(B, N, H, D, dtype=dtypes.float32).contiguous()
with Context(DEBUG=0):
q = base_q.clone().shard(GPUS, axis=0)
k = base_k.clone().shard(GPUS, axis=0)
v = base_v.clone().shard(GPUS, axis=0)
Tensor.realize(q, k, v)
do = base_do.clone().shard(GPUS, axis=0)
Tensor.realize(do)
q_, k_, v_ = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
out = flash_attention(q_, k_, v_, is_causal=True)
out = out.float().transpose(1, 2)
out.backward(do)
Tensor.realize(q.grad, k.grad, v.grad)
with Context(DEBUG=0):
q_ref = base_q.clone()
k_ref = base_k.clone()
v_ref = base_v.clone()
Tensor.realize(q_ref, k_ref, v_ref)
do_ref = base_do.clone()
Tensor.realize(do_ref)
q_ref_, k_ref_, v_ref_ = q_ref.transpose(1, 2), k_ref.transpose(1, 2), v_ref.transpose(1, 2)
ref = flash_attention(q_ref_, k_ref_, v_ref_, is_causal=True)
ref = ref.float().transpose(1, 2)
ref.backward(do_ref)
Tensor.realize(q_ref.grad, k_ref.grad, v_ref.grad)
assert_allclose(q.grad, q_ref.grad, atol=1e-5, rtol=1e-5)
assert_allclose(v.grad, v_ref.grad, atol=1e-5, rtol=1e-5)
assert_allclose(k.grad, k_ref.grad, atol=1e-5, rtol=1e-5)
if __name__ == "__main__":
unittest.main()