From 44f1f45cd53f850ce6323b78f7e18c2727e9d3f2 Mon Sep 17 00:00:00 2001 From: qazal <77887910+Qazalin@users.noreply.github.com> Date: Mon, 10 Aug 2026 15:43:01 +0800 Subject: [PATCH 1/5] llama: custom silu kernels (#17462) * start by copying the C * uop kernel * cleanup tests * estimates is part of SPEC --- examples/mlperf/models/flat_llama.py | 5 +++ extra/llama_kernels/swiglu/__init__.py | 49 ++++++++++++++++++++++++++ test/backend/test_llama_kernels.py | 27 ++++++++++++++ tinygrad/uop/spec.py | 3 +- 4 files changed, 83 insertions(+), 1 deletion(-) create mode 100644 extra/llama_kernels/swiglu/__init__.py diff --git a/examples/mlperf/models/flat_llama.py b/examples/mlperf/models/flat_llama.py index d237823687..4e3a2244ce 100644 --- a/examples/mlperf/models/flat_llama.py +++ b/examples/mlperf/models/flat_llama.py @@ -114,6 +114,11 @@ def silu_w13_quantize_matmul(x_w13:Tensor, w2:Tensor, s_2:Tensor, amax_x2:Tensor|None, next_amax_x2:Tensor|None, grad_amax_xw13:Tensor|None, next_grad_amax_xw13:Tensor|None, grad_amax_xout:Tensor|None, next_grad_amax_xout:Tensor|None): + if FUSED_SILU_W13 and MXFP4: + from extra.llama_kernels.swiglu import swiglu + out, *ret = matmul(swiglu(x_w13), w2, amax_x=amax_x2, w_inv_scale=s_2, grad_amax_state=grad_amax_xout, + next_grad_amax_state=next_grad_amax_xout, next_amax_x=next_amax_x2) + return out, ret if FUSED_SILU_W13 and not MXFP4: from extra.llama_kernels.cast_amax import fused_quantize_fp8_w13 x2_fp8 = fused_quantize_fp8_w13(x_w13, amax_x2, FP8_DTYPE, grad_amax_state=grad_amax_xw13, diff --git a/extra/llama_kernels/swiglu/__init__.py b/extra/llama_kernels/swiglu/__init__.py new file mode 100644 index 0000000000..810cb6d8ec --- /dev/null +++ b/extra/llama_kernels/swiglu/__init__.py @@ -0,0 +1,49 @@ +import functools, math +from tinygrad import Tensor, dtypes +from tinygrad.uop.ops import UOp, KernelInfo +from tinygrad.renderer import Estimates +from extra.llama_kernels import alloc_like + +LOG2E = 1.4426950408889634 + +@functools.cache +def _custom_swiglu(out:UOp, x_w13:UOp) -> UOp: + rows, hidden = math.prod(x_w13.shape[:-1]), x_w13.shape[-1]//2 + n_elems = rows * hidden + out, x_w13 = out.reshape(n_elems), x_w13.reshape(rows, 2*hidden) + i = UOp.range(n_elems, 0) + row, col = i // hidden, i % hidden + act, gate = x_w13[row, col].cast(dtypes.float), x_w13[row, hidden+col].cast(dtypes.float) + sigmoid = (1.0 + (-LOG2E * act).exp2()).reciprocal() + store = out[i].store((act * sigmoid * gate).cast(out.dtype)) + return store.end(i).sink(arg=KernelInfo(f"swiglu_fwd_{n_elems}", estimates=Estimates(ops=5*n_elems, mem=6*n_elems))) + +@functools.cache +def _custom_swiglu_bwd(grad_out:UOp, x_w13:UOp, grad_act:UOp) -> UOp: + rows, hidden = math.prod(x_w13.shape[:-1]), x_w13.shape[-1]//2 + n_elems = rows * hidden + grad_out, x_w13, grad_act = grad_out.reshape(rows, 2*hidden), x_w13.reshape(rows, 2*hidden), grad_act.reshape(n_elems) + i = UOp.range(n_elems, 0) + row, col = i // hidden, i % hidden + act, gate = x_w13[row, col].cast(dtypes.float), x_w13[row, hidden+col].cast(dtypes.float) + grad = grad_act[i].cast(dtypes.float) + sigmoid = (1.0 + (-LOG2E * act).exp2()).reciprocal() + silu = act * sigmoid + dact = grad_out[row, col].store((grad * (sigmoid + silu * (1.0 - sigmoid)) * gate).cast(grad_out.dtype)) + dgate = grad_out.after(dact)[row, hidden+col].store((grad * silu).cast(grad_out.dtype)) + return dgate.end(i).sink(arg=KernelInfo(f"swiglu_bwd_{n_elems}", estimates=Estimates(ops=10*n_elems, mem=10*n_elems))) + +def _swiglu_bwd(gradient:UOp, kernel:UOp): + _, x_w13 = kernel.src[1:] + axis = x_w13.axis if isinstance(x_w13.device, tuple) else None + grad_out = alloc_like(x_w13.shape, dtypes.bfloat16, x_w13.device, axis) + grad_out, *_ = Tensor.custom_kernel(grad_out, Tensor(x_w13, device=x_w13.device), Tensor(gradient, device=x_w13.device), + fxn=_custom_swiglu_bwd) + return (None, grad_out.uop) + +def swiglu(x_w13:Tensor) -> Tensor: + assert x_w13.dtype == dtypes.bfloat16 and x_w13.ndim >= 2 and x_w13.shape[-1] % 32 == 0 + *prefix, two_k = x_w13.shape + axis = x_w13.uop.axis if isinstance(x_w13.device, tuple) else None + out = alloc_like((*prefix, two_k//2), dtypes.bfloat16, x_w13.device, axis) + return Tensor.custom_kernel(out, x_w13, fxn=_custom_swiglu, grad_fxn=_swiglu_bwd)[0] diff --git a/test/backend/test_llama_kernels.py b/test/backend/test_llama_kernels.py index 4542d66a6f..f03a57b393 100644 --- a/test/backend/test_llama_kernels.py +++ b/test/backend/test_llama_kernels.py @@ -5,6 +5,7 @@ from examples.mlperf.models.flat_llama import FP8_DTYPE, quantize_fp8 from extra.llama_kernels.fused_ce import fused_ce_loss from extra.llama_kernels import local_abs_max from extra.llama_kernels.quantize_fp8_delayed import quantize_fp8_delayed, quantize_fp8_scalar +from extra.llama_kernels.swiglu import swiglu from extra.models.llama import apply_rotary_emb, precompute_freqs_cis from extra.thunder.amd.fa import custom_fused_qkv_rope_backward, fused_qkv_rope from test.helpers import needs_second_gpu, assert_kernel_count @@ -161,5 +162,31 @@ class TestFusedQKVRoPE(unittest.TestCase): ref = Tensor.cat(dq_ref, dk_ref, dv_ref, dim=3).reshape(*dx.shape).realize() with Context(DEBUG=0): self.assertTrue(dx.allclose(ref, atol=2e-2, rtol=2e-2).item(), "backward mismatch") +def run_swiglu(test:unittest.TestCase, shape:tuple[int, ...]) -> None: + Tensor.manual_seed(0) + x = (Tensor.randn(*shape) * 2).cast(dtypes.bfloat16).realize() + hidden = x.shape[-1] // 2 + out, ref = swiglu(x), x[..., :hidden].silu() * x[..., hidden:] + Tensor.realize(out, ref) + with Context(DEBUG=0): test.assertTrue(out.allclose(ref, atol=2.5e-1, rtol=3e-2).item(), "SwiGLU forward mismatch") + + grad = (Tensor.randn(*out.shape) * 2).cast(dtypes.bfloat16).realize() + grad_x, grad_ref = out.gradient(x, gradient=grad)[0], ref.gradient(x, gradient=grad)[0] + Tensor.realize(grad_x, grad_ref) + test.assertEqual(grad_x.shape, shape) + test.assertEqual(grad_x.dtype, dtypes.bfloat16) + with Context(DEBUG=0): test.assertTrue(grad_x.allclose(grad_ref, atol=2.5e-1, rtol=3e-2).item(), "SwiGLU backward mismatch") + +class TestSwiGLU(unittest.TestCase): + def setUp(self): + if dtypes.bfloat16 not in Device[Device.DEFAULT].renderer.supported_dtypes(): self.skipTest("need bfloat16") + + def test_simple(self): run_swiglu(self, (2, 32, 64)) + + def test_llama_shape(self): + if Device.DEFAULT != "AMD" or not Device[Device.DEFAULT].renderer.target.arch.startswith("gfx950"): + self.skipTest("only run on real machine for speed") + run_swiglu(self, (2, 8192, 28672)) + if __name__ == '__main__': unittest.main() diff --git a/tinygrad/uop/spec.py b/tinygrad/uop/spec.py index ed28da8553..867ea5f44c 100644 --- a/tinygrad/uop/spec.py +++ b/tinygrad/uop/spec.py @@ -286,10 +286,11 @@ spec_kernel_graph = PatternMatcher([ # late imports to avoid circular import from tinygrad.codegen.opt import Opt, OptOps from tinygrad.schedule.rangeify import BufferizeOpts +from tinygrad.renderer import Estimates glbls:dict[str, Any] = {"inf": math.inf, "nan": math.nan, "KernelInfo": KernelInfo, "Metadata": Metadata, "UOp": UOp, "dtypes": dtypes, "Ops": Ops, "AxisType": AxisType, "Invalid": Invalid, "Opt": Opt, "OptOps": OptOps, "BufferizeOpts": BufferizeOpts, "AddrSpace": AddrSpace, "panic": panic, - "ConstFloat": ConstFloat, "ParamArg": ParamArg} + "ConstFloat": ConstFloat, "ParamArg": ParamArg, "Estimates": Estimates} def eval_pyrender(code:str) -> UOp: lcls:dict[str, Any] = {} exec(code, glbls, lcls) From 2821bd646faada79508313e835456500a37dbd3d Mon Sep 17 00:00:00 2001 From: qazal <77887910+Qazalin@users.noreply.github.com> Date: Mon, 10 Aug 2026 16:31:50 +0800 Subject: [PATCH 2/5] late loss.to("CPU") in llama (#17476) * late loss.to("CPU") in llama * acc = 0 --- examples/mlperf/model_train.py | 21 +++++++++++---------- 1 file changed, 11 insertions(+), 10 deletions(-) diff --git a/examples/mlperf/model_train.py b/examples/mlperf/model_train.py index e9d8086aab..47c4977e39 100644 --- a/examples/mlperf/model_train.py +++ b/examples/mlperf/model_train.py @@ -1458,7 +1458,8 @@ def train_llama3(): # realize everything here if optim.master_params: Tensor.realize(*optim.master_params) - Tensor.realize(*optim.params, *fp8_inv_scales, *fp8_amax, *fp8_next_amax, *fp8_grad_amax, *fp8_next_grad_amax) + loss_acc = Tensor.zeros(1, dtype=dtypes.float32, device=device) + Tensor.realize(loss_acc, *optim.params, *fp8_inv_scales, *fp8_amax, *fp8_next_amax, *fp8_grad_amax, *fp8_next_grad_amax) @TinyJit def minibatch(tokens:Tensor): @@ -1476,8 +1477,8 @@ def train_llama3(): for g, new_g in zip(grads, loss.gradient(*optim.params)): apply_grad(g, new_g.uop) - loss_cpu = loss.flatten().float().to("CPU") - return loss_cpu.realize(*grads, *fp8_amax, *fp8_next_amax, *fp8_grad_amax, *fp8_next_grad_amax) + loss_acc.assign(loss_acc + loss.flatten().float()) + return loss_acc.realize(*grads, *fp8_amax, *fp8_next_amax, *fp8_grad_amax, *fp8_next_grad_amax) @TinyJit def optim_step(): @@ -1490,9 +1491,10 @@ def train_llama3(): lr_cpu = optim.lr.float().to("CPU") grad_norm_cpu = grad_norm.float().to("CPU") - Tensor.realize(lr_cpu, grad_norm_cpu, *grads, *fp8_inv_scales, *fp8_amax, *fp8_grad_amax) + loss_cpu = loss_acc.to("CPU") + Tensor.realize(lr_cpu, grad_norm_cpu, loss_cpu, loss_acc.assign(0), *grads, *fp8_inv_scales, *fp8_amax, *fp8_grad_amax) - return lr_cpu, grad_norm_cpu + return lr_cpu, grad_norm_cpu, loss_cpu @TinyJit @Context(TRAINING=0) @@ -1547,8 +1549,8 @@ def train_llama3(): st = time.perf_counter() stopped = False - losses, data_time, dev_time = [], 0, 0 - for _ in range(grad_acc if i >= 2 else 1): + data_time, dev_time = 0, 0 + for _ in range(accum_steps:=grad_acc if i >= 2 else 1): ist = time.perf_counter() try: tokens = next(train_iter) except StopIteration: @@ -1556,16 +1558,15 @@ def train_llama3(): break mst = time.perf_counter() data_time += mst - ist - losses.append(minibatch(tokens).item()) + minibatch(tokens) dev_time += time.perf_counter() - mst if stopped: break gt = time.perf_counter() ret = optim_step() - lr, grad_norm = ret[0].item(), ret[1].item() + lr, grad_norm, loss = ret[0].item(), ret[1].item(), ret[2].item() / accum_steps et = time.perf_counter() - loss = sum(losses) / len(losses) optim_time = et - gt dev_time += optim_time step_time = et - st From 8611fe22a7fcc7d1928bbde19ded66277cb12f3e Mon Sep 17 00:00:00 2001 From: nimlgen <138685161+nimlgen@users.noreply.github.com> Date: Mon, 10 Aug 2026 13:33:49 +0300 Subject: [PATCH 3/5] fix hevc (#17477) * hevc tests * x --- .github/workflows/platform.yml | 3 ++- test/testextra/test_hevc.py | 25 ++++++++++++++++++++++--- tinygrad/uop/spec.py | 2 +- 3 files changed, 25 insertions(+), 5 deletions(-) diff --git a/.github/workflows/platform.yml b/.github/workflows/platform.yml index 98f15727a3..bb27f446d1 100644 --- a/.github/workflows/platform.yml +++ b/.github/workflows/platform.yml @@ -78,7 +78,8 @@ jobs: # TODO: failing due to library loading error CAPTURE_PROCESS_REPLAY: 0 run: | - python3 -m pytest -n=auto test/device/test_hcq.py test/test_tiny.py --durations=20 + python3 -m pytest -n=auto test/device/test_hcq.py test/test_tiny.py \ + test/testextra/test_hevc.py::TestHevc::test_hevc_decode_compile --durations=20 - name: Run process replay tests uses: ./.github/actions/process-replay diff --git a/test/testextra/test_hevc.py b/test/testextra/test_hevc.py index 058b237f42..174813b8a9 100644 --- a/test/testextra/test_hevc.py +++ b/test/testextra/test_hevc.py @@ -1,7 +1,9 @@ import unittest -from tinygrad import Tensor, Device, dtypes -from tinygrad.helpers import fetch, round_up +from tinygrad import Tensor, Device, Variable, dtypes +from tinygrad.helpers import DEV, fetch, round_up +from tinygrad.engine.realize import compile_linear +from tinygrad.uop.ops import Ops from extra.hevc.hevc import parse_hevc_file_headers, nv_gpu from extra.hevc.decode import hevc_decode @@ -63,7 +65,7 @@ class TestHevc(unittest.TestCase): 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") + @unittest.skipUnless(Device.DEFAULT == "NV" and not DEV.interface.startswith("MOCK"), "real 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() @@ -83,5 +85,22 @@ class TestHevc(unittest.TestCase): self.assertEqual(f.dtype, dtypes.uint8) self.assertEqual(f.device, "NV") + @unittest.skipUnless(Device.DEFAULT == "NV", "NV only") + def test_hevc_decode_compile(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, _, _, luma_w, luma_h, _ = parse_hevc_file_headers(dat) + offset, sz, frame_pos, max_hist, _ = frame_info[1] + out_image_size = luma_h + (luma_h + 1) // 2, round_up(luma_w, 64) + history = [Tensor.empty(*out_image_size, dtype=dtypes.uint8, device="NV") for _ in range(max_hist)] + decoded = Tensor(dat, device="NV")[offset:offset+sz].decode_hevc_frame( + Variable("pos", 0, max_hist + 1).bind(frame_pos), out_image_size, opaque[1], history) + + compiled = compile_linear(decoded.linear_with_vars()[0]) + self.assertTrue(any(call.src[0].op is Ops.PROGRAM for call in compiled.src)) + encdec_calls = [call for call in compiled.src if call.src[0].op is Ops.CUSTOM_FUNCTION and call.src[0].arg == "encdec"] + self.assertEqual(len(encdec_calls), 1) + if __name__ == "__main__": unittest.main() diff --git a/tinygrad/uop/spec.py b/tinygrad/uop/spec.py index 867ea5f44c..654f491cd7 100644 --- a/tinygrad/uop/spec.py +++ b/tinygrad/uop/spec.py @@ -275,7 +275,7 @@ spec_kernel_graph = PatternMatcher([ (UPat(Ops.MSTACK, name="x"), lambda x: all(isinstance(s.device, str) for s in x.src) or (all_same(x.src) and x.src[0].device is None)), (UPat(Ops.MSELECT, name="x"), lambda x: isinstance(x.src[0].device, tuple) and x.arg < len(x.src[0].device)), # all calls are on various sinks - (UPat(Ops.CALL, src=(UPat((Ops.SINK, Ops.LINEAR, Ops.PROGRAM)),), allow_any_len=True), lambda: True), + (UPat(Ops.CALL, src=(UPat((Ops.SINK, Ops.LINEAR, Ops.PROGRAM, Ops.CUSTOM_FUNCTION)),), allow_any_len=True), lambda: True), # after on PARAM or AFTER (UPat(Ops.AFTER, src=(UPat(GroupOp.Movement.union({Ops.PARAM, Ops.AFTER, Ops.BUFFER, Ops.MSTACK, Ops.MSELECT, Ops.BITCAST, Ops.RESHAPE})),), allow_any_len=True, name="x"), lambda x: matches_dtype(x.src[0], x.dtype)), From e29606f07e6cb295dea40a2bc18f8bd0b46ec9d3 Mon Sep 17 00:00:00 2001 From: nimlgen <138685161+nimlgen@users.noreply.github.com> Date: Mon, 10 Aug 2026 17:28:46 +0300 Subject: [PATCH 4/5] hcq2: copy kernel (#17480) * hcq2: copy with kernel * test * x --- .github/workflows/test.yml | 2 ++ extra/hcq2/ops_amd2.py | 6 +++--- test/device/test_hcq2.py | 14 ++++++++++++++ tinygrad/device.py | 2 ++ tinygrad/runtime/support/hcq2.py | 16 ++++++++++++++-- 5 files changed, 35 insertions(+), 5 deletions(-) create mode 100644 test/device/test_hcq2.py diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 45dfa56491..7bd3063e41 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -527,6 +527,8 @@ jobs: TestMultiTensor.test_backward_sum TestMultiTensor.test_matmul_shard_0_0 - name: Run HCQ2 JIT tests run: HCQ_RUNTIME_DEV=PYTHON HCQ2=1 DEV=MOCKKFD+AMD FORWARD_ONLY=1 PYTHONPATH=. python test/unit/test_jit.py + - name: Run HCQ2 unit tests + run: HCQ_RUNTIME_DEV=PYTHON HCQ2=1 DEV=MOCKKFD+AMD FORWARD_ONLY=1 PYTHONPATH=. python -m pytest test/device/test_hcq2.py testmockam: name: Linux (am) diff --git a/extra/hcq2/ops_amd2.py b/extra/hcq2/ops_amd2.py index 0c8e99d693..fbdef9f444 100644 --- a/extra/hcq2/ops_amd2.py +++ b/extra/hcq2/ops_amd2.py @@ -288,10 +288,10 @@ def amd_build_program(prg:UOp) -> UOp: class AMDAllocator(HCQAllocator['AMDDevice']): def __init__(self, dev:AMDDevice): - super().__init__(dev, supports_copy_from_disk=dev.has_sdma_queue, supports_transfer=dev.has_sdma_queue and not dev.is_usb()) + super().__init__(dev, supports_copy_from_disk=dev.has_copy_queue, supports_transfer=dev.has_copy_queue and not dev.is_usb()) def _alloc(self, size:int, options:BufferSpec) -> HCQ2Buffer: - return self.dev.iface.alloc(size, host=options.host, uncached=options.uncached, cpu_access=options.cpu_access or not self.dev.has_sdma_queue) + return self.dev.iface.alloc(size, host=options.host, uncached=options.uncached, cpu_access=options.cpu_access or not self.dev.has_copy_queue) def _do_free(self, opaque, options:BufferSpec): self.dev.iface.free(opaque) @@ -581,7 +581,7 @@ class AMDDevice(HCQ2Compiled): self.max_copy_size = 0x40000000 if self.iface.ip_versions[am.SDMA0_HWIP][0] >= 5 else 0x400000 self.sdma_queues:dict = {} - self.has_sdma_queue = True # self.sdma_queue(0) is not None, TODO: think of this + self.has_copy_queue = not getenv("AMD_DISABLE_SDMA") super().__init__(device, AMDAllocator(self), [HIPRenderer, AMDLLVMRenderer, HIPCCRenderer], None, can_recover=self.is_am(), arch=self.arch) diff --git a/test/device/test_hcq2.py b/test/device/test_hcq2.py new file mode 100644 index 0000000000..19d72d7bb8 --- /dev/null +++ b/test/device/test_hcq2.py @@ -0,0 +1,14 @@ +import unittest, numpy as np +from unittest.mock import patch +from tinygrad import Device, Tensor +from tinygrad.helpers import getenv +from tinygrad.runtime.support.hcq2 import HCQ_DEVS, all_devices_in + +@unittest.skipUnless(getenv("HCQ2") and all_devices_in(Device.DEFAULT, HCQ_DEVS), "hcq2 device required") +class TestHCQ2(unittest.TestCase): + def test_copy_without_copy_queue(self): + with patch.object(Device[Device.DEFAULT], "has_copy_queue", False): + np.testing.assert_equal(Tensor(np.arange(61, dtype=np.float32)).to(Device.DEFAULT).contiguous().realize().numpy(), np.arange(61)) + +if __name__ == "__main__": + unittest.main() diff --git a/tinygrad/device.py b/tinygrad/device.py index 570cd79813..6f796e07a6 100644 --- a/tinygrad/device.py +++ b/tinygrad/device.py @@ -336,6 +336,8 @@ class Compiled: pm_lower:Any = None pm_bufferize:Any = None + has_copy_queue:bool = True + def __init__(self, device:str, allocator:Allocator, renderers:list[type[Renderer]], runtime:type[Program[Self]]|None, graph=None, arch=None): from tinygrad.renderer import Renderer self.device, self.allocator, self.runtime_t, self.graph, self.renderers = device, allocator, runtime, graph, renderers or [Renderer] diff --git a/tinygrad/runtime/support/hcq2.py b/tinygrad/runtime/support/hcq2.py index 9e2429e304..9b0372d035 100644 --- a/tinygrad/runtime/support/hcq2.py +++ b/tinygrad/runtime/support/hcq2.py @@ -96,12 +96,24 @@ pm_replace_buffers = PatternMatcher([(UPat(Ops.CALL, name="call"), replace_call_ def _need_staging(a, b): return all_devices_in(a.device, HCQ_DEVS) and not all_devices_in(b.device, HCQ_P2P_DEVS) +def hcq_call_devs(call:UOp) -> Any|None: return next((b.device for b in call.src[1:] if all_devices_in(b.device, HCQ_DEVS)), None) + def stage_copy(dst:UOp, src:UOp) -> UOp|None: if not (_need_staging(src, dst) or _need_staging(dst, src)): return None stage = UOp.new_buffer("CPU", src.max_numel() * src.dtype.itemsize, dtypes.uint8) return UOp(Ops.LINEAR, src=(src.copy_to_device("CPU").call(stage, src), stage.copy_to_device(dst.device).call(dst, stage))) -pm_insert_copy_staging = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.COPY), UPat(name="dst"), UPat(name="src"))), stage_copy)]) + +def kernel_copy(call:UOp, dst:UOp, src:UOp) -> UOp|None: + if (devs:=hcq_call_devs(call)) is None or Device[(dev:=to_tuple(devs)[0])].has_copy_queue: return None + d, s = (UOp.param(i, dst.dtype, (n:=dst.max_numel(),), device=devs) for i in range(2)) + ast = d.index(r:=UOp.range(n, 0)).store(s.index(r).load()).end(r).sink(arg=KernelInfo(name="copy"), tag=1) + return call.replace(src=(to_program(ast, Device[dev].renderer), dst, src)) + +pm_insert_copy_staging = PatternMatcher([ + (UPat(Ops.CALL, src=(UPat(Ops.COPY), UPat(name="dst"), UPat(name="src"))), stage_copy), + (UPat(Ops.CALL, src=(UPat(Ops.COPY), UPat(name="dst"), UPat(name="src")), name="call"), kernel_copy) +]) # ***************** # 2. deps @@ -217,7 +229,7 @@ def sched_hcq_batches(l:UOp, profile:bool) -> UOp: srcs:list[UOp] = [] batch:list[tuple[UOp, tuple[str, ...]]] = [] for call in l.src: - if (devs:=next((b.device for b in call.src[1:] if all_devices_in(b.device, HCQ_DEVS)), None)) is not None: batch.append((call, to_tuple(devs))) + if (devs:=hcq_call_devs(call)) is not None: batch.append((call, to_tuple(devs))) else: srcs, batch = srcs + _finalize_batch(batch, profile) + [call], [] return l.replace(src=tuple(srcs + _finalize_batch(batch, profile))) From d41ca5e60fc55550b38c0aa0d6c01d9c94397895 Mon Sep 17 00:00:00 2001 From: Raine Date: Mon, 10 Aug 2026 12:39:14 -0300 Subject: [PATCH 5/5] Fix WMMA CI (#17479) * init * split into sub tests * trigger ci --- test/backend/test_linearizer.py | 6 ++---- test/opt/test_tensor_cores.py | 15 ++++++++------- 2 files changed, 10 insertions(+), 11 deletions(-) diff --git a/test/backend/test_linearizer.py b/test/backend/test_linearizer.py index 958a9be840..5a43d68011 100644 --- a/test/backend/test_linearizer.py +++ b/test/backend/test_linearizer.py @@ -437,7 +437,7 @@ def reset_bufs(bufs:list[Buffer]): for buf in bufs: buf.copy_from(Buffer("PYTHON", buf.size, buf.dtype, opaque=memoryview(bytearray(buf.nbytes)))) def _helper_linearizer_opt_ast(realized_ast:UOp, real_bufs:list[Buffer], opts=[], - apply_tc=False, atol=1e-4, rtol=1e-4, color_sizes=[], wanna_output=[]): + apply_tc=False, atol=1e-4, rtol=1e-4, color_sizes=[], wanna_output=[], check_default_opt=True): outbufs = real_bufs[:len(realized_ast.src)] wanna_output = [np.array(x).flatten() for x in wanna_output] buf_uops = [UOp.new_buffer(b.device, b.size, b.dtype) for b in real_bufs] @@ -459,9 +459,7 @@ def _helper_linearizer_opt_ast(realized_ast:UOp, real_bufs:list[Buffer], opts=[] for buf,want in zip(copyout_outputs(outbufs), wanna_output): np.testing.assert_allclose(buf, want, atol=atol, rtol=rtol) # Check correctness of handcoded optimiztions. - reset_bufs(outbufs) - run_prg(opts=None) - for buf,want in zip(copyout_outputs(outbufs), wanna_output): np.testing.assert_allclose(buf, want, atol=atol, rtol=rtol) + if check_default_opt: check_opt(None) for x in opts: # Check custom transformations if any. check_opt(([Opt(OptOps.TC, 0, (TC_SELECT.value, TC_OPT.value, 1))] if apply_tc else [])+x) diff --git a/test/opt/test_tensor_cores.py b/test/opt/test_tensor_cores.py index c25eec469a..0d9b5db3d1 100644 --- a/test/opt/test_tensor_cores.py +++ b/test/opt/test_tensor_cores.py @@ -79,7 +79,8 @@ class TestTensorCores(unittest.TestCase): @unittest.skipUnless(Device[Device.DEFAULT].renderer.tensor_cores, "test requires tensor cores") def test_tensor_cores(self): for tc in Device[Device.DEFAULT].renderer.tensor_cores: - helper_tc_allclose(tc.dims[0], tc.dims[1], tc.dims[2], tc.dtype_in, tc.dtype_out, axis=0, tc_opt=0) + with self.subTest(tc=tc): + helper_tc_allclose(tc.dims[0], tc.dims[1], tc.dims[2], tc.dtype_in, tc.dtype_out, axis=0, tc_opt=0) @unittest.skipUnless(Device[Device.DEFAULT].renderer.tensor_cores, "test requires tensor cores") def test_tensor_cores_nested_reduce(self): @@ -185,10 +186,10 @@ class TestTensorCores(unittest.TestCase): # skip fp8 tcs: the unoptimized ALU baseline quantizes products to fp8 (JAX promotion), which legitimately # differs from the MFMA path (f32 accumulation), so the baseline-vs-TC numerical gate can't hold for fp8. tc = next(tc for tc in Device[Device.DEFAULT].renderer.tensor_cores if tc.dtype_in not in dtypes.fp8s) - x, y = Tensor.rand(64, 64, dtype=tc.dtype_in), Tensor.rand(64, 64, dtype=tc.dtype_in) + x, y = Tensor.rand(16, 64, dtype=tc.dtype_in), Tensor.rand(64, 16, dtype=tc.dtype_in) r = x.matmul(y, dtype=tc.dtype_out) opts = [Opt(OptOps.UNROLL, 0, 2)] - ast = helper_linearizer_opt(r, [opts], apply_tc=True, atol=3e-2, rtol=1e-3) + ast = helper_linearizer_opt(r, [opts], apply_tc=True, atol=3e-2, rtol=1e-3, check_default_opt=False) for u in tuple(to_program(replace_opts(ast, opts), Device[Device.DEFAULT].renderer).src[1].src): if u.op is Ops.WMMA: assert u.src[-1].src[0].op != Ops.STORE @@ -199,10 +200,10 @@ class TestTensorCores(unittest.TestCase): @unittest.skipIf(Device.DEFAULT in {"CPU"}, "CPU does not support using a different type for accumulation") def test_tensor_cores_unroll_casted_phi(self): tc = [tc for tc in Device[Device.DEFAULT].renderer.tensor_cores if tc.dtype_in != tc.dtype_out and tc.dtype_in not in dtypes.fp8s][0] - x, y = Tensor.rand(64, 64, dtype=tc.dtype_in), Tensor.rand(64, 64, dtype=tc.dtype_in) + x, y = Tensor.rand(16, 64, dtype=tc.dtype_in), Tensor.rand(64, 16, dtype=tc.dtype_in) r = x.matmul(y, dtype=tc.dtype_out) opts = [Opt(OptOps.UNROLL, 0, 2)] - ast = helper_linearizer_opt(r, [opts], apply_tc=True, atol=3e-2, rtol=1e-3) + ast = helper_linearizer_opt(r, [opts], apply_tc=True, atol=3e-2, rtol=1e-3, check_default_opt=False) for u in tuple(to_program(replace_opts(ast, opts), Device[Device.DEFAULT].renderer).src[1].src): if u.op is Ops.WMMA: #assert u.src[-1].dtype == dtypes.float.vec(prod(tc.thread_local_sizes[2])) @@ -215,10 +216,10 @@ class TestTensorCores(unittest.TestCase): def test_tensor_cores_unroll_casted_phi_with_children(self): # all STORE children are outside the loop tc = [tc for tc in Device[Device.DEFAULT].renderer.tensor_cores if tc.dtype_in != tc.dtype_out and tc.dtype_in not in dtypes.fp8s][0] - x, y = Tensor.rand(64, 64, dtype=tc.dtype_in), Tensor.rand(64, 64, dtype=tc.dtype_in) + x, y = Tensor.rand(16, 64, dtype=tc.dtype_in), Tensor.rand(64, 16, dtype=tc.dtype_in) r = x.matmul(y, dtype=tc.dtype_out).relu() opts = [Opt(OptOps.UNROLL, 0, 2)] - ast = helper_linearizer_opt(r, [opts], apply_tc=True, atol=3e-2, rtol=1e-3) + ast = helper_linearizer_opt(r, [opts], apply_tc=True, atol=3e-2, rtol=1e-3, check_default_opt=False) for u in tuple(to_program(replace_opts(ast, opts), Device[Device.DEFAULT].renderer).src[1].src): if u.op is Ops.WMMA: #assert u.src[-1].dtype == dtypes.float.vec(prod(tc.thread_local_sizes[2]))