From ee7d44160614dffe7816ebbecdefa2d91b6b56d7 Mon Sep 17 00:00:00 2001 From: George Hotz Date: Tue, 4 Aug 2026 23:23:10 -0700 Subject: [PATCH] fixes --- test/unit/test_assign.py | 34 ++++++++++++++++++++++++++++++++++ test/unit/test_attention.py | 22 ++++++++++++++++++++++ tinygrad/schedule/__init__.py | 17 ++++++++++++++++- 3 files changed, 72 insertions(+), 1 deletion(-) diff --git a/test/unit/test_assign.py b/test/unit/test_assign.py index ee93250480..236d8b38a2 100644 --- a/test/unit/test_assign.py +++ b/test/unit/test_assign.py @@ -697,6 +697,40 @@ class TestAssignOrdering(unittest.TestCase): - Race conditions (concurrent access to same buffer) """ + def test_packed_state_write_not_reordered_before_readers(self): + """A store to buffer B packed in another buffer's AFTER (rec.after(B.store(v))) must not be + reordered before producer kernels that read B. The store is tracked under the AFTER's base buffer + (rec), so without resolving the actual store targets the scheduler generates no WAR dependency for + B's readers and the write-back can run first, corrupting the producers' input.""" + from tinygrad.llm.model import _gated_delta_prefill_kernel + def build(pre_realize:bool) -> np.ndarray: + def Tl(a, b, shape): return Tensor.linspace(a, b, int(np.prod(shape)), dtype=dtypes.float32).reshape(*shape) + x = Tensor.linspace(-1.0, 1.0, 12, dtype=dtypes.float32).reshape(1, 3, 4) + xh = (x / x.square().mean(-1, keepdim=True).sqrt()).half() + qkv = xh @ Tl(-0.15, 0.2, (6, 4)).T + conv_state = Tensor.zeros(1, 1, 6).clone() # read by the conv below + rec_state = Tensor.zeros(1, 1, 2, 2).clone() + window = conv_state.cat(qkv, dim=1) + T = 3 + w_conv = Tl(-0.05, 0.05, (6, 2)) + conv_out = (window[:, 0:T] * w_conv[:, 0] + window[:, 1:T+1] * w_conv[:, 1]).silu() + q, k, v = conv_out.split([2, 2, 2], dim=-1) + q = q.reshape(1, T, 1, 2).normalize(dim=-1) + k = k.reshape(1, T, 1, 2).normalize(dim=-1) + v = v.reshape(1, T, 1, 2) + beta, alpha = Tensor.zeros(1, T, 1) + 0.5, Tensor.zeros(1, T, 1) + 0.9 + q, k, v, beta = [z.transpose(1, 2).float() for z in (q, k, v, beta)] + alpha = alpha.transpose(1, 2).float().exp() + qs, kq = q * 2**-0.5, ((q * 2**-0.5)*k).sum(-1).contiguous() + # conv_state write-back packed into rec_state's AFTER, custom kernel consumes it + new_conv_state = window[:, T:T+1].contiguous() + state = Tensor(rec_state.uop.after(conv_state.uop.store(new_conv_state.uop))) + args = [Tensor.empty_like(v), qs, k, v, beta, alpha, state, kq] + if pre_realize: args = [a.realize() if i != 6 else a for i, a in enumerate(args)] + return Tensor.custom_kernel(*args, fxn=_gated_delta_prefill_kernel)[0].transpose(1, 2).realize().numpy() + # lazy execution (build the whole graph, then realize) must match eager (inputs realized up front) + np.testing.assert_allclose(build(False), build(True), rtol=1e-4, atol=1e-4) + def test_overlapping_slice_assigns(self): """Overlapping slice assigns - later write should win for overlapping elements.""" buf = Tensor.zeros(8).contiguous().realize() diff --git a/test/unit/test_attention.py b/test/unit/test_attention.py index 622712cc6e..1d28d32b84 100644 --- a/test/unit/test_attention.py +++ b/test/unit/test_attention.py @@ -176,6 +176,28 @@ class TestGatedDeltaNetBlock(unittest.TestCase): np.testing.assert_allclose(recurrent_state, expected_recurrent[step], rtol=1e-3, atol=1e-3, err_msg=f"GatedDeltaNet reset recurrent cache mismatch at step {step}") + def test_gatedeltanet_prefill_matches_decode(self): + # chunked prefill (T>1, custom kernel) must produce the same output and state as sequential decode (T=1). + # uses lazy linspace weights, which exercise the scheduler's WAR tracking for the packed conv/recurrent + # state write-backs (a write to one buffer packed in another buffer's AFTER must not be reordered before + # the producers that read the target buffer) + config = self._make_config(max_context=3) + block = self._make_block(config) + x = Tensor.linspace(-1.0, 1.0, 3 * config.dim, dtype=dtypes.float32).reshape(1, 3, config.dim) + + x_norm = block.attn_norm(x) + block._init_state(x_norm) + prefill = block._attention(x_norm, 0).realize().numpy() + prefill_conv, prefill_recurrent = self._cache_views(block) + + block = self._make_block(config) + decode = np.concatenate([self._run_attention(block, x[:, t:t+1], t) for t in range(x.shape[1])], axis=1) + decode_conv, decode_recurrent = self._cache_views(block) + + np.testing.assert_allclose(prefill, decode, rtol=1e-3, atol=1e-3, err_msg="prefill output mismatch") + np.testing.assert_allclose(prefill_conv, decode_conv, rtol=1e-3, atol=1e-3, err_msg="prefill conv cache mismatch") + np.testing.assert_allclose(prefill_recurrent, decode_recurrent, rtol=1e-3, atol=1e-3, err_msg="prefill recurrent cache mismatch") + def test_kda_channel_decay(self): config = self._make_config(n_heads=2, ssm=SSMConfig(conv_kernel=2, state_size=2, group_count=2, time_step_rank=2, inner_size=4, kda=True)) block, x = GatedDeltaNetBlock(config, config.ssm), Tensor([[[1., 2., 0., 0.]]]) diff --git a/tinygrad/schedule/__init__.py b/tinygrad/schedule/__init__.py index f2b9cb0137..0bf7f4e8f8 100644 --- a/tinygrad/schedule/__init__.py +++ b/tinygrad/schedule/__init__.py @@ -26,6 +26,15 @@ def _split_after(after: UOp) -> tuple[tuple[UOp, ...], tuple[UOp, ...]]: raise AssertionError(f"AFTER source should be CALL, END, STORE, or AFTER, not {invalid[0].op}") return tuple(kernels), tuple(deps) +def _kernel_write_targets(k:UOp) -> list[tuple[UOp, UOp]]: + # the buffers a kernel stores to, each with the buffer state that the write supersedes (resolved from the call args) + call = k.src[0] if k.op is Ops.END else k + out: list[tuple[UOp, UOp]] = [] + for s in call.src[0].toposort(): + if s.op is Ops.STORE and s.src[0].buf_uop.op is Ops.PARAM and s.src[0].buf_uop.arg.slot >= 0: + out.extend((st.buf_uop, st) for st in _states(call.src[s.src[0].buf_uop.arg.slot+1])) + return out + def create_schedule(sched_sink:UOp) -> UOp: with cpu_profile(TracingKey("toposort sched_sink")): # build kernel dependency graph: edges from producer kernel to consumer kernels @@ -38,7 +47,13 @@ def create_schedule(sched_sink:UOp) -> UOp: kernels, after_deps = _split_after(u) prev_state = _unwrap_src(u.src[0]) prev_kernels = set(_split_after(prev_state)[0]) if prev_state.op is Ops.AFTER else set() - writes.setdefault(u.buf_uop, []).append((u, prev_state, tuple(k for k in kernels if k not in prev_kernels))) + new_kernels = tuple(k for k in kernels if k not in prev_kernels) + writes.setdefault(u.buf_uop, []).append((u, prev_state, new_kernels)) + # a kernel may store to buffers other than the AFTER's base buffer (e.g. state writes packed into another + # buffer's AFTER); register those writes under the buffer they actually target so readers get WAR deps + for k in new_kernels: + for buf, pstate in _kernel_write_targets(k): + if buf is not u.buf_uop: writes.setdefault(buf, []).append((u, pstate, (k,))) for k in kernels: in_degree.setdefault(k, 0) if k.op is Ops.END: assert k.src[0].op is Ops.CALL, f"END src[0] should be KERNEL, not {k.src[0].op}"