llm kernels: adapt to Ops.BIND removal

Variables are 0-d ALU BUFFERs in the tensor graph and take the ALU PARAM form
inside kernels (UOp.variable(param=True)). Add kernel_var helper for the
conversion, and keep start_pos in bound form at the graph level so function
implicit-input collection and the schedule's binds rename-back line up.

test/unit/test_attention.py

@@ -1,10 +1,12 @@
 import unittest
 import numpy as np
-from tinygrad import Tensor, dtypes
+from tinygrad import Tensor, dtypes, nn
 from tinygrad.llm.model import (
   GatedDeltaNetBlock, SSMConfig, TransformerBlock, TransformerConfig,
   apply_rope as apply_rope_new, precompute_freqs_cis, pairwise_topk,
 )
+from tinygrad.llm.kernels import Linear, gated_delta_prefill
+from tinygrad.llm.gguf import ggml_data_to_tensor
 
 def apply_rope(x:Tensor, start_pos:int):
   B, H, T, Hd = x.shape
@@ -12,6 +14,15 @@ def apply_rope(x:Tensor, start_pos:int):
   freqs_cis = precompute_freqs_cis(Hd, start_pos+T)[start_pos:start_pos+T]
   return apply_rope_new(x, freqs_cis)
 
+class TestLinear(unittest.TestCase):
+  def test_recovers_packed_ggml_weight(self):
+    for ggml_type,packed_size,words in ((13, 176, 44), (14, 210, 210), (23, 136, 34)):
+      packed = Tensor.empty(packed_size+4, dtype=dtypes.uint8, device="CPU")[4:]
+      decoded = ggml_data_to_tensor(packed, 256, ggml_type).reshape(1, 256)
+      linear = Linear(256, 1, bias=False)
+      linear.set_quantized(decoded)
+      self.assertEqual((linear.ggml_type, linear.weight.numel()), (ggml_type, words))
+
 class TestAttention(unittest.TestCase):
   def test_apply_rope(self):
     x = Tensor.randn(1, 2, 4, 8, dtype=dtypes.float32)
@@ -41,14 +52,30 @@ class TestAttention(unittest.TestCase):
     np.testing.assert_allclose(block.cache_kv[0, :, :, :seqlen, :].numpy(), expected.numpy(), rtol=1e-5, atol=1e-5)
 
 class TestGatedDeltaNetBlock(unittest.TestCase):
+  def test_gated_delta_rectangular_state_and_row_decay(self):
+    rng = np.random.default_rng(42)
+    q, k = (rng.normal(size=(1, 1, 3, 32)).astype(np.float32) for _ in range(2))
+    v, beta = rng.normal(size=(1, 1, 3, 4)).astype(np.float32), rng.uniform(size=(1, 1, 3)).astype(np.float32)
+    alpha, initial = rng.uniform(0.8, 1, size=(1, 1, 3, 4)).astype(np.float32), rng.normal(size=(1, 1, 4, 32)).astype(np.float32)
+    expected_state, expected_out = initial.copy(), np.empty_like(v)
+    for t in range(3):
+      previous, av = expected_state.copy(), alpha[:, :, t, :, None]
+      delta = (v[:, :, t] - (previous*k[:, :, t, None]).sum(-1)*alpha[:, :, t]) * beta[:, :, t, None]
+      expected_state = previous*av + delta[..., None]*k[:, :, t, None, :]
+      expected_out[:, :, t] = (previous*q[:, :, t, None]).sum(-1)*alpha[:, :, t] + delta*(q[:, :, t]*k[:, :, t]).sum(-1)
+    state = Tensor(initial).contiguous().realize()
+    out = gated_delta_prefill(Tensor(q), Tensor(k), Tensor(v), Tensor(beta), Tensor(alpha), state).realize()
+    np.testing.assert_allclose(out.numpy(), expected_out, rtol=1e-4, atol=1e-4)
+    np.testing.assert_allclose(state.numpy(), expected_state, rtol=1e-4, atol=1e-4)
+
   def _tensor_linspace(self, start:float, stop:float, shape:tuple[int, ...]) -> Tensor:
     return Tensor.linspace(start, stop, int(np.prod(shape)), dtype=dtypes.float32).reshape(*shape)
 
   def _make_config(self, **kwargs):
-    return TransformerConfig(**({"num_blocks":1, "dim":4, "hidden_dim":8, "n_heads":1, "n_kv_heads":1,
-                                 "norm_eps":1e-5, "vocab_size":32, "head_dim":4, "rope_theta":10000.0,
-                                 "rope_dim":4, "v_head_dim":4, "max_context":4, "ssm_layers":(True,),
-                                 "ssm":SSMConfig(conv_kernel=2, state_size=2, group_count=1, time_step_rank=1, inner_size=2)} | kwargs))
+    return TransformerConfig(**({"num_blocks":1, "dim":32, "hidden_dim":64, "n_heads":1, "n_kv_heads":1,
+                                 "norm_eps":1e-5, "vocab_size":32, "head_dim":32, "rope_theta":10000.0,
+                                 "rope_dim":32, "v_head_dim":32, "max_context":4, "ssm_layers":(True,),
+                                 "ssm":SSMConfig(conv_kernel=2, state_size=32, group_count=1, time_step_rank=1, inner_size=32)} | kwargs))
 
   def _make_block(self, config:TransformerConfig) -> GatedDeltaNetBlock:
     block = GatedDeltaNetBlock(config, config.ssm)
@@ -79,6 +106,10 @@ class TestGatedDeltaNetBlock(unittest.TestCase):
       recurrent_state = cache[:, conv_flat:].reshape(cache.shape[0], block.num_v_heads, block.head_v_dim, block.head_v_dim)
       return conv_state, recurrent_state
 
+  def _reset_state(self, block:GatedDeltaNetBlock):
+    Tensor.realize(block.conv_state.assign(block.conv_state.const_like(0)),
+                   block.recurrent_state.assign(block.recurrent_state.const_like(0)))
+
   def _linear_np(self, x:np.ndarray, weight:np.ndarray) -> np.ndarray:
     return x.astype(np.float32) @ weight.T.astype(np.float32)
 
@@ -86,7 +117,7 @@ class TestGatedDeltaNetBlock(unittest.TestCase):
     x_float = x.astype(np.float32)
     return (x_float / np.sqrt((x_float * x_float).mean(axis=-1, keepdims=True) + eps)) * weight.astype(np.float32)
 
-  def _normalize_np(self, x:np.ndarray, eps:float=1e-12) -> np.ndarray:
+  def _normalize_np(self, x:np.ndarray, eps:float=1e-6) -> np.ndarray:
     return x / np.maximum(np.sqrt((x * x).sum(axis=-1, keepdims=True)), eps)
 
   def _softplus_np(self, x:np.ndarray) -> np.ndarray:
@@ -148,6 +179,12 @@ class TestGatedDeltaNetBlock(unittest.TestCase):
     x = Tensor.linspace(-1.0, 1.0, 3 * config.dim, dtype=dtypes.float32).reshape(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)
+    self._reset_state(block)
 
     for step in range(x.shape[1]):
       out = self._run_attention(block, x[:, step:step+1], step)
@@ -163,7 +200,7 @@ class TestGatedDeltaNetBlock(unittest.TestCase):
     prompt = Tensor.linspace(0.75, -0.75, 2 * config.dim, dtype=dtypes.float32).reshape(1, 2, config.dim)
 
     for i in range(warmup.shape[1]): self._run_attention(block, warmup[:, i:i+1], i)
-    Tensor.realize(*block._state_reset_ops())
+    self._reset_state(block)
     expected_outs, expected_conv, expected_recurrent = self._naive_attention(block, prompt)
 
     for step in range(prompt.shape[1]):
@@ -177,18 +214,41 @@ class TestGatedDeltaNetBlock(unittest.TestCase):
                                  err_msg=f"GatedDeltaNet reset recurrent cache mismatch at step {step}")
 
   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.]]])
-    # f_b(f_a(x)) = [1, 2, 3, 4]
+    config = self._make_config(dim=4, hidden_dim=8, n_heads=2, head_dim=4, rope_dim=4, v_head_dim=4,
+      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.], [2., 1., 0., 0.]]])
     block.ssm_f_a.weight = Tensor([[1., 0., 0., 0.], [0., 1., 0., 0.]])
     block.ssm_f_b.weight = Tensor([[1., 0.], [0., 1.], [1., 1.], [2., 1.]])
     block._init_state(x)
     initial_state = Tensor.arange(8, dtype=dtypes.float32).reshape(1, 2, 2, 2)
     block.recurrent_state.assign(initial_state).realize()
     block.ssm_a = Tensor([[-1.], [-1.]])
