fixes

extra/benchmark_qwen_llm.py

@@ -0,0 +1,56 @@
+import argparse, time
+from tinygrad.llm.cli import models
+from tinygrad.llm.model import Transformer
+from tinygrad.helpers import fetch
+
+if __name__ == "__main__":
+  parser = argparse.ArgumentParser()
+  parser.add_argument("--model", default="qwen3.5:0.8b", help=f"Model choice ({', '.join(models.keys())}) or path to a local GGUF file")
+  parser.add_argument("--max-context", type=int, default=32768)
+  parser.add_argument("--prompt-tokens", type=int, default=3072)
+  parser.add_argument("--decode-tokens", type=int, default=16)
+  parser.add_argument("--chunk-size", type=int, default=256)
+  parser.add_argument("--expect-output", type=int, nargs="+", default=None, help="expected output tokens to assert on")
+  parser.add_argument("--skip-resume-check", action="store_true")
+  args = parser.parse_args()
+
+  startup_st = st = time.perf_counter()
+  model, _ = Transformer.from_gguf(fetch(models.get(args.model, args.model)), args.max_context)
+  print(f"load {time.perf_counter()-st:.3f}s", flush=True)
+  st = time.perf_counter()
+  model.warmup()
+  # warm up the chunked prefill JIT too, then invalidate the prompt cache (forces the state reset path)
+  if model.has_recurrent_block:
+    for _ in range(2):
+      warm = model.generate([0]*args.chunk_size, chunk_size=args.chunk_size)
+      next(warm), next(warm)
+      model._cached_tokens = [-1]
+  print(f"warm {time.perf_counter()-st:.3f}s", flush=True)
+  states = [getattr(block, name) for block in model.blk for name in ("cache_kv", "cache_k", "cache_v", "conv_state", "recurrent_state")
+            if hasattr(block, name)]
+  device = str(model.token_embd.weight.device)
+  assert all(str(state.device) == device and state.uop.is_realized for state in states)
+  assert all(block.cache_kv.shape[3] >= args.max_context for block in model.blk if hasattr(block, "cache_kv"))
+  assert model.prefill_jit.cnt >= 2 and model.rollout_jit.cnt >= 2
+  print(f"preallocated {sum(state.nbytes() for state in states)/2**30:.3f} GiB state on {device}", flush=True)
+
+  prompt = [257] + [1000+i%1000 for i in range(args.prompt_tokens-1)]
+  gen, st = model.generate(prompt, chunk_size=args.chunk_size), time.perf_counter()
+  output = [next(gen)]
+  pt = time.perf_counter()
+  print(f"prefill {args.prompt_tokens/(pt-st):.3f} tok/s", flush=True)
+  for _ in range(args.decode_tokens): output.append(next(gen))
+  print(f"decode {args.decode_tokens/(time.perf_counter()-pt):.3f} tok/s output {output}", flush=True)
+  if args.expect_output is not None: assert output == args.expect_output, f"expected {args.expect_output}, got {output}"
+
+  if not args.skip_resume_check:
+    follow = model._cached_tokens + [1234+i for i in range(8)]
+    full_prompt = list(follow)
+    resume_pos, gen, st = model.get_start_pos(follow), model.generate(follow, chunk_size=args.chunk_size), time.perf_counter()
+    resumed_token = next(gen)
+    print(f"resume {len(follow)-1} tokens from {resume_pos} in {time.perf_counter()-st:.3f}s token {resumed_token}", flush=True)
+    model._cached_tokens = [-1]
+    st = time.perf_counter()
+    full_token = next(model.generate(full_prompt, chunk_size=args.chunk_size))
+    print(f"full {time.perf_counter()-st:.3f}s token {full_token} match {resumed_token == full_token}", flush=True)
+    assert resumed_token == full_token

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/llm/model.py

@@ -1,10 +1,12 @@
 from __future__ import annotations
 import functools, itertools, pathlib
 from dataclasses import dataclass, replace
-from tinygrad import Tensor, nn, UOp, TinyJit, getenv, function
+from typing import cast
+from tinygrad import Tensor, nn, UOp, TinyJit, getenv, function, dtypes
 from tinygrad.nn import Linear
