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

400 lines

Browse Source
This commit is contained in:
geohot
2026-08-03 16:28:20 +00:00
parent 79ef0a16c0
commit f26128e8d2
8 changed files with 157 additions and 250 deletions
Download Patch File Download Diff File Expand all files Collapse all files
extra/benchmark_qwen_amd.py
+11 -4
View File
@@ -12,27 +12,33 @@ if __name__ == "__main__":
parser.add_argument("--skip-resume-check", action="store_true")
args = parser.parse_args()
st = time.perf_counter()
startup_st = st = time.perf_counter()
model, _ = Transformer.from_gguf(args.model, args.max_context)
print(f"load {time.perf_counter()-st:.3f}s", flush=True)
st = time.perf_counter()
with Context(BEAM=0): model.warmup(args.chunk_size)
print(f"warm {time.perf_counter()-st:.3f}s", flush=True)
assert time.perf_counter()-startup_st < 60
states = [getattr(block, name) for block in model.blk for name in ("cache_kv", "cache_kv_scale", "conv_state", "recurrent_state")
if hasattr(block, name)]
assert all(str(state.device).startswith("AMD") and state.uop.is_realized for state in states)
assert all(block.cache_kv.shape[3] >= args.max_context for block in model.blk if hasattr(block, "cache_kv"))
assert model.flash_prefill_jit.cnt >= 2 and model.recurrent_rollout_jit.cnt >= 2
assert model.prefill_jit.cnt >= 2 and model.recurrent_rollout_jit.cnt >= 2
print(f"preallocated {sum(state.nbytes() for state in states)/2**30:.3f} GiB state on AMD", flush=True)
prompt = [257] + [1000+i%1000 for i in range(args.prompt_tokens-1)]
gen, st = model.generate(prompt, chunk_size=args.chunk_size), time.perf_counter()
output = [next(gen)]
pt = time.perf_counter()
print(f"prefill {args.prompt_tokens/(pt-st):.3f} tok/s", flush=True)
prefill = args.prompt_tokens/(pt-st)
print(f"prefill {prefill:.3f} tok/s", flush=True)
for _ in range(args.decode_tokens): output.append(next(gen))
et = time.perf_counter()
print(f"decode {args.decode_tokens/(et-pt):.3f} tok/s output {output}", flush=True)
decode = args.decode_tokens/(et-pt)
print(f"decode {decode:.3f} tok/s output {output}", flush=True)
if args.prompt_tokens == 3000 and args.decode_tokens == 16:
assert prefill > 800 and decode > 40
assert output == [13, 271, 248068, 198, 8160, 579, 264, 7047, 1817, 25, 271, 16, 13, 220, 2972, 2014, 53983]
if not args.skip_resume_check:
follow = model._cached_tokens + [1234+i for i in range(8)]
@@ -44,3 +50,4 @@ if __name__ == "__main__":
st = time.perf_counter()
full_token = next(model.generate(full_prompt, chunk_size=args.chunk_size))
print(f"full {time.perf_counter()-st:.3f}s token {full_token} match {resumed_token == full_token}", flush=True)
assert resumed_token == full_token
test/unit/test_attention.py
+4 -5
View File
@@ -16,10 +16,10 @@ def apply_rope(x:Tensor, start_pos:int):
class TestLinear(unittest.TestCase):
def test_recovers_packed_ggml_weight(self):
for ggml_type,packed_size in ((13, 176), (14, 210), (23, 136)):
packed = Tensor.empty(packed_size+1, dtype=dtypes.uint8, device="CPU")[1:]
packed = Tensor.empty(packed_size+4, dtype=dtypes.uint8, device="CPU")[4:]
decoded = ggml_data_to_tensor(packed, 256, ggml_type).reshape(1, 256)
linear = Linear(256, 1, bias=False)
self.assertTrue(linear.set_quantized(decoded))
self.assertIsNotNone(linear.set_quantized(decoded))
self.assertEqual((linear.ggml_type, linear.weight.shape), (ggml_type, (packed_size,)))
class TestAttention(unittest.TestCase):
@@ -77,9 +77,8 @@ class TestGatedDeltaNetBlock(unittest.TestCase):
def _run_attention(self, block:GatedDeltaNetBlock, x:Tensor, start_pos:int):
x_norm = block.attn_norm(x)
block._init_state(x_norm)
out = block._attention(x_norm, start_pos).realize()
if block.pending_state is not None:
Tensor.realize(block.conv_state.assign(block.pending_state[0]), block.recurrent_state.assign(block.pending_state[1]))
out, conv_state, recurrent_state = block._attention(x_norm, start_pos)
Tensor.realize(out, block.conv_state.assign(conv_state), block.recurrent_state.assign(recurrent_state))
return out.numpy()
def _cache_views(self, block:GatedDeltaNetBlock) -> tuple[np.ndarray, np.ndarray]:
test/unit/test_function.py
+3 -2
View File
@@ -67,8 +67,9 @@ class TestFunction(unittest.TestCase):
a, b = Tensor.zeros(8).contiguous().realize(), Tensor.zeros(8).contiguous().realize()
@function(precompile=True, allow_implicit=True)
def f(x:Tensor, start:UOp):
stores = (a[start:start+2].uop.store(x.uop), b[start:start+2].uop.store((x+1).uop))
return Tensor(a.uop.after(*stores)) + Tensor(b.uop.after(*stores))
a[start:start+2].assign(x)
b[start:start+2].assign(x+1)
return a+b
out = f(Tensor([2., 3.]).realize(), UOp.variable("start", 0, 6).bind(1))
np.testing.assert_equal(out.numpy(), [0, 5, 7, 0, 0, 0, 0, 0])
tinygrad/llm/cli.py
+6 -9
View File
@@ -16,24 +16,21 @@ class SimpleTokenizer:
raise ValueError(f"Invalid tokenizer preset '{preset}'")
# https://github.com/openai/gpt-2/blob/9b63575ef42771a015060c964af2c3da4cf7c8ab/src/encoder.py#L9
bs = [*range(33, 127), *range(161, 173), *range(174, 256)] # bytes that map to themselves
byte_decoder = str.maketrans({chr(b): b for b in bs} | {chr(256+i): b for i,b in enumerate(b for b in range(256) if b not in bs)})
self._byte_decoder = {chr(b): b for b in bs} | {chr(256+i): b for i,b in enumerate(b for b in range(256) if b not in bs)}
# https://github.com/ggml-org/llama.cpp/blob/94933c8c2eeaa9a7983e3f6c08af76bd86724094/src/llama-vocab.cpp#L286
# 0x323b0 is one past the max codepoint in unicode categories L/N/Z (0x323af is max L)
# compact adjacent codepoints into ranges: listing them all makes re spend seconds on large prompts
categories:dict[str, list[int]] = {category:[] for category in "LNZ"}
for cp in range(0x323b0):
if (category:=unicodedata.category(chr(cp))[0]) in categories: categories[category].append(cp)
def ucat_range(cps:list[int]) -> str:
entries = enumerate(cps)
runs = [list(g) for _, g in itertools.groupby(entries, lambda e: e[1]-e[0])]
def ucat_range(pre:str) -> str:
cps = enumerate(cp for cp in range(0x323b0) if unicodedata.category(chr(cp)).startswith(pre))
runs = [list(g) for _, g in itertools.groupby(cps, lambda e: e[1]-e[0])]
return "".join(re.escape(chr(g[0][1])) + (f"-{re.escape(chr(g[-1][1]))}" if len(g) > 1 else "") for g in runs)
r_ws, r_p_N, r_p_L = r"\t\n\x0b\x0c\r\x85" + ucat_range(categories["Z"]), ucat_range(categories["N"]), ucat_range(categories["L"])
r_ws, r_p_N, r_p_L = r"\t\n\x0b\x0c\r\x85" + ucat_range("Z"), ucat_range("N"), ucat_range("L")
self._split_to_word = re.compile("(?i:'s|'t|'re|'ve|'m|'ll|'d)|" + \
f"[^\\r\\n{r_p_N}{r_p_L}]?[{r_p_L}]+|[{r_p_N}]{{1,3}}| ?[^{r_ws}{r_p_N}{r_p_L}]+[\\r\\n]*|[{r_ws}]*[\\r\\n]+|[{r_ws}]+(?![^{r_ws}])|[{r_ws}]+")
self._split_to_sentence = re.compile("|".join(re.escape(tok) for tok in special_tokens.keys()) if special_tokens else r"(?!)")
self._normal_tokens = {tok.translate(byte_decoder).encode("latin1"): tid for tok, tid in normal_tokens.items()}
self._normal_tokens = {bytes(self._byte_decoder[c] for c in tok): tid for tok, tid in normal_tokens.items()}
self._special_tokens = special_tokens
self._tok2bytes = {tid: tok for tok, tid in self._normal_tokens.items()} | {tid: tok.encode() for tok, tid in self._special_tokens.items()}
self.preset = preset
tinygrad/llm/kernels/amd.py
+69 -117
View File
@@ -4,18 +4,14 @@ from typing import TYPE_CHECKING, Callable, cast
from tinygrad import Tensor, UOp
from tinygrad.uop.ops import AxisType, KernelInfo, Ops
from tinygrad.dtype import AddrSpace, dtypes
from tinygrad.llm.gguf import _GGML_QUANT
if TYPE_CHECKING:
from tinygrad.llm.model import Linear
if TYPE_CHECKING: from tinygrad.llm.model import Linear
BLOCK_M, BLOCK_N, DECODE_HEAD_TILE, WARP_SIZE = 32, 32, 8, 32
WMMA_M, WMMA_N, WMMA_K = 16, 16, 16
WAVES_M, WAVES_N, LANES_PER_WAVE_M, LANES_PER_WAVE_N = 2, 2, 2, 16
WMMA_ACC = WMMA_M // LANES_PER_WAVE_M
THREADS_PER_BLOCK = WARP_SIZE * WAVES_M * WAVES_N
LDS_PAD = 4 # pad LDS rows to reduce bank conflicts
WMMA_ARG, LOG2E = ((WMMA_M, WMMA_N, WMMA_K), 'AMD', 32), math.log2(math.e)
Q5_K, Q6_K, IQ4_XS, GGML_BLOCK_SIZE, Q8_GROUP_SIZE = 13, 14, 23, 256, 32
WMMA_ACC, THREADS_PER_BLOCK = WMMA_M // LANES_PER_WAVE_M, WARP_SIZE * WAVES_M * WAVES_N
LDS_PAD, WMMA_ARG, LOG2E = 4, ((WMMA_M, WMMA_N, WMMA_K), 'AMD', 32), math.log2(math.e)
Q5_K, Q6_K, IQ4_XS, GGML_BLOCK_SIZE, Q8_GROUP_SIZE, Q5_WORDS, Q6_BYTES, IQ4_WORDS = 13, 14, 23, 256, 32, 44, 210, 34
def warp_shfl_xor(val, offset, lane):
idx = ((lane ^ offset) * 4).int()
@@ -33,8 +29,7 @@ def _reg(shape:tuple[int, ...], slot:int, value:float, dep:UOp|None=None) -> UOp
return ret.after((ret if dep is None else ret.after(dep)).store(ret.const_like(value)))
@functools.cache
def _amd_flash_attention_decode_partial(out:UOp, stats:UOp, q:UOp, cache_kv:UOp, cache_scale:UOp,
valid_kv_len:int|UOp, max_kv_len:int, block_n:int) -> UOp:
def _amd_flash_attention_decode_partial(out, stats, q, cache_kv, cache_scale, valid_kv_len, max_kv_len, block_n):
_, B, H_KV, N, D = cast(tuple[int, int, int, int, int], cache_kv.shape)
_, H, M, _ = cast(tuple[int, int, int, int], q.shape)
assert M == 1 and H % H_KV == 0 and D % WARP_SIZE == 0 and max_kv_len <= N and max_kv_len % block_n == 0
@@ -77,37 +72,31 @@ def _amd_flash_attention_decode_partial(out:UOp, stats:UOp, q:UOp, cache_kv:UOp,
stores = [store for head in range(heads_per_wave) for store in head_stores(head)]
return UOp.group(*stores).end(lane, wave, block_n, block_bhkv).sink(arg=KernelInfo(name="flash_decode_partial", opts_to_apply=()))
def amd_flash_attention_decode(q:Tensor, cache_kv:Tensor, valid_kv_len:int|UOp, cache_scale:Tensor,
max_kv_len:int|None=None) -> Tensor:
_, B, H_KV, N, D = cast(tuple[int, int, int, int, int], cache_kv.shape)
_, H, M, _ = cast(tuple[int, int, int, int], q.shape)
max_kv_len = N if max_kv_len is None else max_kv_len
def amd_flash_attention_decode(q:Tensor, cache_kv:Tensor, valid_kv_len:int|UOp, cache_scale:Tensor, max_kv_len:int) -> Tensor:
B, H, D = cache_kv.shape[1], q.shape[1], cache_kv.shape[4]
block_n = 128
assert M == 1 and max_kv_len <= N and max_kv_len % block_n == 0
chunks = min(64, max_kv_len // block_n)
partial = Tensor.empty(B, H, chunks, D, dtype="float32", device=q.device)
stats = Tensor.empty(B, H, chunks, 2, dtype="float32", device=q.device)
def decode_partial(out:UOp, stat:UOp, query:UOp, cache:UOp, scale:UOp) -> UOp:
return _amd_flash_attention_decode_partial(out, stat, query, cache, scale, valid_kv_len, max_kv_len, block_n)
decode_partial = functools.partial(_amd_flash_attention_decode_partial, valid_kv_len=valid_kv_len, max_kv_len=max_kv_len, block_n=block_n)
partial, stats = Tensor.custom_kernel(partial, stats, q, cache_kv, cache_scale, fxn=decode_partial)[:2]
live_chunks = (valid_kv_len+block_n-1)//block_n
live_chunks = min(live_chunks, chunks) if isinstance(live_chunks, int) else live_chunks.minimum(chunks)
partial, stats = partial[:, :, :live_chunks], stats[:, :, :live_chunks]
row_max = stats[..., 0].max(2, keepdim=True)
weights = ((stats[..., 0]-row_max)*LOG2E).exp2()
weights = ((stats[..., 0]-stats[..., 0].max(2, keepdim=True))*LOG2E).exp2()
return ((partial*weights.unsqueeze(-1)).sum(2) / (stats[..., 1]*weights).sum(2, keepdim=True)).unsqueeze(2)
@functools.cache
def _amd_flash_attention(o:UOp, q:UOp, k:UOp, v:UOp, kv_scale:UOp, valid_kv_len:int|UOp, key_limit:int|UOp|None=None) -> UOp:
def _amd_flash_attention(o:UOp, q:UOp, cache:UOp, kv_scale:UOp, valid_kv_len:int|UOp, key_limit:int|UOp|None=None) -> UOp:
BH, M, D = q.shape
physical_n = k.shape[1]
N = valid_kv_len
_, B, H_KV, physical_n, cache_dim = cache.shape
k, v = cache[0].reshape(B*H_KV, physical_n, cache_dim), cache[1].reshape(B*H_KV, physical_n, cache_dim)
kv_scale = kv_scale.reshape(2, B*H_KV, physical_n)
assert k.shape == v.shape and BH % k.shape[0] == 0 and k.shape[2] == D
gqa_group = BH // k.shape[0]
if isinstance(M, int) and isinstance(N, int):
assert M % BLOCK_M == 0 and N % BLOCK_N == 0, f"M={M} and N={N} must be divisible by BLOCK_M={BLOCK_M} and BLOCK_N={BLOCK_N}"
if isinstance(M, int) and isinstance(valid_kv_len, int):
assert M % BLOCK_M == 0 and valid_kv_len % BLOCK_N == 0
assert isinstance(D, int) and D % WMMA_K == 0 and D % LANES_PER_WAVE_N == 0
assert BLOCK_M % (WAVES_M * WMMA_M) == 0 and BLOCK_N % LANES_PER_WAVE_N == 0
TM, TN, TD, SCALE = BLOCK_M//(WAVES_M*LANES_PER_WAVE_M), BLOCK_N//LANES_PER_WAVE_N, D//(WAVES_N*LANES_PER_WAVE_N), 1/math.sqrt(D)
block_bh, block_m = UOp.range(BH, 0, AxisType.GLOBAL), UOp.range(M // BLOCK_M, 1, AxisType.GLOBAL)
q = q.reshape(BH, M//BLOCK_M, BLOCK_M, D)[block_bh, block_m]
@@ -121,7 +110,7 @@ def _amd_flash_attention(o:UOp, q:UOp, k:UOp, v:UOp, kv_scale:UOp, valid_kv_len:
QP_lds = UOp.placeholder((BLOCK_M, D + LDS_PAD), dtypes.half, slot=0, addrspace=AddrSpace.LOCAL)
KV_lds = UOp.placeholder((BLOCK_N, D + LDS_PAD), dtypes.half, slot=1, addrspace=AddrSpace.LOCAL)[:, :D]
acc, m_i, l_i = _reg((TM, TD), 2, 0), _reg((TM,), 3, -math.inf), _reg((TM,), 4, 0)
n_tiles = (N - M + (block_m + 1) * BLOCK_M + BLOCK_N - 1) // BLOCK_N
n_tiles = (valid_kv_len - M + (block_m + 1) * BLOCK_M + BLOCK_N - 1) // BLOCK_N
n_tile = UOp.range(n_tiles, 100, AxisType.REDUCE)
Q_lds = QP_lds[:, :D]
Q_store = Q_lds.after(n_tile).reshape(THREADS_PER_BLOCK, Q_ELEMS_PER_THREAD)[tid].store(q.reshape(THREADS_PER_BLOCK, Q_ELEMS_PER_THREAD)[tid])
@@ -140,7 +129,7 @@ def _amd_flash_attention(o:UOp, q:UOp, k:UOp, v:UOp, kv_scale:UOp, valid_kv_len:
S_reg = S_reg.after(qk_done)
S_reg = S_reg.after(S_reg.store(S_reg * SCALE))
rm, rn = UOp.range(TM, 250, AxisType.WEAK), UOp.range(TN, 251, AxisType.WEAK)
q_idx = N - M + block_m * BLOCK_M + wave_m * WMMA_M + rm * LANES_PER_WAVE_M + lane_m
q_idx = valid_kv_len - M + block_m * BLOCK_M + wave_m * WMMA_M + rm * LANES_PER_WAVE_M + lane_m
k_idx = n_tile * BLOCK_N + rn * LANES_PER_WAVE_N + lane_n
valid = k_idx <= q_idx
if key_limit is not None: valid = valid & (k_idx < key_limit)
@@ -188,35 +177,28 @@ def _amd_flash_attention(o:UOp, q:UOp, k:UOp, v:UOp, kv_scale:UOp, valid_kv_len:
o = o.permute((0, 4, 2, 6, 1, 3, 5, 7)).reshape(THREADS_PER_BLOCK, TM, TD)
return o[tid].store(acc).end(wave_m, wave_n, lane).end(block_m, block_bh).sink(arg=KernelInfo(opts_to_apply=()))
def flash_attention_causal_cached(q:Tensor, cache_kv:Tensor, valid_kv_len:int|UOp, cache_scale:Tensor,
key_limit:int|UOp|None=None) -> Tensor:
def flash_attention_causal_cached(q:Tensor, cache_kv:Tensor, valid_kv_len:int|UOp, cache_scale:Tensor, key_limit:int|UOp|None=None) -> Tensor:
B, H, T, D = cast(tuple[int, int, int, int], q.shape)
q_flat = q.reshape(B*H, T, D)
out = Tensor.empty(B*H, T, D, dtype="float32", device=q.device)
def flash_cached(*uops:UOp) -> UOp:
output, query, cache, scale = uops
_, b, h_kv, n, d = cast(tuple[int, int, int, int, int], cache.shape)
return _amd_flash_attention(output, query, cache[0].reshape(b*h_kv, n, d), cache[1].reshape(b*h_kv, n, d),
scale.reshape(2, b*h_kv, n), valid_kv_len, key_limit)
return Tensor.custom_kernel(out, q_flat, cache_kv, cache_scale, fxn=flash_cached)[0].reshape(B, H, T, D)
flash_cached = functools.partial(_amd_flash_attention, valid_kv_len=valid_kv_len, key_limit=key_limit)
return Tensor.custom_kernel(out, q.reshape(B*H, T, D), cache_kv, cache_scale, fxn=flash_cached)[0].reshape(B, H, T, D)
def _amd_dp4a(a:UOp, b:UOp, c:UOp) -> UOp: return UOp(Ops.CUSTOMI, dtypes.int32, (a.int(), b.int(), c),
arg="__builtin_amdgcn_sudot4(true, {}, true, {}, {}, false)")
def _amd_dp4a(a:UOp, b:UOp, c:UOp) -> UOp:
return UOp(Ops.CUSTOMI, dtypes.int32, (a.int(), b.int(), c), arg="__builtin_amdgcn_sudot4(true, {}, true, {}, {}, false)")
def _amd_byte_perm(a:UOp, b:UOp, selectors:UOp) -> UOp: return UOp(Ops.CUSTOMI, dtypes.uint32,
(a.cast(dtypes.uint32), b.cast(dtypes.uint32), selectors.cast(dtypes.uint32)), arg="__builtin_amdgcn_perm({}, {}, {})")
def _amd_byte_perm(a:UOp, b:UOp, selectors:UOp) -> UOp:
return UOp(Ops.CUSTOMI, dtypes.uint32, (a.cast(dtypes.uint32), b.cast(dtypes.uint32), selectors.cast(dtypes.uint32)),
arg="__builtin_amdgcn_perm({}, {}, {})")
def _amd_vector_load(ptr:UOp, lanes:int) -> UOp:
def _amd_load(ptr:UOp, lanes:int|None=None) -> UOp:
assert ptr.op is Ops.INDEX
if lanes is None: return UOp(Ops.CUSTOMI, ptr.dtype, (ptr,), arg="__builtin_nontemporal_load({0})")
buf, coords = ptr.src[0], ptr.src[1:]
idx = sum((coord*math.prod(buf.shape[i+1:]) for i,coord in enumerate(coords)), UOp.const(0, dtypes.weakint))
return UOp(Ops.SHRINK, src=(buf.flatten(), idx, UOp.const(lanes, dtypes.weakint))).load(dtype=ptr.dtype)
def _amd_stream_load(ptr:UOp) -> UOp:
assert ptr.op is Ops.INDEX
return UOp(Ops.CUSTOMI, ptr.dtype, (ptr,), arg="__builtin_nontemporal_load({0})")
def _load_byte(raw:UOp, base:UOp, offset:UOp) -> UOp: return (raw[base + offset//4] >> ((offset&3)*8).cast(dtypes.uint32)) & 255
def _half(value:UOp) -> UOp: return value.cast(dtypes.uint16).bitcast(dtypes.float16).float()
def _iq4_bytes(packed:UOp, shift:int) -> UOp:
selectors = (packed >> shift) & 0x0f0f0f0f
@@ -230,8 +212,7 @@ def q8_quantize(x:Tensor, tokens:int, in_features:int) -> tuple[Tensor, Tensor]:
return (groups/scale.unsqueeze(-1)).round().clip(-127, 127).cast(dtypes.int8).contiguous().bitcast(dtypes.uint32), scale
@functools.cache
def _gated_delta_prefill_kernel(core:UOp, next_state:UOp, q:UOp, k:UOp, v:UOp, beta:UOp, alpha:UOp, state:UOp,
kq:UOp) -> UOp:
def _gated_delta_prefill_kernel(core:UOp, next_state: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))
@@ -258,9 +239,9 @@ def _gated_delta_prefill_kernel(core:UOp, next_state:UOp, q:UOp, k:UOp, v:UOp, b
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)
return UOp.group(*(next_state[bh, row, col].store(current.after(step)[row_idx*dim//32+i].load().cast(next_state.dtype))
for row_idx,row in enumerate(rows) for i,col in enumerate(cols))).end(lane, bh_row).sink(
arg=KernelInfo(name="gated_delta_prefill", opts_to_apply=()))
state_stores = (next_state[bh, row, col].store(current.after(step)[row_idx*dim//32+i].load().cast(next_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) -> tuple[Tensor, Tensor]:
batch, heads, tokens, dim = q.shape
@@ -280,8 +261,7 @@ def _wmma_layout(out:UOp, out_features:int, token_tile:int, output_tiles:int):
tokens = tuple(tuple(token_block*token_tile + tile*16 + half*8 + i for i in range(8)) for tile in range(token_tile//16))
return output_waves, token_block, output_block, lane, wave, half, outputs, inputs, tokens
def _wmma_stores(out:UOp, outputs:tuple[UOp, ...], tokens:tuple[tuple[UOp, ...], ...], accs:tuple[tuple[UOp, ...], ...],
update:UOp, half:UOp) -> list[UOp]:
def _wmma_stores(out, outputs, tokens, accs, update, half):
def values(acc:UOp) -> tuple[UOp, ...]:
vals = tuple(acc.after(update)[i].load() for i in range(8))
swapped = tuple(UOp(Ops.CUSTOM, dtypes.float32, (value,),
@@ -307,19 +287,16 @@ def _q5_scales(raw:UOp, base:UOp, subgroup:UOp) -> tuple[UOp, UOp, UOp, UOp]:
minimum = (subgroup < 4).where(_load_byte(raw, base, 8 + subgroup) & 63,
(_load_byte(raw, base, 8 + subgroup) >> 4) | ((_load_byte(raw, base, 4 + subgroup) >> 6) << 4))
d, dmin = (raw[base] & 0xffff).cast(dtypes.uint16), (raw[base] >> 16).cast(dtypes.uint16)
return d.bitcast(dtypes.float16).float(), dmin.bitcast(dtypes.float16).float(), scale.float(), minimum.float()
return _half(d), _half(dmin), scale.float(), minimum.float()
@functools.cache
def _q5_linear_kernel(out:UOp, raw:UOp, xq:UOp, xd:UOp, raw_offset:UOp,
out_features:int, in_features:int) -> UOp:
raw_offset = raw_offset.cast(dtypes.uint64)
group_count, type_words = in_features // Q8_GROUP_SIZE, _GGML_QUANT[Q5_K][1] // 4
output_words = in_features // GGML_BLOCK_SIZE * type_words
def _q5_linear_kernel(out:UOp, raw:UOp, xq:UOp, xd:UOp, raw_offset:UOp, out_features:int, in_features:int) -> UOp:
group_count, type_words = in_features // Q8_GROUP_SIZE, Q5_WORDS
def group_dot(output:UOp, group:UOp) -> UOp:
block, subgroup = group // 8, group % 8
base = raw_offset + output * output_words + block * type_words
base = raw_offset + (output * in_features//GGML_BLOCK_SIZE + block) * type_words
qs_base = base + 12 + (subgroup // 2) * 8
xwords = _amd_vector_load(xq[0, group, 0], 8)
xwords = _amd_load(xq[0, group, 0], 8)
dot, qsum = UOp.const(0, dtypes.int32), UOp.const(0, dtypes.int32)
for word_idx in range(8):
word = (raw[qs_base + word_idx] >> ((subgroup & 1) * 4).cast(dtypes.uint32)) & 0x0f0f0f0f
@@ -330,8 +307,7 @@ def _q5_linear_kernel(out:UOp, raw:UOp, xq:UOp, xd:UOp, raw_offset:UOp,
return (dot.float()*d*scale - qsum.float()*dmin*minimum) * xd[0, group]
return _decode_linear(out, out_features, group_count, group_dot, "linear_q5_k")
def _quant_linear_wmma(out:UOp, x:UOp, out_features:int, in_features:int, type_words:int, layout:tuple,
dequant:Callable[[UOp, UOp, int], tuple[UOp, ...]], name:str) -> UOp:
def _quant_linear_wmma(out, x, out_features, in_features, type_words, layout, dequant, name):
x = x.reshape(out.shape[0], in_features)
_, token_block, output_block, lane, wave, physical_half, outputs, input_tokens, tokens = layout
token_tile, output_tiles = len(tokens)*16, len(outputs)
@@ -357,7 +333,6 @@ def _quant_linear_wmma(out:UOp, x:UOp, out_features:int, in_features:int, type_w
@functools.cache
def _q5_linear_f16_wmma_kernel(out:UOp, raw:UOp, x:UOp, raw_offset:UOp, out_features:int, in_features:int) -> UOp:
raw_offset = raw_offset.cast(dtypes.uint64)
token_tile, output_tiles = (64, 1) if (out_features <= 1024 or out_features > 6144) and out.shape[0] % 64 == 0 else \
(64, 2) if out.shape[0] % 64 == 0 else (32 if out.shape[0] % 32 == 0 else 16, 2)
def dequant(base:UOp, subgroup:UOp, half:int) -> tuple[UOp, ...]:
@@ -367,42 +342,37 @@ def _q5_linear_f16_wmma_kernel(out:UOp, raw:UOp, x:UOp, raw_offset:UOp, out_feat
words = tuple((raw[qs_base+i] >> ((subgroup&1)*4).cast(dtypes.uint32) & 0x0f0f0f0f) |
((raw[base+4+half*4+i] >> subgroup.cast(dtypes.uint32) & 0x01010101) << 4) for i in range(4))
return tuple(((word >> (byte*8) & 255).float()*d*scale-dmin*minimum).cast(dtypes.float16) for word in words for byte in range(4))
return _quant_linear_wmma(out, x, out_features, in_features, _GGML_QUANT[Q5_K][1]//4,
return _quant_linear_wmma(out, x, out_features, in_features, Q5_WORDS,
_wmma_layout(out, out_features, token_tile, output_tiles), dequant, "linear_q5_k_f16_wmma")
def _iq4_scales(raw:UOp, base:UOp, subgroup:UOp) -> tuple[UOp, UOp]:
low = _load_byte(raw, base, 4 + subgroup//2)
scale = ((low >> (4*(subgroup%2)).cast(dtypes.uint32)) & 15) | ((((raw[base] >> 16) >> (2*subgroup).cast(dtypes.uint32)) & 3) << 4)
return (raw[base] & 0xffff).cast(dtypes.uint16).bitcast(dtypes.float16).float(), (scale.cast(dtypes.uint8).bitcast(dtypes.int8)-32).float()
return _half(raw[base] & 0xffff), (scale.cast(dtypes.uint8).bitcast(dtypes.int8)-32).float()
@functools.cache
def _iq4_linear_kernel(out:UOp, raw:UOp, xq:UOp, xd:UOp, raw_offset:UOp, out_features:int, in_features:int) -> UOp:
raw_offset = raw_offset.cast(dtypes.uint64)
group_count, type_words = in_features // Q8_GROUP_SIZE, _GGML_QUANT[IQ4_XS][1] // 4
output_words = in_features // GGML_BLOCK_SIZE * type_words
group_count, type_words = in_features // Q8_GROUP_SIZE, IQ4_WORDS
def group_dot(output:UOp, group:UOp) -> UOp:
block, subgroup = group // 8, group % 8
base = raw_offset + output * output_words + block * type_words
xwords = _amd_vector_load(xq[0, group, 0], 8)
base = raw_offset + (output * in_features//GGML_BLOCK_SIZE + block) * type_words
xwords = _amd_load(xq[0, group, 0], 8)
dot = UOp.const(0, dtypes.int32)
for word_idx in range(8):
packed, shift = _amd_stream_load(raw[base + 2 + subgroup*4 + word_idx % 4]), 4 * (word_idx // 4)
packed, shift = _amd_load(raw[base + 2 + subgroup*4 + word_idx % 4]), 4 * (word_idx // 4)
dot = _amd_dp4a(_iq4_bytes(packed, shift), xwords[word_idx], dot)
d, scale = _iq4_scales(raw, base, subgroup)
return dot.float() * xd[0, group] * d * scale
return _decode_linear(out, out_features, group_count, group_dot, "linear_iq4_xs")
@functools.cache
def _iq4_linear_f16_wmma_kernel(out:UOp, raw:UOp, x:UOp, lut:UOp, raw_offset:UOp,
out_features:int, in_features:int) -> UOp:
raw_offset = raw_offset.cast(dtypes.uint64)
def _iq4_linear_f16_wmma_kernel(out:UOp, raw:UOp, x:UOp, lut:UOp, raw_offset:UOp, out_features:int, in_features:int) -> UOp:
token_tile = 32 if out_features <= 1024 and out.shape[0] % 32 == 0 else 64 if out.shape[0] % 64 == 0 and \
(out_features <= 6144 or out_features == 5120 and in_features > 8192) else 128 if out.shape[0] % 128 == 0 else \
32 if out.shape[0] % 32 == 0 else 16
output_tiles = 1 if out_features <= 1024 else 2 if out_features <= 6144 else 1 if out_features < 8192 else 2
layout = _wmma_layout(out, out_features, token_tile, output_tiles)
output_waves, _, _, lane, wave, _, _, _, _ = layout
assert out_features % (16*output_tiles*output_waves) == 0
local_lut = UOp.placeholder((256,), dtypes.uint32, slot=32, addrspace=AddrSpace.LOCAL)
tid, lut_items = wave*32+lane, 256//(32*output_waves)
lut = local_lut.after(UOp.group(*(local_lut[tid*lut_items+i].store(lut[tid*lut_items+i]) for i in range(lut_items))).barrier())
@@ -413,25 +383,20 @@ def _iq4_linear_f16_wmma_kernel(out:UOp, raw:UOp, x:UOp, lut:UOp, raw_offset:UOp
if out_features <= 6144:
pairs = tuple(lut[((raw[base + 2 + subgroup*4 + word] >> (byte*8)) & 255).cast(dtypes.weakint)]
for word in range(4) for byte in range(4))
return tuple((((pair >> (half*16)) & 0xffff).cast(dtypes.uint16).bitcast(dtypes.float16).float()*scale).cast(dtypes.float16)
for pair in pairs)
values = []
for packed in (raw[base+2+subgroup*4+i] for i in range(4)):
nibbles = tuple((packed >> (8*i+4*half)) & 15 for i in range(4))
pairs = tuple(lut[(nibbles[i] | nibbles[i+1]<<4).cast(dtypes.weakint)] for i in (0, 2))
values += [(((pair >> (i*16)) & 0xffff).cast(dtypes.uint16).bitcast(dtypes.float16).float()*scale).cast(dtypes.float16)
for pair in pairs for i in range(2)]
return tuple(values)
return _quant_linear_wmma(out, x, out_features, in_features, _GGML_QUANT[IQ4_XS][1]//4, layout, dequant, "linear_iq4_xs_f16_wmma")
return tuple((_half((pair >> (half*16)) & 0xffff)*scale).cast(dtypes.float16) for pair in pairs)
def nibble(packed:UOp, index:int): return (packed >> (8*index+4*half)) & 15
lut_pairs = (lut[(nibble(packed, i) | nibble(packed, i+1)<<4).cast(dtypes.weakint)]
for packed in (raw[base+2+subgroup*4+i] for i in range(4)) for i in (0, 2))
return tuple((_half((pair >> (i*16)) & 0xffff)*scale).cast(dtypes.float16) for pair in lut_pairs for i in range(2))
return _quant_linear_wmma(out, x, out_features, in_features, IQ4_WORDS, layout, dequant, "linear_iq4_xs_f16_wmma")
@functools.cache
def _q6_linear_kernel(out:UOp, raw:UOp, xq:UOp, xd:UOp, raw_offset:UOp, out_features:int, in_features:int) -> UOp:
raw_offset = raw_offset.cast(dtypes.uint64) * 4
group_count, type_size = in_features // Q8_GROUP_SIZE, _GGML_QUANT[Q6_K][1]
output_size = in_features // GGML_BLOCK_SIZE * type_size
raw_offset = raw_offset * 4
group_count, type_size = in_features // Q8_GROUP_SIZE, Q6_BYTES
def group_dot(output:UOp, group:UOp) -> UOp:
block, subgroup = group//8, group%8
xwords, base = _amd_vector_load(xq[0, group, 0], 8), raw_offset + output*output_size + block*type_size
xwords, base = _amd_load(xq[0, group, 0], 8), raw_offset + (output*in_features//GGML_BLOCK_SIZE+block)*type_size
dots = [UOp.const(0, dtypes.int32), UOp.const(0, dtypes.int32)]
for word_idx in range(8):
word = UOp.const(0, dtypes.uint32)
@@ -444,16 +409,9 @@ def _q6_linear_kernel(out:UOp, raw:UOp, xq:UOp, xd:UOp, raw_offset:UOp, out_feat
dots[word_idx//4] = _amd_dp4a(word, xwords[word_idx], dots[word_idx//4])
scales = [raw[base + 192 + subgroup*2+i].cast(dtypes.uint8).bitcast(dtypes.int8).float() for i in range(2)]
dbits = raw[base+208].cast(dtypes.uint16) | (raw[base+209].cast(dtypes.uint16) << 8)
return (dots[0].float()*scales[0] + dots[1].float()*scales[1]) * xd[0, group] * dbits.bitcast(dtypes.float16).float()
return (dots[0].float()*scales[0] + dots[1].float()*scales[1]) * xd[0, group] * _half(dbits)
return _decode_linear(out, out_features, group_count, group_dot, "linear_q6")
def _linear_result(layer:Linear, x:Tensor, srcs:tuple[UOp, ...], fxn:Callable[..., UOp]) -> Tensor:
params = tuple(UOp.placeholder_like(src, slot=i) for i,src in enumerate(srcs))
*args, offset = params
kernel = fxn(*args, offset[0]).call(*srcs)
out = Tensor(srcs[0].after(kernel)).reshape(*x.shape[:-1], layer.out_features)
return out if layer.bias is None else out + layer.bias
def q8_linear(layer:Linear, x:Tensor) -> Tensor:
assert layer.ggml_type in (Q5_K, Q6_K, IQ4_XS) and layer._raw_uop is not None and layer._raw_offset_uop is not None
tokens = int(x.numel()) // layer.in_features
@@ -461,27 +419,21 @@ def q8_linear(layer:Linear, x:Tensor) -> Tensor:
raw, offset = layer._raw_uop, layer._raw_offset_uop
out_features, in_features = layer.out_features, layer.in_features
use_wmma = tokens % 16 == 0 and layer.out_features % 16 == 0
srcs:tuple[UOp, ...]
def run(srcs:tuple[UOp, ...], fxn:Callable[..., UOp]) -> Tensor:
return _linear_result(layer, x, srcs, functools.partial(fxn, out_features=out_features, in_features=in_features))
if layer.ggml_type == Q5_K:
if use_wmma:
srcs = (out, raw.bitcast(dtypes.uint32), x.cast(dtypes.float16).contiguous().uop, offset)
return run(srcs, _q5_linear_f16_wmma_kernel)
xq, xd = q8_quantize(x, tokens, layer.in_features)
srcs = (out, raw.bitcast(dtypes.uint32), xq.uop, xd.uop, offset)
return run(srcs, _q5_linear_kernel)
def run(fxn:Callable[..., UOp], *srcs:UOp) -> Tensor:
params = tuple(UOp.placeholder_like(src, slot=i) for i,src in enumerate(srcs))
kernel = fxn(*params[:-1], params[-1][0], out_features=out_features, in_features=in_features).call(*srcs)
result = Tensor(srcs[0].after(kernel)).reshape(*x.shape[:-1], layer.out_features)
return result if layer.bias is None else result + layer.bias
if layer.ggml_type == Q5_K and use_wmma:
return run(_q5_linear_f16_wmma_kernel, out, raw.bitcast(dtypes.uint32), x.cast(dtypes.float16).contiguous().uop, offset)
xq, xd = q8_quantize(x, tokens, layer.in_features)
if layer.ggml_type == IQ4_XS:
if use_wmma: srcs = (out, raw.bitcast(dtypes.uint32), x.cast(dtypes.float16).contiguous().uop,
iq4_half_lut(str(x.device)).uop, offset)
else: srcs = (out, raw.bitcast(dtypes.uint32), xq.uop, xd.uop, offset)
fxn = _iq4_linear_f16_wmma_kernel if use_wmma else _iq4_linear_kernel
return run(srcs, fxn)
return run((out, raw, xq.uop, xd.uop, offset), _q6_linear_kernel)
if layer.ggml_type == Q5_K: return run(_q5_linear_kernel, out, raw.bitcast(dtypes.uint32), xq.uop, xd.uop, offset)
if layer.ggml_type == IQ4_XS and use_wmma:
return run(_iq4_linear_f16_wmma_kernel, out, raw.bitcast(dtypes.uint32), x.cast(dtypes.float16).contiguous().uop,
iq4_half_lut(str(x.device)).uop, offset)
if layer.ggml_type == IQ4_XS: return run(_iq4_linear_kernel, out, raw.bitcast(dtypes.uint32), xq.uop, xd.uop, offset)
return run(_q6_linear_kernel, out, raw, xq.uop, xd.uop, offset)
@functools.cache
def iq4_half_lut(device:str) -> Tensor:
tinygrad/llm/model.py
+59 -107
View File
@@ -1,7 +1,7 @@
from __future__ import annotations
import functools, itertools, pathlib
from dataclasses import dataclass, replace
from typing import Any, cast
from typing import cast
from tinygrad import Tensor, nn, UOp, TinyJit, getenv, function, dtypes, Context
from tinygrad.llm.kernels import amd as llm_amd
from tinygrad.llm.gguf import gguf_load
@@ -15,16 +15,12 @@ class Linear(nn.Linear):
ggml_type:int|None
_raw_uop:UOp|None
_raw_offset_uop:UOp|None
def set_quantized(self, decoded:Tensor) -> bool:
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 False
if raw is None: return None
packed = Tensor(raw)
self.weight, self.ggml_type = packed.flatten(), packed_sizes[prod(raw.shape)]
self._raw_uop = self._raw_offset_uop = None
if self.ggml_type == 23 and str(packed.device).startswith("AMD"): llm_amd.iq4_half_lut(str(packed.device))
return True
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:
@@ -32,6 +28,7 @@ class Linear(nn.Linear):
raw = raw.src[0]
assert raw_offset % 4 == 0 and raw.dtype == dtypes.uint8
self._raw_uop = raw
if self.ggml_type == 23 and str(packed.device).startswith("AMD"): llm_amd.iq4_half_lut(str(packed.device))
return Tensor([raw_offset // 4], dtype=dtypes.uint64, device=self.weight.device)
def __init__(self, in_features:int, out_features:int, bias=True):
if LLM_EMPTY_WEIGHTS:
@@ -114,8 +111,6 @@ class TransformerConfig:
class FFNBlock:
def __init__(self, config:TransformerConfig):
self.config = config
self.pending_state:tuple[Tensor, Tensor]|None = None
# --- RMSNorms --------------------------------------------------------
self.attn_norm = nn.RMSNorm(config.dim, config.norm_eps)
self.ffn_norm = nn.RMSNorm(config.dim, config.norm_eps)
@@ -164,25 +159,23 @@ class FFNBlock:
# 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, use_flash:bool=False, kv_len:int|UOp|None=None, valid_len:int|UOp|None=None) -> Tensor:
def _attention(self, x:Tensor, start_pos:int|UOp, kv_len:int|UOp|None=None, valid_len:int|UOp|None=None) -> Tensor:
raise NotImplementedError
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):
def __call__(self, x: Tensor, start_pos: int|UOp, kv_len:int|UOp|None=None, valid_len:int|UOp|None=None):
self._init_state(x)
if hasattr(self, 'attn_gate'):
self.pending_state = None
@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._feed_forward(self.ffn_norm(h))).contiguous()
assert self.pending_state is not None
return (out, *self.pending_state)
attn, conv_state, recurrent_state = cast(tuple[Tensor, Tensor, Tensor], self._attention(self.attn_norm(x), start_pos, kv_len, valid_len))
h = x + attn
return (h + self._feed_forward(self.ffn_norm(h))).contiguous(), conv_state, recurrent_state
out, conv_state, next_recurrent_state = _run_stateful(x, start_pos, valid_len)
recurrent_state = getattr(self, "recurrent_state")
stores = (getattr(self, "conv_state").uop.store(conv_state.uop), recurrent_state.uop.store(next_recurrent_state.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, use_flash, kv_len)
h = x + self._attention(self.attn_norm(x), start_pos, kv_len)
return (h + self._feed_forward(self.ffn_norm(h))).contiguous()
return function(precompile=True, allow_implicit=True)(_run)(x, start_pos)
@@ -200,7 +193,7 @@ class TransformerBlock(FFNBlock):
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, use_flash:bool=False, kv_len:int|UOp|None=None, valid_len:int|UOp|None=None) -> Tensor:
def _attention(self, x:Tensor, start_pos:int|UOp, kv_len:int|UOp|None=None, valid_len:int|UOp|None=None):
q, k, v = self.attn_q(x), self.attn_k(x), self.attn_v(x)
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)
@@ -233,25 +226,17 @@ class TransformerBlock(FFNBlock):
# NOTE: this mask is causal_lower_right, not the causal_upper_left generated by is_casual = True
# TODO: this if statement should be removed and it shouldn't generate extra kernels
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:
assert assigned_scale is not None
decode_pos = cast(int|UOp, unwrap_var(start_pos)) + 1
decode_len = kv_len if isinstance(kv_len, int) else self.config.max_context
attn = llm_amd.amd_flash_attention_decode(q.half(), assigned_kv, decode_pos, assigned_scale, decode_len)
elif use_flash:
if resolve(T == 1) and kv_len is not None and str(x.device).startswith("AMD") and self.config.head_dim == 256:
assert assigned_scale is not None and isinstance(kv_len, int)
attn = llm_amd.amd_flash_attention_decode(q.half(), assigned_kv, cast(int|UOp, unwrap_var(start_pos))+1, assigned_scale, kv_len)
elif self.config.ssm is not None and resolve(T != 1):
assert assigned_scale is not None
start, valid = cast(int|UOp, unwrap_var(start_pos)), unwrap_var(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, assigned_scale, key_limit)
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
causal = resolve(T != 1) or kv_len is not None and self.config.ssm is None
mask = Tensor.full((1, 1, T, cache_len), float("-inf"), dtype=x.dtype, device=x.device, buffer=False).triu(start_pos+1) if causal 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()))
@@ -281,7 +266,7 @@ class MLATransformerBlock(FFNBlock):
self.attn_v_b = {"weight": Tensor.zeros(config.n_heads, config.v_head_dim, config.kv_lora_rank)}
self.attn_output = Linear(config.n_heads * config.v_head_dim, config.dim, bias=False)
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) -> Tensor:
def _attention(self, x:Tensor, start_pos:int|UOp, kv_len:int|UOp|None=None, valid_len:int|UOp|None=None) -> Tensor:
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)
@@ -333,13 +318,7 @@ class GatedDeltaNetBlock(FFNBlock):
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), Linear(ssm.inner_size, config.dim, bias=False)
def _qwen_projections(self, x:Tensor) -> tuple[Tensor, Tensor, Tensor, Tensor]:
out_gate, qkv, (beta, alpha) = self.attn_gate(x), self.attn_qkv(x), \
((x @ self.ssm_beta_alpha_weight.T).split(self.num_v_heads, dim=-1) if self.ssm_beta_alpha_weight is not None else
(self.ssm_beta(x), self.ssm_alpha(x)))
return out_gate, qkv, beta, alpha
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) -> Tensor:
def _attention(self, x:Tensor, start_pos:int|UOp, kv_len:int|UOp|None=None, valid_len:int|UOp|None=None):
if not hasattr(self, "ssm_g_a"): return self._qwen_attention(x, valid_len)
B, T, _ = x.shape
assert T == 1, "GatedDeltaNetBlock currently only supports T=1"
@@ -376,32 +355,33 @@ class GatedDeltaNetBlock(FFNBlock):
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))
def _qwen_attention(self, x:Tensor, valid_len:int|UOp|None) -> Tensor:
def _qwen_attention(self, x:Tensor, valid_len:int|UOp|None):
B, T, _ = x.shape
conv_state, initial_state = self.conv_state, self.recurrent_state
x = x.half()
out_gate, qkv, beta, alpha = self._qwen_projections(x)
out_gate, qkv = self.attn_gate(x), self.attn_qkv(x)
beta, alpha = (x @ self.ssm_beta_alpha_weight.T).split(self.num_v_heads, dim=-1) if self.ssm_beta_alpha_weight is not None else \
(self.ssm_beta(x), self.ssm_alpha(x))
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, T, self.num_k_heads, self.head_k_dim).normalize(dim=-1, eps=1e-6).repeat(1, 1, self.num_v_heads//self.num_k_heads, 1)
k = k.reshape(B, T, self.num_k_heads, self.head_k_dim).normalize(dim=-1, eps=1e-6).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)
if T == 1:
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()
q, k = q.reshape(B, self.num_k_heads, self.head_k_dim), k.reshape(B, self.num_k_heads, self.head_k_dim)
q, k = q.normalize(dim=-1, eps=1e-6), k.normalize(dim=-1, eps=1e-6)
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
q, k, v = q[:, 0] * self.head_k_dim**-0.5, k[:, 0], v[:, 0]
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)
state_k, state_q = (initial_state @ kv.cat(qv, dim=-1)).split(1, dim=-1)
delta = (v.unsqueeze(-1) - state_k * alpha4) * beta4
self.pending_state = (conv_window[:, 1:, :].cast(self.conv_state.dtype).contiguous(),
(initial_state * alpha4 + delta @ kv.transpose(-1, -2)).cast(self.recurrent_state.dtype).contiguous())
recurrent_state = self.pending_state[1]
conv_state = conv_window[:, 1:, :].cast(self.conv_state.dtype).contiguous()
recurrent_state = (initial_state * alpha4 + delta @ kv.transpose(-1, -2)).cast(self.recurrent_state.dtype).contiguous()
core = (state_q * alpha4 + delta * (kv.transpose(-1, -2) @ 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))
return self.ssm_out((core * out_gate.silu()).reshape(B, 1, -1).cast(x.dtype)), conv_state, recurrent_state
assert str(x.device).startswith("AMD"), "batched GatedDeltaNet prefill currently requires AMD"
out_gate = out_gate.reshape(B, T, self.num_v_heads, self.head_v_dim)
@@ -410,22 +390,15 @@ class GatedDeltaNetBlock(FFNBlock):
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 = q.reshape(B, T, self.num_k_heads, self.head_k_dim).normalize(dim=-1, eps=1e-6).repeat(1, 1, self.num_v_heads//self.num_k_heads, 1)
k = k.reshape(B, T, self.num_k_heads, self.head_k_dim).normalize(dim=-1, eps=1e-6).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)
out = self.ssm_out((self.ssm_norm(core.transpose(1, 2)) * out_gate.silu()).reshape(B, T, -1).cast(x.dtype)).contiguous()
state_pos = T if valid_len is None else valid_len
self.pending_state = (conv_window[:, state_pos:state_pos+self.ssm_conv_kernel-1, :].cast(self.conv_state.dtype).contiguous(),
recurrent_state)
return out
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 []
def _reusable_prefix_len(self, prefix_len:int, cached_len:int) -> int: return 0 if prefix_len != cached_len else prefix_len
conv_state = conv_window[:, state_pos:state_pos+self.ssm_conv_kernel-1, :].cast(self.conv_state.dtype).contiguous()
return out, conv_state, recurrent_state
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 []
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()
@@ -448,38 +421,25 @@ class Transformer:
# 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.recurrent_prefill_jits:dict[tuple[int, bool], Any] = {}
self.recurrent_rollout_jit = TinyJit(functools.partial(self.forward_recurrent_decode, decode_len=self.max_context))
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) -> Tensor:
def forward(self, tokens:Tensor, start_pos:int|UOp, temperature:Tensor, kv_len:int|UOp|None=None, valid_len:int|UOp|None=None) -> Tensor:
x = self.token_embd(tokens).float() # (B, T, D)
for block in self.blk: x = block(x, start_pos, use_flash, kv_len, valid_len)
for block in self.blk: x = block(x, start_pos, 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]
logits = self.output(self.output_norm(last))[:, -1, :]
# Gumbel-max trick: argmax(logits/temp - log(-log(uniform))) is equivalent to sampling from softmax(logits/temp)
if self.has_recurrent_block: return logits.argmax(-1, keepdim=True)
return (logits / temperature.maximum(1e-12) - (Tensor.rand_like(logits).maximum(1e-12).log().neg()).log()).argmax(-1, keepdim=True)
def forward_recurrent_decode(self, tokens:Tensor, start_pos:int|UOp, temperature:Tensor, decode_len:int,
valid_len:int|UOp|None=None) -> Tensor:
def forward_recurrent_decode(self, tokens:Tensor, start_pos:int|UOp, temperature:Tensor, decode_len:int, valid_len:int|UOp|None=None) -> Tensor:
return tokens.assign(self.forward(tokens, start_pos, temperature, kv_len=decode_len, valid_len=valid_len))
def __call__(self, tokens:Tensor, start_pos:int|UOp, temperature:Tensor, use_flash:bool=False,
valid_len:int|UOp|None=None) -> Tensor:
jit:Any
def __call__(self, tokens:Tensor, start_pos:int|UOp, temperature:Tensor, valid_len:int|UOp|None=None) -> Tensor:
if not self.has_recurrent_block:
return (self.prefill_jit if resolve(tokens.shape[1] != 1) else self.rollout_jit)(tokens.contiguous(), start_pos, temperature)
if resolve(tokens.shape[1] == 1):
jit = self.recurrent_rollout_jit
else:
prefill_key = (int(tokens.shape[1]), use_flash)
if prefill_key not in self.recurrent_prefill_jits:
self.recurrent_prefill_jits[prefill_key] = TinyJit(functools.partial(self.forward, use_flash=use_flash))
jit = self.recurrent_prefill_jits[prefill_key]
ret = jit(tokens.contiguous(), start_pos, temperature, valid_len=valid_len)
return ret[0] if isinstance(ret, tuple) else ret
jit = self.recurrent_rollout_jit if resolve(tokens.shape[1] == 1) else self.prefill_jit
return jit(tokens.contiguous(), start_pos, temperature, valid_len=valid_len)
@staticmethod
def from_gguf(gguf:Tensor|str|pathlib.Path, max_context:int|None=None,
@@ -555,27 +515,23 @@ class Transformer:
with Context(LLM_EMPTY_WEIGHTS=1): 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_layers:list[Linear] = []
packed_layers:list[tuple[Linear, Tensor]] = []
def resolve_owner(path:list[str]):
return functools.reduce(lambda obj, part: obj[int(part)] if isinstance(obj, list) else getattr(obj, part), path, model)
for name, weight in state_dict.items():
parts = name.split('.')
owner = resolve_owner(parts[:-1]) if parts[-1] == "weight" else None
if str(load_device).startswith("AMD") and isinstance(owner, Linear) and owner.set_quantized(weight):
packed_layers.append(owner)
if str(load_device).startswith("AMD") and isinstance(owner, Linear) and (offset:=owner.set_quantized(weight)) is not None:
packed_layers.append((owner, offset))
state_dict[name] = owner.weight
elif getenv("HALF", 1): state_dict[name] = weight.cast('float16')
nn.state.load_state_dict(model, state_dict, verbose=False, consume=True, realize=False) # NOTE: rope_freqs.weight (32,) is unused
recurrent_weights:list[Tensor] = []
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)
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
if isinstance(block, GatedDeltaNetBlock) and hasattr(block, "ssm_alpha") and \
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()
if packed_layers: Tensor.realize(*(offset for _,offset in packed_layers))
for layer,offset in packed_layers: 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())
@@ -589,19 +545,16 @@ class Transformer:
def warmup(self, chunk_size:int=256):
assert self.has_recurrent_block
device = self.token_embd.weight.device
warm_len = min(chunk_size, 256, self.max_context - 1)
if warm_len > 0:
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.flash_prefill_jit.cnt = self.recurrent_rollout_jit.cnt = 1
self.recurrent_prefill_jits[(warm_len, True)] = self.flash_prefill_jit
warm = self.generate([0] * warm_len, chunk_size=chunk_size)
with Context(JIT_BATCH_SIZE=getenv("PREFILL_JIT_BATCH_SIZE", 512)): next(warm)
with Context(JIT_BATCH_SIZE=0): next(warm)
x = Tensor.zeros(1, 1, self.blk[0].config.dim, device=self.token_embd.weight.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.prefill_jit.cnt = self.recurrent_rollout_jit.cnt = 1
warm = self.generate([0] * warm_len, chunk_size=chunk_size)
with Context(JIT_BATCH_SIZE=getenv("PREFILL_JIT_BATCH_SIZE", 512)): next(warm)
with Context(JIT_BATCH_SIZE=0): next(warm)
if resets := [r for block in self.blk for r in block._state_reset_ops()]: Tensor.realize(*resets)
self._cached_tokens = []
@@ -621,7 +574,6 @@ class Transformer:
out, prompt_len = None, len(tokens)
while len(tokens) < self.max_context:
recurrent_prefill = self.has_recurrent_block and start_pos < prompt_len and not decode_resume
use_flash = recurrent_prefill and bool(getenv("AMD_FLASH_ATTENTION", 1)) and chunk_size % 64 == 0
sp = v_start_pos.bind(start_pos)
actual_nt = min(1 if decode_resume and start_pos < prompt_len else chunk_size, len(tokens)-start_pos)
nt = chunk_size if recurrent_prefill else 1 if self.has_recurrent_block else v_toks.bind(actual_nt)
@@ -633,8 +585,8 @@ class Transformer:
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
out = self(inp, sp, temp, use_flash=use_flash, valid_len=valid_len).realize()
valid_len = v_toks.bind(actual_nt) if recurrent_prefill else None
out = self(inp, sp, temp, valid_len=valid_len).realize()
start_pos += actual_nt
# chunked prefill: keep processing until all prompt tokens are consumed
if start_pos < len(tokens): continue
tinygrad/llm/serve.py
+1 -2
View File
@@ -1,7 +1,7 @@
from __future__ import annotations
import json, pathlib, re, time, typing, uuid
from typing import TYPE_CHECKING
from tinygrad.helpers import DEBUG, START_TIME, colored, stderr_log
from tinygrad.helpers import DEBUG, colored, stderr_log
from tinygrad.viz.serve import TCPServerWithReuse, HTTPRequestHandler
if TYPE_CHECKING:
from tinygrad.llm.cli import SimpleTokenizer
@@ -164,5 +164,4 @@ class Handler(HTTPRequestHandler):
class LLMServer(TCPServerWithReuse):
def __init__(self, server_address:tuple, model:Transformer, model_name:str, tok:SimpleTokenizer, template:typing.Any):
self.model, self.model_name, self.tok, self.template = model, model_name, tok, template
print(f"*** started server on http://127.0.0.1:{server_address[1]} at {time.perf_counter()-START_TIME:.2f} s")
super().__init__(server_address, Handler)
tinygrad/schedule/rangeify.py
+4 -4
View File
@@ -412,13 +412,13 @@ def remove_noop_afters(x:UOp) -> UOp|None:
if len(src) != len(x.src): return src[0] if len(src) == 1 else x.replace(src=src)
return None
def close_after_stores(after:UOp) -> UOp|None:
def close_after_stores(after:UOp):
def close(store:UOp) -> UOp:
if store.op is not Ops.STORE: return store
kernel = store.end(*sorted([r for r in store.ranges if r.arg[-1] is not AxisType.DEVICE], key=lambda r:r.arg))
return kernel if store.src[0].buf_uop is after.buf_uop else store.src[0].buf_uop.after(kernel)
src = (after.src[0], *map(close, after.src[1:]))
return after.replace(src=src) if src != after.src else None
target = store.src[0].src[0].base
return kernel if target.buf_uop is after.buf_uop else target.after(kernel)
if (src := (after.src[0], *map(close, after.src[1:]))) != after.src: return after.replace(src=src)
pm_add_buffers = pm_mops+pm_flatten_bufferize+PatternMatcher([
(UPat(Ops.AFTER, name="after"), close_after_stores),