-    block._attention(x, 0).realize()
-    alpha = np.exp(-self._softplus_np(np.arange(1, 5)).reshape(1, 2, 1, 2))
-    np.testing.assert_allclose(block.recurrent_state.numpy(), initial_state.numpy() * alpha, rtol=1e-5, atol=1e-5)
+    block._attention(x, x.shape[1]).realize()
+    alpha = np.exp(-self._softplus_np(np.array([[1, 2, 3, 4], [2, 1, 3, 5]])).reshape(2, 2, 2)).prod(0)
+    np.testing.assert_allclose(block.recurrent_state.numpy(), initial_state.numpy() * alpha[..., None], rtol=1e-5, atol=1e-5)
+
+  def test_kda_prefill_matches_decode(self):
+    config = self._make_config(ssm=SSMConfig(conv_kernel=2, state_size=32, group_count=1, time_step_rank=1, inner_size=32, kda=True))
+    block = GatedDeltaNetBlock(config, config.ssm)
+    for p in nn.state.get_parameters(block):
+      p.replace(self._tensor_linspace(-0.05, 0.05, p.shape) if len(p.shape) > 1 else self._tensor_linspace(0.05, 0.1, p.shape))
+    x = self._tensor_linspace(-0.5, 0.5, (1, 3, config.dim))
+    prefill = self._run_attention(block, x, 0)
+    prefill_conv, prefill_recurrent = self._cache_views(block)
+    self._reset_state(block)
+    decode = np.concatenate([self._run_attention(block, x[:, i:i+1], i) for i in range(3)], axis=1)
+    decode_conv, decode_recurrent = self._cache_views(block)
+    np.testing.assert_allclose(prefill, decode, rtol=1e-3, atol=1e-3)
+    np.testing.assert_allclose(prefill_conv, decode_conv, rtol=1e-3, atol=1e-3)
+    np.testing.assert_allclose(prefill_recurrent, decode_recurrent, rtol=1e-3, atol=1e-3)
+
+  def test_start_zero_resets_realized_state(self):
+    config, x = self._make_config(max_context=3), self._tensor_linspace(-1, 1, (1, 3, 32))
+    block = self._make_block(config)
+    self._run_attention(block, x, 0)
+    restarted = self._run_attention(block, x[:, :2], 0)
+    fresh = self._run_attention(self._make_block(config), x[:, :2], 0)
+    np.testing.assert_allclose(restarted, fresh, rtol=1e-3, atol=1e-3)
 
 class TestPairwiseTopk(unittest.TestCase):
   def test_basic_topk(self):

test/unit/test_llm_amd.py

