sirhcm/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests 2 Actions 8 Activity

cleanups

Browse Source
This commit is contained in:
geohot
2026-08-04 01:08:49 +00:00
parent 70ed0a70e0
commit ca78469398
5 changed files with 143 additions and 115 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/unit/test_attention.py
+38 -6
View File
@@ -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):
tinygrad/llm/kernels/__init__.py
+25 -33
View File
@@ -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
tinygrad/llm/kernels/amd.py
+23 -36
View File
@@ -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]
tinygrad/llm/kernels/generic.py
+15 -10
View File
@@ -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]
tinygrad/llm/model.py
+42 -30
View File
@@ -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)