clean

test/unit/test_function.py

@@ -557,6 +557,15 @@ class TestFunctionTuple(unittest.TestCase):
     def f(a:Tensor): return Tensor.custom_kernel(Tensor.empty(*a.shape, dtype=a.dtype, device=a.device), a, fxn=inplace_add)[0]
     with self.assertRaisesRegex(RuntimeError, "implicit buffer"): f(Tensor([1., 2., 3., 4.]).contiguous().realize())
 
+  def test_custom_kernel_bound_scalar(self):
+    def add_scalar(out:UOp, x:UOp, value:UOp):
+      i = UOp.range(x.shape[0], 0)
+      return out[i].store(x[i] + value).end(i).sink(arg=KernelInfo(name="add_scalar"))
+    value = Tensor(UOp.variable("value", 0, 10).bind(3))
+    x = Tensor([0, 1, 2, 3])
+    out = Tensor.custom_kernel(Tensor.empty_like(x), x, value, fxn=add_scalar)[0]
+    np.testing.assert_equal(out.numpy(), [3, 4, 5, 6])
+
   def test_custom_kernel_write_only_persistent_output_is_implicit(self):
     # a write-only custom_kernel output that is a realized buffer must be captured
     def write(C:UOp, A:UOp) -> UOp:

test/unit/test_llm_server.py

@@ -13,11 +13,11 @@ V_TOKS = UOp.variable("toks", 1, 32)  # 32 is the default chunk_size in generate
 class TestTransformerGenerate(unittest.TestCase):
   def test_warmup(self):
     model, calls = Transformer(TEST_CONFIG), []
-    def generate(tokens):
+    def generate(tokens, **kwargs):
       calls.append(tokens)
       yield from (1, 2)
     with patch.object(model, "generate", generate): model.warmup()
-    self.assertEqual(calls, [[0], [0]])
+    self.assertEqual(calls, [[0]])
 
   def test_first_recurrent_generate_before_state_init(self):
     model = Transformer(TEST_CONFIG)

tinygrad/llm/kernels/__init__.py

@@ -43,7 +43,7 @@ def cached_attention(q:Tensor, stacked_kv:Tensor, cache_kv:Tensor, cache_scale:T
   return q.scaled_dot_product_attention(k, v, attn_mask=mask, enable_gqa=True)
 
 @functools.cache
-def _gated_delta_prefill_kernel(core:UOp, q:UOp, k:UOp, v:UOp, beta:UOp, alpha:UOp, state:UOp, kq:UOp) -> UOp:
+def _gated_delta_prefill_kernel(core:UOp, q:UOp, k:UOp, v:UOp, beta:UOp, alpha:UOp, state:UOp, kq:UOp, start_pos:UOp|None=None) -> 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))
@@ -52,7 +52,8 @@ def _gated_delta_prefill_kernel(core:UOp, q:UOp, k:UOp, v:UOp, beta:UOp, alpha:U
   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)))
