From ca78469398b7233d4a9829edbe2005e9d7fade6b Mon Sep 17 00:00:00 2001 From: George Hotz Date: Tue, 4 Aug 2026 01:08:49 +0000 Subject: [PATCH] cleanups --- test/unit/test_attention.py | 44 ++++++++++++++++--- tinygrad/llm/kernels/__init__.py | 58 +++++++++++-------------- tinygrad/llm/kernels/amd.py | 59 ++++++++++---------------- tinygrad/llm/kernels/generic.py | 25 ++++++----- tinygrad/llm/model.py | 72 +++++++++++++++++++------------- 5 files changed, 143 insertions(+), 115 deletions(-) diff --git a/test/unit/test_attention.py b/test/unit/test_attention.py index 250f6c9197..de694a6265 100644 --- a/test/unit/test_attention.py +++ b/test/unit/test_attention.py @@ -1,10 +1,11 @@ import unittest import numpy as np -from tinygrad import Tensor, dtypes +from tinygrad import Tensor, dtypes, nn from tinygrad.llm.model import ( GatedDeltaNetBlock, Linear, SSMConfig, TransformerBlock, TransformerConfig, apply_rope as apply_rope_new, precompute_freqs_cis, pairwise_topk, ) +from tinygrad.llm.kernels import gated_delta_prefill from tinygrad.llm.gguf import ggml_data_to_tensor def apply_rope(x:Tensor, start_pos:int): @@ -51,6 +52,22 @@ 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(np.linspace(start, stop, int(np.prod(shape)), dtype=np.float32).reshape(shape), device=Tensor.empty(1).device).realize() @@ -195,17 +212,32 @@ class TestGatedDeltaNetBlock(unittest.TestCase): def test_kda_channel_decay(self): 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.]]]) - # f_b(f_a(x)) = [1, 2, 3, 4] + 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) + ret = block._attention(x, 0) + (block._update_state(*ret) if isinstance(ret, tuple) else ret).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) + Tensor.realize(*block._state_reset_ops()) + 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) class TestPairwiseTopk(unittest.TestCase): def test_basic_topk(self): diff --git a/tinygrad/llm/kernels/__init__.py b/tinygrad/llm/kernels/__init__.py index 3e8f57fa28..835cace61c 100644 --- a/tinygrad/llm/kernels/__init__.py +++ b/tinygrad/llm/kernels/__init__.py @@ -1,17 +1,29 @@ -from typing import Any, cast from tinygrad import Tensor, UOp, nn, dtypes -from tinygrad.llm.kernels.amd import Linear -from tinygrad.uop.ops import resolve +from tinygrad.helpers import prod +from tinygrad.uop.ops import Ops, resolve -def select_cache_dtype(device:str|tuple[str, ...]|None, recurrent:bool, max_context:int): - return dtypes.int8 if recurrent and max_context > 8192 and str(device).startswith("AMD") else dtypes.default_float - -def make_attention_cache(batch:int|UOp, n_kv_heads:int, max_context:int, head_dim:int, device, recurrent:bool): - dtype, cache_len = select_cache_dtype(device, recurrent, max_context), (max_context+255)//256*256 if recurrent else max_context - shape = (2, batch, n_kv_heads, cache_len, head_dim) - cache = Tensor.empty(*shape, dtype=dtype, device=device).contiguous() - scale = Tensor.empty(*shape[:-1], dtype=dtypes.float16, device=device).contiguous() if dtype == dtypes.int8 else None - return cache, scale, cache_len +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._raw_offset_uop:UOp|None = None + def set_quantized(self, decoded:Tensor) -> Tensor|None: + 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 None + self.weight, self.ggml_type = Tensor(raw).flatten(), packed_sizes[prod(raw.shape)] + raw_offset = self.weight.uop.contiguous_view_offset() + assert raw_offset is not None and raw_offset % 4 == 0 and self.weight.uop.buf_uop.dtype == dtypes.uint8 + if self.ggml_type == 23 and str(self.weight.device).startswith("AMD"): + from tinygrad.llm.kernels.amd import iq4_half_lut + iq4_half_lut(str(self.weight.device)) + return Tensor([raw_offset // 4], dtype=dtypes.uint64, device=self.weight.device) + def __call__(self, x:Tensor) -> Tensor: + if self.ggml_type in (13, 14, 23) and str(self.weight.device).startswith("AMD"): + from tinygrad.llm.kernels.amd import q8_linear + return q8_linear(self, x) + return super().__call__(x) def cached_attention(q:Tensor, stacked_kv:Tensor, cache_kv:Tensor, cache_scale:Tensor|None, start_pos:int|UOp, max_context:int) -> Tensor: @@ -27,28 +39,8 @@ def cached_attention(q:Tensor, stacked_kv:Tensor, cache_kv:Tensor, cache_scale:T return q.scaled_dot_product_attention(k, v, attn_mask=mask, enable_gqa=True) def gated_delta_prefill(q:Tensor, k:Tensor, v:Tensor, beta:Tensor, alpha:Tensor, state:Tensor) -> Tensor: - if str(q.device).startswith("AMD"): + if str(q.device).startswith("AMD") and q.shape[-1] % 32 == 0 and v.shape[-1] % 4 == 0: from tinygrad.llm.kernels.amd import gated_delta_prefill as kernel else: from tinygrad.llm.kernels.generic import gated_delta_prefill as kernel return kernel(q, k, v, beta, alpha, state) - -def _prepare_quantized_weights(model:Any, state_dict:dict[str, Tensor]) -> list[tuple[Linear, Tensor]]: - packed:list[tuple[Linear, Tensor]] = [] - layers = cast(dict[str, Linear], nn.state.get_state_dict(model, tensor_type=Linear)) - for name,owner in layers.items(): - key, weight = f"{name}.weight", state_dict[f"{name}.weight"] - if str(weight.device).startswith("AMD") and (offset:=owner.set_quantized(weight)) is not None: - packed.append((owner, offset)) - state_dict[key] = owner.weight - return packed - -def load_state_dict(model:Any, state_dict:dict[str, Tensor]): - for key in nn.state.get_state_dict(model): - if key.endswith(".ssm_beta_alpha.weight") and key not in state_dict: - prefix = key.removesuffix("beta_alpha.weight") - state_dict[key] = state_dict.pop(prefix+"beta.weight").cat(state_dict.pop(prefix+"alpha.weight"), dim=0).contiguous() - packed = _prepare_quantized_weights(model, state_dict) - nn.state.load_state_dict(model, state_dict, verbose=False, consume=True, realize=False) - if packed: Tensor.realize(*(offset for _,offset in packed)) - for layer,offset in packed: layer._raw_offset_uop = offset.uop diff --git a/tinygrad/llm/kernels/amd.py b/tinygrad/llm/kernels/amd.py index e904d49a40..cad9f06314 100644 --- a/tinygrad/llm/kernels/amd.py +++ b/tinygrad/llm/kernels/amd.py @@ -1,8 +1,8 @@ from __future__ import annotations import functools, math from typing import Callable, cast -from tinygrad import Tensor, UOp, nn -from tinygrad.helpers import prod +from tinygrad import Tensor, UOp +from tinygrad.llm.kernels import Linear from tinygrad.uop.ops import AxisType, KernelInfo, Ops, resolve from tinygrad.dtype import AddrSpace, dtypes @@ -13,24 +13,6 @@ WMMA_ACC, THREADS_PER_BLOCK = WMMA_M // LANES_PER_WAVE_M, WARP_SIZE * WAVES_M * 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 -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._raw_offset_uop:UOp|None = None - def set_quantized(self, decoded:Tensor) -> Tensor|None: - packed_sizes = {decoded.numel() // 256 * type_size:typ for typ,type_size in ((Q5_K, 176), (Q6_K, 210), (IQ4_XS, 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 None - self.weight, self.ggml_type = Tensor(raw).flatten(), packed_sizes[prod(raw.shape)] - raw_offset = self.weight.uop.contiguous_view_offset() - assert raw_offset is not None and raw_offset % 4 == 0 and self.weight.uop.buf_uop.dtype == dtypes.uint8 - if self.ggml_type == IQ4_XS and str(self.weight.device).startswith("AMD"): iq4_half_lut(str(self.weight.device)) - return Tensor([raw_offset // 4], dtype=dtypes.uint64, device=self.weight.device) - def __call__(self, x:Tensor) -> Tensor: - return q8_linear(self, x) if self.ggml_type in (Q5_K, Q6_K, IQ4_XS) and str(self.weight.device).startswith("AMD") else super().__call__(x) - def warp_reduce(val:UOp, lane: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]): idx = ((lane ^ offset) * 4).int() @@ -238,40 +220,45 @@ def q8_quantize(x:Tensor, tokens:int, in_features:int) -> tuple[Tensor, Tensor]: @functools.cache def _gated_delta_prefill_kernel(core:UOp, q:UOp, k:UOp, v:UOp, beta:UOp, alpha:UOp, state:UOp, kq:UOp) -> UOp: - batch, heads, tokens, dim, row_tile = *core.shape, 4 - assert all(isinstance(x, int) for x in (batch, heads, tokens, dim)) and dim % 32 == 0 and dim % row_tile == 0 - batch, heads, tokens, dim = cast(tuple[int, int, int, int], (batch, heads, tokens, dim)) - core, q, k, v = (x.reshape(batch*heads, tokens, dim) for x in (core, q, k, v)) - beta, alpha, kq = (x.reshape(batch*heads, tokens) for x in (beta, alpha, kq)) - state = state.reshape(batch*heads, dim, dim) - bh_row, lane = UOp.range(batch*heads*dim//row_tile, 0), UOp.range(32, 1, axis_type=AxisType.LOCAL) - bh, row_base = bh_row // (dim//row_tile), (bh_row % (dim//row_tile))*row_tile + 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(dim//32)) - current = UOp.placeholder((row_tile*dim//32,), dtypes.float32, slot=0, addrspace=AddrSpace.REG) + cols = tuple(lane + i*32 for i in range(key_dim//32)) + current = UOp.placeholder((row_tile*key_dim//32,), dtypes.float32, slot=0, addrspace=AddrSpace.REG) current = current.after(current.store(UOp.stack(*(state[bh, row, col].float() for row in rows for col in cols)))) 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) - av, bv = alpha[bh, token].load(), beta[bh, token].load() updates:list[UOp] = [] stores:list[UOp] = [] for row_idx,row in enumerate(rows): - previous = tuple(current.after(token)[row_idx*dim//32+i].load() for i in range(dim//32)) + 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)), lane, full_wave=True) state_q = warp_reduce(sum((x*y for x,y in zip(previous, queries)), UOp.const(0, dtypes.float32)), lane, 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*dim//32+i].load().cast(state.dtype)) + 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 gated_delta_prefill(q:Tensor, k:Tensor, v:Tensor, beta:Tensor, alpha:Tensor, state:Tensor) -> Tensor: - batch, heads, tokens, dim = q.shape - assert q.shape == k.shape == v.shape and beta.shape == alpha.shape == (batch, heads, tokens) and state.shape == (batch, heads, dim, dim) - core, kq = Tensor.empty_like(q), (q*k).sum(-1).contiguous() + 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) + core, kq = Tensor.empty_like(v), (q*k).sum(-1).contiguous() return Tensor.custom_kernel(core, q.contiguous(), k.contiguous(), v.contiguous(), beta.contiguous(), alpha.contiguous(), state, kq, fxn=_gated_delta_prefill_kernel)[0] diff --git a/tinygrad/llm/kernels/generic.py b/tinygrad/llm/kernels/generic.py index a8947183d9..3f4d036313 100644 --- a/tinygrad/llm/kernels/generic.py +++ b/tinygrad/llm/kernels/generic.py @@ -6,17 +6,19 @@ from tinygrad.uop.ops import AxisType, KernelInfo @functools.cache def _gated_delta_prefill_kernel(core:UOp, q:UOp, k:UOp, v:UOp, beta:UOp, alpha:UOp, state:UOp, kq:UOp) -> UOp: - batch, heads, tokens, dim = cast(tuple[int, int, int, int], core.shape) - core, q, k, v = (x.reshape(batch*heads, tokens, dim) for x in (core, q, k, v)) - beta, alpha, kq = (x.reshape(batch*heads, tokens) for x in (beta, alpha, kq)) - state = state.reshape(batch*heads, dim, dim) - bh, row, cols = UOp.range(batch*heads, 0, AxisType.GLOBAL), UOp.range(dim, 2), tuple(range(dim)) - current = UOp.placeholder((dim,), dtypes.float32, slot=0, addrspace=AddrSpace.REG) + batch, heads, tokens, value_dim = cast(tuple[int, int, int, int], core.shape) + key_dim, alpha_dim = cast(int, q.shape[-1]), cast(int, alpha.shape[-1]) if len(alpha.shape) == 4 else 1 + core, v = (x.reshape(batch*heads, tokens, value_dim) for x in (core, v)) + q, k = (x.reshape(batch*heads, tokens, key_dim) for x in (q, k)) + beta, kq = (x.reshape(batch*heads, tokens) for x in (beta, kq)) + alpha, state = alpha.reshape(batch*heads, tokens, alpha_dim), state.reshape(batch*heads, value_dim, key_dim) + bh, row, cols = UOp.range(batch*heads, 0, AxisType.GLOBAL), UOp.range(value_dim, 2), tuple(range(key_dim)) + current = UOp.placeholder((key_dim,), dtypes.float32, slot=0, addrspace=AddrSpace.REG) current = current.after(UOp.group(*(current[col].store(state[bh, row, col].float()) for col in cols))) token = UOp.range(tokens, 1, AxisType.REDUCE) previous = tuple(current.after(token)[col].load() for col in cols) keys, queries = (tuple(x[bh, token, col].load() for col in cols) for x in (k, q)) - av, bv = alpha[bh, token].load(), beta[bh, token].load() + 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 @@ -26,8 +28,11 @@ def _gated_delta_prefill_kernel(core:UOp, q:UOp, k:UOp, v:UOp, beta:UOp, alpha:U return UOp.group(*stores).end(row, bh).sink(arg=KernelInfo(name="gated_delta_prefill", opts_to_apply=())) def gated_delta_prefill(q:Tensor, k:Tensor, v:Tensor, beta:Tensor, alpha:Tensor, state:Tensor) -> Tensor: - batch, heads, tokens, dim = q.shape - assert q.shape == k.shape == v.shape and beta.shape == alpha.shape == (batch, heads, tokens) and state.shape == (batch, heads, dim, dim) - core, kq = Tensor.empty_like(q), (q*k).sum(-1).contiguous() + 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) + core, kq = Tensor.empty_like(v), (q*k).sum(-1).contiguous() return Tensor.custom_kernel(core, q.contiguous(), k.contiguous(), v.contiguous(), beta.contiguous(), alpha.contiguous(), state, kq, fxn=_gated_delta_prefill_kernel)[0] diff --git a/tinygrad/llm/model.py b/tinygrad/llm/model.py index 7dad162056..0203dba087 100644 --- a/tinygrad/llm/model.py +++ b/tinygrad/llm/model.py @@ -1,11 +1,28 @@ from __future__ import annotations import functools, itertools, pathlib from dataclasses import dataclass, replace -from tinygrad import Tensor, nn, UOp, TinyJit, getenv, function, Context -from tinygrad.llm.kernels import Linear, cached_attention, gated_delta_prefill, load_state_dict, make_attention_cache +from typing import Any, cast +from tinygrad import Tensor, nn, UOp, TinyJit, getenv, function, Context, dtypes +from tinygrad.llm.kernels import Linear, cached_attention, gated_delta_prefill from tinygrad.llm.gguf import gguf_load from tinygrad.uop.ops import resolve +def _prepare_quantized_weights(model:Any, state_dict:dict[str, Tensor]) -> list[tuple[Linear, Tensor]]: + packed:list[tuple[Linear, Tensor]] = [] + layers = cast(dict[str, Linear], nn.state.get_state_dict(model, tensor_type=Linear)) + for name,owner in layers.items(): + key, weight = f"{name}.weight", state_dict[f"{name}.weight"] + if str(weight.device).startswith("AMD") and (offset:=owner.set_quantized(weight)) is not None: + packed.append((owner, offset)) + state_dict[key] = owner.weight + return packed + +def load_state_dict(model:Any, state_dict:dict[str, Tensor]): + packed = _prepare_quantized_weights(model, state_dict) + nn.state.load_state_dict(model, state_dict, verbose=False, consume=True, realize=False) + if packed: Tensor.realize(*(offset for _,offset in packed)) + for layer,offset in packed: layer._raw_offset_uop = offset.uop + @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)) @@ -179,9 +196,12 @@ class TransformerBlock(FFNBlock): def _init_state(self, x:Tensor): if not hasattr(self, "cache_kv"): - self.cache_kv, cache_scale, cache_len = make_attention_cache( - x.shape[0], self.config.n_kv_heads, self.config.max_context, self.config.head_dim, x.device, self.config.ssm is not None) - if cache_scale is not None: self.cache_kv_scale = cache_scale + recurrent = self.config.ssm is not None + cache_len, cache_dtype = ((self.config.max_context+255)//256*256 if recurrent else self.config.max_context), \ + (dtypes.int8 if recurrent and self.config.max_context > 8192 and str(x.device).startswith("AMD") else dtypes.default_float) + shape = (2, x.shape[0], self.config.n_kv_heads, cache_len, self.config.head_dim) + self.cache_kv = Tensor.empty(*shape, dtype=cache_dtype, device=x.device).contiguous() + if cache_dtype == dtypes.int8: self.cache_kv_scale = Tensor.empty(*shape[:-1], dtype=dtypes.float16, device=x.device).contiguous() self.freqs_cis = precompute_freqs_cis(self.config.rope_dim, cache_len, self.config.rope_theta, device=x.device) class MLATransformerBlock(FFNBlock): @@ -258,36 +278,25 @@ class GatedDeltaNetBlock(FFNBlock): beta, alpha = (self.ssm_beta(x), self.ssm_f_b(self.ssm_f_a(x))) if is_kda else \ self.ssm_beta_alpha(x).split(self.num_v_heads, dim=-1) 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) if is_kda else functools.reduce(lambda a,b: a+b, + 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, k = (z.reshape(B, T, self.num_k_heads, self.head_k_dim).normalize(dim=-1, eps=1e-12 if is_kda else 1e-6).repeat( 1, 1, self.num_v_heads//self.num_k_heads, 1) for z in (q, k)) v = v.reshape(B, T, self.num_v_heads, self.head_v_dim) - state_pos:int|UOp - if is_kda: - assert T == 1, "channel-wise gated delta prefill is not supported" - beta = beta.sigmoid().reshape(B, self.num_v_heads, 1, 1) - 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) - q, k, v = q[:, 0].mul(self.head_k_dim**-0.5).unsqueeze(-1), k[:, 0].unsqueeze(-1), v[:, 0].unsqueeze(-1) - recurrent_state = self.recurrent_state * alpha - recurrent_state = recurrent_state + ((v - recurrent_state@k) * beta)@k.transpose(-1, -2) - recurrent_state = Tensor(self.recurrent_state.uop.after( - self.recurrent_state.uop.store(recurrent_state.cast(self.recurrent_state.dtype).uop))) - core, gate, state_pos = (recurrent_state@q).squeeze(-1).reshape(B, 1, self.num_v_heads, self.head_v_dim), out_gate.sigmoid(), 1 - else: - beta, log_alpha = beta.sigmoid().reshape(B, T, self.num_v_heads), \ - ((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 - q, k, v, beta, log_alpha = [z.transpose(1, 2).float() for z in (q, k, v, beta, log_alpha)] - core = gated_delta_prefill(q * self.head_k_dim**-0.5, k, v, beta, log_alpha.exp(), self.recurrent_state).transpose(1, 2) - gate, state_pos = out_gate.silu(), T if valid_len is None else valid_len + beta = beta.sigmoid().reshape(B, T, self.num_v_heads) + 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) + 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.unsqueeze(-1) if is_kda else active) + q, k, v, beta = [z.transpose(1, 2).float() for z in (q, k, v, beta)] + alpha = log_alpha.transpose(1, 2).float().exp() + core = gated_delta_prefill(q * self.head_k_dim**-0.5, k, v, beta, alpha, self.recurrent_state).transpose(1, 2) + gate, state_pos = (out_gate.sigmoid() if is_kda else out_gate.silu()), T if valid_len is None else valid_len out = self.ssm_out((self.ssm_norm(core) * gate).reshape(B, T, -1).cast(x.dtype)).contiguous() conv_state = conv_window[:, state_pos:state_pos+self.ssm_conv_kernel-1].cast(self.conv_state.dtype).contiguous() - return Tensor(out.uop.after(self.conv_state.uop.after(self.conv_state.uop.store(conv_state.uop)))) if is_kda else (out, conv_state) + return out, conv_state def _update_state(self, out:Tensor, *state:Tensor) -> Tensor: conv_state, = state @@ -302,7 +311,7 @@ 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() + self.recurrent_state = Tensor.zeros(x.shape[0], self.num_v_heads, self.head_v_dim, self.head_k_dim, device=x.device).clone() class Transformer: def __init__(self, config:TransformerConfig): @@ -317,7 +326,7 @@ class Transformer: 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.has_recurrent_prefill = config.ssm is not None and not config.ssm.kda + self.has_recurrent_prefill = config.ssm is not None self._cached_tokens: list[int] = [] # we specialize the JIT for prefill and rollout self.prefill_jit = TinyJit(self.forward) @@ -361,6 +370,9 @@ class Transformer: 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'])) + for i,is_ssm in enumerate(ssm_layers): + if is_ssm: state_dict[f"blk.{i}.ssm_beta_alpha.weight"] = state_dict.pop(f"blk.{i}.ssm_beta.weight").cat( + state_dict.pop(f"blk.{i}.ssm_alpha.weight"), dim=0).contiguous() elif arch == 'kimi-linear': ssm_layers = tuple(x == 0 for x in n_kv_heads) n_kv_heads = max(n_kv_heads)