@@ -0,0 +1,91 @@
+import unittest
+import numpy as np
+from tinygrad import Tensor, UOp, dtypes, nn
+from tinygrad.llm.kernels import Linear, amd_custom_kernels_supported
+from tinygrad.llm.kernels.amd import q8_quantize, quantized_attention
+from tinygrad.llm.gguf import ggml_data_to_tensor
+
+class TestQ8Quantize(unittest.TestCase):
+  def test_word_quant_weights_use_typed_buffer_view(self):
+    for ggml_type, type_size in ((13, 176), (23, 136)):
+      with self.subTest(ggml_type=ggml_type):
+        raw = Tensor(np.zeros(type_size + 4, dtype=np.uint8), device="CPU").contiguous().realize()[4:]
+        decoded = ggml_data_to_tensor(raw, 256, ggml_type).reshape(1, 256)
+        linear = Linear(256, 1, bias=False)
+        linear.set_quantized(decoded)
+        self.assertEqual(linear.ggml_type, ggml_type)
+        self.assertEqual(linear.weight.dtype, dtypes.uint32)
+        self.assertEqual(linear.weight.nbytes(), type_size)
+        self.assertEqual(linear.weight.uop.buf_uop.buffer.offset, 4)
+
+  def test_values_and_scales(self):
+    if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
+    x = np.linspace(-3.1, 2.7, 64, dtype=np.float32).reshape(2, 32)
+    quant, scale = q8_quantize(Tensor(x), 2, 32)
+    scale_np = np.maximum(np.max(np.abs(x), axis=-1, keepdims=True) / 127, 1e-8)
+    expected = np.clip(np.rint(x / scale_np), -127, 127).astype(np.int8)
+    np.testing.assert_array_equal(quant.bitcast(dtypes.int8).reshape(2, 32).numpy(), expected)
+    np.testing.assert_allclose(scale.numpy(), scale_np, rtol=1e-6)
+
+  def test_q6_linear_compiles(self):
+    if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
+    rng = np.random.default_rng(42)
+    packed = rng.integers(0, 256, 210, dtype=np.uint8)
+    packed[-2:] = np.array([0.01], dtype=np.float16).view(np.uint8)
+    raw = Tensor(np.pad(packed, (4, 0))).contiguous().realize()[4:]
+    decoded = ggml_data_to_tensor(raw, 256, 14).reshape(1, 256)
+    linear = Linear(256, 1, bias=False)
+    nn.state.load_state_dict(linear, {"weight":decoded}, verbose=False, realize=False)
+    self.assertTrue(np.isfinite(linear(Tensor.randn(1, 256)).realize().item()))
+    self.assertEqual(linear.weight.uop.buf_uop.buffer.offset, 4)
+
+  def test_q6_linear_multiple_tokens(self):
+    if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
+    rng = np.random.default_rng(42)
+    in_features, blocks = 2048, 16*2048//256
+    packed = rng.integers(0, 256, blocks*210, dtype=np.uint8)
+    for i in range(blocks): packed[i*210+208:i*210+210] = np.array([0.01], dtype=np.float16).view(np.uint8)
+    raw = Tensor(np.pad(packed, (4, 0))).contiguous().realize()[4:]
+    decoded = ggml_data_to_tensor(raw, 16*in_features, 14).reshape(16, in_features)
+    weight = decoded.numpy()
+    linear = Linear(in_features, 16, bias=False)
+    nn.state.load_state_dict(linear, {"weight":decoded}, verbose=False, realize=False)
+    x = rng.normal(size=(3, in_features)).astype(np.float32)
+    scale = np.maximum(np.abs(x).reshape(3, in_features//32, 32).max(-1, keepdims=True) / 127, 1e-8)
+    xq = np.clip(np.rint(x.reshape(3, in_features//32, 32) / scale), -127, 127) * scale
+    np.testing.assert_allclose(linear(Tensor(x)).numpy(), xq.reshape(3, in_features) @ weight.T, rtol=2e-3, atol=2e-2)
+    self.assertEqual(linear.ggml_type, 14)
+
+    generic = Linear(in_features, 16, bias=False)
+    nn.state.load_state_dict(generic, {"weight":decoded}, verbose=False, realize=False)
+    generic(Tensor.randn(4, in_features)[:UOp.variable("tokens", 1, 4).bind(2)])
+    self.assertFalse(generic.use_custom_quant)
+    self.assertIsNone(generic.ggml_type)
+
+  def test_attention_uses_physical_cache_length(self):
+    if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
+    q, k, v = Tensor.zeros(1, 2, 1, 32), Tensor.randn(1, 1, 1, 32), Tensor.randn(1, 1, 1, 32)
+    cache = Tensor.empty(2, 1, 1, 256, 32, dtype=dtypes.int8).contiguous()
+    scale = Tensor.empty(2, 1, 1, 256, dtype=dtypes.float16).contiguous()
+    out = quantized_attention(q, Tensor.stack(k, v), cache, scale, 0).realize()
+    np.testing.assert_allclose(out.numpy(), v.expand(1, 2, 1, 32).numpy(), rtol=2e-2, atol=2e-2)
+
+  def test_prefill_attention_unaligned_start(self):
+    if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
+    rng = np.random.default_rng(42)
+    start_pos = 1718
+    q = Tensor.zeros(1, 8, 32, 128)
+    old_kv = rng.normal(size=(2, 1, 1, start_pos, 128)).astype(np.float32)
+    new_kv = rng.normal(size=(2, 1, 1, 32, 128)).astype(np.float32)
+    cache = Tensor.empty(2, 1, 1, 2048, 128, dtype=dtypes.int8).contiguous()
+    scale = Tensor.zeros(2, 1, 1, 2048, dtype=dtypes.float16).contiguous()
+    old_scale = np.maximum(np.max(np.abs(old_kv), axis=-1, keepdims=True) / 127, 1e-8).astype(np.float16)
+    Tensor.realize(cache[:, :, :, :start_pos].assign(Tensor(np.rint(old_kv / old_scale).astype(np.int8))),
+                   scale[:, :, :, :start_pos].assign(Tensor(old_scale.squeeze(-1))))
+    out = quantized_attention(q, Tensor(new_kv), cache, scale, UOp.variable("start_pos", 0, 2047).bind(start_pos)).realize()
+    values = cache[1, 0, 0, :start_pos+32].numpy().astype(np.float32) * \
+      scale[1, 0, 0, :start_pos+32].numpy().astype(np.float32)[:, None]
+    expected = np.stack([values[:start_pos+i+1].mean(0) for i in range(32)])[None, None].repeat(8, axis=1)
+    np.testing.assert_allclose(out.numpy(), expected, rtol=2e-3, atol=2e-3)
+
+if __name__ == "__main__": unittest.main()

tinygrad/llm/kernels/__init__.py

@@ -0,0 +1,92 @@
+import functools
+from typing import cast
+from tinygrad import Tensor, UOp, nn, dtypes, Device, Context
+from tinygrad.device import Buffer
+from tinygrad.dtype import AddrSpace
+from tinygrad.helpers import prod
+from tinygrad.uop.ops import AxisType, KernelInfo, Ops
+
+def kernel_var(x:UOp) -> UOp:
+  # a Variable is a 0-d ALU BUFFER in the tensor graph; inside kernels it takes the ALU PARAM form (same name keeps the value binding)
+  return x.substitute({v: UOp.variable(v.expr, v.vmin, v.vmax, dtype=v.dtype, multiple_of=v.arg.multiple_of, param=True)
+                       for v in x.toposort() if v.is_variable})
+
+def amd_custom_kernels_supported(device:str|tuple[str, ...]|None) -> bool:
+  # the custom kernels are tuned for RDNA3 (gfx11): the WMMA register layouts don't match gfx12 (RDNA4)
+  # or CDNA (MFMA-only, wave64), and the dp4a builtins and 32-lane wave ops aren't portable either.
+  if isinstance(device, tuple): device = device[0]
+  if device is None or device.split(":")[0] != "AMD": return False
+  # Device[...] trips ALLOW_DEVICE_USAGE=0 in function contexts, the device is always open here anyway
+  with Context(ALLOW_DEVICE_USAGE=1):
+    return (t:=getattr(Device[device], "target", None)) is not None and t[0] == 11
+
+class Linear(nn.Linear):
+  ggml_type:int|None = None
+  def __init__(self, in_features:int, out_features:int, bias=True):
+    super().__init__(in_features, out_features, bias)
+    self.in_features, self.out_features = in_features, out_features
+    self.use_custom_quant = True
+  def set_quantized(self, decoded:Tensor):
+    packed_sizes = {decoded.numel() // 256 * type_size:typ for typ,type_size in ((13, 176), (14, 210), (23, 136))}
+    raw = next((u for u in decoded.uop.toposort() if u.op is Ops.SHRINK and u.dtype == dtypes.uint8 and prod(u.shape) in packed_sizes), None)
+    if raw is None: return
+    raw_offset = raw.contiguous_view_offset()
+    assert raw_offset is not None and raw_offset % 4 == 0 and raw.buf_uop.dtype == dtypes.uint8
+    self.ggml_type = packed_sizes[prod(raw.shape)]
+    # Q5_K and IQ4_XS kernels consume words. Store a typed buffer view directly: a lazy BITCAST is decomposed into
+    # byte-combining ALU before custom-kernel scheduling and would copy the entire packed weight on every JIT graph.
+    packed_dtype = dtypes.uint8 if self.ggml_type == 14 else dtypes.uint32
+    self.weight = Tensor(UOp.from_buffer(cast(Buffer, raw.buf_uop.buffer).view(raw.max_numel() * raw.dtype.itemsize // packed_dtype.itemsize,
+                                                                              packed_dtype, raw_offset)))
+  def __call__(self, x:Tensor) -> Tensor:
+    static = isinstance(x.numel(), int)
+    supported = self.use_custom_quant and amd_custom_kernels_supported(self.weight.device)
+    if self.ggml_type is None and not static: supported = self.use_custom_quant = False
+    if self.ggml_type is None and supported: self.set_quantized(self.weight)
+    if self.ggml_type in (13, 14, 23) and supported:
+      from tinygrad.llm.kernels.amd import q8_linear
+      return q8_linear(self, x)
+    return super().__call__(x)
+
+@functools.cache
+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))
+  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)
+  initial = None if start_pos is None else start_pos.eq(0)
+  current = current.after(UOp.group(*(current[col].store(state[bh, row, col].float() if initial is None else
+    initial.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))
+  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=()))
+
+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)
+  assert alpha.shape in ((batch, heads, tokens), (batch, heads, tokens, value_dim))
+  assert state.shape == (batch, heads, value_dim, key_dim)
+  kernel = _gated_delta_prefill_kernel
+  if amd_custom_kernels_supported(q.device) 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()
+  srcs = (core, q.contiguous(), k.contiguous(), v.contiguous(), beta.contiguous(), alpha.contiguous(), state, kq)
+  if start_pos is None: return Tensor.custom_kernel(*srcs, fxn=kernel)[0]
+  contig = tuple(x.uop if x.uop.op is Ops.AFTER else x.uop.contiguous() for x in srcs)
+  params = tuple(UOp.placeholder_like(x, slot=i) for i,x in enumerate(contig))
+  assert start_pos.uop.is_bound_var
+  call = kernel(*params, kernel_var(start_pos.uop.src[0])).call(*contig, start_pos.uop)
+  return Tensor(contig[0].after(call))

tinygrad/llm/kernels/amd.py

@@ -0,0 +1,451 @@
+from __future__ import annotations
+import functools, math
+from typing import Callable, cast
+from tinygrad import Tensor, UOp
+from tinygrad.llm.kernels import Linear, kernel_var
+from tinygrad.uop.ops import AxisType, KernelInfo, Ops, resolve
+from tinygrad.dtype import AddrSpace, dtypes
+
+BLOCK_M, BLOCK_N, DECODE_HEAD_TILE, WARP_SIZE = 32, 32, 8, 32
+WMMA_M, WMMA_N, WMMA_K = 16, 16, 16
+WAVES_M, WAVES_N, LANES_PER_WAVE_M, LANES_PER_WAVE_N = 2, 2, 2, 16
+WMMA_ACC, THREADS_PER_BLOCK = WMMA_M // LANES_PER_WAVE_M, WARP_SIZE * WAVES_M * WAVES_N
+LDS_PAD, WMMA_ARG, LOG2E = 4, ((WMMA_M, WMMA_N, WMMA_K), 'AMD', 32), math.log2(math.e)
+Q5_K, Q6_K, IQ4_XS, GGML_BLOCK_SIZE, Q8_GROUP_SIZE, Q5_WORDS, Q6_BYTES, IQ4_WORDS = 13, 14, 23, 256, 32, 44, 210, 34
+
+def warp_reduce(val:UOp, maximum:bool=False, full_wave:bool=False) -> UOp:
+  for offset in ((16, 8, 4, 2, 1) if full_wave else (8, 4, 2, 1)):
+    if val.op is Ops.INDEX and val.addrspace == AddrSpace.REG: val = val.load()
+    other = UOp(Ops.CUSTOM, dtypes.float, (val,), arg=
+      f"__builtin_bit_cast(float, __builtin_amdgcn_ds_swizzle(__builtin_bit_cast(int, {{0}}), {0x1f | offset<<10}))")
+    val = val.maximum(other) if maximum else val + other
+  return val
+
+def _reg(shape:tuple[int, ...], slot:int, value:float, dep:UOp|None=None) -> UOp:
+  ret = UOp.placeholder(shape, dtypes.float, slot=slot, addrspace=AddrSpace.REG)
+  return ret.after((ret if dep is None else ret.after(dep)).store(ret.const_like(value)))
+
+@functools.cache
+def _amd_flash_attention_decode_partial(out, stats, q, cache_kv, cache_scale, valid_kv_len, max_kv_len, block_n):
+  if isinstance(valid_kv_len, UOp): valid_kv_len = kernel_var(valid_kv_len.unbind_all()[0])
+  _, B, H_KV, N, D = cast(tuple[int, int, int, int, int], cache_kv.shape)
+  _, H, M, _ = cast(tuple[int, int, int, int], q.shape)
+  assert M == 1 and H % H_KV == 0 and D % WARP_SIZE == 0 and max_kv_len <= N and max_kv_len % block_n == 0
+  G, CHUNK, DV, heads_per_wave = H // H_KV, block_n, D // WARP_SIZE, 2
+  head_tile = min(DECODE_HEAD_TILE, G)  # share each KV stream across two GQA heads per wave
+  assert G % head_tile == 0 and head_tile % heads_per_wave == 0
+  decode_waves, decode_group = head_tile // heads_per_wave, 4
+  block_bhkv = UOp.range(B*H_KV*(G//head_tile), 0, AxisType.GLOBAL)
+  valid_chunks = (valid_kv_len+CHUNK-1)//CHUNK
+  group_count = min(valid_chunks, out.shape[2]) if isinstance(valid_chunks, int) else valid_chunks.minimum(out.shape[2])
+  block_n = UOp.range(group_count, 1, AxisType.GLOBAL)
+  lane, wave = UOp.range(WARP_SIZE, 2, AxisType.LOCAL), UOp.range(decode_waves, 3, AxisType.LOCAL)
+  head_group, bhkv = block_bhkv % (G//head_tile), block_bhkv // (G//head_tile)
+  b, kv_head = bhkv // H_KV, bhkv % H_KV
+  dims = tuple(lane + i*WARP_SIZE for i in range(DV))
+  acc, row_max, row_sum = _reg((heads_per_wave, DV), 0, 0), _reg((heads_per_wave,), 1, -math.inf), _reg((heads_per_wave,), 2, 0)
+  groups_per_chunk, offset = CHUNK // decode_group, UOp.range(((valid_chunks+group_count-1)//group_count)*(CHUNK//decode_group), 100, AxisType.REDUCE)
+  chunk = block_n + (offset // groups_per_chunk) * group_count
+  keys = tuple(chunk*CHUNK + (offset % groups_per_chunk)*decode_group + i for i in range(decode_group))
+  valid = tuple(key < valid_kv_len for key in keys)
+  kvals, vvals = (tuple(tuple(cache_kv[kv, b, kv_head, key, d].float() *
+    is_valid.where(cache_scale[kv, b, kv_head, key].float(), UOp.const(0, dtypes.float)) for d in dims)
+    for key,is_valid in zip(keys, valid)) for kv in range(2))
+  q_heads = tuple(kv_head*G + head_group*head_tile + wave*heads_per_wave + head for head in range(heads_per_wave))
+  updates:list[UOp] = []
+  for head,q_head in enumerate(q_heads):
+    scores = tuple(warp_reduce(sum((q[b, q_head, 0, d].float()*k for d,k in zip(dims, key_kvals)),
+                                   UOp.const(0, dtypes.float)), full_wave=True) / math.sqrt(D) for key_kvals in kvals)
+    prev_acc, prev_max, prev_sum = acc.after(offset)[head], row_max.after(offset)[head], row_sum.after(offset)[head]
+    new_max = functools.reduce(lambda a,vs:a.maximum(vs[0].where(vs[1], UOp.const(-math.inf, dtypes.float))), zip(valid, scores), prev_max)
+    alpha = ((prev_max-new_max)*LOG2E).exp2()
+    betas = tuple(is_valid.where(((score-new_max)*LOG2E).exp2(), UOp.const(0, dtypes.float)) for is_valid,score in zip(valid, scores))
+    updates += [acc[head].store(prev_acc*alpha + sum((UOp.stack(*value)*beta for value,beta in zip(vvals, betas)), acc[head].const_like(0))),
+                row_sum[head].store(prev_sum*alpha + sum(betas, UOp.const(0, dtypes.float))), row_max[head].store(new_max)]
+  update = UOp.group(*updates).end(offset)
+  acc, row_max, row_sum = acc.after(update), row_max.after(update), row_sum.after(update)
+  stores = [out[b, q_head, block_n, d].store(acc[head, i]) for head,q_head in enumerate(q_heads) for i,d in enumerate(dims)] + \
+    [stats[b, q_head.valid(lane.eq(0)), block_n, i].store(x[head])
+     for head,q_head in enumerate(q_heads) for i,x in enumerate((row_max, row_sum))]
+  return UOp.group(*stores).end(lane, wave, block_n, block_bhkv).sink(arg=KernelInfo(name="flash_decode_partial", opts_to_apply=()))
+
+def amd_flash_attention_decode(q:Tensor, cache_kv:Tensor, valid_kv_len:int|UOp, cache_scale:Tensor, max_kv_len:int) -> Tensor:
+  B, H, D = cache_kv.shape[1], q.shape[1], cache_kv.shape[4]
+  block_n = 128
+  chunks = min(64, max_kv_len // block_n)
+  partial = Tensor.empty(B, H, chunks, D, dtype="float32", device=q.device)
+  stats = Tensor.empty(B, H, chunks, 2, dtype="float32", device=q.device)
+  decode_partial = functools.partial(_amd_flash_attention_decode_partial, valid_kv_len=valid_kv_len, max_kv_len=max_kv_len, block_n=block_n)
+  partial, stats = Tensor.custom_kernel(partial, stats, q, cache_kv, cache_scale, fxn=decode_partial)[:2]
+  live_chunks = (valid_kv_len+block_n-1)//block_n
+  live_chunks = min(live_chunks, chunks) if isinstance(live_chunks, int) else live_chunks.minimum(chunks)
+  partial, stats = partial[:, :, :live_chunks], stats[:, :, :live_chunks]
+  weights = ((stats[..., 0]-stats[..., 0].max(2, keepdim=True))*LOG2E).exp2()
+  return ((partial*weights.unsqueeze(-1)).sum(2) / (stats[..., 1]*weights).sum(2, keepdim=True)).unsqueeze(2)
+
+@functools.cache
+def _amd_flash_attention(o:UOp, q:UOp, cache:UOp, kv_scale:UOp, valid_kv_len:int|UOp) -> UOp:
+  if isinstance(valid_kv_len, UOp): valid_kv_len = kernel_var(valid_kv_len.unbind_all()[0])
+  BH, M, D = q.shape
+  _, B, H_KV, physical_n, cache_dim = cache.shape
+  k, v = cache[0].reshape(B*H_KV, physical_n, cache_dim), cache[1].reshape(B*H_KV, physical_n, cache_dim)
+  kv_scale = kv_scale.reshape(2, B*H_KV, physical_n)
+  assert k.shape == v.shape and BH % k.shape[0] == 0 and k.shape[2] == D
+  gqa_group = BH // k.shape[0]
+  if isinstance(M, int) and isinstance(valid_kv_len, int):
+    assert M % BLOCK_M == 0 and valid_kv_len % BLOCK_N == 0
+  assert isinstance(D, int) and D % WMMA_K == 0 and D % LANES_PER_WAVE_N == 0
+  TM, TN, TD, SCALE = BLOCK_M//(WAVES_M*LANES_PER_WAVE_M), BLOCK_N//LANES_PER_WAVE_N, D//(WAVES_N*LANES_PER_WAVE_N), 1/math.sqrt(D)
+  block_bh, block_m = UOp.range(BH, 0, AxisType.GLOBAL), UOp.range(M // BLOCK_M, 1, AxisType.GLOBAL)
+  q = q.reshape(BH, M//BLOCK_M, BLOCK_M, D)[block_bh, block_m]
+  kv_head = block_bh // gqa_group
+  k, v = k[kv_head], v[kv_head]
+  o = o.reshape(BH, M//BLOCK_M, BLOCK_M, D)[block_bh, block_m]
+  wave_m, wave_n, lane = UOp.range(WAVES_M, 2, AxisType.LOCAL), UOp.range(WAVES_N, 3, AxisType.LOCAL), UOp.range(WARP_SIZE, -1, AxisType.WARP)
+  tid = (wave_m * WAVES_N + wave_n) * WARP_SIZE + lane
+  lane_m, lane_n = lane // LANES_PER_WAVE_N, lane % LANES_PER_WAVE_N
+  Q_ELEMS_PER_THREAD, KV_ELEMS_PER_THREAD = BLOCK_M * D // THREADS_PER_BLOCK, BLOCK_N * D // THREADS_PER_BLOCK
+  QP_lds = UOp.placeholder((BLOCK_M, D + LDS_PAD), dtypes.half, slot=0, addrspace=AddrSpace.LOCAL)
+  KV_lds = UOp.placeholder((BLOCK_N, D + LDS_PAD), dtypes.half, slot=1, addrspace=AddrSpace.LOCAL)[:, :D]
+  acc, m_i, l_i = _reg((TM, TD), 2, 0), _reg((TM,), 3, -math.inf), _reg((TM,), 4, 0)
+  n_tiles = (valid_kv_len - M + (block_m + 1) * BLOCK_M + BLOCK_N - 1) // BLOCK_N
+  n_tile = UOp.range(n_tiles, 100, AxisType.REDUCE)
+  Q_lds = QP_lds[:, :D]
+  Q_store = Q_lds.after(n_tile).reshape(THREADS_PER_BLOCK, Q_ELEMS_PER_THREAD)[tid].store(q.reshape(THREADS_PER_BLOCK, Q_ELEMS_PER_THREAD)[tid])
+  load_k = UOp.range(KV_ELEMS_PER_THREAD, 90, AxisType.WEAK)
+  kidx = n_tile*BLOCK_N*D + tid*KV_ELEMS_PER_THREAD + load_k
+  kval = k.reshape(physical_n*D)[kidx].float() * kv_scale[0, kv_head, kidx // D].float()
+  K_store = KV_lds.reshape(THREADS_PER_BLOCK, KV_ELEMS_PER_THREAD)[tid, load_k].store(kval).end(load_k)
+  qk_load_barrier = UOp.barrier(UOp.group(Q_store, K_store))
+  Q_lds, KV_lds_k = Q_lds.after(qk_load_barrier), KV_lds.after(qk_load_barrier)
+  S_reg = _reg((TM, TN), 6, 0, n_tile)
+  k_qk, tm1, tn1 = UOp.range(D//WMMA_K, 101, AxisType.REDUCE), UOp.range(TM//WMMA_ACC, 200), UOp.range(TN, 201)
+  S_frag = S_reg.reshape(TM // WMMA_ACC, WMMA_ACC, TN).permute(0, 2, 1)[tm1, tn1]
+  q_frag = Q_lds.reshape(WAVES_M, TM // WMMA_ACC, WMMA_M, D // WMMA_K, WMMA_K)[wave_m, tm1, lane_n, k_qk]
+  k_frag = KV_lds_k.reshape(TN, WMMA_N, D // WMMA_K, WMMA_K)[tn1, lane_n, k_qk]
+  qk_done = S_frag.store(UOp.wmma(q_frag, k_frag, S_frag.after(k_qk), *WMMA_ARG)).end(tm1, tn1).end(k_qk)
+  S_reg = S_reg.after(qk_done)
+  S_reg = S_reg.after(S_reg.store(S_reg * SCALE))
+  rm, rn = UOp.range(TM, 250, AxisType.WEAK), UOp.range(TN, 251, AxisType.WEAK)
+  q_idx = valid_kv_len - M + block_m * BLOCK_M + wave_m * WMMA_M + rm * LANES_PER_WAVE_M + lane_m
+  k_idx = n_tile * BLOCK_N + rn * LANES_PER_WAVE_N + lane_n
+  valid = k_idx <= q_idx
+  S_reg = S_reg.after(S_reg[rm, rn].store(valid.where(S_reg[rm, rn], S_reg[rm, rn].const_like(-math.inf))).end(rm, rn))
+  m_ij, rm2 = _reg((TM,), 7, -math.inf, n_tile), UOp.range(TN, 261, AxisType.REDUCE)
+  m_ij = m_ij.after(m_ij.store(m_ij.after(rm2).maximum(S_reg[:, rm2])).end(rm2))
+  ri_w = UOp.range(TM, 270)
+  m_ij = m_ij.after(m_ij[ri_w].store(warp_reduce(m_ij[ri_w], maximum=True)).end(ri_w))
+  tile_max = m_ij.reshape(TM, 1).expand(TM, TN).maximum(-1e30)
+  S_reg = S_reg.after(S_reg.store(((S_reg - tile_max) * LOG2E).exp2()))
+  p_local, ri_ws = _reg((TM,), 8, 0, n_tile), UOp.range(TM, 295, AxisType.WEAK)
+  p_sum = p_local.after(p_local[ri_ws].store(sum((warp_reduce(S_reg[ri_ws, rn]) for rn in range(TN)), S_reg.const_like(0))).end(ri_ws))
+  P_lds = QP_lds.flatten()[:WAVES_N * BLOCK_M * BLOCK_N].reshape(WAVES_N, BLOCK_M, BLOCK_N)
+  P_write = P_lds.reshape(WAVES_N, WAVES_M, TM, LANES_PER_WAVE_M, 1, TN, LANES_PER_WAVE_N, 1)
+  P_write = P_write.permute((1, 0, 3, 6, 2, 4, 5, 7)).reshape(THREADS_PER_BLOCK, TM, TN)
+  P_store = P_write[tid].store(S_reg.cast(dtypes.half))
+  beta_i, ri4 = UOp.placeholder((TM,), dtypes.float, slot=9, addrspace=AddrSpace.REG), UOp.range(TM, 330, AxisType.WEAK)
+  m_new_val = m_i[ri4].maximum(m_ij[ri4])
+  alpha_val = ((m_i[ri4] - m_new_val) * LOG2E).exp2()
+  beta_val = ((m_ij[ri4] - m_new_val) * LOG2E).exp2()
+  rj4 = UOp.range(TD, 331)
+  correction = UOp.group(acc[ri4, rj4].store(alpha_val * acc[ri4, rj4]).end(rj4),
+                         l_i[ri4].store(alpha_val * l_i[ri4] + beta_val * p_sum[ri4]),
+                         m_i[ri4].store(m_new_val), beta_i[ri4].store(beta_val)).end(ri4)
+  acc, l_i, m_i, beta_i = acc.after(correction), l_i.after(correction), m_i.after(correction), beta_i.after(correction)
+  V_lds = UOp.placeholder((D, BLOCK_N + LDS_PAD), dtypes.half, slot=1, addrspace=AddrSpace.LOCAL)[:, :BLOCK_N]
+  V_copy, load_v = V_lds.after(qk_done).permute(1, 0), UOp.range(KV_ELEMS_PER_THREAD, 390, AxisType.WEAK)
+  vidx = n_tile*BLOCK_N*D + tid*KV_ELEMS_PER_THREAD + load_v
+  vval = v.reshape(physical_n*D)[vidx].float() * kv_scale[1, kv_head, vidx // D].float()
+  V_store = V_copy.reshape(THREADS_PER_BLOCK, KV_ELEMS_PER_THREAD)[tid, load_v].store(vval).end(load_v)
+  pv_barrier = UOp.barrier(UOp.group(P_store, V_store))
+  P_lds, V_lds = P_lds.after(pv_barrier), V_lds.after(pv_barrier)
+  pv_acc = _reg((TM, TD), 10, 0, n_tile).after(pv_barrier)
+  k_pv, tm2, tn2 = UOp.range(BLOCK_N//WMMA_K, 400, AxisType.REDUCE), UOp.range(TM//WMMA_ACC, 401, AxisType.WEAK), UOp.range(TD, 402, AxisType.WEAK)
+  pv_frag = pv_acc.reshape(TM // WMMA_ACC, WMMA_ACC, TD).permute(0, 2, 1)[tm2, tn2]
+  p_frag = P_lds[wave_n].reshape(WAVES_M, TM // WMMA_ACC, WMMA_M, BLOCK_N // WMMA_K, WMMA_K)[wave_m, tm2, lane_n, k_pv]
+  v_frag = V_lds.reshape(WAVES_N, TD, WMMA_N, BLOCK_N // WMMA_K, WMMA_K)[wave_n, tn2, lane_n, k_pv]
+  pv_done = pv_frag.store(UOp.wmma(p_frag, v_frag, pv_frag.after(k_pv), *WMMA_ARG)).end(tm2, tn2).end(k_pv)
+  pv_acc = pv_acc.after(pv_done)
+  ri5, rj5 = UOp.range(TM, 410, AxisType.WEAK), UOp.range(TD, 411, AxisType.WEAK)
+  accumulate = acc[ri5, rj5].store(acc[ri5, rj5] + beta_i[ri5] * pv_acc[ri5, rj5]).end(ri5, rj5)
+  n_tile_end = accumulate.barrier().end(n_tile)
+  acc, l_i, m_i = acc.after(n_tile_end), l_i.after(n_tile_end), m_i.after(n_tile_end)
+  acc = acc.after(acc.store(acc * (1 / l_i).reshape(TM, 1).expand(TM, TD)))
+  o = o.reshape(WAVES_M, TM, LANES_PER_WAVE_M, 1, WAVES_N, TD, LANES_PER_WAVE_N, 1)
+  o = o.permute((0, 4, 2, 6, 1, 3, 5, 7)).reshape(THREADS_PER_BLOCK, TM, TD)
+  return o[tid].store(acc).end(wave_m, wave_n, lane).end(block_m, block_bh).sink(arg=KernelInfo(opts_to_apply=()))
+
+def flash_attention_causal_cached(q:Tensor, cache_kv:Tensor, valid_kv_len:int|UOp, cache_scale:Tensor) -> Tensor:
+  B, H, T, D = cast(tuple[int, int, int, int], q.shape)
+  out = Tensor.empty(B*H, T, D, dtype="float32", device=q.device)
+  flash_cached = functools.partial(_amd_flash_attention, valid_kv_len=valid_kv_len)
+  return Tensor.custom_kernel(out, q.reshape(B*H, T, D), cache_kv, cache_scale, fxn=flash_cached)[0].reshape(B, H, T, D)
+
+def quantized_attention(q:Tensor, stacked_kv:Tensor, cache_kv:Tensor, cache_scale:Tensor, start_pos:int|UOp) -> Tensor:
+  T = q.shape[2]
+  scale = (stacked_kv.float().abs().max(axis=-1, keepdim=True) / 127).maximum(1e-8).half()
+  packed_kv = (stacked_kv.float() / scale).round().clip(-127, 127).cast(dtypes.int8)
+  store_kv = cache_kv[:, :, :, start_pos:start_pos+T, :].uop.store(packed_kv.uop)
+  store_scale = cache_scale[:, :, :, start_pos:start_pos+T].uop.store(scale.squeeze(-1).uop)
+  # each store goes on its own buffer's AFTER: sharing both stores across both AFTERs leaves
+  # un-ended stores with open ranges in the kernel graph
+  assigned_kv, assigned_scale = Tensor(cache_kv.uop.after(store_kv)), Tensor(cache_scale.uop.after(store_scale))
+  # keep start_pos in its bound form at the graph level, the kernel builders unbind it to the kernel-side PARAM form
+  valid_end = start_pos+T
+  return amd_flash_attention_decode(q.half(), assigned_kv, valid_end, assigned_scale, cast(int, cache_kv.shape[3])) if resolve(T == 1) else \
+    flash_attention_causal_cached(q.half(), assigned_kv, valid_end, assigned_scale)
+
+def _amd_dp4a(a:UOp, b:UOp, c:UOp) -> UOp:
+  return UOp(Ops.CUSTOMI, dtypes.int32, (a.int(), b.int(), c), arg="__builtin_amdgcn_sudot4(true, {}, true, {}, {}, false)")
+
+def _amd_byte_perm(a:UOp, b:UOp, selectors:UOp) -> UOp:
+  return UOp(Ops.CUSTOMI, dtypes.uint32, tuple(x.cast(dtypes.uint32) for x in (a, b, selectors)), arg="__builtin_amdgcn_perm({}, {}, {})")
+
+def _amd_load(ptr:UOp, lanes:int|None=None) -> UOp:
+  assert ptr.op is Ops.INDEX
+  if lanes is None: return UOp(Ops.CUSTOMI, ptr.dtype, (ptr,), arg="__builtin_nontemporal_load({0})")
+  buf, coords = ptr.src[0], ptr.src[1:]
+  idx = sum((coord*math.prod(buf.shape[i+1:]) for i,coord in enumerate(coords)), UOp.const(0, dtypes.weakint))
+  return UOp(Ops.SHRINK, src=(buf.flatten(), idx, UOp.const(lanes, dtypes.weakint))).load(dtype=ptr.dtype)
+
+def _load_byte(raw:UOp, base:UOp, offset:UOp) -> UOp: return (raw[base + offset//4] >> ((offset&3)*8).cast(dtypes.uint32)) & 255
+def _half(value:UOp) -> UOp: return value.cast(dtypes.uint16).bitcast(dtypes.float16).float()
+
+def _iq4_bytes(packed:UOp, shift:int) -> UOp:
+  selectors = (packed >> shift) & 0x0f0f0f0f
+  low = _amd_byte_perm(UOp.const(0xf6eaddcf, dtypes.uint32), UOp.const(0xbfad9881, dtypes.uint32), selectors)
+  high = _amd_byte_perm(UOp.const(0x71594535, dtypes.uint32), UOp.const(0x26190d01, dtypes.uint32), selectors & 0x07070707)
+  return _amd_byte_perm(high, low, 0x03020100 | ((selectors & 0x08080808) >> 1))
+
+@functools.cache
+def _q8_quantize_kernel(q:UOp, scale:UOp, x:UOp, tokens:int, in_features:int) -> UOp:
+  groups = in_features//Q8_GROUP_SIZE
+  token_group, lane = UOp.range(tokens*groups, 0), UOp.range(32, 1, axis_type=AxisType.LOCAL)
+  token, group = token_group//groups, token_group%groups
+  x = x.reshape(tokens, groups, 32)
+  group_scale = warp_reduce(x[token, group, lane].float().abs(), maximum=True, full_wave=True) / 127
+  group_scale = group_scale.maximum(1e-8)
+  word_lane = lane.minimum(7)
+  xs = tuple(x[token, group, word_lane*4+i].float() for i in range(4))
+  word = sum(((v/group_scale).round().clip(-127, 127).cast(dtypes.int8).cast(dtypes.uint8).cast(dtypes.uint32) << (i*8)
+              for i,v in enumerate(xs)), UOp.const(0, dtypes.uint32))
+  stores = (q[token, group, lane.valid(lane < 8)].store(word), scale[token, group.valid(lane.eq(0))].store(group_scale))
+  return UOp.group(*stores).end(token_group, lane).sink(arg=KernelInfo(name="q8_quantize", opts_to_apply=()))
+
+def q8_quantize(x:Tensor, tokens:int, in_features:int) -> tuple[Tensor, Tensor]:
+  groups = in_features//Q8_GROUP_SIZE
+  q, scale = Tensor.empty(tokens, groups, 8, dtype=dtypes.uint32, device=x.device), \
+    Tensor.empty(tokens, groups, dtype=dtypes.float32, device=x.device)
+  q, scale = Tensor.custom_kernel(q, scale, x, fxn=functools.partial(_q8_quantize_kernel, tokens=tokens, in_features=in_features))[:2]
+  return q, scale
+
+@functools.cache
+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
+  batch, heads, tokens, value_dim, key_dim = cast(tuple[int, int, int, int, int], (batch, heads, tokens, value_dim, key_dim))
+  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, lane = UOp.range(batch*heads*value_dim//row_tile, 0), UOp.range(32, 1, axis_type=AxisType.LOCAL)
+  bh, row_base = bh_row // (value_dim//row_tile), (bh_row % (value_dim//row_tile))*row_tile
+  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)
+  initial = None if start_pos is None else start_pos.eq(0)
+  current = current.after(current.store(UOp.stack(*(state[bh, row, col].float() if initial is None else
+    initial.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)
+  updates:list[UOp] = []
+  stores:list[UOp] = []
+  for row_idx,row in enumerate(rows):
+    previous = tuple(current.after(token)[row_idx*key_dim//32+i].load() for i in range(key_dim//32))
+    av, bv = alpha[bh, token, row if alpha_dim > 1 else 0].load(), beta[bh, token].load()
+    state_k = warp_reduce(sum((x*y for x,y in zip(previous, keys)), UOp.const(0, dtypes.float32)), full_wave=True)
+    state_q = warp_reduce(sum((x*y for x,y in zip(previous, queries)), UOp.const(0, dtypes.float32)), full_wave=True)
+    delta = (v[bh, token, row].load() - state_k*av) * bv
+    updates += [x*av + delta*y for x,y in zip(previous, keys)]
+    stores.append(core[bh, token, row.valid(lane.eq(0))].store(state_q*av + delta*kq[bh, token]))
+  step = UOp.group(*stores, current.store(UOp.stack(*updates))).end(token)
+  state_stores = (state[bh, row, col].store(current.after(step)[row_idx*key_dim//32+i].load().cast(state.dtype))
+                  for row_idx,row in enumerate(rows) for i,col in enumerate(cols))
+  return UOp.group(*state_stores).end(lane, bh_row).sink(arg=KernelInfo(name="gated_delta_prefill", opts_to_apply=()))
+
+def _wmma_layout(out:UOp, out_features:int, token_tile:int, output_tiles:int):
+  output_waves = 2 if out_features % (32*output_tiles) == 0 else 1
+  token_block, output_block = UOp.range(out.shape[0]//token_tile, 0), UOp.range(out_features//(16*output_tiles*output_waves), 1)
+  lane, wave = UOp.range(WARP_SIZE, 2, axis_type=AxisType.LOCAL), UOp.range(output_waves, 3, axis_type=AxisType.LOCAL)
+  hw_lane = UOp(Ops.CUSTOM, dtypes.int32, (lane.int(),), arg="__builtin_amdgcn_mbcnt_lo(-1, 0)").cast(dtypes.weakint)
+  col, half = hw_lane % 16, hw_lane // 16
+  outputs = tuple((output_block*output_waves+wave)*(16*output_tiles) + tile*16 + col for tile in range(output_tiles))
+  inputs = tuple(token_block*token_tile + tile*16 + col for tile in range(token_tile//16))
+  tokens = tuple(tuple(token_block*token_tile + tile*16 + half*8 + i for i in range(8)) for tile in range(token_tile//16))
+  return output_waves, token_block, output_block, lane, wave, half, outputs, inputs, tokens
+
+def _wmma_stores(out, outputs, tokens, accs, update, half):
+  def values(acc:UOp) -> tuple[UOp, ...]:
+    vals = tuple(acc.after(update)[i].load() for i in range(8))
+    swapped = tuple(UOp(Ops.CUSTOM, dtypes.float32, (value,),
+      arg="__builtin_bit_cast(float, __builtin_amdgcn_ds_swizzle(__builtin_bit_cast(int, {0}), 50688))") for value in vals)
+    low = half.eq(0)
+    return tuple(low.where(vals[i], swapped[i+4]) if j == 0 else low.where(swapped[i], vals[i+4]) for i in range(4) for j in range(2))
+  return [out[token, output].store(value) for output,output_accs in zip(outputs, accs)
+          for tile_tokens,acc in zip(tokens, output_accs) for token,value in zip(tile_tokens, values(acc))]
+
+def _decode_linear(out:UOp, out_features:int, group_count:int, group_dot, name:str) -> UOp:
+  chunks = (group_count+31)//32
+  token_output_chunk, lane = UOp.range(out.shape[0]*out_features*chunks, 0), UOp.range(32, 1, axis_type=AxisType.LOCAL)
+  token, output, chunk = token_output_chunk // (out_features*chunks), (token_output_chunk//chunks) % out_features, token_output_chunk % chunks
+  group = lane+chunk*32
+  value = group_dot(token, output, group) if group_count % 32 == 0 else \
+    (group < group_count).where(group_dot(token, output, group.minimum(group_count-1)), UOp.const(0, dtypes.float32))
+  total = warp_reduce(value, full_wave=True)
+  return out[token, output, chunk.valid(lane.eq(0))].store(total.cast(out.dtype)).end(token_output_chunk, lane).sink(
+    arg=KernelInfo(name=name, opts_to_apply=()))
+
+def _q5_scales(raw:UOp, base:UOp, subgroup:UOp) -> tuple[UOp, UOp, UOp, UOp]:
+  scale = (subgroup < 4).where(_load_byte(raw, base, 4 + subgroup) & 63,
+    (_load_byte(raw, base, 8 + subgroup) & 15) | ((_load_byte(raw, base, subgroup) >> 6) << 4))
+  minimum = (subgroup < 4).where(_load_byte(raw, base, 8 + subgroup) & 63,
+    (_load_byte(raw, base, 8 + subgroup) >> 4) | ((_load_byte(raw, base, 4 + subgroup) >> 6) << 4))
+  d, dmin = (raw[base] & 0xffff).cast(dtypes.uint16), (raw[base] >> 16).cast(dtypes.uint16)
+  return _half(d), _half(dmin), scale.float(), minimum.float()
+
+@functools.cache
+def _quant_decode_kernel(out:UOp, raw:UOp, xq:UOp, xd:UOp, out_features:int, in_features:int, ggml_type:int) -> UOp:
+  group_count = in_features // Q8_GROUP_SIZE
+  def group_dot(token:UOp, output:UOp, group:UOp) -> UOp:
+    block, subgroup = group // 8, group % 8
+    xwords = _amd_load(xq[token, group, 0], 8)
+    if ggml_type == Q5_K:
+      base = (output * in_features//GGML_BLOCK_SIZE + block) * Q5_WORDS
+      qs_base, dot, qsum = base + 12 + (subgroup//2)*8, UOp.const(0, dtypes.int32), UOp.const(0, dtypes.int32)
+      for word_idx in range(8):
+        word = (raw[qs_base+word_idx] >> ((subgroup&1)*4).cast(dtypes.uint32)) & 0x0f0f0f0f
+        word |= ((raw[base+4+word_idx] >> subgroup.cast(dtypes.uint32)) & 0x01010101) << 4
+        dot, qsum = _amd_dp4a(word, xwords[word_idx], dot), _amd_dp4a(UOp.const(0x01010101, dtypes.uint32), xwords[word_idx], qsum)
+      d, dmin, scale, minimum = _q5_scales(raw, base, subgroup)
+      return (dot.float()*d*scale - qsum.float()*dmin*minimum) * xd[token, group]
+    if ggml_type == IQ4_XS:
+      base = (output * in_features//GGML_BLOCK_SIZE + block) * IQ4_WORDS
+      dot = UOp.const(0, dtypes.int32)
+      for word_idx in range(8):
+        packed = _amd_load(raw[base + 2 + subgroup*4 + word_idx%4])
+        dot = _amd_dp4a(_iq4_bytes(packed, 4*(word_idx//4)), xwords[word_idx], dot)
+      d, scale = _iq4_scales(raw, base, subgroup)
+      return dot.float() * xd[token, group] * d * scale
+    base = (output*in_features//GGML_BLOCK_SIZE+block)*Q6_BYTES
+    dots = [UOp.const(0, dtypes.int32), UOp.const(0, dtypes.int32)]
+    for word_idx in range(8):
+      pos, within = subgroup*32 + word_idx*4, (subgroup*32 + word_idx*4)%128
+      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 = 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)
+    return (dots[0].float()*scales[0] + dots[1].float()*scales[1]) * xd[token, group] * _half(dbits)
+  return _decode_linear(out, out_features, group_count, group_dot, {Q5_K:"linear_q5_k", IQ4_XS:"linear_iq4_xs", Q6_K:"linear_q6"}[ggml_type])
+
+def _quant_linear_wmma(out, x, out_features, in_features, type_words, layout, dequant, name):
+  x = x.reshape(out.shape[0], in_features)
+  _, token_block, output_block, lane, wave, physical_half, outputs, input_tokens, tokens = layout
+  token_tile, output_tiles = len(tokens)*16, len(outputs)
+  output_words = in_features // GGML_BLOCK_SIZE * type_words
+  accs = tuple(tuple(UOp.placeholder((8,), dtypes.float32, slot=ot*(token_tile//16)+tile, addrspace=AddrSpace.REG)
+                     for tile in range(token_tile // 16)) for ot in range(output_tiles))
+  accs = tuple(tuple(acc.after(acc.store(acc.const_like(0))) for acc in output_accs) for output_accs in accs)
+  group = UOp.range(in_features // Q8_GROUP_SIZE, 4, AxisType.REDUCE)
+  block, subgroup = group // 8, group % 8
+  wmma_accs = [list(output_accs) for output_accs in accs]
+  for half in range(2):
+    afrags = tuple(UOp.stack(*(x[input_token, group*32 + half*16 + i].cast(dtypes.float16) for i in range(16)))
+                   for input_token in input_tokens)
+    for output_tile,output in enumerate(outputs):
+      bfrag = UOp.stack(*dequant(output*output_words + block*type_words, subgroup, half))
+      for tile,afrag in enumerate(afrags):
+        previous = accs[output_tile][tile].after(group) if half == 0 else wmma_accs[output_tile][tile]
+        wmma_accs[output_tile][tile] = UOp.wmma(afrag, bfrag, previous, *WMMA_ARG)
+  update = UOp.group(*(acc.store(value) for output_accs,output_values in zip(accs, wmma_accs)
+                       for acc,value in zip(output_accs, output_values))).end(group)
+  return UOp.group(*_wmma_stores(out, outputs, tokens, accs, update, physical_half)).end(token_block, output_block, lane, wave).sink(
+    arg=KernelInfo(name=name, opts_to_apply=()))
+
+@functools.cache
+def _q5_linear_f16_wmma_kernel(out:UOp, raw:UOp, x:UOp, out_features:int, in_features:int) -> UOp:
+  token_tile, output_tiles = (64, 1) if out_features <= 1024 and out.shape[0] % 64 == 0 else \
+    (64, 2) if out.shape[0] % 64 == 0 else (32 if out.shape[0] % 32 == 0 else 16, 2)
+  def dequant(base:UOp, subgroup:UOp, half:int) -> tuple[UOp, ...]:
+    d, dmin, scale, minimum = _q5_scales(raw, base, subgroup)
+    qs_base = base + 12 + (subgroup // 2)*8 + half*4
+    words = tuple((raw[qs_base+i] >> ((subgroup&1)*4).cast(dtypes.uint32) & 0x0f0f0f0f) |
+      ((raw[base+4+half*4+i] >> subgroup.cast(dtypes.uint32) & 0x01010101) << 4) for i in range(4))
+    return tuple(((word >> (byte*8) & 255).float()*d*scale-dmin*minimum).cast(dtypes.float16) for word in words for byte in range(4))
+  return _quant_linear_wmma(out, x, out_features, in_features, Q5_WORDS,
+                            _wmma_layout(out, out_features, token_tile, output_tiles), dequant, "linear_q5_k_f16_wmma")
+
+def _iq4_scales(raw:UOp, base:UOp, subgroup:UOp) -> tuple[UOp, UOp]:
+  low = _load_byte(raw, base, 4 + subgroup//2)
+  scale = ((low >> (4*(subgroup%2)).cast(dtypes.uint32)) & 15) | ((((raw[base] >> 16) >> (2*subgroup).cast(dtypes.uint32)) & 3) << 4)
+  return _half(raw[base] & 0xffff), (scale.cast(dtypes.uint8).bitcast(dtypes.int8)-32).float()
+
+@functools.cache
+def _iq4_linear_f16_wmma_kernel(out:UOp, raw:UOp, x:UOp, lut:UOp, out_features:int, in_features:int) -> UOp:
+  token_tile = 32 if out_features <= 1024 and out.shape[0] % 32 == 0 else 64 if out.shape[0] % 64 == 0 and \
+    (out_features <= 6144 or out_features == 5120 and in_features > 8192) else 128 if out.shape[0] % 128 == 0 else \
+    32 if out.shape[0] % 32 == 0 else 16
+  output_tiles = 1 if out_features <= 1024 else 2 if out_features <= 6144 else 1 if out_features < 8192 else 2
+  layout = _wmma_layout(out, out_features, token_tile, output_tiles)
+  output_waves, _, _, lane, wave, _, _, _, _ = layout
+  local_lut = UOp.placeholder((256,), dtypes.uint32, slot=32, addrspace=AddrSpace.LOCAL)
+  tid, lut_items = wave*32+lane, 256//(32*output_waves)
+  lut = local_lut.after(UOp.group(*(local_lut[tid*lut_items+i].store(lut[tid*lut_items+i]) for i in range(lut_items))).barrier())
+  def dequant(base:UOp, subgroup:UOp, half:int) -> tuple[UOp, ...]:
+    d, scale = _iq4_scales(raw, base, subgroup)
+    scale = scale * d
+    if out_features <= 6144:
+      pairs = tuple(lut[((raw[base + 2 + subgroup*4 + word] >> (byte*8)) & 255).cast(dtypes.weakint)]
+                    for word in range(4) for byte in range(4))
+      return tuple((_half((pair >> (half*16)) & 0xffff)*scale).cast(dtypes.float16) for pair in pairs)
+    def nibble(packed:UOp, index:int): return (packed >> (8*index+4*half)) & 15
+    lut_pairs = (lut[(nibble(packed, i) | nibble(packed, i+1)<<4).cast(dtypes.weakint)]
+                 for packed in (raw[base+2+subgroup*4+i] for i in range(4)) for i in (0, 2))
+    return tuple((_half((pair >> (i*16)) & 0xffff)*scale).cast(dtypes.float16) for pair in lut_pairs for i in range(2))
+  return _quant_linear_wmma(out, x, out_features, in_features, IQ4_WORDS, layout, dequant, "linear_iq4_xs_f16_wmma")
+
+def q8_linear(layer:Linear, x:Tensor) -> Tensor:
+  assert layer.ggml_type in (Q5_K, Q6_K, IQ4_XS)
+  tokens = int(x.numel()) // layer.in_features
+  raw = layer.weight.uop.buf_uop
+  out_features, in_features = layer.out_features, layer.in_features
+  use_wmma = tokens % 16 == 0 and layer.out_features % 16 == 0
+  def run(fxn:Callable[..., UOp], out:UOp, *srcs:UOp) -> Tensor:
+    all_srcs = (out,)+srcs
+    params = tuple(UOp.placeholder_like(src, slot=i) for i,src in enumerate(all_srcs))
+    kernel = fxn(*params, out_features=out_features, in_features=in_features).call(*all_srcs)
+    result = Tensor(out.after(kernel))
+    if len(result.shape) == 3: result = result.sum(-1)
+    result = result.reshape(*x.shape[:-1], layer.out_features)
+    return result if layer.bias is None else result + layer.bias
+
+  out = Tensor.empty(tokens, layer.out_features, dtype=dtypes.float32, device=x.device).uop
+  if layer.ggml_type == Q5_K and use_wmma:
+    return run(_q5_linear_f16_wmma_kernel, out, raw, x.cast(dtypes.float16).contiguous().uop)
+  if layer.ggml_type == IQ4_XS and use_wmma:
+    return run(_iq4_linear_f16_wmma_kernel, out, raw, x.cast(dtypes.float16).contiguous().uop,
+               iq4_half_lut(str(x.device)).uop)
+  xq, xd = q8_quantize(x, tokens, layer.in_features)
+  decode = functools.partial(_quant_decode_kernel, ggml_type=layer.ggml_type)
+  out = Tensor.empty(tokens, layer.out_features, (layer.in_features+1023)//1024, dtype=dtypes.float32, device=x.device).uop
+  return run(decode, out, raw, xq.uop, xd.uop)
+
+@functools.cache
+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()

tinygrad/llm/model.py

@@ -1,8 +1,8 @@
 from __future__ import annotations
 import functools, itertools, pathlib
 from dataclasses import dataclass, replace
-from tinygrad import Tensor, nn, UOp, TinyJit, getenv, function, dtypes
-from tinygrad.nn import Linear
+from tinygrad import Tensor, nn, UOp, TinyJit, getenv, function, Context, dtypes
+from tinygrad.llm.kernels import Linear, gated_delta_prefill, amd_custom_kernels_supported
 from tinygrad.llm.gguf import gguf_load
 from tinygrad.uop.ops import resolve
 
@@ -124,8 +124,6 @@ class FFNBlock:
 
   # given the token-prefix match, return how much cached state this block can still reuse
   def _reusable_prefix_len(self, prefix_len:int, cached_len:int) -> int: return prefix_len
-  # return writes that reset this block's state after a cache mismatch
-  def _state_reset_ops(self) -> list[Tensor]: return []
   def _init_state(self, x:Tensor): raise NotImplementedError
   def _attention(self, x:Tensor, start_pos:int|UOp) -> Tensor: raise NotImplementedError
 
@@ -169,26 +167,38 @@ class TransformerBlock(FFNBlock):
     k = apply_rope(k[..., :self.config.rope_dim], self.freqs_cis[start_pos:start_pos+T]).cat(k[..., self.config.rope_dim:], dim=-1)
 
     # NOTE: we don't want to change self.cache_kv, the function API doesn't support this well
-    assigned_kv = Tensor(self.cache_kv.uop.after(self.cache_kv[:, :, :, start_pos:start_pos+T, :].uop.store(Tensor.stack(k, v).uop)))
-    k = assigned_kv[0, :, :, 0:start_pos+T, :]
-    v = assigned_kv[1, :, :, 0:start_pos+T, :]
+    stacked_kv = Tensor.stack(k, v)
+    if hasattr(self, "cache_kv_scale"):
+      from tinygrad.llm.kernels.amd import quantized_attention
+      attn = quantized_attention(q, stacked_kv, self.cache_kv, self.cache_kv_scale, start_pos)
+    else:
+      assigned_kv = Tensor(self.cache_kv.uop.after(self.cache_kv[:, :, :, start_pos:start_pos+T, :].uop.store(stacked_kv.uop)))
+      k = assigned_kv[0, :, :, 0:start_pos+T, :]
+      v = assigned_kv[1, :, :, 0:start_pos+T, :]
 
-    #self.cache_kv[:, :, :, start_pos:start_pos+T, :].assign(Tensor.stack(k, v))
-    #k = self.cache_kv[0, :, :, 0:start_pos+T, :]
-    #v = self.cache_kv[1, :, :, 0:start_pos+T, :]
+      #self.cache_kv[:, :, :, start_pos:start_pos+T, :].assign(stacked_kv)
+      #k = self.cache_kv[0, :, :, 0:start_pos+T, :]
+      #v = self.cache_kv[1, :, :, 0:start_pos+T, :]
 
-    # NOTE: this mask is causal_lower_right, not the causal_upper_left generated by is_casual = True
-    # TODO: this if statement should be removed and it shouldn't generate extra kernels
-    mask = Tensor.full((1, 1, T, start_pos+T), float("-inf"), dtype=x.dtype, buffer=False).triu(start_pos+1) \
-      if resolve(T != 1) else None
-    attn = q.scaled_dot_product_attention(k, v, attn_mask=mask, enable_gqa=True)     # (B,H,T,Hd)
+      # NOTE: this mask is causal_lower_right, not the causal_upper_left generated by is_casual = True
+      # TODO: this if statement should be removed and it shouldn't generate extra kernels
+      mask = Tensor.full((1, 1, T, start_pos+T), float("-inf"), dtype=x.dtype, buffer=False).triu(start_pos+1) \
+        if resolve(T != 1) else None
+      attn = q.scaled_dot_product_attention(k, v, attn_mask=mask, enable_gqa=True)     # (B,H,T,Hd)
     attn = attn.transpose(1, 2).reshape(B, T, -1)                                    # back to (B,T,D)
     return self.attn_output(attn if not self.config.attn_output_gate else (attn * gate.sigmoid()))
 
   def _init_state(self, x:Tensor):
     if not hasattr(self, "cache_kv"):
+      # hybrid models use a quantized KV cache on AMD, sized in flash decode blocks of 256
+      quantize = amd_custom_kernels_supported(x.device) and self.config.ssm is not None
+      assert not quantize or self.config.max_context % 256 == 0, \
+        f"quantized KV cache needs max_context to be a multiple of 256, got {self.config.max_context}"
       self.cache_kv = Tensor.empty(2, x.shape[0], self.config.n_kv_heads, self.config.max_context, self.config.head_dim,
-                                   dtype=dtypes.default_float, device=x.device)
+                                   dtype=dtypes.int8 if quantize else dtypes.default_float, device=x.device)
+      if quantize:
+        self.cache_kv_scale = Tensor.zeros(2, x.shape[0], self.config.n_kv_heads, self.config.max_context,
+                                           dtype=dtypes.float16, device=x.device)
       self.freqs_cis = precompute_freqs_cis(self.config.rope_dim, self.config.max_context, self.config.rope_theta, device=x.device)
 
 class MLATransformerBlock(FFNBlock):
@@ -260,50 +270,45 @@ 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"
+    start_pos = start_pos if isinstance(start_pos, UOp) else UOp.variable("start_pos", 0, self.config.max_context-1).bind(start_pos)
     is_kda = hasattr(self, "ssm_g_a")
 
     # input processing
     x = x.half()
     out_gate = self.ssm_g_b(self.ssm_g_a(x)) if is_kda 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)
+    out_gate = out_gate.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, self.num_v_heads, -1) *
-             self.ssm_a.reshape(1, self.num_v_heads, -1)).exp().unsqueeze(-2)
+    log_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()
+    # qkv conv, conv_state is reset when starting from position 0
+    conv_state = Tensor(start_pos).eq(0).where(0, 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)
-    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 = log_alpha.transpose(1, 2).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, the conv and recurrent states are updated in place
+    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)))
+    core = gated_delta_prefill(q * self.head_k_dim**-0.5, k, v, beta, alpha, state, Tensor(start_pos)).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))
+    core_attn_out = self.ssm_norm(core)
     out_gate = out_gate.sigmoid() if is_kda else out_gate.silu()
-    return self.ssm_out((core_attn_out * out_gate).reshape(B, 1, -1).cast(x.dtype))
-
-  # recurrent state can't be partially reused after divergence, force a full rebuild
-  def _state_reset_ops(self):
-    return [self.conv_state.assign(self.conv_state.const_like(0)),
-            self.recurrent_state.assign(self.recurrent_state.const_like(0))] if hasattr(self, "conv_state") else []
+    return self.ssm_out((core_attn_out * out_gate).reshape(B, T, -1).cast(x.dtype)).contiguous()
 
   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)
+      self.recurrent_state = Tensor.empty(x.shape[0], self.num_v_heads, self.head_v_dim, self.head_k_dim, device=x.device)
 
 class Transformer:
   def __init__(self, config:TransformerConfig):
@@ -417,9 +422,6 @@ class Transformer:
       Tensor.realize(*params)
     return model, kv
 
-  def warmup(self):
-    for _ in range(2): list(zip(range(2), self.generate([0])))
-
   def get_start_pos(self, tokens:list[int]) -> int:
     # recurrent state can't be partially reused after divergence: reuse it only when tokens extend the cached prefix
     if self.has_recurrent_block:
@@ -428,8 +430,19 @@ class Transformer:
     prefix_len = sum(1 for _ in itertools.takewhile(lambda ab: ab[0] == ab[1], zip(tokens[:-1], self._cached_tokens)))
     return min(block._reusable_prefix_len(prefix_len, len(self._cached_tokens)) for block in self.blk)
 
+  def warmup(self, chunk_size:int=32):
+    prompt = [0] * (min(chunk_size, 256, self.max_context-1) if self.has_recurrent_block else 1)
+    if self.has_recurrent_block:
+      x = Tensor.empty(1, 1, self.blk[0].config.dim, device=self.token_embd.weight.device)
+      for block in self.blk: block._init_state(x)
+    for _ in range(2):
+      # NOTE: chunk_size must match what generate uses at serve time, otherwise the captured JIT rejects the new toks range
+      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)
+      self._cached_tokens = []
+
   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
@@ -438,11 +451,12 @@ class Transformer:
     t = Tensor(tokens + [0] * (self.max_context - len(tokens)), dtype="int32").reshape(1, self.max_context)
     # recompute start_pos from what's currently valid in the caches
     start_pos = self.get_start_pos(tokens)
-    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)
-      sp, nt = v_start_pos.bind(start_pos), v_toks.bind(n_toks)
+      # recurrent blocks prefill full chunks with a static shape, the tail of the prompt goes through the decode graph
+      remaining = len(tokens)-start_pos
+      n_toks = 1 if self.has_recurrent_block and remaining < chunk_size else min(chunk_size, remaining)
+      sp, nt = v_start_pos.bind(start_pos), n_toks if self.has_recurrent_block else 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()
       start_pos += n_toks
       # chunked prefill: keep processing until all prompt tokens are consumed