+  current = current.after(UOp.group(*(current[col].store(state[bh, row, col].float() if start_pos is None else
+    start_pos.eq(0).where(0, 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))
@@ -65,7 +66,7 @@ def _gated_delta_prefill_kernel(core:UOp, q:UOp, k:UOp, v:UOp, beta:UOp, alpha:U
   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=()))
 
-def gated_delta_prefill(q:Tensor, k:Tensor, v:Tensor, beta:Tensor, alpha:Tensor, state:Tensor) -> Tensor:
+def gated_delta_prefill(q:Tensor, k:Tensor, v:Tensor, beta:Tensor, alpha:Tensor, state:Tensor, start_pos:Tensor|None=None) -> Tensor:
   batch, heads, tokens, key_dim = q.shape
   value_dim = v.shape[-1]
   assert q.shape == k.shape and v.shape[:3] == q.shape[:3] and beta.shape == (batch, heads, tokens)
@@ -75,5 +76,5 @@ def gated_delta_prefill(q:Tensor, k:Tensor, v:Tensor, beta:Tensor, alpha:Tensor,
   if str(q.device).startswith("AMD") and key_dim % 32 == 0 and value_dim % 4 == 0:
     from tinygrad.llm.kernels.amd import _gated_delta_prefill_kernel as kernel
   core, kq = Tensor.empty_like(v), (q*k).sum(-1).contiguous()
-  return Tensor.custom_kernel(core, q.contiguous(), k.contiguous(), v.contiguous(), beta.contiguous(), alpha.contiguous(), state, kq,
-                              fxn=kernel)[0]
+  srcs = (core, q.contiguous(), k.contiguous(), v.contiguous(), beta.contiguous(), alpha.contiguous(), state, kq)
+  return Tensor.custom_kernel(*srcs, *((start_pos,) if start_pos is not None else ()), fxn=kernel)[0]

tinygrad/llm/kernels/amd.py

@@ -218,7 +218,7 @@ def q8_quantize(x:Tensor, tokens:int, in_features:int) -> tuple[Tensor, Tensor]:
   return (groups/scale.unsqueeze(-1)).round().clip(-127, 127).cast(dtypes.int8).contiguous().bitcast(dtypes.uint32), scale
 
 @functools.cache
-def _gated_delta_prefill_kernel(core:UOp, q:UOp, k:UOp, v:UOp, beta:UOp, alpha:UOp, state:UOp, kq:UOp) -> UOp:
+def _gated_delta_prefill_kernel(core:UOp, q:UOp, k:UOp, v:UOp, beta:UOp, alpha:UOp, state:UOp, kq:UOp, start_pos:UOp|None=None) -> UOp:
   batch, heads, tokens, value_dim, row_tile = *core.shape, 4
   key_dim, alpha_dim = q.shape[-1], alpha.shape[-1] if len(alpha.shape) == 4 else 1
   assert all(isinstance(x, int) for x in (batch, heads, tokens, value_dim, key_dim)) and key_dim % 32 == 0 and value_dim % row_tile == 0
@@ -232,7 +232,8 @@ def _gated_delta_prefill_kernel(core:UOp, q:UOp, k:UOp, v:UOp, beta:UOp, alpha:U
   rows = tuple(row_base+i for i in range(row_tile))
   cols = tuple(lane + i*32 for i in range(key_dim//32))
   current = UOp.placeholder((row_tile*key_dim//32,), dtypes.float32, slot=0, addrspace=AddrSpace.REG)
-  current = current.after(current.store(UOp.stack(*(state[bh, row, col].float() for row in rows for col in cols))))
+  current = current.after(current.store(UOp.stack(*(state[bh, row, col].float() if start_pos is None else
+    start_pos.eq(0).where(0, state[bh, row, col].float()) for row in rows for col in cols))))
   token = UOp.range(tokens, 2, AxisType.REDUCE)
   keys = tuple(k[bh, token, col].load() for col in cols)
   queries = tuple(q[bh, token, col].load() for col in cols)
@@ -418,4 +419,4 @@ def q8_linear(layer:Linear, x:Tensor) -> Tensor:
 def iq4_half_lut(device:str) -> Tensor:
   from tinygrad.runtime.autogen.ggml_common import kvalues_iq4nl
   return Tensor([x for j in range(16) for i in range(16) for x in (kvalues_iq4nl[i], kvalues_iq4nl[j])],
-                dtype=dtypes.float16, device=device).bitcast(dtypes.uint32).contiguous().realize()
+                dtype=dtypes.float16, device=device).bitcast(dtypes.uint32).contiguous()

tinygrad/llm/model.py

@@ -256,7 +256,10 @@ class GatedDeltaNetBlock(FFNBlock):
     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)
     alpha = self.ssm_f_b(self.ssm_f_a(x)) if is_kda else self.ssm_alpha(x)
-    conv_window = self.conv_state.cat(self.attn_qkv(x), dim=1)
+    start_pos = start_pos if isinstance(start_pos, UOp) else UOp.variable("start_pos", 0, self.config.max_context-1).bind(start_pos)
+    initial = Tensor(start_pos).eq(0)
+    conv_state = initial.where(0, self.conv_state) if not self.conv_state.uop.is_realized else self.conv_state
+    conv_window = conv_state.cat(self.attn_qkv(x), dim=1)
     conv_out = ((conv_window * self.ssm_conv1d["weight"].T.unsqueeze(0)).sum(1) if is_kda and resolve(T == 1) else 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)
@@ -275,7 +278,8 @@ class GatedDeltaNetBlock(FFNBlock):
     state_pos = T if valid_len is None else valid_len
     conv_state = conv_window[:, state_pos:state_pos+self.ssm_conv_kernel-1].cast(self.conv_state.dtype).contiguous()
     state = Tensor(self.recurrent_state.uop.after(self.conv_state.uop.after(self.conv_state.uop.store(conv_state.uop))))
-    core = gated_delta_prefill(q * self.head_k_dim**-0.5, k, v, beta, alpha, state).transpose(1, 2)
+    reset_pos = Tensor(start_pos) if not self.recurrent_state.uop.is_realized else None
+    core = gated_delta_prefill(q * self.head_k_dim**-0.5, k, v, beta, alpha, state, reset_pos).transpose(1, 2)
     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()
 
@@ -287,8 +291,8 @@ class GatedDeltaNetBlock(FFNBlock):
 
   def _init_state(self, x):
     if not hasattr(self, "conv_state"):
-      self.conv_state = Tensor.zeros(x.shape[0], self.ssm_conv_kernel-1, self.conv_channels, device=x.device).clone()
-      self.recurrent_state = Tensor.zeros(x.shape[0], self.num_v_heads, self.head_v_dim, self.head_k_dim, device=x.device).clone()
+      self.conv_state = Tensor.empty(x.shape[0], self.ssm_conv_kernel-1, self.conv_channels, device=x.device).contiguous()
+      self.recurrent_state = Tensor.empty(x.shape[0], self.num_v_heads, self.head_v_dim, self.head_k_dim, device=x.device).contiguous()
 
 class Transformer:
   def __init__(self, config:TransformerConfig):
@@ -410,19 +414,14 @@ class Transformer:
     return min(block._reusable_prefix_len(prefix_len, len(self._cached_tokens)) for block in self.blk)
 
   def warmup(self, chunk_size:int=256):
-    if not self.has_recurrent_block:
-      for _ in range(2): list(zip(range(2), self.generate([0])))
-      return
-    x = Tensor.zeros(1, 1, self.blk[0].config.dim, device=self.token_embd.weight.device)
-    for block in self.blk: block._init_state(x)
-    Tensor.realize(*(getattr(block, name) for block in self.blk
-      for name in ("cache_kv", "cache_kv_scale", "freqs_cis", "conv_state", "recurrent_state") if hasattr(block, name)))
-    self.prefill_jit.cnt = self.rollout_jit.cnt = 1
-    warm = self.generate([0] * min(chunk_size, 256, self.max_context-1), chunk_size=chunk_size)
-    with Context(JIT_BATCH_SIZE=getenv("PREFILL_JIT_BATCH_SIZE", 512)): next(warm)
+    prompt = [0] * (min(chunk_size, 256, self.max_context-1) if self.has_recurrent_block else 1)
+    if self.has_recurrent_block:
+      x = Tensor.zeros(1, 1, self.blk[0].config.dim, device=self.token_embd.weight.device)
+      for block in self.blk: block._init_state(x)
+      self.prefill_jit.cnt = self.rollout_jit.cnt = 1
+    warm = self.generate(prompt, chunk_size=chunk_size)
+    with Context(JIT_BATCH_SIZE=getenv("PREFILL_JIT_BATCH_SIZE", 512) if self.has_recurrent_block else 0): next(warm)
     with Context(JIT_BATCH_SIZE=0): next(warm)
-
-    if resets := [r for block in self.blk for r in block._state_reset_ops()]: Tensor.realize(*resets)
     self._cached_tokens = []
 
   def generate(self, tokens:list[int], chunk_size:int|None=None, temperature:float=0.0):
@@ -436,7 +435,6 @@ class Transformer:
     # recompute start_pos from what's currently valid in the caches
     start_pos = self.get_start_pos(tokens)
     decode_resume = self.has_recurrent_block and bool(self._cached_tokens) and start_pos > 0
-    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:
       padded_prefill = self.has_recurrent_block and start_pos < prompt_len and not decode_resume

tinygrad/uop/ops.py

@@ -1157,7 +1157,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
     return UOp(Ops.PARAM, src=src, arg=ParamArg(slot, dtype, vmin_vmax, multiple_of, name, addrspace, axis, device, volatile))
   def param_like(self, slot:int):
     addrspace = self.addrspace if self.addrspace is not None else AddrSpace.GLOBAL
-    if self.op is Ops.BIND: return self.src[0].replace(arg=replace(self.src[0].arg, slot=slot, addrspace=addrspace))
+    if self.op is Ops.BIND: return self.src[0].replace(arg=replace(self.src[0].arg, slot=slot))
     return UOp.param(slot, self.dtype, self.shard_shape if self.axis is not None else self._shape, self.device, addrspace=addrspace, axis=self.axis)
 
   @staticmethod
@@ -1177,8 +1177,8 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
     body = self if self.op is Ops.TUPLE else UOp.maketuple(self)
     return UOp(Ops.FUNCTION, src=(body,)+srcs, arg=CallInfo(grad_fxn, name, precompile, precompile_backward, aux))
   def custom_kernel(*srcs:UOp, fxn:Callable, grad_fxn:Callable|None=None) -> list[UOp]:
-    contig_srcs = tuple(x.contiguous() if x.op is not Ops.AFTER else x for x in srcs)
-    placeholders = [UOp.placeholder_like(s, slot=i) for i,s in enumerate(contig_srcs)]
+    contig_srcs = tuple(x if x.op in (Ops.AFTER, Ops.BIND) else x.contiguous() for x in srcs)
+    placeholders = [s.param_like(i) if s.op is Ops.BIND else UOp.placeholder_like(s, slot=i) for i,s in enumerate(contig_srcs)]
     kernel = fxn(*placeholders).call(*contig_srcs, grad_fxn=grad_fxn)
     return [s.after(kernel) for s in contig_srcs]