+from tinygrad.dtype import AddrSpace
 from tinygrad.llm.gguf import gguf_load
-from tinygrad.uop.ops import resolve
+from tinygrad.uop.ops import AxisType, KernelInfo, resolve
 
 @functools.cache
 def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0, device:str|None=None) -> Tensor:
@@ -238,6 +240,30 @@ class MLATransformerBlock(FFNBlock):
       self.cache_k = Tensor.empty(x.shape[0], 1, self.config.max_context, self.config.kv_lora_rank + self.config.rope_dim, device=x.device)
       self.freqs_cis = precompute_freqs_cis(self.config.rope_dim, self.config.max_context, self.config.rope_theta, device=x.device)
 
+# NOTE: this basic kernel works on any backend (CPU, AMD, ...); it's a sequential scan over all T tokens in a single kernel
+@functools.cache
+def _gated_delta_prefill_kernel(core:UOp, q:UOp, k:UOp, v:UOp, beta:UOp, alpha:UOp, state:UOp, kq:UOp) -> UOp:
+  batch, heads, tokens, value_dim = cast(tuple[int, int, int, int], core.shape)
+  key_dim, alpha_dim = cast(int, q.shape[-1]), cast(int, alpha.shape[-1]) if len(alpha.shape) == 4 else 1
+  core, v = (x.reshape(batch*heads, tokens, value_dim) for x in (core, v))
+  q, k = (x.reshape(batch*heads, tokens, key_dim) for x in (q, k))
+  beta, kq = (x.reshape(batch*heads, tokens) for x in (beta, kq))
+  alpha, state = alpha.reshape(batch*heads, tokens, alpha_dim), state.reshape(batch*heads, value_dim, key_dim)
+  bh, row, cols = UOp.range(batch*heads, 0, AxisType.GLOBAL), UOp.range(value_dim, 2), tuple(range(key_dim))
+  current = UOp.placeholder((key_dim,), dtypes.float32, slot=0, addrspace=AddrSpace.REG)
+  current = current.after(UOp.group(*(current[col].store(state[bh, row, col].float()) for col in cols)))
+  token = UOp.range(tokens, 1, AxisType.REDUCE)
+  previous = tuple(current.after(token)[col].load() for col in cols)
+  keys, queries = (tuple(x[bh, token, col].load() for col in cols) for x in (k, q))
+  av, bv = alpha[bh, token, row if alpha_dim > 1 else 0].load(), beta[bh, token].load()
+  state_k = sum((x*y for x,y in zip(previous, keys)), UOp.const(0, dtypes.float32))
+  state_q = sum((x*y for x,y in zip(previous, queries)), UOp.const(0, dtypes.float32))
+  delta = (v[bh, token, row].load() - state_k*av) * bv
+  step = UOp.group(core[bh, token, row].store(state_q*av + delta*kq[bh, token]),
+                   *(current[col].store(x*av + delta*y) for col,x,y in zip(cols, previous, keys))).end(token)
+  stores = (state[bh, row, col].store(current.after(step)[col].load().cast(state.dtype)) for col in cols)
+  return UOp.group(*stores).end(row, bh).sink(arg=KernelInfo(name="gated_delta_prefill", opts_to_apply=()))
+
 class GatedDeltaNetBlock(FFNBlock):
   def __init__(self, config:TransformerConfig, ssm:SSMConfig):
     super().__init__(config)
@@ -260,39 +286,71 @@ class GatedDeltaNetBlock(FFNBlock):
 
   def _attention(self, x:Tensor, start_pos:int|UOp) -> Tensor:
     B, T, _ = x.shape
