llm: keep recurrent prefill portable

test/unit/test_attention.py

@@ -155,12 +155,17 @@ class TestGatedDeltaNetBlock(unittest.TestCase):
     return outputs, conv_states, recurrent_states
 
   def test_gatedeltanet_reference_and_reset(self):
-    if not str(Tensor.empty(1).device).startswith("AMD"): self.skipTest("AMD required")
     config = self._make_config(max_context=3)
     block = self._make_block(config)
     x = self._tensor_linspace(-1.0, 1.0, (1, 3, config.dim))
 
     expected_outs, expected_conv, expected_recurrent = self._naive_attention(block, x)
+    out = self._run_attention(block, x, 0)
+    conv_state, recurrent_state = self._cache_views(block)
+    np.testing.assert_allclose(out, np.concatenate(expected_outs, axis=1), rtol=1e-3, atol=1e-3)
+    np.testing.assert_allclose(conv_state, expected_conv[-1], rtol=1e-3, atol=1e-3)
+    np.testing.assert_allclose(recurrent_state, expected_recurrent[-1], rtol=1e-3, atol=1e-3)
+    Tensor.realize(*block._state_reset_ops())
 
     for step in range(x.shape[1]):
       out = self._run_attention(block, x[:, step:step+1], step)

tinygrad/llm/kernels/amd.py

@@ -313,7 +313,7 @@ def _quant_decode_kernel(out:UOp, raw:UOp, xq:UOp, xd:UOp, raw_offset:UOp, out_f
       low = _amd_load(raw[base + (pos//128)*64 + within%64], 4) >> ((within//64)*4).cast(dtypes.uint8)
       high = _amd_load(raw[base + 128 + (pos//128)*32 + within%32], 4) >> ((within//32)*2).cast(dtypes.uint8)
       quant = ((low & 15) | ((high & 3) << 4)).bitcast(dtypes.int8) - 32
-      word = quant.cast(dtypes.int8).bitcast(dtypes.uint8).bitcast(dtypes.uint32).squeeze(0)
+      word = sum((quant[i].cast(dtypes.uint8).cast(dtypes.uint32) << (i*8) for i in range(4)), UOp.const(0, dtypes.uint32))
       dots[word_idx//4] = _amd_dp4a(word, xwords[word_idx], dots[word_idx//4])
     scales = [raw[base + 192 + subgroup*2+i].cast(dtypes.uint8).bitcast(dtypes.int8).float() for i in range(2)]
     dbits = raw[base+208].cast(dtypes.uint16) | (raw[base+209].cast(dtypes.uint16) << 8)

tinygrad/llm/kernels/generic.py

@@ -0,0 +1,33 @@
+import functools
+from typing import cast
+from tinygrad import Tensor, UOp, dtypes
+from tinygrad.dtype import AddrSpace
+from tinygrad.uop.ops import AxisType, KernelInfo
+
+@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, dim = cast(tuple[int, int, int, int], core.shape)
+  core, q, k, v = (x.reshape(batch*heads, tokens, dim) for x in (core, q, k, v))
+  beta, alpha, kq = (x.reshape(batch*heads, tokens) for x in (beta, alpha, kq))
+  state = state.reshape(batch*heads, dim, dim)
+  bh, row, cols = UOp.range(batch*heads, 0, AxisType.GLOBAL), UOp.range(dim, 2), tuple(range(dim))
+  current = UOp.placeholder((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].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=()))
+
+def gated_delta_prefill(q:Tensor, k:Tensor, v:Tensor, beta:Tensor, alpha:Tensor, state:Tensor) -> Tensor:
+  batch, heads, tokens, dim = q.shape
+  assert q.shape == k.shape == v.shape and beta.shape == alpha.shape == (batch, heads, tokens) and state.shape == (batch, heads, dim, dim)
+  core, kq = Tensor.empty_like(q), (q*k).sum(-1).contiguous()
+  return Tensor.custom_kernel(core, q.contiguous(), k.contiguous(), v.contiguous(), beta.contiguous(), alpha.contiguous(), state, kq,
+                              fxn=_gated_delta_prefill_kernel)[0]

tinygrad/llm/model.py

@@ -4,6 +4,7 @@ from dataclasses import dataclass, replace
 from typing import cast
 from tinygrad import Tensor, nn, UOp, TinyJit, getenv, function, dtypes, Context
 from tinygrad.llm.kernels import amd as llm_amd
+from tinygrad.llm.kernels import generic as llm_generic
 from tinygrad.llm.gguf import gguf_load
 from tinygrad.helpers import prod
 from tinygrad.uop.ops import resolve, Ops
@@ -334,7 +335,8 @@ class GatedDeltaNetBlock(FFNBlock):
       active = (Tensor.arange(T).to(x.device) < Tensor(valid_len, device=x.device)).reshape(1, T, 1)
       beta, log_alpha = beta * active, log_alpha * active
     q, k, v, beta, log_alpha = [z.transpose(1, 2).float() for z in (q, k, v, beta, log_alpha)]
-    core = llm_amd.gated_delta_prefill(q * self.head_k_dim**-0.5, k, v, beta, log_alpha.exp(), initial_state)
+    gated_delta = llm_amd.gated_delta_prefill if str(x.device).startswith("AMD") else llm_generic.gated_delta_prefill
+    core = gated_delta(q * self.head_k_dim**-0.5, k, v, beta, log_alpha.exp(), initial_state)
     out = self.ssm_out((self.ssm_norm(core.transpose(1, 2)) * out_gate.silu()).reshape(B, T, -1).cast(x.dtype)).contiguous()
     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()