llm: fast packed Qwen3.6 inference on AMD

test/backend/test_amd_flash_attention.py

@@ -0,0 +1,66 @@
+import unittest
+import numpy as np
+
+from tinygrad import Device, Tensor, TinyJit
+from tinygrad.llm.kernels.amd import amd_flash_attention_decode, flash_attention_causal_cached
+
+
+@unittest.skipUnless(Device.DEFAULT.startswith("AMD"), "AMD flash attention required")
+class TestAMDFlashAttention(unittest.TestCase):
+  def _test_decode(self, max_kv_len:int, valid_kv_len:int, n_heads:int=16, n_kv_heads:int=2, quantized:bool=False):
+    rng = np.random.default_rng(1)
+    q_np = rng.standard_normal((1, n_heads, 1, 256)).astype(np.float16)
+    kv_np = rng.standard_normal((2, 1, n_kv_heads, max_kv_len, 256)).astype(np.float16)
+    scale_np = np.maximum(np.max(np.abs(kv_np.astype(np.float32)), axis=-1), 1e-8) / 127
+    if quantized:
+      kv_np = np.clip(np.rint(kv_np.astype(np.float32) / scale_np[..., None]), -127, 127).astype(np.int8)
+    q, kv = Tensor(q_np).realize(), Tensor(kv_np).realize()
+    scale = Tensor(scale_np.astype(np.float16)).realize() if quantized else None
+
+    @TinyJit
+    def decode(q:Tensor, kv:Tensor): return amd_flash_attention_decode(q, kv, valid_kv_len, max_kv_len, scale).realize()
+
+    out = None
+    for _ in range(3): out = decode(q, kv).numpy()
+    assert out is not None
+    q_ref = q_np[0, :, 0].astype(np.float32)
+    kv_ref = kv_np.astype(np.float32) * scale_np[..., None] if quantized else kv_np.astype(np.float32)
+    k_ref, v_ref = kv_ref[:, 0, :, :valid_kv_len]
+    expected = np.empty((n_heads, 256), dtype=np.float32)
+    for head in range(n_heads):
+      scores = q_ref[head] @ k_ref[head // (n_heads // n_kv_heads)].T / np.sqrt(256)
+      probs = np.exp(scores - scores.max())
+      expected[head] = probs @ v_ref[head // (n_heads // n_kv_heads)] / probs.sum()
+
+    self.assertTrue(np.isfinite(out).all())
+    np.testing.assert_allclose(out[0, :, 0], expected, rtol=2e-3, atol=2e-3)
+
+  def test_short_decode_is_finite_and_matches_reference(self): self._test_decode(8192, 25)
+
+  def test_q8_cache_matches_dequantized_reference(self): self._test_decode(8192, 25, quantized=True)
+
+  def test_q8_cached_prefill_matches_dequantized_reference(self):
+    rng = np.random.default_rng(2)
+    heads, kv_heads, tokens, dim = 16, 2, 32, 256
+    q = rng.standard_normal((1, heads, tokens, dim)).astype(np.float16)
+    kv = rng.standard_normal((2, 1, kv_heads, tokens, dim)).astype(np.float16)
+    scale = np.maximum(np.max(np.abs(kv.astype(np.float32)), axis=-1), 1e-8) / 127
+    packed = np.clip(np.rint(kv.astype(np.float32) / scale[..., None]), -127, 127).astype(np.int8)
+    got = flash_attention_causal_cached(Tensor(q).realize(), Tensor(packed).realize(), tokens, tokens,
+                                        Tensor(scale.astype(np.float16)).realize()).numpy()
+    dequant = packed.astype(np.float32) * scale.astype(np.float16).astype(np.float32)[..., None]
+    expected = np.empty_like(got)
+    for head in range(heads):
+      scores = q[0, head].astype(np.float32) @ dequant[0, 0, head // (heads // kv_heads)].T / np.sqrt(dim)
+      scores[np.triu_indices(tokens, 1)] = -np.inf
+      probs = np.exp(scores - scores.max(axis=-1, keepdims=True))
+      expected[0, head] = probs @ dequant[1, 0, head // (heads // kv_heads)] / probs.sum(axis=-1, keepdims=True)
+    np.testing.assert_allclose(got, expected, rtol=2e-3, atol=2e-3)
+
+  def test_six_query_heads_per_kv_head(self): self._test_decode(8192, 25, n_heads=12, n_kv_heads=2)
+
+  def test_hierarchical_decode_matches_reference(self): self._test_decode(16384, 4097)
+
+
+if __name__ == "__main__": unittest.main()
+

test/null/test_hcq_iface.py

@@ -1,9 +1,26 @@
 import unittest, array, time
 from tinygrad.helpers import mv_address
 from tinygrad.runtime.support.hcq import MMIOInterface
+from tinygrad.runtime.support.memory import VirtMapping
+from tinygrad.runtime.support.system import PCIIfaceBase
 from tinygrad.runtime.support.usb import USBMMIOInterface
 from test.mockgpu.usb import MockUSB
 
+class TestPCIIface(unittest.TestCase):
+  def test_sysmem_mapping_respects_uncached(self):
+    class MM:
+      def alloc_vaddr(self, size, align): return 0x10000
+      def map_range(self, vaddr, size, paddrs, aspace, uncached=False, snooped=False):
+        return VirtMapping(vaddr, size, paddrs, aspace, uncached, snooped)
+    class PCI:
+      def bar_info(self, bar): return 0, 256 << 20
+      def alloc_sysmem(self, size, **kwargs): return memoryview(bytearray(size)), [0x20000]
+    iface = PCIIfaceBase.__new__(PCIIfaceBase)
+    iface.dev, iface.vram_bar, iface.pci_dev = None, 0, PCI()
+    iface.dev_impl = type("DevImpl", (), {"mm": MM()})()
+    for uncached in (False, True):
+      with self.subTest(uncached=uncached): self.assertEqual(iface.alloc(4096, host=True, uncached=uncached).meta.mapping.uncached, uncached)
+
 class TestHCQIface(unittest.TestCase):
   def setUp(self):
     self.size = 4 << 10

test/unit/test_llm_quant_amd.py

@@ -0,0 +1,105 @@
+import unittest
+import numpy as np
+
+from tinygrad import Device, Tensor, TinyJit, dtypes
+from tinygrad.llm.gguf import _GGML_QUANT, ggml_data_to_tensor
+from tinygrad.llm.kernels import amd as llm_amd
+from tinygrad.llm.model import Embedding, Linear
+
+def q8_activation(x:np.ndarray) -> np.ndarray:
+  grouped = x.reshape(*x.shape[:-1], -1, 32)
+  scale = np.maximum(np.max(np.abs(grouped), axis=-1, keepdims=True) / 127, 1e-8)
+  return (np.clip(np.rint(grouped / scale), -127, 127) * scale).reshape(x.shape)
+
+
+def random_packed(rng:np.random.Generator, ggml_type:int, elements:int) -> np.ndarray:
+  block_size, type_size = _GGML_QUANT[ggml_type]
+  blocks = rng.integers(0, 256, size=(elements // block_size, type_size), dtype=np.uint8)
+  scales = rng.uniform(0.001, 0.02, size=len(blocks)).astype(np.float16).view(np.uint8).reshape(-1, 2)
+  blocks[:, :2] = scales
+  if ggml_type in (12, 13): blocks[:, 2:4] = scales
+  if ggml_type == 14: blocks[:, -2:] = scales
+  return blocks.flatten()
+
+
+@unittest.skipUnless(Device.DEFAULT == "AMD", "requires DEV=AMD")
+class TestLLMQuantAMD(unittest.TestCase):
+  @staticmethod
+  def assert_q8_equal(result:tuple[Tensor, Tensor, Tensor], expected:np.ndarray):
+    grouped = expected.reshape(expected.shape[0], -1, 32)
+    scale = np.maximum(np.max(np.abs(grouped), axis=-1) / 127, 1e-8)
+    quant = np.clip(np.rint(grouped / scale[..., None]), -127, 127).astype(np.int8)
+    np.testing.assert_equal(result[0].numpy().view(np.int8).reshape(grouped.shape), quant)
+    np.testing.assert_allclose(result[1].numpy(), scale, rtol=1e-6, atol=1e-8)
+    np.testing.assert_equal(result[2].numpy(), quant.astype(np.int32).sum(-1))
+
+  def test_gated_delta_prefill_matches_sequential_reference(self):
+    rng = np.random.default_rng(36)
+    batch, heads, tokens, dim = 1, 2, 5, 128
+    q, k, v = [rng.standard_normal((batch, heads, tokens, dim), dtype=np.float32) for _ in range(3)]
+    q = q / np.linalg.norm(q, axis=-1, keepdims=True) / np.float32(np.sqrt(dim))
+    k = k / np.linalg.norm(k, axis=-1, keepdims=True)
+    beta, alpha = rng.random((batch, heads, tokens), dtype=np.float32), rng.uniform(0.9, 1, (batch, heads, tokens)).astype(np.float32)
+    state = rng.standard_normal((batch, heads, dim, dim), dtype=np.float32).astype(np.float16)
+    expected_core, expected_state = np.empty_like(q), state.astype(np.float32)
+    for token in range(tokens):
+      state_k = np.einsum("bhij,bhj->bhi", expected_state, k[:, :, token])
+      state_q = np.einsum("bhij,bhj->bhi", expected_state, q[:, :, token])
+      delta = (v[:, :, token] - state_k * alpha[:, :, token, None]) * beta[:, :, token, None]
+      expected_core[:, :, token] = state_q * alpha[:, :, token, None] + delta * np.sum(k[:, :, token] * q[:, :, token], axis=-1)[..., None]
+      expected_state = expected_state * alpha[:, :, token, None, None] + delta[..., None] * k[:, :, token, None, :]
+    core, next_state = llm_amd.gated_delta_prefill(
+      *(Tensor(x, device="AMD") for x in (q, k, v, beta, alpha)), Tensor(state, device="AMD"))
+    np.testing.assert_allclose(core.numpy(), expected_core, rtol=2e-4, atol=1e-3)
+    np.testing.assert_allclose(next_state.numpy(), expected_state.astype(np.float16), rtol=2e-4, atol=1e-3)
+
+  def test_q8_quantize_matches_reference(self):
+    rng, tokens, in_features = np.random.default_rng(35), 3, 256
+    x = rng.standard_normal((tokens, in_features), dtype=np.float32)
+    grouped = x.reshape(tokens, -1, 32)
+    expected_scale = np.maximum(np.max(np.abs(grouped), axis=-1) / 127, 1e-8)
+    expected_quant = np.clip(np.rint(grouped / expected_scale[..., None]), -127, 127).astype(np.int8)
+    quant, scale, group_sum = llm_amd.q8_quantize_sum(Tensor(x, device="AMD"), tokens, in_features)
+    np.testing.assert_equal(quant.numpy().view(np.int8).reshape(grouped.shape), expected_quant)
+    np.testing.assert_allclose(scale.numpy(), expected_scale, rtol=1e-7, atol=0)
+    np.testing.assert_equal(group_sum.numpy(), expected_quant.astype(np.int32).sum(-1))
+
+  def test_q4_embedding_matches_reference(self):
+    rng, vocab_size, embed_size = np.random.default_rng(34), 16, 256
+    raw = random_packed(rng, 12, vocab_size * embed_size)
+    expected = ggml_data_to_tensor(Tensor(raw), vocab_size * embed_size, 12).reshape(vocab_size, embed_size).half()
+    storage = Tensor(np.concatenate((np.zeros(68, dtype=np.uint8), raw)), dtype=dtypes.uint8, device="AMD").realize()
+    embedding = Embedding(vocab_size, embed_size)
+    embedding.set_quantized(storage[68:], 12)
+    idx = np.array([[7, 1, 15], [0, 4, 7]], dtype=np.int32)
+    np.testing.assert_equal(embedding(Tensor(idx, device="AMD")).numpy(), expected.numpy()[idx])
+
+  def test_iq4_lut_is_ready_for_jit_capture(self):
+    rng, in_features, out_features = np.random.default_rng(33), 256, 16
+    raw = random_packed(rng, 23, out_features * in_features)
+    weight = ggml_data_to_tensor(Tensor(raw), out_features * in_features, 23).numpy().reshape(out_features, in_features)
+    llm_amd.iq4_half_lut.cache_clear()
+    layer = Linear(in_features, out_features, bias=False)
+    layer.set_quantized(Tensor(raw, dtype=dtypes.uint8, device="AMD").realize(), 23)
+    @TinyJit
+    def run(x:Tensor): return layer(x).realize()
+    x = rng.standard_normal((16, in_features), dtype=np.float32)
+    expected = x.astype(np.float16).astype(np.float32) @ weight.astype(np.float16).astype(np.float32).T
+    for _ in range(2): np.testing.assert_allclose(run(Tensor(x, device="AMD")).numpy(), expected, rtol=1e-5, atol=2e-3)
+
+  def test_packed_linear_offset_matches_reference(self):
+    rng = np.random.default_rng(32)
+    for ggml_type,in_features in ((8, 256), (12, 256), (13, 256), (14, 256), (23, 256)):
+      for tokens in ((1, 16, 32, 64, 128) if ggml_type == 23 else (1, 16, 128) if ggml_type in (12, 13) else
+                     (1, 16) if ggml_type == 14 else (1,)):
+        raw, out_features = random_packed(rng, ggml_type, 64 * in_features), 64
+        weight = ggml_data_to_tensor(Tensor(raw), out_features * in_features, ggml_type).numpy().reshape(out_features, in_features)
+        storage = Tensor(np.concatenate((np.zeros(68, dtype=np.uint8), raw)), dtype=dtypes.uint8, device="AMD").realize()
+        layer = Linear(in_features, out_features, bias=False)
+        layer.set_quantized(storage[68:], ggml_type)
+        x = rng.standard_normal((tokens, in_features), dtype=np.float32)
+        expected = x.astype(np.float16).astype(np.float32) @ weight.astype(np.float16).astype(np.float32).T \
+          if ggml_type in (12, 13, 23) and tokens > 1 else q8_activation(x) @ weight.T
+        np.testing.assert_allclose(layer(Tensor(x, device="AMD")).numpy(), expected, rtol=1e-5, atol=2e-3)
+
+

tinygrad/llm/gguf.py

@@ -1,7 +1,8 @@
-import functools, io, pathlib, re, struct
+import functools, io, pathlib, re, struct, weakref
 from typing import Any, Callable
 
 from tinygrad.tensor import Tensor
+from tinygrad.uop.ops import UOp
 from tinygrad.dtype import dtypes
 from tinygrad.helpers import prod, round_up
 from tinygrad.nn.state import TensorIO
@@ -20,7 +21,14 @@ _GGML_NATIVE = {0: dtypes.float32, 1: dtypes.float16, 24: dtypes.int8, 25: dtype
 _GGML_QUANT = {2:(32,18), 3:(32,20), 6:(32,22), 7:(32,24), 8:(32,34),
                12:(256,144), 13:(256,176), 14:(256,210), 18:(256,98), 21:(256,110), 22:(256,82), 23:(256,136), 39:(32,17), 41:(128,18)}
 
-def ggml_data_to_tensor(t: Tensor, n: int, ggml_type: int) -> Tensor:
+_quantized_tensors:weakref.WeakKeyDictionary[UOp, tuple[UOp, int]] = weakref.WeakKeyDictionary()
+
+def get_ggml_quantization(tensor:Tensor) -> tuple[Tensor, int]|None:
+  if (meta:=_quantized_tensors.get(tensor.uop)) is None: return None
+  packed, ggml_type = meta
+  return Tensor(packed), ggml_type
+
+def ggml_data_to_tensor(t:Tensor, n:int, ggml_type:int, contiguous:bool=True) -> Tensor:
   """
   Converts ggml tensor data to a tinygrad tensor.
 
@@ -42,7 +50,8 @@ def ggml_data_to_tensor(t: Tensor, n: int, ggml_type: int) -> Tensor:
 
   if (nelements_nbytes := _GGML_QUANT.get(ggml_type)) is not None:
     from tinygrad.runtime.autogen import ggml_common as _ggml
-    blocks = t[:(n//nelements_nbytes[0])*nelements_nbytes[1]].reshape((-1, nelements_nbytes[1])).contiguous()
+    blocks = t[:(n//nelements_nbytes[0])*nelements_nbytes[1]].reshape((-1, nelements_nbytes[1]))
+    if contiguous: blocks = blocks.contiguous()
     if ggml_type == 2: return (q_to_uint8(blocks[:,2:], 4).bitcast(dtypes.int8) - 8) * blocks[:,:2].bitcast(dtypes.float16).cast(dtypes.float32)
     if ggml_type == 3:
       d, m = (blocks[:,s:s+2].bitcast(dtypes.float16).cast(dtypes.float32) for s in [ 0, 2 ])
@@ -145,7 +154,14 @@ def _gguf_parse(tensor: Tensor) -> tuple[dict, dict[str, Tensor]]:
   alignment, pos = kv_data.get("general.alignment", 32), r.tell()
   data_start = round_up(pos, alignment)
 
-  state_dict = {name: ggml_data_to_tensor(tensor[data_start + off:], prod(dims), typ).reshape(*reversed(dims)) for name, dims, typ, off in t_infos}
+  state_dict = {}
+  for name, dims, typ, off in t_infos:
+    n, shape = prod(dims), tuple(reversed(dims))
+    decoded = ggml_data_to_tensor(data:=tensor[data_start + off:], n, typ).reshape(*shape)
+    if typ in _GGML_QUANT:
+      block_size, type_size = _GGML_QUANT[typ]
+      _quantized_tensors[decoded.uop] = (data[:n//block_size*type_size].uop, typ)
+    state_dict[name] = decoded
   return kv_data, state_dict
 
 def _gguf_split_paths(path: pathlib.Path, kv: dict) -> list[pathlib.Path]:

tinygrad/llm/kernels/__init__.py

@@ -0,0 +1 @@
+"""Custom kernels used by tinygrad.llm."""

tinygrad/llm/model.py

@@ -1,15 +1,82 @@
 from __future__ import annotations
 import functools, itertools, pathlib
 from dataclasses import dataclass, replace
-from tinygrad import Tensor, nn, UOp, TinyJit, getenv, function
-from tinygrad.llm.gguf import gguf_load
-from tinygrad.uop.ops import resolve
+from typing import Any
+from tinygrad import Device, Tensor, nn, UOp, TinyJit, getenv, function, dtypes, Context
+from tinygrad.device import Buffer, BufferSpec
+from tinygrad.llm.kernels import amd as llm_amd
+from tinygrad.llm.gguf import get_ggml_quantization, ggml_data_to_tensor, gguf_load
+from tinygrad.uop.ops import resolve, Ops
+
+class PackedWeight:
+  weight:Tensor
+  def _init_packed(self):
+    self.ggml_type:int|None = None
+    self._raw_uop:UOp|None = None
+    self._raw_offset_uop:UOp|None = None
+    self._raw_offset_words:int|None = None
+  def set_quantized(self, packed:Tensor, ggml_type:int):
+    self.weight, self.ggml_type = packed.flatten(), ggml_type
+    self._raw_uop = self._raw_offset_uop = self._raw_offset_words = None
+    # IQ4 prefill uses a device LUT. Build it with the weights, before this linear can be captured by TinyJit.
+    if ggml_type == 23 and str(packed.device).startswith("AMD"): llm_amd.iq4_half_lut(str(packed.device))
+  def _packed_offset(self) -> Tensor:
+    raw, raw_offset = self.weight.uop, 0
+    while raw.op in (Ops.BITCAST, Ops.RESHAPE): raw = raw.src[0]
+    while raw.op is Ops.SHRINK:
+      raw_offset += raw.src[1].arg * raw.dtype.itemsize
+      raw = raw.src[0]
+    assert raw_offset % 4 == 0 and raw.dtype == dtypes.uint8
+    self._raw_uop = raw
+    self._raw_offset_words = raw_offset // 4
+    return Tensor([self._raw_offset_words], dtype=dtypes.uint64, device=self.weight.device)
+  def _prepare_packed(self):
+    self._raw_offset_uop = self._packed_offset().realize().uop
+
+class Linear(nn.Linear, PackedWeight):
+  def __init__(self, in_features:int, out_features:int, bias=True):
+    # GGUF loading replaces every LLM weight. Lazy zeros avoid constructing hundreds of random-init graphs first,
+    # while keeping directly-created test models deterministic and valid.
+    self.weight = Tensor.zeros(out_features, in_features)
+    self.bias = Tensor.zeros(out_features) if bias else None
+    self.in_features, self.out_features = in_features, out_features
+    self._init_packed()
+  def prepare(self, x:Tensor, with_sum:bool=False) -> tuple[Tensor, ...]|None:
+    if (with_sum or self.ggml_type in (12, 13)) and self.ggml_type in (8, 12, 13, 14, 23) and \
+       str(self.weight.device).startswith("AMD"):
+      return llm_amd.q8_quantize_sum(x, int(x.numel()) // self.in_features, self.in_features)
+    return llm_amd.q8_quantize(x, int(x.numel()) // self.in_features, self.in_features) \
+      if self.ggml_type in (8, 14, 23) and str(self.weight.device).startswith("AMD") else None
+  def __call__(self, x:Tensor, prepared:tuple[Tensor, ...]|None=None) -> Tensor:
+    if self.ggml_type in (8, 12, 13, 14, 23) and str(self.weight.device).startswith("AMD"):
+      return llm_amd.q8_linear(self, x, prepared)
+    if self.ggml_type is not None:
+      weight = ggml_data_to_tensor(self.weight, self.out_features * self.in_features, self.ggml_type,
+                                   contiguous=False).reshape(self.out_features, self.in_features)
+      if getenv("HALF", 1): weight = weight.cast('float16')
+      return x.linear(weight.transpose(), self.bias)
+    return super().__call__(x)
+
+class Embedding(nn.Embedding, PackedWeight):
+  def __init__(self, vocab_size:int, embed_size:int):
+    self.weight = Tensor.zeros(vocab_size, embed_size)
+    self.vocab_size, self.embed_size = vocab_size, embed_size
+    self._init_packed()
+  def __call__(self, idx:Tensor) -> Tensor:
+    if self.ggml_type == 12 and str(self.weight.device).startswith("AMD"): return llm_amd.q4_embedding(self, idx)
+    return super().__call__(idx)
 
 @functools.cache
 def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0, device:str|None=None) -> Tensor:
-  freqs = 1.0 / (theta ** (Tensor.arange(0, dim, 2)[:(dim // 2)] / dim))
-  freqs = Tensor.arange(end).unsqueeze(dim=1) * freqs.unsqueeze(dim=0)
-  return freqs.cos().cat(freqs.sin(), dim=-1).clone(device)
+  freqs = 1.0 / (theta ** (Tensor.arange(0, dim, 2).to(device)[:(dim // 2)] / dim))
+  freqs = Tensor.arange(end).to(device).unsqueeze(dim=1) * freqs.unsqueeze(dim=0)
+  table = freqs.cos().cat(freqs.sin(), dim=-1)
+  if device is not None and str(device).startswith("AMD") and end > 8192:
+    size = table.numel()
+    assert isinstance(size, int)
+    storage = Buffer(str(device), size, table.dtype, options=BufferSpec(host=True))
+    return Tensor(UOp.from_buffer(storage).reshape(table.shape)).assign(table).realize()
+  return table.clone(device)
 
 class ExpertWeights:
   """Like nn.Linear but with num_experts dimension. Weight shape: (num_experts, out_features, in_features)."""
@@ -27,12 +94,14 @@ def apply_rope(x:Tensor, freqs_cis:Tensor) -> Tensor:
 
 def pairwise_topk(x: Tensor, k: int) -> tuple[Tensor, Tensor]:
   n = x.shape[-1]
-  vals = Tensor.arange(n).reshape(1,1,n).cast(x.dtype).expand(x.shape)
+  vals = Tensor.arange(n).to(x.device).reshape(1,1,n).cast(x.dtype).expand(x.shape)
   cmp = (x.unsqueeze(-1) > x.unsqueeze(-2)) | ((x.unsqueeze(-1) == x.unsqueeze(-2)) & \
-    (Tensor.arange(n).reshape(1,1,n,1) < Tensor.arange(n).reshape(1,1,1,n)))
+    (Tensor.arange(n).to(x.device).reshape(1,1,n,1) < Tensor.arange(n).to(x.device).reshape(1,1,1,n)))
   sel = x.const_like(0).scatter(-1, cmp.sum(axis=-1).cast('int32'), vals)[:,:,n-k:].cast('int32')
   return x.gather(-1, sel), sel
 
+def l2norm(x:Tensor) -> Tensor: return x * (x.square().sum(-1, keepdim=True) + 1e-6).rsqrt()
+
 @dataclass(frozen=True)
 class SSMConfig:
   conv_kernel: int
@@ -76,32 +145,26 @@ class TransformerConfig:
 class FFNBlock:
   def __init__(self, config:TransformerConfig):
     self.config = config
-
-    # --- RMSNorms --------------------------------------------------------
-    self.attn_norm   = nn.RMSNorm(config.dim, config.norm_eps)
-    self.ffn_norm    = nn.RMSNorm(config.dim, config.norm_eps)
-
-    # --- feed-forward (MoE or dense) -------------------------------------
+    self.pending_state:tuple[Tensor, Tensor]|None = None
+    self.attn_norm, self.ffn_norm = nn.RMSNorm(config.dim, config.norm_eps), nn.RMSNorm(config.dim, config.norm_eps)
     if config.num_experts > 0:
-      self.ffn_gate_inp = nn.Linear(config.dim, config.num_experts, bias=False)  # router
+      self.ffn_gate_inp = Linear(config.dim, config.num_experts, bias=False)
       if config.expert_bias: self.exp_probs_b = {"bias": Tensor.zeros(config.num_experts)}
       self.ffn_gate_exps = ExpertWeights(config.num_experts, config.dim, config.hidden_dim)
       self.ffn_up_exps = ExpertWeights(config.num_experts, config.dim, config.hidden_dim)
       self.ffn_down_exps = ExpertWeights(config.num_experts, config.hidden_dim, config.dim)
       if config.shared_expert_dim > 0:
-        self.ffn_gate_shexp = nn.Linear(config.dim, config.shared_expert_dim, bias=False)
-        self.ffn_up_shexp = nn.Linear(config.dim, config.shared_expert_dim, bias=False)
-        self.ffn_down_shexp = nn.Linear(config.shared_expert_dim, config.dim, bias=False)
+        self.ffn_gate_shexp = Linear(config.dim, config.shared_expert_dim, bias=False)
+        self.ffn_up_shexp = Linear(config.dim, config.shared_expert_dim, bias=False)
+        self.ffn_down_shexp = Linear(config.shared_expert_dim, config.dim, bias=False)
         if config.shared_expert_gate: self.ffn_gate_inp_shexp = {"weight": Tensor.zeros(config.dim)}
     else:
-      self.ffn_gate    = nn.Linear(config.dim, config.hidden_dim, bias=False)
-      self.ffn_up      = nn.Linear(config.dim, config.hidden_dim, bias=False)
-      self.ffn_down    = nn.Linear(config.hidden_dim, config.dim, bias=False)
+      self.ffn_gate, self.ffn_up = Linear(config.dim, config.hidden_dim, bias=False), Linear(config.dim, config.hidden_dim, bias=False)
+      self.ffn_down = Linear(config.hidden_dim, config.dim, bias=False)
 
   def _feed_forward(self, x:Tensor) -> Tensor:
     if hasattr(self, 'ffn_gate_exps'):
-      h = x.unsqueeze(2)  # (B, T, 1, D) - add expert dim for broadcasting
-      logits = self.ffn_gate_inp(x)
+      h, logits = x.unsqueeze(2), self.ffn_gate_inp(x)
       if hasattr(self, 'exp_probs_b'):
         probs = logits.sigmoid()
         _, sel = pairwise_topk(probs + self.exp_probs_b["bias"], self.config.num_experts_per_tok)
@@ -111,48 +174,70 @@ class FFNBlock:
         vals, sel = pairwise_topk(logits, self.config.num_experts_per_tok)
         probs = vals.softmax(-1) if self.config.norm_topk_prob else logits.softmax(-1).gather(-1, sel)
       probs = probs * self.config.routed_scaling_factor
-      x_down = self.ffn_down_exps(sel, (self.ffn_gate_exps(sel, h).silu() * self.ffn_up_exps(sel, h)).contiguous())  # (B, T, k, D)
-      out = (x_down * probs.unsqueeze(-1)).sum(axis=2)  # (B, T, D)
+      x_down = self.ffn_down_exps(sel, (self.ffn_gate_exps(sel, h).silu() * self.ffn_up_exps(sel, h)).contiguous())
+      out = (x_down * probs.unsqueeze(-1)).sum(axis=2)
       if hasattr(self, 'ffn_gate_shexp'):
         shexp = self.ffn_down_shexp(self.ffn_gate_shexp(x).silu().contiguous() * self.ffn_up_shexp(x))
         if hasattr(self, 'ffn_gate_inp_shexp'): shexp = shexp * (x * self.ffn_gate_inp_shexp["weight"]).sum(axis=-1, keepdim=True).sigmoid()
         out = out + shexp
       return out
-    # TODO: remove the need for this contiguous
-    return self.ffn_down(self.ffn_gate(x).silu().contiguous() * self.ffn_up(x))
+    prepared = self.ffn_gate.prepare(x)
+    gate, up = self.ffn_gate(x, prepared), self.ffn_up(x, prepared)
+    return self.ffn_down(gate.silu() * up)
+
+  def _normalized_feed_forward(self, x:Tensor) -> Tensor:
+    return self._feed_forward(self.ffn_norm(x))
 
-  # 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
+  def _attention(self, x:Tensor, start_pos:int|UOp, use_flash:bool=False, kv_len:int|UOp|None=None,
+                 valid_len:int|UOp|None=None, input_norm:nn.RMSNorm|None=None) -> Tensor: raise NotImplementedError
 
-  def __call__(self, x: Tensor, start_pos: int|UOp):
+  def __call__(self, x:Tensor, start_pos:int|UOp, use_flash:bool=False, kv_len:int|UOp|None=None, valid_len:int|UOp|None=None):
     self._init_state(x)
-    # we pass in the weights implicitly so we unpack the GGUF on the fly
-    @function(precompile=True, allow_implicit=True)
+    if hasattr(self, 'ssm_a'):
+      self.pending_state, self.pending_recurrent_inplace = None, False
+      @function(precompile=True, allow_implicit=True)
+      def _run_stateful(x:Tensor, start_pos:int|UOp, valid_len:int|UOp|None):
+        h = x + self._attention(self.attn_norm(x), start_pos, use_flash, kv_len, valid_len)
+        out = (h + self._normalized_feed_forward(h)).contiguous()
+        assert self.pending_state is not None
+        return (out, self.pending_state[0]) if self.pending_recurrent_inplace else (out, *self.pending_state)
+      stateful_out = _run_stateful(x, start_pos, valid_len)
+      out, conv_state = stateful_out[:2]
+      recurrent_state = getattr(self, "recurrent_state")
+      stores = [getattr(self, "conv_state").uop.store(conv_state.uop)]
+      if not self.pending_recurrent_inplace: stores.append(recurrent_state.uop.store(stateful_out[2].uop))
+      return Tensor(out.uop.after(recurrent_state.uop.after(*stores)))
     def _run(x:Tensor, start_pos:int|UOp):
-      h =     x + self._attention(self.attn_norm(x), start_pos)
-      return (h + self._feed_forward(self.ffn_norm(h))).contiguous()
-    return _run(x, start_pos)
+      h = x + self._attention(self.attn_norm(x), start_pos, use_flash, kv_len)
+      return (h + self._normalized_feed_forward(h)).contiguous()
+    return function(precompile=True, allow_implicit=True)(_run)(x, start_pos)
 
 class TransformerBlock(FFNBlock):
   def __init__(self, config:TransformerConfig):
     super().__init__(config)
+    self.kv_cache_host = False
     assert config.v_head_dim == config.head_dim, "TransformerBlock requires v_head_dim == head_dim"
 
     # --- attention projections (all linear, bias-free) ------------------
     q_proj_out       = config.head_dim * config.n_heads * (2 if config.attn_output_gate else 1)
     kv_proj_out      = config.head_dim * config.n_kv_heads
-    self.attn_q      = nn.Linear(config.dim, q_proj_out,  bias=config.qkv_bias)
-    self.attn_k      = nn.Linear(config.dim, kv_proj_out, bias=config.qkv_bias)
-    self.attn_v      = nn.Linear(config.dim, kv_proj_out, bias=config.qkv_bias)
-    self.attn_output = nn.Linear(config.head_dim * config.n_heads, config.dim, bias=False)
+    self.attn_q      = Linear(config.dim, q_proj_out,  bias=config.qkv_bias)
+    self.attn_k      = Linear(config.dim, kv_proj_out, bias=config.qkv_bias)
+    self.attn_v      = Linear(config.dim, kv_proj_out, bias=config.qkv_bias)
+    self.attn_output = Linear(config.head_dim * config.n_heads, config.dim, bias=False)
     if config.qk_norm: self.attn_q_norm, self.attn_k_norm = nn.RMSNorm(config.qk_norm, config.norm_eps), nn.RMSNorm(config.qk_norm, config.norm_eps)
 
-  def _attention(self, x:Tensor, start_pos:int|UOp) -> Tensor:
-    q, k, v = self.attn_q(x), self.attn_k(x), self.attn_v(x)
+  def _attention(self, x:Tensor, start_pos:int|UOp, use_flash:bool=False, kv_len:int|UOp|None=None,
+                 valid_len:int|UOp|None=None, input_norm:nn.RMSNorm|None=None) -> Tensor:
+    prepared:tuple[Tensor, ...]|None
+    prepared = self.attn_q.prepare(x, any(layer.ggml_type in (12, 13) for layer in (self.attn_q, self.attn_k, self.attn_v)))
+    q = self.attn_q(x, prepared)
+    if prepared is not None and self.attn_k.ggml_type == self.attn_v.ggml_type == 8:
+      k, v = self.attn_k(x, prepared), self.attn_v(x, prepared)
+    else: k, v = self.attn_k(x, prepared), self.attn_v(x, prepared)
     if self.config.qk_norm and self.config.qk_norm != self.config.head_dim: q, k = self.attn_q_norm(q), self.attn_k_norm(k)
 
     B, T, _ = x.shape
@@ -168,45 +253,90 @@ 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 self.cache_kv.dtype == dtypes.int8:
+      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)
+      stores = (self.cache_kv[:, :, :, start_pos:start_pos+T, :].uop.store(packed_kv.uop),
+                self.cache_kv_scale[:, :, :, start_pos:start_pos+T].uop.store(scale.squeeze(-1).uop))
+      assigned_kv, assigned_scale = Tensor(self.cache_kv.uop.after(*stores)), Tensor(self.cache_kv_scale.uop.after(*stores))
+    else:
+      assigned_kv = Tensor(self.cache_kv.uop.after(
+        self.cache_kv[:, :, :, start_pos:start_pos+T, :].uop.store(stacked_kv.cast(self.cache_kv.dtype).uop)))
+      assigned_scale = None
+    cache_len = start_pos + T if kv_len is None else kv_len
+    k, v = assigned_kv[0, :, :, 0:cache_len, :], assigned_kv[1, :, :, 0:cache_len, :]
+    if assigned_scale is not None:
+      k, v = k.float() * assigned_scale[0, :, :, 0:cache_len, None], v.float() * assigned_scale[1, :, :, 0:cache_len, None]
 
     #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, :]
 
-    # NOTE: this mask is causal_lower_right, not the causal_upper_left generated by is_casual = True
+    # NOTE: this mask is causal_lower_right, not the causal_upper_left generated by is_causal = 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)
+    flash_decode = resolve(T == 1) and kv_len is not None and str(x.device).startswith("AMD") and self.config.head_dim == 256
+    if flash_decode:
+      decode_len = kv_len if isinstance(kv_len, int) else self.config.max_context
+      decode_pos = (start_pos.unbind()[0] if isinstance(start_pos, UOp) else start_pos) + 1
+      attn = llm_amd.amd_flash_attention_decode(q.half(), assigned_kv, decode_pos, decode_len, assigned_scale)
+    elif use_flash:
+      start = start_pos.unbind()[0] if isinstance(start_pos, UOp) else start_pos
+      valid = valid_len.unbind()[0] if isinstance(valid_len, UOp) else valid_len
+      valid_kv_len, key_limit = start + T, start + valid if valid is not None else None
+      attn = llm_amd.flash_attention_causal_cached(q.half(), assigned_kv, valid_kv_len, key_limit, assigned_scale)
+    else:
+      mask:Tensor|None
+      if kv_len is not None:
+        mask = None if resolve(T == 1) and self.config.ssm is not None else \
+          Tensor.full((1, 1, 1, kv_len), float("-inf"), dtype=x.dtype, device=x.device, buffer=False).triu(start_pos+1)
+      else:
+        mask = Tensor.full((1, 1, T, start_pos+T), float("-inf"), dtype=x.dtype, device=x.device, buffer=False).triu(start_pos+1) \
+          if resolve(T != 1) else None
+      attn = q.float().scaled_dot_product_attention(k.float(), v.float(), 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"):
       # TODO: how is the dtype of this determined?
-      self.cache_kv = Tensor.empty(2, x.shape[0], self.config.n_kv_heads, self.config.max_context, self.config.head_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)
+      # Decode uses fixed-size KV buckets. Unwritten entries must be zero: masking happens after QK, so values left
+      # uninitialized by Tensor.empty can inject NaNs before the mask is applied.
+      cache_dtype = dtypes.int8 if self.config.max_context > 8192 and str(x.device).startswith("AMD") else dtypes.float16
+      cache_shape = (2, x.shape[0], self.config.n_kv_heads, self.config.max_context+192, self.config.head_dim)
+      if self.kv_cache_host and str(x.device).startswith("AMD"):
+        cache_size = 1
+        for dim in cache_shape:
+          assert isinstance(dim, int)
+          cache_size *= dim
+        storage = Buffer(str(x.device), cache_size, cache_dtype, options=BufferSpec(host=True))
+        self.cache_kv = Tensor(UOp.from_buffer(storage).reshape(cache_shape))
+        self.cache_kv.assign(self.cache_kv.const_like(0)).realize()
+      else: self.cache_kv = Tensor.zeros(*cache_shape, dtype=cache_dtype, device=x.device).contiguous()
+      if cache_dtype == dtypes.int8:
+        self.cache_kv_scale = Tensor.zeros(2, x.shape[0], self.config.n_kv_heads, self.config.max_context+192,
+                                           dtype=dtypes.float16, device=x.device).contiguous()
+      self.freqs_cis = precompute_freqs_cis(self.config.rope_dim, self.config.max_context+192, self.config.rope_theta, device=x.device)
 
 class MLATransformerBlock(FFNBlock):
   def __init__(self, config:TransformerConfig):
     super().__init__(config)
     qk_nope_head_dim = config.head_dim - config.rope_dim
     if config.q_lora_rank > 0:
-      self.attn_q_a = nn.Linear(config.dim, config.q_lora_rank, bias=False)
+      self.attn_q_a = Linear(config.dim, config.q_lora_rank, bias=False)
       self.attn_q_a_norm = nn.RMSNorm(config.q_lora_rank, config.norm_eps)
-      self.attn_q_b = nn.Linear(config.q_lora_rank, config.n_heads * config.head_dim, bias=False)
+      self.attn_q_b = Linear(config.q_lora_rank, config.n_heads * config.head_dim, bias=False)
     else:
-      self.attn_q = nn.Linear(config.dim, config.n_heads * config.head_dim, bias=False)
-    self.attn_kv_a_mqa = nn.Linear(config.dim, config.kv_lora_rank + config.rope_dim, bias=False)
+      self.attn_q = Linear(config.dim, config.n_heads * config.head_dim, bias=False)
+    self.attn_kv_a_mqa = Linear(config.dim, config.kv_lora_rank + config.rope_dim, bias=False)
     self.attn_kv_a_norm = nn.RMSNorm(config.kv_lora_rank, config.norm_eps)
     self.attn_k_b = {"weight": Tensor.zeros(config.n_heads, config.kv_lora_rank, qk_nope_head_dim)}
     self.attn_v_b = {"weight": Tensor.zeros(config.n_heads, config.v_head_dim, config.kv_lora_rank)}
-    self.attn_output = nn.Linear(config.n_heads * config.v_head_dim, config.dim, bias=False)
+    self.attn_output = Linear(config.n_heads * config.v_head_dim, config.dim, bias=False)
 
-  def _attention(self, x:Tensor, start_pos:int|UOp) -> Tensor:
+  def _attention(self, x:Tensor, start_pos:int|UOp, use_flash:bool=False, kv_len:int|UOp|None=None,
+                 valid_len:int|UOp|None=None, input_norm:nn.RMSNorm|None=None) -> Tensor:
+    assert input_norm is None
     B, T, _ = x.shape
     q_nope_head_dim = self.config.head_dim - self.config.rope_dim
     q_proj = self.attn_q_b(self.attn_q_a_norm(self.attn_q_a(x))) if self.config.q_lora_rank > 0 else self.attn_q(x)
@@ -224,7 +354,7 @@ class MLATransformerBlock(FFNBlock):
     k = Tensor(self.cache_k.uop.after(self.cache_k[:, :, start_pos:start_pos+T, :].uop.store(k_store.uop)))[:, :, 0:start_pos+T, :]
     v = k[..., :self.config.kv_lora_rank]
 
-    mask = Tensor.full((1, 1, T, start_pos+T), float("-inf"), dtype=x.dtype, buffer=False).triu(start_pos+1) \
+    mask = Tensor.full((1, 1, T, start_pos+T), float("-inf"), dtype=x.dtype, device=x.device, buffer=False).triu(start_pos+1) \
       if resolve(T != 1) else None
     attn = q @ k.transpose(-1, -2) * (1.0 / self.config.head_dim ** 0.5)
     if mask is not None: attn = attn + mask
@@ -234,8 +364,8 @@ class MLATransformerBlock(FFNBlock):
 
   def _init_state(self, x:Tensor):
     if not hasattr(self, "cache_k"):
-      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)
+      self.cache_k = Tensor.empty(x.shape[0], 1, self.config.max_context+192, 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+192, self.config.rope_theta, device=x.device)
 
 class GatedDeltaNetBlock(FFNBlock):
   def __init__(self, config:TransformerConfig, ssm:SSMConfig):
@@ -244,56 +374,116 @@ class GatedDeltaNetBlock(FFNBlock):
     assert self.num_v_heads % self.num_k_heads == 0
     self.head_v_dim, self.ssm_conv_kernel = ssm.inner_size // ssm.time_step_rank, ssm.conv_kernel
     self.conv_channels, self.q_dim = ssm.inner_size + 2*ssm.group_count*ssm.state_size, ssm.state_size*ssm.group_count
-    self.attn_qkv = nn.Linear(config.dim, self.conv_channels, bias=False)
+    self.attn_qkv = Linear(config.dim, self.conv_channels, bias=False)
     if ssm.kda:
-      self.ssm_g_a, self.ssm_g_b = nn.Linear(config.dim, self.head_v_dim, bias=False), nn.Linear(self.head_v_dim, ssm.inner_size, bias=False)
-      self.ssm_f_a, self.ssm_f_b = nn.Linear(config.dim, self.head_k_dim, bias=False), nn.Linear(self.head_k_dim, ssm.inner_size, bias=False)
+      self.ssm_g_a, self.ssm_g_b = Linear(config.dim, self.head_v_dim, bias=False), Linear(self.head_v_dim, ssm.inner_size, bias=False)
+      self.ssm_f_a, self.ssm_f_b = Linear(config.dim, self.head_k_dim, bias=False), Linear(self.head_k_dim, ssm.inner_size, bias=False)
     else:
-      self.attn_gate = nn.Linear(config.dim, ssm.inner_size, bias=False)
-      self.ssm_alpha = nn.Linear(config.dim, self.num_v_heads, bias=False)
-    self.ssm_beta = nn.Linear(config.dim, self.num_v_heads, bias=False)
+      self.attn_gate = Linear(config.dim, ssm.inner_size, bias=False)
+      self.ssm_alpha = Linear(config.dim, self.num_v_heads, bias=False)
+    self.ssm_beta = Linear(config.dim, self.num_v_heads, bias=False)
+    self.ssm_beta_alpha_weight:Tensor|None = None
     self.ssm_conv1d = {"weight": Tensor.zeros(self.conv_channels, self.ssm_conv_kernel)}
     self.ssm_dt = {"bias": Tensor.zeros(ssm.inner_size if ssm.kda else self.num_v_heads)}
     self.ssm_a = Tensor.zeros(self.num_v_heads, 1) if ssm.kda else Tensor.zeros(self.num_v_heads)
-    self.ssm_norm, self.ssm_out = nn.RMSNorm(self.head_v_dim, config.norm_eps), nn.Linear(ssm.inner_size, config.dim, bias=False)
+    self.ssm_norm, self.ssm_out = nn.RMSNorm(self.head_v_dim, config.norm_eps), Linear(ssm.inner_size, config.dim, bias=False)
 
-  def _attention(self, x:Tensor, start_pos:int|UOp) -> Tensor:
+  def _finish_state(self, conv_state:Tensor, recurrent_state:Tensor) -> Tensor:
+    self.pending_state = (conv_state.contiguous(), recurrent_state.contiguous())
+    if str(recurrent_state.device).startswith("AMD"): return recurrent_state
+    stores = (self.conv_state.uop.store(self.pending_state[0].uop), self.recurrent_state.uop.store(self.pending_state[1].uop))
+    self.pending_recurrent_inplace = True
+    return Tensor(self.recurrent_state.uop.after(*stores))
+
+  def _attention(self, x:Tensor, start_pos:int|UOp, use_flash:bool=False, kv_len:int|UOp|None=None,
+                 valid_len:int|UOp|None=None, input_norm:nn.RMSNorm|None=None) -> Tensor:
     B, T, _ = x.shape
-    assert T == 1, "GatedDeltaNetBlock currently only supports T=1"
+    conv_state, initial_state = self.conv_state, self.recurrent_state
+    if hasattr(self, "ssm_g_a"):
+      assert T == 1
+      x = x.half()
+      out_gate, qkv = self.ssm_g_b(self.ssm_g_a(x)), self.attn_qkv(x)
+      beta = self.ssm_beta(x).sigmoid().reshape(B, self.num_v_heads, 1, 1)
+      alpha = self.ssm_f_b(self.ssm_f_a(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)
+      conv_window = conv_state.cat(qkv, 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 = l2norm(q.reshape(B, self.num_k_heads, self.head_k_dim)).repeat(1, self.num_v_heads//self.num_k_heads, 1)
+      k = l2norm(k.reshape(B, self.num_k_heads, self.head_k_dim)).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_state = initial_state * alpha
+      recurrent_state = recurrent_state + ((v - recurrent_state @ k) * beta) @ k.transpose(-1, -2)
+      recurrent_state = self._finish_state(conv_window[:, 1:, :].cast(self.conv_state.dtype),
+                                           recurrent_state.cast(self.recurrent_state.dtype))
+      core = self.ssm_norm((recurrent_state @ q).squeeze(-1).reshape(B, 1, self.num_v_heads, self.head_v_dim))
+      gate = out_gate.reshape(B, 1, self.num_v_heads, self.head_v_dim).sigmoid()
+      return self.ssm_out((core * gate).reshape(B, 1, -1).cast(x.dtype))
+    if T == 1:
+      if input_norm is None: x = x.half()
+      prepared = self.attn_gate.prepare(x, self.attn_qkv.ggml_type in (12, 13))
+      out_gate, qkv = self.attn_gate(x, prepared), self.attn_qkv(x, prepared)
+      if self.ssm_beta_alpha_weight is not None:
+        beta_alpha = x @ self.ssm_beta_alpha_weight.T
+        beta, alpha = beta_alpha.reshape(B, 1, -1).split(self.num_v_heads, dim=-1)
+      else: beta, alpha = self.ssm_beta(x, prepared), self.ssm_alpha(x, prepared)
+      out_gate = out_gate.reshape(B, 1, self.num_v_heads, self.head_v_dim)
+      beta, alpha = beta.sigmoid(), ((alpha.float() + self.ssm_dt["bias"]).softplus() * self.ssm_a).exp()
+      conv_window = conv_state.cat(qkv, 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)
+      k = k.reshape(B, self.num_k_heads, self.head_k_dim)
+      q, k = (l2norm(q), l2norm(k)) if str(x.device).startswith("AMD") else (q.normalize(dim=-1), k.normalize(dim=-1))
+      q, k = q.repeat(1, self.num_v_heads//self.num_k_heads, 1), k.repeat(1, self.num_v_heads//self.num_k_heads, 1)
+      v, q = v.reshape(B, self.num_v_heads, self.head_v_dim), q * self.head_k_dim**-0.5
+      qv, kv = q.unsqueeze(-1), k.unsqueeze(-1)
+      alpha4, beta4 = alpha.reshape(B, self.num_v_heads, 1, 1), beta.reshape(B, self.num_v_heads, 1, 1)
+      if str(x.device).startswith("AMD"):
+        state_k, state_q = (initial_state @ kv.cat(qv, dim=-1)).split(1, dim=-1)
+        delta = (v.unsqueeze(-1) - state_k * alpha4) * beta4
+        recurrent_state = initial_state * alpha4 + delta @ kv.transpose(-1, -2)
+        core = (state_q * alpha4 + delta * (kv.transpose(-1, -2) @ qv)).squeeze(-1)
+      else:
+        recurrent_state = initial_state * alpha4
+        recurrent_state = recurrent_state + ((v.unsqueeze(-1) - recurrent_state @ kv) * beta4) @ kv.transpose(-1, -2)
+        core = (recurrent_state @ qv).squeeze(-1)
+      recurrent_state = self._finish_state(conv_window[:, 1:, :].cast(self.conv_state.dtype),
+                                           recurrent_state.cast(self.recurrent_state.dtype))
+      if not str(x.device).startswith("AMD"): core = (recurrent_state @ qv).squeeze(-1)
+      core = self.ssm_norm(core.reshape(B, 1, self.num_v_heads, self.head_v_dim))
+      return self.ssm_out((core * out_gate.silu()).reshape(B, 1, -1).cast(x.dtype))
 
-    # input processing
+    assert str(x.device).startswith("AMD"), "batched GatedDeltaNet prefill currently requires AMD"
     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()
+    prepared = self.attn_gate.prepare(x, self.attn_qkv.ggml_type in (12, 13))
+    out_gate, qkv = self.attn_gate(x, prepared), self.attn_qkv(x, prepared)
+    if self.ssm_beta_alpha_weight is not None: beta, alpha = (x @ self.ssm_beta_alpha_weight.T).split(self.num_v_heads, dim=-1)
+    else: beta, alpha = self.ssm_beta(x, prepared), self.ssm_alpha(x, prepared)
+    out_gate = out_gate.reshape(B, T, self.num_v_heads, self.head_v_dim)
+    beta = beta.sigmoid().reshape(B, T, self.num_v_heads)
+    log_alpha = ((alpha.float() + self.ssm_dt["bias"]).softplus() * self.ssm_a).reshape(B, T, self.num_v_heads)
+    if valid_len is not None:
+      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
+    conv_window = conv_state.cat(qkv, dim=1)
+    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 = l2norm(q.reshape(B, T, self.num_k_heads, self.head_k_dim)).repeat(1, 1, self.num_v_heads//self.num_k_heads, 1)
+    k = l2norm(k.reshape(B, T, self.num_k_heads, self.head_k_dim)).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, log_alpha = [z.transpose(1, 2).float() for z in (q, k, v, beta, log_alpha)]
+    core, recurrent_state = llm_amd.gated_delta_prefill(q * self.head_k_dim**-0.5, k, v, beta, log_alpha.exp(), initial_state)
+    core = self.ssm_norm(core.transpose(1, 2))
+    out = self.ssm_out((core * out_gate.silu()).reshape(B, T, -1).cast(x.dtype)).contiguous()
+    state_pos = T if valid_len is None else valid_len
+    self.pending_state = (conv_window[:, state_pos:state_pos+self.ssm_conv_kernel-1, :].cast(self.conv_state.dtype).contiguous(),
+                          recurrent_state)
+    return out
 
-    # 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))
-
-  # 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 []
@@ -302,44 +492,85 @@ 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_v_dim, device=x.device).clone()
+      state_dtype = dtypes.float16 if str(x.device).startswith("AMD") else x.dtype
+      self.recurrent_state = Tensor.zeros(x.shape[0], self.num_v_heads, self.head_v_dim, self.head_v_dim,
+                                          dtype=state_dtype, device=x.device).clone()
 
 class Transformer:
   def __init__(self, config:TransformerConfig):
     dense_config = replace(config, num_experts=0, num_experts_per_tok=0, shared_expert_dim=0, hidden_dim=config.dense_hidden_dim or config.hidden_dim)
     if config.ssm: config = replace(config, qk_norm=config.head_dim)
     block_cls = MLATransformerBlock if config.kv_lora_rank > 0 else TransformerBlock
-    self.blk:list[FFNBlock] = [GatedDeltaNetBlock(dense_config if i < config.leading_dense_blocks else config, config.ssm)
-                               if config.ssm and config.ssm_layers[i] else
+    self.blk:list[FFNBlock] = [GatedDeltaNetBlock(config, config.ssm) if config.ssm and config.ssm_layers[i] else
                                block_cls(dense_config if i < config.leading_dense_blocks else config) for i in range(config.num_blocks)]
-    self.token_embd  = nn.Embedding(config.vocab_size, config.dim)
+    if config.max_context > 8192:
+      # A full Q8 cache for 262k Qwen leaves no graph workspace on a 24 GB card. Keep two fixed layers host-mapped;
+      # this avoids runtime cache growth while leaving the other attention layers resident in VRAM.
+      for block in [block for block in self.blk if isinstance(block, TransformerBlock)][-2:]: block.kv_cache_host = True
+    self.token_embd  = Embedding(config.vocab_size, config.dim)
     self.output_norm = nn.RMSNorm(config.dim, config.norm_eps)
-    self.output = nn.Linear(config.dim, config.vocab_size, bias=False)
+    self.output = Linear(config.dim, config.vocab_size, bias=False)
     self.max_context = config.max_context
     self.has_recurrent_block = any(isinstance(b, GatedDeltaNetBlock) for b in self.blk)
     self._cached_tokens: list[int] = []
+    self._state_checkpoints: list[Tensor] = []
+    self._state_checkpoint_pos = 0
+    self._save_state_jit:Any = None
+    self._restore_state_jit:Any = None
+    self._warming_up = False
     # we specialize the JIT for prefill and rollout
     self.prefill_jit = TinyJit(self.forward)
-    self.rollout_jit = TinyJit(self.forward)
+    self.flash_prefill_jit = TinyJit(functools.partial(self.forward, use_flash=True))
+    self.sample_prefill_jit = TinyJit(functools.partial(self.forward, sample=True))
+    self.recurrent_prefill_jits:dict[tuple[int, bool, bool], Any] = {}
+    self.rollout_jits:dict[int, Any] = {}
+    self.sample_rollout_jits:dict[int, Any] = {}
 
-  def forward(self, tokens:Tensor, start_pos:int|UOp, temperature:Tensor) -> Tensor:
+  def forward(self, tokens:Tensor, start_pos:int|UOp, temperature:Tensor, use_flash:bool=False, kv_len:int|UOp|None=None,
+              valid_len:int|UOp|None=None, sample:bool=False) -> Tensor:
     x = self.token_embd(tokens).float()                   # (B, T, D)
-    for block in self.blk: x = block(x, start_pos)
-    logits = self.output(self.output_norm(x))[:, -1, :]
+    for block in self.blk: x = block(x, start_pos, use_flash, kv_len, valid_len)
+    last = x[:, tokens.shape[1]-1:tokens.shape[1]] if valid_len is None else x[:, valid_len-1:valid_len]
+    normalized = self.output_norm(last)
+    logits = self.output(normalized)[:, -1, :]
     # Gumbel-max trick: argmax(logits/temp - log(-log(uniform))) is equivalent to sampling from softmax(logits/temp)
-    return (logits / temperature.maximum(1e-12) - (Tensor.rand_like(logits).maximum(1e-12).log().neg()).log()).argmax(-1, keepdim=True)
+    if not sample: return logits.argmax(-1, keepdim=True)
+    return (logits / temperature - (Tensor.rand_like(logits).maximum(1e-12).log().neg()).log()).argmax(-1, keepdim=True)
 
-  def __call__(self, tokens:Tensor, start_pos:int|UOp, temperature:Tensor) -> Tensor:
-    return (self.prefill_jit if resolve(tokens.shape[1] != 1) else self.rollout_jit)(tokens.contiguous(), start_pos, temperature)
+  def forward_recurrent_decode(self, tokens:Tensor, start_pos:int|UOp, temperature:Tensor, decode_len:int,
+                               valid_len:int|UOp|None=None, sample:bool=False) -> Tensor:
+    return tokens.assign(self.forward(tokens, start_pos, temperature, kv_len=decode_len, valid_len=valid_len, sample=sample))
+
+  def __call__(self, tokens:Tensor, start_pos:int|UOp, temperature:Tensor, use_flash:bool=False,
+               valid_len:int|UOp|None=None, sample:bool|None=None) -> Tensor:
+    if sample is None: sample = getattr(self, "_sample", False)
+    jit_kwargs = {"valid_len":valid_len}
+    if resolve(tokens.shape[1] == 1):
+      if self.has_recurrent_block:
+        key = self.max_context
+      else:
+        pos = start_pos.unbind()[1] if isinstance(start_pos, UOp) else start_pos
+        min_bucket = max(1, getenv("DECODE_BUCKET", 256))
+        kv_len = key = min(self.max_context, max(min_bucket, 1 << pos.bit_length()))
+      rollout_jits = self.sample_rollout_jits if sample else self.rollout_jits
+      if key not in rollout_jits:
+        rollout_jits[key] = TinyJit(functools.partial(self.forward_recurrent_decode, decode_len=key, sample=sample) if self.has_recurrent_block else
+                                    functools.partial(self.forward, kv_len=kv_len, sample=sample))
+      jit = rollout_jits[key]
+    elif self.has_recurrent_block:
+      prefill_key = (int(tokens.shape[1]), use_flash, sample)
+      if prefill_key not in self.recurrent_prefill_jits:
+        self.recurrent_prefill_jits[prefill_key] = TinyJit(functools.partial(self.forward, use_flash=use_flash, sample=sample))
+      jit = self.recurrent_prefill_jits[prefill_key]
+    else:
+      jit = self.sample_prefill_jit if sample else self.flash_prefill_jit if use_flash else self.prefill_jit
+    ret = jit(tokens.contiguous(), start_pos, temperature, **jit_kwargs)
+    return ret[0] if isinstance(ret, tuple) else ret
 
   @staticmethod
   def from_gguf(gguf:Tensor|str|pathlib.Path, max_context:int|None=None,
                 realize=bool(getenv("REALIZE", 0))) -> tuple[Transformer, dict]:
-    # TODO: remove the need for copy to default device
-    kv, state_dict = gguf_load(gguf.to(None).realize() if isinstance(gguf, Tensor) else gguf)
-
-    # all state items should be float16, not float32
-    state_dict = {k:v.cast('float16') if getenv("HALF", 1) else v for k,v in state_dict.items()}
+    kv, state_dict = gguf_load(gguf)
 
     # some models like Llama 3.2 don't have an output.weight, they just tie to the token_embd.weight
     if 'output.weight' not in state_dict: state_dict['output.weight'] = state_dict['token_embd.weight']
@@ -349,14 +580,15 @@ class Transformer:
     n_heads, n_kv_heads = kv[f'{arch}.attention.head_count'], kv[f'{arch}.attention.head_count_kv']
 
     ssm = None
-    ssm_layers: tuple[bool, ...] = ()
+    ssm_layers:tuple[bool, ...] = ()
     if arch in ('qwen35', 'qwen35moe'):
       ssm = SSMConfig(**{k: kv[f'{arch}.ssm.{k}'] for k in ('conv_kernel','state_size','group_count','time_step_rank','inner_size')})
       ssm_layers = tuple((i+1) % kv[f'{arch}.full_attention_interval'] != 0 for i in range(kv[f'{arch}.block_count']))
     elif arch == 'kimi-linear':
       ssm_layers = tuple(x == 0 for x in n_kv_heads)
       n_kv_heads = max(n_kv_heads)
-      ssm = SSMConfig(kv[f'{arch}.ssm.conv_kernel'], kv[f'{arch}.kda.head_dim'], n_heads, n_heads, n_heads*kv[f'{arch}.kda.head_dim'], kda=True)
+      ssm = SSMConfig(kv[f'{arch}.ssm.conv_kernel'], kv[f'{arch}.kda.head_dim'], n_heads, n_heads,
+                      n_heads*kv[f'{arch}.kda.head_dim'], kda=True)
       for i, is_ssm in enumerate(ssm_layers):
         if not is_ssm: continue
         state_dict[f"blk.{i}.attn_qkv.weight"] = state_dict.pop(f"blk.{i}.attn_q.weight").cat(
@@ -408,7 +640,40 @@ class Transformer:
       qkv_bias='blk.0.attn_q.bias' in state_dict,
       expert_bias=f"blk.{kv.get(f'{arch}.leading_dense_block_count', 0)}.exp_probs_b.bias" in state_dict)
     model = Transformer(config)
+    load_device = next(iter(state_dict.values())).device
+    for param in nn.state.get_parameters(model): param.replace(param.to(load_device))
+    packed_weights:set[str] = set()
+    packed_layers:list[PackedWeight] = []
+    def resolve_owner(path:list[str]):
+      obj = model
+      for part in path: obj = obj[int(part)] if isinstance(obj, list) else getattr(obj, part)
+      return obj
+    for name, weight in state_dict.items():
+      parts = name.split('.')
+      quantization = get_ggml_quantization(weight)
+      owner = resolve_owner(parts[:-1]) if parts[-1] == "weight" else None
+      packed = quantization is not None and str(load_device).startswith("AMD") and \
+        (isinstance(owner, Linear) and quantization[1] in (8, 12, 13, 14, 23) or
+        isinstance(owner, Embedding) and quantization[1] == 12 and str(load_device).startswith("AMD"))
+      if packed:
+        assert quantization is not None and isinstance(owner, PackedWeight)
+        owner.set_quantized(*quantization)
+        packed_layers.append(owner)
+        state_dict[name], packed_weights = owner.weight, packed_weights | {name}
+
+    state_dict = {k:v if k in packed_weights else v.cast('float16') if getenv("HALF", 1) else v for k,v in state_dict.items()}
     nn.state.load_state_dict(model, state_dict, verbose=False, consume=True, realize=False)  # NOTE: rope_freqs.weight (32,) is unused
+    recurrent_weights = []
+    for block in model.blk:
+      if isinstance(block, GatedDeltaNetBlock) and hasattr(block, "ssm_alpha"):
+        if block.ssm_beta.ggml_type is None and block.ssm_alpha.ggml_type is None:
+          block.ssm_beta_alpha_weight = block.ssm_beta.weight.cat(block.ssm_alpha.weight).contiguous()
+          recurrent_weights.append(block.ssm_beta_alpha_weight)
+    if recurrent_weights: Tensor.realize(*recurrent_weights)
+    # Custom kernels need the shared GGUF buffer and byte offset before function tracing disables device access.
+    packed_offsets = [layer._packed_offset() for layer in packed_layers]
+    if packed_offsets: Tensor.realize(*packed_offsets)
+    for layer,offset in zip(packed_layers, packed_offsets): layer._raw_offset_uop = offset.uop
     # NOTE: without this contiguous, it unpacks the weights from the model every time. we shouldn't need this, but for now it's faster
     if realize:
       for s in (params:=nn.state.get_parameters(model)): s.replace(s.contiguous())
@@ -417,25 +682,168 @@ class Transformer:
 
   def get_start_pos(self, tokens:list[int]) -> int:
     prefix_len = sum(1 for _ in itertools.takewhile(lambda ab: ab[0] == ab[1], zip(tokens[:-1], self._cached_tokens)))
+    # Recurrent state has no token dimension to slice. Roll back to its latest aligned checkpoint so resumed flash
+    # prefill uses the same global chunk boundaries as a full prompt.
+    if self.has_recurrent_block:
+      return self._state_checkpoint_pos if prefix_len >= self._state_checkpoint_pos else 0
     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
+  def _recurrent_states(self) -> list[Tensor]:
+    return [getattr(block, name) for block in self.blk for name in ("conv_state", "recurrent_state") if hasattr(block, name)]
+
+  def _init_state_checkpoints(self):
+    if not self._state_checkpoints:
+      states = self._recurrent_states()
+      if not states: return
+      self._state_checkpoints = [Tensor.zeros_like(state).contiguous().realize() for state in states]
+      def copy_jit(pairs:list[tuple[Tensor, Tensor]]) -> Any:
+        def copy_states() -> Tensor:
+          copies = [dest.assign(src) for dest,src in pairs]
+          Tensor.realize(*copies)
+          return copies[-1]
+        jit = TinyJit(copy_states)
+        jit()
+        jit()
+        return jit
+      self._save_state_jit = copy_jit(list(zip(self._state_checkpoints, states)))
+      self._restore_state_jit = copy_jit(list(zip(states, self._state_checkpoints)))
+
+  def _save_state_checkpoint(self, pos:int):
+    self._init_state_checkpoints()
+    if self._save_state_jit is None: return
+    self._save_state_jit()
+    self._state_checkpoint_pos = pos
+
+  def _restore_state_checkpoint(self):
+    assert self._restore_state_jit is not None
+    self._restore_state_jit()
+
+  def warmup(self, chunk_size:int=256):
+    device = self.token_embd.weight.device
+    direct_capture = not self.has_recurrent_block and all(isinstance(block, TransformerBlock) for block in self.blk)
+    if direct_capture:
+      device = str(self.token_embd.weight.device)
+      direct_capture = device.startswith("AMD") and Device[device].renderer.target.arch.startswith("gfx11")
+
+    # Recurrent prefill has one fixed padded shape with symbolic valid length.
+    recurrent_chunk = min(chunk_size, 256)
+    warm_len = min(recurrent_chunk if self.has_recurrent_block else chunk_size * 2, self.max_context - 1)
+    if warm_len > 0:
+      if direct_capture:
+        x = Tensor.zeros(1, 1, self.blk[0].config.dim, device=device)
+        for block in self.blk: block._init_state(x)
+        Tensor.realize(*[state for block in self.blk for state in (getattr(block, "cache_kv"), getattr(block, "freqs_cis"))])
+        self.flash_prefill_jit.cnt = 1
+        next(self.generate([0] * warm_len, chunk_size=chunk_size))
+      elif self.has_recurrent_block:
+        # State creation must happen outside JIT capture: capturing _init_state makes the first real request
+        # reuse initialization buffers instead of the persistent recurrent/KV state.
+        x = Tensor.zeros(1, 1, self.blk[0].config.dim, device=device)
+        for block in self.blk: block._init_state(x)
+        states = [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)]
+        Tensor.realize(*states)
+        self._init_state_checkpoints()
+        self.prefill_jit.cnt = self.flash_prefill_jit.cnt = 1
+        self.recurrent_prefill_jits[(warm_len, False, False)] = self.prefill_jit
+        self.recurrent_prefill_jits[(warm_len, True, False)] = self.flash_prefill_jit
+        decode_len = self.max_context
+        self.rollout_jits[decode_len] = TinyJit(
+          functools.partial(self.forward_recurrent_decode, decode_len=decode_len, sample=False))
+        self.rollout_jits[decode_len].cnt = 1
+        self._warming_up = True
+        warm = self.generate([0] * warm_len, chunk_size=chunk_size)
+        prefill_batch = getenv("PREFILL_JIT_BATCH_SIZE", 128)
+        with Context(JIT_BATCH_SIZE=prefill_batch): next(warm)
+        with Context(JIT_BATCH_SIZE=0): next(warm)
+        self._warming_up = False
+      else:
+        for salt in range(2): next(self.generate([salt] + [0] * (warm_len - 1), chunk_size=chunk_size))
+
+    # Rollout uses fixed power-of-two KV shapes. Capture every shape up front so requests never pay a JIT transition.
+    if not self.has_recurrent_block:
+      min_bucket = max(1, getenv("DECODE_BUCKET", 256))
+      bucket_positions:dict[int, int] = {}
+      for pos in [0] + [1 << i for i in range(self.max_context.bit_length())]:
+        bucket = min(self.max_context, max(min_bucket, 1 << pos.bit_length()))
+        bucket_positions.setdefault(bucket, pos)
+      v_start_pos = UOp.variable("start_pos", 0, self.max_context-1)
+      token, temperature = Tensor([[0]], dtype="int32", device=device), Tensor([0.0], device=device)
+      for bucket, pos in sorted(bucket_positions.items()):
+        if direct_capture:
+          self.rollout_jits[bucket] = TinyJit(functools.partial(self.forward, kv_len=bucket))
+          self.rollout_jits[bucket].cnt = 1
+        for _ in range(1 if direct_capture else 2):
+          result = self(token, v_start_pos.bind(pos), temperature)
+          assert isinstance(result, Tensor)
+          result.realize()
+
+    if self._state_checkpoints:
+      # Recurrent warmup starts from the zero checkpoint. Restore it directly instead of scheduling state clears and
+      # then copying the same zeros back into the checkpoint.
+      self._restore_state_checkpoint()
+      self._state_checkpoint_pos = 0
+    elif 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):
+    start_pos = self.get_start_pos(tokens)
+    self._sample = temperature > 0
+    chunk_size = min(chunk_size or 256, 256) if self.has_recurrent_block else chunk_size or 32
     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
-    temp = Tensor([temperature])
-    # assign all input tokens once, then slice from start_pos for the model call
-    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)
+    device = self.token_embd.weight.device
+    temp = Tensor([temperature], device=device)
+    # Dense attention needs a symbolic slice into one input buffer. Recurrent prefill instead creates fixed-size
+    # chunk tensors below; allocating its input at max_context makes short requests spend most of their time converting zeros.
+    t = None if self.has_recurrent_block else \
+      Tensor(tokens + [0] * (self.max_context + chunk_size - len(tokens)), dtype="int32", device=device).reshape(1, self.max_context + chunk_size)
+    # start_pos describes what's currently valid in the caches
+    if start_pos < len(self._cached_tokens):
+      if self.has_recurrent_block and self._state_checkpoints and start_pos == self._state_checkpoint_pos:
+        self._restore_state_checkpoint()
+      elif resets := [r for b in self.blk for r in b._state_reset_ops()]:
+        Tensor.realize(*resets)
+        if self._state_checkpoints: self._save_state_checkpoint(0)
     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)
-      out = self(t[:, sp:sp+nt] if start_pos < prompt_len or out is None else out, sp, temp).realize()
-      start_pos += n_toks
+      remaining = len(tokens) - start_pos
+      recurrent_prefill = self.has_recurrent_block and start_pos < prompt_len
+      can_flash = bool(getenv("AMD_FLASH_ATTENTION", 1)) and chunk_size % 64 == 0 and \
+                  (recurrent_prefill if self.has_recurrent_block else start_pos > 0)
+      if can_flash:
+        device = str(self.token_embd.weight.device)
+        can_flash = device.startswith("AMD") and Device[device].renderer.target.arch.startswith("gfx11")
+      use_flash = can_flash and (self.has_recurrent_block or start_pos % 64 == 0)
+      sp = v_start_pos.bind(start_pos)
+      # Dense attention aligns cache reuse to a 64-token flash tile. Recurrent prefill always uses its fixed,
+      # padded chunk shape, and key_limit excludes the padding from attention.
+      actual_nt = min(chunk_size, remaining)
+      nt = chunk_size if use_flash or self.has_recurrent_block and start_pos < prompt_len else 1 if self.has_recurrent_block else \
+           v_toks.bind(min(64 - start_pos % 64, remaining) if can_flash else actual_nt)
+      if self.has_recurrent_block and (start_pos < prompt_len or out is None):
+        assert isinstance(nt, int)
+        inp = Tensor(tokens[start_pos:start_pos+actual_nt] + [0] * (nt-actual_nt), dtype="int32", device=device).reshape(1, nt)
+      elif start_pos < prompt_len or out is None:
+        assert t is not None
+        inp = t[:, sp:sp+nt]
+      else: inp = out
+      valid_len = v_toks.bind(actual_nt) if recurrent_prefill or use_flash and actual_nt < chunk_size else None
+      # Save once immediately before a short final chunk. This is the nearest globally aligned state that can be
+      # reused without changing flash-attention's numerical tile layout.
+      if not self._warming_up and recurrent_prefill and remaining < chunk_size and start_pos % chunk_size == 0:
+        self._save_state_checkpoint(start_pos)
+      if use_flash: result = self(inp, sp, temp, use_flash=True, valid_len=valid_len)
+      elif valid_len is not None: result = self(inp, sp, temp, valid_len=valid_len)
+      else: result = self(inp, sp, temp)
+      out = result.realize()
+      start_pos += actual_nt
+      # Generated tool calls are reconstructed by clients and are not guaranteed token-identical on the next request.
+      # Keep the reusable checkpoint at the stable prompt boundary instead of overwriting it inside generated output.
+      if not self._warming_up and self.has_recurrent_block and start_pos == prompt_len and start_pos % chunk_size == 0:
+        self._save_state_checkpoint(start_pos)
       # chunked prefill: keep processing until all prompt tokens are consumed
       if start_pos < len(tokens): continue
       tokens.append(int(out.item()))

tinygrad/runtime/support/system.py

@@ -269,7 +269,8 @@ class PCIIfaceBase:
     if should_use_sysmem:
       vaddr = self.dev_impl.mm.alloc_vaddr(size:=round_up(size, mmap.PAGESIZE), align=mmap.PAGESIZE)
       memview, paddrs = self.pci_dev.alloc_sysmem(size, vaddr=vaddr, contiguous=contiguous)
-      mapping = self.dev_impl.mm.map_range(vaddr, size, [(paddr, 0x1000) for paddr in paddrs], aspace=AddrSpace.SYS, snooped=True, uncached=True)
+      mapping = self.dev_impl.mm.map_range(vaddr, size, [(paddr, 0x1000) for paddr in paddrs], aspace=AddrSpace.SYS,
+                                           snooped=True, uncached=uncached)
       return HCQBuffer(vaddr, size, meta=PCIAllocationMeta(mapping, has_cpu_mapping=True, hMemory=paddrs[0]), view=memview, owner=self.dev)
 
     mapping = self.dev_impl.mm.valloc(size:=round_up(size, 0x1000), uncached=uncached, contiguous=cpu_access)