fixes

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()

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.]]])

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}"