-    assert T == 1, "GatedDeltaNetBlock currently only supports T=1"
+    if resolve(T == 1):
+      # input processing
+      x = x.half()
+      out_gate = self.ssm_g_b(self.ssm_g_a(x)) if hasattr(self, "ssm_g_a") else self.attn_gate(x)
+      out_gate = out_gate.reshape(B, 1, self.num_v_heads, self.head_v_dim)
+      beta = self.ssm_beta(x).sigmoid().reshape(B, self.num_v_heads, 1, 1)
+      alpha = self.ssm_f_b(self.ssm_f_a(x)) if hasattr(self, "ssm_f_a") else self.ssm_alpha(x)
+      alpha = ((alpha.float() + self.ssm_dt["bias"]).softplus().reshape(B, self.num_v_heads, -1) *
+               self.ssm_a.reshape(1, self.num_v_heads, -1)).exp().unsqueeze(-2)
+
+      # qkv conv
+      conv_window = self.conv_state.cat(self.attn_qkv(x), dim=1)
+      conv_out = (conv_window * self.ssm_conv1d["weight"].T.unsqueeze(0)).sum(1).silu()
+      q, k, v = conv_out.split([self.q_dim, self.q_dim, self.conv_channels - 2*self.q_dim], dim=-1)
+      q = q.reshape(B, self.num_k_heads, self.head_k_dim).normalize(dim=-1).repeat(1, self.num_v_heads//self.num_k_heads, 1)
+      k = k.reshape(B, self.num_k_heads, self.head_k_dim).normalize(dim=-1).repeat(1, self.num_v_heads//self.num_k_heads, 1)
+      v = v.reshape(B, self.num_v_heads, self.head_v_dim)
+      q, k, v = q.mul(self.head_k_dim**-0.5).unsqueeze(-1), k.unsqueeze(-1), v.unsqueeze(-1)
+
+      # recurrent
+      recurrent_state = self.recurrent_state * alpha
+      recurrent_state = recurrent_state + ((v - recurrent_state@k) * beta)@k.transpose(-1, -2)
+
+      # store the updated state
+      conv_state_store = self.conv_state.uop.store(conv_window[:, 1:, :].cast(self.conv_state.dtype).uop)
+      recurrent_state_store = self.recurrent_state.uop.store(recurrent_state.cast(self.recurrent_state.dtype).uop)
+      recurrent_state = Tensor(self.recurrent_state.uop.after(recurrent_state_store, conv_state_store))
+
+      # output
+      core_attn_out = self.ssm_norm((recurrent_state@q).squeeze(-1).reshape(B, 1, self.num_v_heads, self.head_v_dim))
+      out_gate = out_gate.sigmoid() if hasattr(self, "ssm_g_a") else out_gate.silu()
+      return self.ssm_out((core_attn_out * out_gate).reshape(B, 1, -1).cast(x.dtype))
+
+    is_kda = hasattr(self, "ssm_g_a")
 
     # input processing
     x = x.half()
-    out_gate = self.ssm_g_b(self.ssm_g_a(x)) if hasattr(self, "ssm_g_a") else self.attn_gate(x)
-    out_gate = out_gate.reshape(B, 1, self.num_v_heads, self.head_v_dim)
-    beta = self.ssm_beta(x).sigmoid().reshape(B, self.num_v_heads, 1, 1)
-    alpha = self.ssm_f_b(self.ssm_f_a(x)) if hasattr(self, "ssm_f_a") else self.ssm_alpha(x)
-    alpha = ((alpha.float() + self.ssm_dt["bias"]).softplus().reshape(B, self.num_v_heads, -1) *
-             self.ssm_a.reshape(1, self.num_v_heads, -1)).exp().unsqueeze(-2)
+    out_gate = (self.ssm_g_b(self.ssm_g_a(x)) if is_kda else self.attn_gate(x)).reshape(B, T, self.num_v_heads, self.head_v_dim)
+    beta = self.ssm_beta(x).sigmoid().reshape(B, T, self.num_v_heads)
+    alpha = self.ssm_f_b(self.ssm_f_a(x)) if is_kda else self.ssm_alpha(x)
+    alpha = (((alpha.float() + self.ssm_dt["bias"]).softplus().reshape(B, T, self.num_v_heads, -1) * self.ssm_a).squeeze(-1)
+             if is_kda else ((alpha.float() + self.ssm_dt["bias"]).softplus() * self.ssm_a).reshape(B, T, self.num_v_heads))
 
     # qkv conv
     conv_window = self.conv_state.cat(self.attn_qkv(x), dim=1)
-    conv_out = (conv_window * self.ssm_conv1d["weight"].T.unsqueeze(0)).sum(1).silu()
+    conv_out = (functools.reduce(lambda a,b: a+b,
+      (conv_window[:, i:i+T] * self.ssm_conv1d["weight"][:, i] for i in range(self.ssm_conv_kernel)))).silu()
     q, k, v = conv_out.split([self.q_dim, self.q_dim, self.conv_channels - 2*self.q_dim], dim=-1)
-    q = q.reshape(B, self.num_k_heads, self.head_k_dim).normalize(dim=-1).repeat(1, self.num_v_heads//self.num_k_heads, 1)
-    k = k.reshape(B, self.num_k_heads, self.head_k_dim).normalize(dim=-1).repeat(1, self.num_v_heads//self.num_k_heads, 1)
-    v = v.reshape(B, self.num_v_heads, self.head_v_dim)
-    q, k, v = q.mul(self.head_k_dim**-0.5).unsqueeze(-1), k.unsqueeze(-1), v.unsqueeze(-1)
+    q = q.reshape(B, T, self.num_k_heads, self.head_k_dim).normalize(dim=-1, eps=1e-12 if is_kda else 1e-6)
+    k = k.reshape(B, T, self.num_k_heads, self.head_k_dim).normalize(dim=-1, eps=1e-12 if is_kda else 1e-6)
+    q, k = q.repeat(1, 1, self.num_v_heads//self.num_k_heads, 1), k.repeat(1, 1, self.num_v_heads//self.num_k_heads, 1)
+    v = v.reshape(B, T, self.num_v_heads, self.head_v_dim)
+    q, k, v, beta = [z.transpose(1, 2).float() for z in (q, k, v, beta)]
+    alpha = alpha.transpose(1, 2).float().exp()
 
-    # recurrent
-    recurrent_state = self.recurrent_state * alpha
-    recurrent_state = recurrent_state + ((v - recurrent_state@k) * beta)@k.transpose(-1, -2)
-
-    # store the updated state
-    conv_state_store = self.conv_state.uop.store(conv_window[:, 1:, :].cast(self.conv_state.dtype).uop)
-    recurrent_state_store = self.recurrent_state.uop.store(recurrent_state.cast(self.recurrent_state.dtype).uop)
-    recurrent_state = Tensor(self.recurrent_state.uop.after(recurrent_state_store, conv_state_store))
+    # recurrent: run the gated delta rule over all T tokens in a single custom kernel, writing back the updated state
+    conv_state = conv_window[:, T:T+self.ssm_conv_kernel-1].cast(self.conv_state.dtype).contiguous()
+    state = Tensor(self.recurrent_state.uop.after(self.conv_state.uop.store(conv_state.uop)))
+    q, kq = q * self.head_k_dim**-0.5, ((q * self.head_k_dim**-0.5)*k).sum(-1).contiguous()
+    core = Tensor.custom_kernel(Tensor.empty_like(v), q, k, v, beta, alpha, state, kq,
+                                fxn=_gated_delta_prefill_kernel)[0].transpose(1, 2)
 
     # output
-    core_attn_out = self.ssm_norm((recurrent_state@q).squeeze(-1).reshape(B, 1, self.num_v_heads, self.head_v_dim))
-    out_gate = out_gate.sigmoid() if hasattr(self, "ssm_g_a") else out_gate.silu()
-    return self.ssm_out((core_attn_out * out_gate).reshape(B, 1, -1).cast(x.dtype))
+    gate = out_gate.sigmoid() if is_kda else out_gate.silu()
+    return self.ssm_out((self.ssm_norm(core) * gate).reshape(B, T, -1).cast(x.dtype)).contiguous()
 
   # recurrent state can't be partially reused after divergence, force a full rebuild
   def _state_reset_ops(self):
@@ -424,7 +482,6 @@ class Transformer:
     return min(block._reusable_prefix_len(prefix_len, len(self._cached_tokens)) for block in self.blk)
 
   def generate(self, tokens:list[int], chunk_size:int=32, temperature:float=0.0):
-    if self.has_recurrent_block: chunk_size = 1
     v_start_pos = UOp.variable("start_pos", 0, self.max_context-1)
     v_toks = UOp.variable("toks", 1, chunk_size)
     # TODO: use UOp.variable for temperature once float variables are supported
@@ -436,9 +493,12 @@ class Transformer:
     if start_pos < len(self._cached_tokens) and (resets := [r for b in self.blk for r in b._state_reset_ops()]): Tensor.realize(*resets)
     out, prompt_len = None, len(tokens)
     while len(tokens) < self.max_context:
-      n_toks = min(chunk_size, len(tokens) - start_pos)
+      # NOTE: Tensor.custom_kernel requires all-int input shapes (gated delta prefill), so recurrent blocks run the prompt
+      # in fixed-size chunks and roll out T=1, giving the JIT one prefill graph (per chunk size) and one rollout graph
+      n_toks = 1 if self.has_recurrent_block and len(tokens) - start_pos < chunk_size else min(chunk_size, len(tokens) - start_pos)
       sp, nt = v_start_pos.bind(start_pos), v_toks.bind(n_toks)
-      out = self(t[:, sp:sp+nt] if start_pos < prompt_len or out is None else out, sp, temp).realize()
+      sl = t[:, sp:sp+(n_toks if self.has_recurrent_block else nt)]
+      out = self(sl if start_pos < prompt_len or out is None else out, sp, temp).realize()
       start_pos += n_toks
       # chunked prefill: keep processing until all prompt tokens are consumed
       if start_pos < len(tokens): continue

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}"
@@ -196,4 +211,6 @@ def create_linear_with_vars(big_sink:UOp) -> tuple[UOp, dict[str, int]]:
     return UOp(Ops.LINEAR, src=()), var_vals
 
   held_bufs = ({b for b in linear_call.src[1:] if b.op is Ops.BUFFER} if linear_call.op is Ops.CALL else set())
+  # buffers that already hold data can't be suballocated by the memory planner, custom kernels write them in place
+  held_bufs |= {b for b in big_sink.toposort(gate_kernel_sink) if b.op is Ops.BUFFER and b.buffer.is_allocated()}
   return memory_plan_rewrite(linear, held_bufs), var_vals

tinygrad/schedule/indexing.py

@@ -25,11 +25,21 @@ def realize_store_after_src(ctx:dict[UOp, None], dest:UOp, src:UOp):
   # you don't usually have to do this for assign unless there's a WAR hazard like TestAssign.test_assign_double_diamond_reduce
   if dest.base in src.backward_slice_with_self: ctx[src] = None
 
+# the inputs of a custom kernel resolve to whole buffers (one per PARAM slot), so they have to materialize.
+# buffer states (AFTER/BUFFER/PARAM) and views of buffers already materialize, anything else must be realized.
+def realize_custom_kernel_srcs(ctx:dict[UOp, None], c:UOp) -> None:
+  for s in c.src[1:]:
+    t = s
+    while t.op in GroupOp.Movement or t.op is Ops.SLICE: t = t.src[0]
+    if t.op not in {Ops.AFTER, Ops.BUFFER, Ops.PARAM, Ops.MSELECT, Ops.MSTACK}: ctx[s] = None
+
 pm_generate_realize_map = PatternMatcher([
   # always realize
   (UPat({Ops.CONTIGUOUS, Ops.STORE}, name="tr"), realize),
   # realize srcs of these
   (UPat((Ops.MSELECT, Ops.MSTACK), name="rb"), realize_srcs),
+  # realize the inputs of custom kernel calls
+  (UPat(Ops.CALL, src=(UPat(Ops.SINK),), name="c", allow_any_len=True), realize_custom_kernel_srcs),
   # sometimes we need to realize the src of STORE if there's a self-access
   (UPat(Ops.STORE, src=(UPat.var("dest"), UPat.var("src"))), realize_store_after_src),
 ])