diff --git a/test/unit/test_attention.py b/test/unit/test_attention.py index 2731f3277a..5e00807d17 100644 --- a/test/unit/test_attention.py +++ b/test/unit/test_attention.py @@ -52,21 +52,22 @@ class TestAttention(unittest.TestCase): class TestGatedDeltaNetBlock(unittest.TestCase): def _tensor_linspace(self, start:float, stop:float, shape:tuple[int, ...]) -> Tensor: - return Tensor.linspace(start, stop, int(np.prod(shape)), dtype=dtypes.float32).reshape(*shape) + return Tensor(np.linspace(start, stop, int(np.prod(shape)), dtype=np.float32).reshape(shape), device=Tensor.empty(1).device).realize() def _make_config(self, **kwargs): - return TransformerConfig(**({"num_blocks":1, "dim":4, "hidden_dim":8, "n_heads":1, "n_kv_heads":1, - "norm_eps":1e-5, "vocab_size":32, "head_dim":4, "rope_theta":10000.0, - "rope_dim":4, "v_head_dim":4, "max_context":4, "ssm_layers":(True,), - "ssm":SSMConfig(conv_kernel=2, state_size=2, group_count=1, time_step_rank=1, inner_size=2)} | kwargs)) + return TransformerConfig(**({"num_blocks":1, "dim":32, "hidden_dim":64, "n_heads":1, "n_kv_heads":1, + "norm_eps":1e-5, "vocab_size":32, "head_dim":32, "rope_theta":10000.0, + "rope_dim":32, "v_head_dim":32, "max_context":4, "ssm_layers":(True,), + "ssm":SSMConfig(conv_kernel=2, state_size=32, group_count=1, time_step_rank=1, inner_size=32)} | kwargs)) def _make_block(self, config:TransformerConfig) -> GatedDeltaNetBlock: block = GatedDeltaNetBlock(config, config.ssm) block.attn_norm.weight = self._tensor_linspace(0.8, 1.2, (config.dim,)) block.attn_qkv.weight = self._tensor_linspace(-0.15, 0.2, (block.conv_channels, config.dim)) block.attn_gate.weight = self._tensor_linspace(-0.1, 0.15, (config.ssm.inner_size, config.dim)) - block.ssm_alpha.weight = self._tensor_linspace(-0.08, 0.12, (block.num_v_heads, config.dim)) - block.ssm_beta.weight = self._tensor_linspace(-0.12, 0.07, (block.num_v_heads, config.dim)) + beta = self._tensor_linspace(-0.12, 0.07, (block.num_v_heads, config.dim)) + alpha = self._tensor_linspace(-0.08, 0.12, (block.num_v_heads, config.dim)) + block.ssm_beta_alpha.weight = beta.cat(alpha) block.ssm_conv1d["weight"] = self._tensor_linspace(-0.05, 0.05, (block.conv_channels, block.ssm_conv_kernel)) block.ssm_dt["bias"] = self._tensor_linspace(-0.1, 0.1, (block.num_v_heads,)) block.ssm_a = self._tensor_linspace(-0.1, -0.05, (block.num_v_heads,)) @@ -77,8 +78,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, 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)) + out, conv_state = block._attention(x_norm, start_pos) + Tensor.realize(out, block.conv_state.assign(conv_state)) return out.numpy() def _cache_views(self, block:GatedDeltaNetBlock) -> tuple[np.ndarray, np.ndarray]: @@ -115,8 +116,7 @@ class TestGatedDeltaNetBlock(unittest.TestCase): conv_weight = block.ssm_conv1d["weight"].numpy().astype(np.float32).T[None, :, :] qkv_weight = block.attn_qkv.weight.numpy().astype(np.float32) gate_weight = block.attn_gate.weight.numpy().astype(np.float32) - alpha_weight = block.ssm_alpha.weight.numpy().astype(np.float32) - beta_weight = block.ssm_beta.weight.numpy().astype(np.float32) + beta_weight, alpha_weight = np.split(block.ssm_beta_alpha.weight.numpy().astype(np.float32), 2) out_weight = block.ssm_out.weight.numpy().astype(np.float32) dt_bias = block.ssm_dt["bias"].numpy().astype(np.float32) ssm_a = block.ssm_a.numpy().astype(np.float32) @@ -155,9 +155,10 @@ class TestGatedDeltaNetBlock(unittest.TestCase): return outputs, conv_states, recurrent_states def test_gatedeltanet_reference_and_reset(self): + if not str(Tensor.empty(1).device).startswith("AMD"): self.skipTest("AMD required") config = self._make_config(max_context=3) block = self._make_block(config) - x = Tensor.linspace(-1.0, 1.0, 3 * config.dim, dtype=dtypes.float32).reshape(1, 3, config.dim) + x = self._tensor_linspace(-1.0, 1.0, (1, 3, config.dim)) expected_outs, expected_conv, expected_recurrent = self._naive_attention(block, x) @@ -171,8 +172,7 @@ class TestGatedDeltaNetBlock(unittest.TestCase): np.testing.assert_allclose(recurrent_state, expected_recurrent[step], rtol=1e-3, atol=1e-3, err_msg=f"GatedDeltaNet recurrent cache mismatch at step {step}") - warmup = Tensor.linspace(-0.5, 0.5, 2 * config.dim, dtype=dtypes.float32).reshape(1, 2, config.dim) - prompt = Tensor.linspace(0.75, -0.75, 2 * config.dim, dtype=dtypes.float32).reshape(1, 2, config.dim) + warmup, prompt = self._tensor_linspace(-0.5, 0.5, (1, 2, config.dim)), self._tensor_linspace(0.75, -0.75, (1, 2, config.dim)) for i in range(warmup.shape[1]): self._run_attention(block, warmup[:, i:i+1], i) Tensor.realize(*block._state_reset_ops()) @@ -189,7 +189,8 @@ class TestGatedDeltaNetBlock(unittest.TestCase): err_msg=f"GatedDeltaNet reset recurrent cache mismatch at step {step}") def test_kda_channel_decay(self): - config = self._make_config(n_heads=2, ssm=SSMConfig(conv_kernel=2, state_size=2, group_count=2, time_step_rank=2, inner_size=4, kda=True)) + 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.ssm_f_a.weight = Tensor([[1., 0., 0., 0.], [0., 1., 0., 0.]]) diff --git a/test/unit/test_function.py b/test/unit/test_function.py index 3db67a9fa1..e333d8f9a6 100644 --- a/test/unit/test_function.py +++ b/test/unit/test_function.py @@ -63,16 +63,6 @@ class TestFunction(unittest.TestCase): np.testing.assert_equal(c.numpy(), [12,15,18]) np.testing.assert_equal(d.numpy(), [12,15,19]) - def test_precompile_cross_buffer_stores(self): - a, b = Tensor.zeros(8).contiguous().realize(), Tensor.zeros(8).contiguous().realize() - @function(precompile=True, allow_implicit=True) - def f(x:Tensor, start:UOp): - 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]) - def test_implicit_unrealized(self): inp = Tensor([1,2,3]) + Tensor([4,5,6]) @function(allow_implicit=True) diff --git a/test/unit/test_llm_amd.py b/test/unit/test_llm_amd.py index 9dd552b436..9959f3eee3 100644 --- a/test/unit/test_llm_amd.py +++ b/test/unit/test_llm_amd.py @@ -2,6 +2,8 @@ import unittest import numpy as np from tinygrad import Tensor, dtypes from tinygrad.llm.kernels.amd import q8_quantize +from tinygrad.llm.gguf import ggml_data_to_tensor +from tinygrad.llm.model import Linear class TestQ8Quantize(unittest.TestCase): def test_values_and_scales(self): @@ -12,4 +14,17 @@ class TestQ8Quantize(unittest.TestCase): np.testing.assert_array_equal(quant.bitcast(dtypes.int8).reshape(2, 32).numpy(), expected) np.testing.assert_allclose(scale.numpy(), scale_np, rtol=1e-6) + def test_q6_linear_compiles(self): + if not str(Tensor.empty(1).device).startswith("AMD"): self.skipTest("AMD required") + rng = np.random.default_rng(42) + packed = rng.integers(0, 256, 210, dtype=np.uint8) + packed[-2:] = np.array([0.01], dtype=np.float16).view(np.uint8) + raw = Tensor(np.pad(packed, (4, 0))).contiguous().realize()[4:] + decoded = ggml_data_to_tensor(raw, 256, 14).reshape(1, 256) + linear = Linear(256, 1, bias=False) + offset = linear.set_quantized(decoded) + assert offset is not None + linear._raw_offset_uop = offset.realize().uop + self.assertTrue(np.isfinite(linear(Tensor.randn(1, 256)).realize().item())) + if __name__ == "__main__": unittest.main() diff --git a/tinygrad/llm/kernels/amd.py b/tinygrad/llm/kernels/amd.py index 7cf0e8587d..463004234b 100644 --- a/tinygrad/llm/kernels/amd.py +++ b/tinygrad/llm/kernels/amd.py @@ -13,14 +13,12 @@ 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 -def warp_shfl_xor(val, offset, lane): - idx = ((lane ^ offset) * 4).int() - if val.op is Ops.INDEX and val.addrspace == AddrSpace.REG: val = val.load() - return UOp(Ops.CUSTOM, dtypes.float, (idx, val), arg="__builtin_bit_cast(float, __builtin_amdgcn_ds_bpermute({0}, __builtin_bit_cast(int, {1})))") - 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]): - other = warp_shfl_xor(val, offset, lane) + idx = ((lane ^ offset) * 4).int() + if val.op is Ops.INDEX and val.addrspace == AddrSpace.REG: val = val.load() + other = UOp(Ops.CUSTOM, dtypes.float, (idx, val), + arg="__builtin_bit_cast(float, __builtin_amdgcn_ds_bpermute({0}, __builtin_bit_cast(int, {1})))") val = val.maximum(other) if maximum else val + other return val @@ -49,27 +47,25 @@ def _amd_flash_attention_decode_partial(out, stats, q, cache_kv, cache_scale, va groups_per_chunk, offset = CHUNK // decode_group, UOp.range(((valid_chunks+group_count-1)//group_count)*(CHUNK//decode_group), 100, AxisType.REDUCE) chunk = block_n + (offset // groups_per_chunk) * group_count keys = tuple(chunk*CHUNK + (offset % groups_per_chunk)*decode_group + i for i in range(decode_group)) - valid = tuple(key < valid_kv_len for key in keys) kvals, vvals = (tuple(tuple(cache_kv[kv, b, kv_head, key, d].float() * cache_scale[kv, b, kv_head, key].float() for d in dims) for key in keys) for kv in range(2)) + q_heads = tuple(kv_head*G + head_group*head_tile + wave*heads_per_wave + head for head in range(heads_per_wave)) updates:list[UOp] = [] - for head in range(heads_per_wave): - q_head = kv_head*G + head_group*head_tile + wave*heads_per_wave + head + for head,q_head in enumerate(q_heads): scores = tuple(warp_reduce(sum((q[b, q_head, 0, d].float()*k for d,k in zip(dims, key_kvals)), UOp.const(0, dtypes.float)), lane + wave*WARP_SIZE, full_wave=True) / math.sqrt(D) for key_kvals in kvals) prev_acc, prev_max, prev_sum = acc.after(offset)[head], row_max.after(offset)[head], row_sum.after(offset)[head] - new_max = functools.reduce(lambda a,vs:a.maximum(vs[0].where(vs[1], UOp.const(-math.inf, dtypes.float))), zip(valid, scores), prev_max) + new_max = functools.reduce(lambda a,ks:a.maximum((ks[0] < valid_kv_len).where(ks[1], UOp.const(-math.inf, dtypes.float))), + zip(keys, scores), prev_max) alpha = ((prev_max-new_max)*LOG2E).exp2() - betas = tuple(is_valid.where(((score-new_max)*LOG2E).exp2(), UOp.const(0, dtypes.float)) for is_valid,score in zip(valid, scores)) + betas = tuple((key < valid_kv_len).where(((score-new_max)*LOG2E).exp2(), UOp.const(0, dtypes.float)) for key,score in zip(keys, scores)) updates += [acc[head].store(prev_acc*alpha + sum((UOp.stack(*value)*beta for value,beta in zip(vvals, betas)), acc[head].const_like(0))), row_sum[head].store(prev_sum*alpha + sum(betas, UOp.const(0, dtypes.float))), row_max[head].store(new_max)] update = UOp.group(*updates).end(offset) acc, row_max, row_sum = acc.after(update), row_max.after(update), row_sum.after(update) - def head_stores(head:int) -> list[UOp]: - q_head = kv_head*G + head_group*head_tile + wave*heads_per_wave + head - return [out[b, q_head, block_n, d].store(acc[head, i]) for i,d in enumerate(dims)] + \ - [stats[b, q_head.valid(lane.eq(0)), block_n, i].store(x[head]) for i,x in enumerate((row_max, row_sum))] - stores = [store for head in range(heads_per_wave) for store in head_stores(head)] + stores = [out[b, q_head, block_n, d].store(acc[head, i]) for head,q_head in enumerate(q_heads) for i,d in enumerate(dims)] + \ + [stats[b, q_head.valid(lane.eq(0)), block_n, i].store(x[head]) + for head,q_head in enumerate(q_heads) for i,x in enumerate((row_max, row_sum))] return UOp.group(*stores).end(lane, wave, block_n, block_bhkv).sink(arg=KernelInfo(name="flash_decode_partial", opts_to_apply=())) def amd_flash_attention_decode(q:Tensor, cache_kv:Tensor, valid_kv_len:int|UOp, cache_scale:Tensor, max_kv_len:int) -> Tensor: @@ -187,8 +183,7 @@ 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({}, {}, {})") + return UOp(Ops.CUSTOMI, dtypes.uint32, tuple(x.cast(dtypes.uint32) for x in (a, b, selectors)), arg="__builtin_amdgcn_perm({}, {}, {})") def _amd_load(ptr:UOp, lanes:int|None=None) -> UOp: assert ptr.op is Ops.INDEX @@ -212,13 +207,13 @@ 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, 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, next_state = (x.reshape(batch*heads, dim, dim) for x in (state, next_state)) + 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 rows = tuple(row_base+i for i in range(row_tile)) @@ -239,16 +234,16 @@ 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) - state_stores = (next_state[bh, row, col].store(current.after(step)[row_idx*dim//32+i].load().cast(next_state.dtype)) + state_stores = (state[bh, row, col].store(current.after(step)[row_idx*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) -> tuple[Tensor, Tensor]: +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, next_state, kq = Tensor.empty_like(q), Tensor.empty_like(state), (q*k).sum(-1).contiguous() - return tuple(Tensor.custom_kernel(core, next_state, q.contiguous(), k.contiguous(), v.contiguous(), beta.contiguous(), - alpha.contiguous(), state, kq, fxn=_gated_delta_prefill_kernel)[:2]) # type: ignore[return-value] + core, kq = Tensor.empty_like(q), (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] def _wmma_layout(out:UOp, out_features:int, token_tile:int, output_tiles:int): output_waves = 2 if out_features % (32*output_tiles) == 0 else 1 @@ -273,12 +268,11 @@ def _wmma_stores(out, outputs, tokens, accs, update, half): def _decode_linear(out:UOp, out_features:int, group_count:int, group_dot, name:str) -> UOp: output, lane = UOp.range(out_features, 0), UOp.range(32, 1, axis_type=AxisType.LOCAL) - acc = UOp.placeholder((1,), dtypes.float32, slot=0, addrspace=AddrSpace.REG) - acc = acc.after(acc.store(acc.const_like(0))) + acc = UOp.placeholder((1,), dtypes.float32, slot=0, addrspace=AddrSpace.REG)[0].set(0) chunk = UOp.range((group_count+31)//32, 2, AxisType.REDUCE) group = (lane+chunk*32).valid(lane+chunk*32 < group_count) - update = acc.store(acc.after(chunk) + group_dot(output, group)).end(chunk) - total = warp_reduce(acc.after(update)[0].load(), lane, full_wave=True) + acc = acc[0].set(acc.after(chunk)[0] + group_dot(output, group), end=chunk) + total = warp_reduce(acc[0].load(), lane, full_wave=True) return out[0, output.valid(lane.eq(0))].store(total.cast(out.dtype)).end(output, lane).sink(arg=KernelInfo(name=name, opts_to_apply=())) def _q5_scales(raw:UOp, base:UOp, subgroup:UOp) -> tuple[UOp, UOp, UOp, UOp]: @@ -290,22 +284,41 @@ def _q5_scales(raw:UOp, base:UOp, subgroup:UOp) -> tuple[UOp, UOp, UOp, UOp]: 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: - group_count, type_words = in_features // Q8_GROUP_SIZE, Q5_WORDS +def _quant_decode_kernel(out:UOp, raw:UOp, xq:UOp, xd:UOp, raw_offset:UOp, out_features:int, in_features:int, ggml_type:int) -> UOp: + group_count = in_features // Q8_GROUP_SIZE def group_dot(output:UOp, group:UOp) -> UOp: block, subgroup = group // 8, group % 8 - base = raw_offset + (output * in_features//GGML_BLOCK_SIZE + block) * type_words - qs_base = base + 12 + (subgroup // 2) * 8 xwords = _amd_load(xq[0, group, 0], 8) - dot, qsum = UOp.const(0, dtypes.int32), UOp.const(0, dtypes.int32) + if ggml_type == Q5_K: + base = raw_offset + (output * in_features//GGML_BLOCK_SIZE + block) * Q5_WORDS + qs_base, dot, qsum = base + 12 + (subgroup//2)*8, UOp.const(0, dtypes.int32), UOp.const(0, dtypes.int32) + for word_idx in range(8): + word = (raw[qs_base+word_idx] >> ((subgroup&1)*4).cast(dtypes.uint32)) & 0x0f0f0f0f + word |= ((raw[base+4+word_idx] >> subgroup.cast(dtypes.uint32)) & 0x01010101) << 4 + dot, qsum = _amd_dp4a(word, xwords[word_idx], dot), _amd_dp4a(UOp.const(0x01010101, dtypes.uint32), xwords[word_idx], qsum) + d, dmin, scale, minimum = _q5_scales(raw, base, subgroup) + return (dot.float()*d*scale - qsum.float()*dmin*minimum) * xd[0, group] + if ggml_type == IQ4_XS: + base = raw_offset + (output * in_features//GGML_BLOCK_SIZE + block) * IQ4_WORDS + dot = UOp.const(0, dtypes.int32) + for word_idx in range(8): + packed = _amd_load(raw[base + 2 + subgroup*4 + word_idx%4]) + dot = _amd_dp4a(_iq4_bytes(packed, 4*(word_idx//4)), xwords[word_idx], dot) + d, scale = _iq4_scales(raw, base, subgroup) + return dot.float() * xd[0, group] * d * scale + base = raw_offset*4 + (output*in_features//GGML_BLOCK_SIZE+block)*Q6_BYTES + dots = [UOp.const(0, dtypes.int32), UOp.const(0, dtypes.int32)] for word_idx in range(8): - word = (raw[qs_base + word_idx] >> ((subgroup & 1) * 4).cast(dtypes.uint32)) & 0x0f0f0f0f - word = word | (((raw[base + 4 + word_idx] >> subgroup.cast(dtypes.uint32)) & 0x01010101) << 4) - dot = _amd_dp4a(word, xwords[word_idx], dot) - qsum = _amd_dp4a(UOp.const(0x01010101, dtypes.uint32), xwords[word_idx], qsum) - d, dmin, scale, minimum = _q5_scales(raw, base, subgroup) - return (dot.float()*d*scale - qsum.float()*dmin*minimum) * xd[0, group] - return _decode_linear(out, out_features, group_count, group_dot, "linear_q5_k") + pos, within = subgroup*32 + word_idx*4, (subgroup*32 + word_idx*4)%128 + low = _amd_load(raw[base + (pos//128)*64 + within%64], 4) >> ((within//64)*4).cast(dtypes.uint8) + high = _amd_load(raw[base + 128 + (pos//128)*32 + within%32], 4) >> ((within//32)*2).cast(dtypes.uint8) + quant = ((low & 15) | ((high & 3) << 4)).bitcast(dtypes.int8) - 32 + word = quant.cast(dtypes.int8).bitcast(dtypes.uint8).bitcast(dtypes.uint32).squeeze(0) + 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] * _half(dbits) + return _decode_linear(out, out_features, group_count, group_dot, {Q5_K:"linear_q5_k", IQ4_XS:"linear_iq4_xs", Q6_K:"linear_q6"}[ggml_type]) def _quant_linear_wmma(out, x, out_features, in_features, type_words, layout, dequant, name): x = x.reshape(out.shape[0], in_features) @@ -333,7 +346,7 @@ def _quant_linear_wmma(out, x, out_features, in_features, type_words, layout, de @functools.cache def _q5_linear_f16_wmma_kernel(out:UOp, raw:UOp, x:UOp, raw_offset:UOp, out_features:int, in_features:int) -> UOp: - token_tile, output_tiles = (64, 1) if (out_features <= 1024 or out_features > 6144) and out.shape[0] % 64 == 0 else \ + token_tile, output_tiles = (64, 1) if out_features <= 1024 and out.shape[0] % 64 == 0 else \ (64, 2) if out.shape[0] % 64 == 0 else (32 if out.shape[0] % 32 == 0 else 16, 2) def dequant(base:UOp, subgroup:UOp, half:int) -> tuple[UOp, ...]: base = raw_offset + base @@ -350,21 +363,6 @@ def _iq4_scales(raw:UOp, base:UOp, subgroup:UOp) -> tuple[UOp, UOp]: scale = ((low >> (4*(subgroup%2)).cast(dtypes.uint32)) & 15) | ((((raw[base] >> 16) >> (2*subgroup).cast(dtypes.uint32)) & 3) << 4) return _half(raw[base] & 0xffff), (scale.cast(dtypes.uint8).bitcast(dtypes.int8)-32).float() -@functools.cache -def _iq4_linear_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, IQ4_WORDS - def group_dot(output:UOp, group:UOp) -> UOp: - block, subgroup = group // 8, group % 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_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: token_tile = 32 if out_features <= 1024 and out.shape[0] % 32 == 0 else 64 if out.shape[0] % 64 == 0 and \ @@ -390,33 +388,11 @@ def _iq4_linear_f16_wmma_kernel(out:UOp, raw:UOp, x:UOp, lut:UOp, raw_offset:UOp 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 * 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_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) - for byte_idx in range(4): - pos, within = subgroup*32 + word_idx*4 + byte_idx, (subgroup*32 + word_idx*4 + byte_idx)%128 - low = (raw[base + (pos//128)*64 + within%64] >> ((within//64)*4).cast(dtypes.uint8)) & 15 - high = (raw[base + 128 + (pos//128)*32 + within%32] >> ((within//32)*2).cast(dtypes.uint8)) & 3 - q = (low | (high << 4)).cast(dtypes.uint8).bitcast(dtypes.int8) - 32 - word = word | (q.cast(dtypes.int8).bitcast(dtypes.uint8).cast(dtypes.uint32) << (8*byte_idx)) - 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] * _half(dbits) - return _decode_linear(out, out_features, group_count, group_dot, "linear_q6") - 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 + assert layer.ggml_type in (Q5_K, Q6_K, IQ4_XS) and layer._raw_offset_uop is not None tokens = int(x.numel()) // layer.in_features out = Tensor.empty(tokens, layer.out_features, dtype=dtypes.float32, device=x.device).uop - raw, offset = layer._raw_uop, layer._raw_offset_uop + raw, offset = layer.weight.uop.buf_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 def run(fxn:Callable[..., UOp], *srcs:UOp) -> Tensor: @@ -427,13 +403,12 @@ def q8_linear(layer:Linear, x:Tensor) -> Tensor: 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 == 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) + xq, xd = q8_quantize(x, tokens, layer.in_features) + decode = functools.partial(_quant_decode_kernel, ggml_type=layer.ggml_type) + return run(decode, out, raw if layer.ggml_type == Q6_K else raw.bitcast(dtypes.uint32), xq.uop, xd.uop, offset) @functools.cache def iq4_half_lut(device:str) -> Tensor: diff --git a/tinygrad/llm/model.py b/tinygrad/llm/model.py index 7faa66ea74..a7ef4be64f 100644 --- a/tinygrad/llm/model.py +++ b/tinygrad/llm/model.py @@ -5,37 +5,26 @@ 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 -from tinygrad.helpers import prod, ContextVar +from tinygrad.helpers import prod from tinygrad.uop.ops import resolve, Ops def unwrap_var(x:int|UOp|None): return x.unbind()[0] if isinstance(x, UOp) else x -LLM_EMPTY_WEIGHTS = ContextVar("LLM_EMPTY_WEIGHTS", 0) class Linear(nn.Linear): - ggml_type:int|None - _raw_uop:UOp|None - _raw_offset_uop:UOp|None + ggml_type:int|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 - packed = Tensor(raw) - self.weight, self.ggml_type = packed.flatten(), packed_sizes[prod(raw.shape)] - raw, raw_offset = self.weight.uop, 0 - while raw.op in (Ops.BITCAST, Ops.RESHAPE): raw = raw.src[0] - while raw.op is Ops.SHRINK: - raw_offset += raw.src[1].arg * raw.dtype.itemsize - raw = raw.src[0] - assert raw_offset % 4 == 0 and raw.dtype == dtypes.uint8 - self._raw_uop = raw - if self.ggml_type == 23 and str(packed.device).startswith("AMD"): llm_amd.iq4_half_lut(str(packed.device)) + 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"): llm_amd.iq4_half_lut(str(self.weight.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: - self.weight, self.bias = Tensor.empty(out_features, in_features), Tensor.empty(out_features) if bias else None - else: super().__init__(in_features, out_features, bias) + super().__init__(in_features, out_features, bias) self.in_features, self.out_features = in_features, out_features - self.ggml_type, self._raw_uop, self._raw_offset_uop = None, None, None + self._raw_offset_uop:UOp|None = None def __call__(self, x:Tensor) -> Tensor: return llm_amd.q8_linear(self, x) if self.ggml_type in (13, 14, 23) and str(self.weight.device).startswith("AMD") else \ super().__call__(x) @@ -154,28 +143,23 @@ class FFNBlock: return out return self.ffn_down(self.ffn_gate(x).silu() * self.ffn_up(x)) - # given the token-prefix match, return how much cached state this block can still reuse - def _reusable_prefix_len(self, prefix_len:int, cached_len:int) -> int: return prefix_len - # return writes that reset this block's state after a cache mismatch - def _state_reset_ops(self) -> list[Tensor]: return [] def _init_state(self, x:Tensor): raise NotImplementedError - def _attention(self, x:Tensor, start_pos:int|UOp, kv_len:int|UOp|None=None, valid_len:int|UOp|None=None) -> Tensor: + def _attention(self, x:Tensor, start_pos:int|UOp, valid_len:int|UOp|None=None) -> Tensor: raise NotImplementedError - def __call__(self, x: Tensor, start_pos: int|UOp, kv_len:int|UOp|None=None, valid_len:int|UOp|None=None): + def __call__(self, x: Tensor, start_pos: int|UOp, valid_len:int|UOp|None=None): self._init_state(x) if hasattr(self, 'attn_gate'): @function(precompile=True, allow_implicit=True) def _run_stateful(x:Tensor, start_pos:int|UOp, valid_len:int|UOp|None): - attn, conv_state, recurrent_state = cast(tuple[Tensor, Tensor, Tensor], self._attention(self.attn_norm(x), start_pos, kv_len, valid_len)) + attn, conv_state = cast(tuple[Tensor, Tensor], self._attention(self.attn_norm(x), start_pos, 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))) + return (h + self._feed_forward(self.ffn_norm(h))).contiguous(), conv_state + out, conv_state = _run_stateful(x, start_pos, valid_len) + state = getattr(self, "conv_state") + return Tensor(out.uop.after(state.uop.after(state.uop.store(conv_state.uop)))) def _run(x:Tensor, start_pos:int|UOp): - h = x + self._attention(self.attn_norm(x), start_pos, kv_len) + h = x + self._attention(self.attn_norm(x), start_pos) return (h + self._feed_forward(self.ffn_norm(h))).contiguous() return function(precompile=True, allow_implicit=True)(_run)(x, start_pos) @@ -193,7 +177,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, kv_len:int|UOp|None=None, valid_len:int|UOp|None=None): + def _attention(self, x:Tensor, start_pos:int|UOp, 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) @@ -210,33 +194,27 @@ class TransformerBlock(FFNBlock): k = apply_rope(k[..., :self.config.rope_dim], self.freqs_cis[start_pos:start_pos+T]).cat(k[..., self.config.rope_dim:], dim=-1) # NOTE: we don't want to change self.cache_kv, the function API doesn't support this well - stacked_kv, assigned_scale = Tensor.stack(k, v), None + stacked_kv = Tensor.stack(k, v) if self.cache_kv.dtype == dtypes.int8: scale = (stacked_kv.float().abs().max(axis=-1, keepdim=True) / 127).maximum(1e-8).half() packed_kv = (stacked_kv.float() / scale).round().clip(-127, 127).cast(dtypes.int8) stores = (self.cache_kv[:, :, :, start_pos:start_pos+T, :].uop.store(packed_kv.uop), self.cache_kv_scale[:, :, :, start_pos:start_pos+T].uop.store(scale.squeeze(-1).uop)) assigned_kv, assigned_scale = Tensor(self.cache_kv.uop.after(*stores)), Tensor(self.cache_kv_scale.uop.after(*stores)) - else: assigned_kv = Tensor(self.cache_kv.uop.after( - self.cache_kv[:, :, :, start_pos:start_pos+T, :].uop.store(stacked_kv.cast(self.cache_kv.dtype).uop))) - cache_len = start_pos + T if kv_len is None else kv_len - k, v = assigned_kv[0, :, :, 0:cache_len, :], assigned_kv[1, :, :, 0:cache_len, :] - if assigned_scale is not None: - k, v = k.float() * assigned_scale[0, :, :, 0:cache_len, None], v.float() * assigned_scale[1, :, :, 0:cache_len, None] - - # 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 - 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) + if resolve(T == 1): + attn = llm_amd.amd_flash_attention_decode(q.half(), assigned_kv, cast(int|UOp, unwrap_var(start_pos))+1, + assigned_scale, self.config.max_context) + else: + start, valid = cast(int|UOp, unwrap_var(start_pos)), unwrap_var(valid_len) + attn = llm_amd.flash_attention_causal_cached(q.half(), assigned_kv, start+T, assigned_scale, + start+valid if valid is not None else None) else: - 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 + assigned_kv = Tensor(self.cache_kv.uop.after( + self.cache_kv[:, :, :, start_pos:start_pos+T, :].uop.store(stacked_kv.cast(self.cache_kv.dtype).uop))) + k, v = assigned_kv[0, :, :, 0:start_pos+T, :], assigned_kv[1, :, :, 0:start_pos+T, :] + # NOTE: this mask is causal_lower_right, not the causal_upper_left generated by is_casual = True + mask = Tensor.full((1, 1, T, k.shape[-2]), float("-inf"), dtype=x.dtype, device=x.device, buffer=False).triu(start_pos+1) \ + if resolve(T != 1) else None attn = q.float().scaled_dot_product_attention(k.float(), v.float(), attn_mask=mask, enable_gqa=True) # (B,H,T,Hd) attn = attn.transpose(1, 2).reshape(B, T, -1) # back to (B,T,D) return self.attn_output(attn if not self.config.attn_output_gate else (attn * gate.sigmoid())) @@ -244,7 +222,8 @@ class TransformerBlock(FFNBlock): def _init_state(self, x:Tensor): if not hasattr(self, "cache_kv"): # Zero padding prevents masked, unwritten cache entries from injecting NaNs before the mask. - cache_dtype = dtypes.int8 if self.config.max_context > 8192 and str(x.device).startswith("AMD") else dtypes.float16 + cache_dtype = dtypes.int8 if self.config.ssm is not None and self.config.max_context > 8192 and \ + str(x.device).startswith("AMD") else dtypes.float16 cache_shape = (2, x.shape[0], self.config.n_kv_heads, self.config.max_context+192, self.config.head_dim) self.cache_kv = Tensor.zeros(*cache_shape, dtype=cache_dtype, device=x.device).contiguous() if cache_dtype == dtypes.int8: self.cache_kv_scale = Tensor.zeros(*cache_shape[:-1], dtype=dtypes.float16, device=x.device).contiguous() @@ -266,7 +245,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, kv_len:int|UOp|None=None, valid_len:int|UOp|None=None) -> Tensor: + def _attention(self, x:Tensor, start_pos:int|UOp, 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) @@ -308,82 +287,46 @@ class GatedDeltaNetBlock(FFNBlock): if ssm.kda: self.ssm_g_a, self.ssm_g_b = Linear(config.dim, self.head_v_dim, bias=False), Linear(self.head_v_dim, ssm.inner_size, bias=False) self.ssm_f_a, self.ssm_f_b = Linear(config.dim, self.head_k_dim, bias=False), Linear(self.head_k_dim, ssm.inner_size, bias=False) + self.ssm_beta = Linear(config.dim, self.num_v_heads, bias=False) else: self.attn_gate = Linear(config.dim, ssm.inner_size, bias=False) - self.ssm_alpha = Linear(config.dim, self.num_v_heads, bias=False) - self.ssm_beta = Linear(config.dim, self.num_v_heads, bias=False) - self.ssm_beta_alpha_weight:Tensor|None = None + self.ssm_beta_alpha = Linear(config.dim, 2*self.num_v_heads, bias=False) self.ssm_conv1d = {"weight": Tensor.zeros(self.conv_channels, self.ssm_conv_kernel)} self.ssm_dt = {"bias": Tensor.zeros(ssm.inner_size if ssm.kda else self.num_v_heads)} self.ssm_a = Tensor.zeros(self.num_v_heads, 1) if ssm.kda else Tensor.zeros(self.num_v_heads) self.ssm_norm, self.ssm_out = nn.RMSNorm(self.head_v_dim, config.norm_eps), Linear(ssm.inner_size, config.dim, bias=False) - 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) + def _attention(self, x:Tensor, start_pos:int|UOp, valid_len:int|UOp|None=None): B, T, _ = x.shape - assert T == 1, "GatedDeltaNetBlock currently only supports T=1" - - # input processing + is_qwen = hasattr(self, "attn_gate") + assert is_qwen or T == 1, "Kimi GatedDeltaNetBlock currently only supports T=1" x = x.half() - out_gate = self.ssm_g_b(self.ssm_g_a(x)) if hasattr(self, "ssm_g_a") else self.attn_gate(x) - out_gate = out_gate.reshape(B, 1, self.num_v_heads, self.head_v_dim) - beta = self.ssm_beta(x).sigmoid().reshape(B, self.num_v_heads, 1, 1) - alpha = self.ssm_f_b(self.ssm_f_a(x)) if hasattr(self, "ssm_f_a") else self.ssm_alpha(x) - alpha = ((alpha.float() + self.ssm_dt["bias"]).softplus().reshape(B, self.num_v_heads, -1) * - self.ssm_a.reshape(1, self.num_v_heads, -1)).exp().unsqueeze(-2) - - # qkv conv + out_gate = self.attn_gate(x) if is_qwen else self.ssm_g_b(self.ssm_g_a(x)) + if is_qwen: + beta, alpha = self.ssm_beta_alpha(x).split(self.num_v_heads, dim=-1) + else: beta, alpha = self.ssm_beta(x), self.ssm_f_b(self.ssm_f_a(x)) conv_window = self.conv_state.cat(self.attn_qkv(x), dim=1) - conv_out = (conv_window * self.ssm_conv1d["weight"].T.unsqueeze(0)).sum(1).silu() - q, k, v = conv_out.split([self.q_dim, self.q_dim, self.conv_channels - 2*self.q_dim], dim=-1) - q = q.reshape(B, self.num_k_heads, self.head_k_dim).normalize(dim=-1).repeat(1, self.num_v_heads//self.num_k_heads, 1) - k = k.reshape(B, self.num_k_heads, self.head_k_dim).normalize(dim=-1).repeat(1, self.num_v_heads//self.num_k_heads, 1) - v = v.reshape(B, self.num_v_heads, self.head_v_dim) - q, k, v = q.mul(self.head_k_dim**-0.5).unsqueeze(-1), k.unsqueeze(-1), v.unsqueeze(-1) - - # recurrent - recurrent_state = self.recurrent_state * alpha - recurrent_state = recurrent_state + ((v - recurrent_state@k) * beta)@k.transpose(-1, -2) - - # store the updated state - conv_state_store = self.conv_state.uop.store(conv_window[:, 1:, :].cast(self.conv_state.dtype).uop) - recurrent_state_store = self.recurrent_state.uop.store(recurrent_state.cast(self.recurrent_state.dtype).uop) - recurrent_state = Tensor(self.recurrent_state.uop.after(recurrent_state_store, conv_state_store)) - - # output - core_attn_out = self.ssm_norm((recurrent_state@q).squeeze(-1).reshape(B, 1, self.num_v_heads, self.head_v_dim)) - out_gate = out_gate.sigmoid() if hasattr(self, "ssm_g_a") else out_gate.silu() - return self.ssm_out((core_attn_out * out_gate).reshape(B, 1, -1).cast(x.dtype)) - - 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 = 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) + q = q.reshape(B, T, self.num_k_heads, self.head_k_dim).normalize(dim=-1, eps=1e-6 if is_qwen else 1e-12).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 if is_qwen else 1e-12).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, 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 - 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)), conv_state, recurrent_state - - assert str(x.device).startswith("AMD"), "batched GatedDeltaNet prefill currently requires AMD" + initial_state = self.recurrent_state + if not is_qwen: + out_gate = out_gate.reshape(B, 1, self.num_v_heads, self.head_v_dim).sigmoid() + 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 = initial_state*alpha + ((v-initial_state*alpha@k)*beta)@k.transpose(-1, -2) + stores = (self.recurrent_state.uop.store(recurrent_state.cast(self.recurrent_state.dtype).uop), + self.conv_state.uop.store(conv_window[:, 1:, :].cast(self.conv_state.dtype).uop)) + state = Tensor(self.recurrent_state.uop.after(*stores)) + core = self.ssm_norm((state@q).squeeze(-1).reshape(B, 1, self.num_v_heads, self.head_v_dim)) + return self.ssm_out((core*out_gate).reshape(B, 1, -1).cast(x.dtype)) out_gate = out_gate.reshape(B, T, self.num_v_heads, self.head_v_dim) 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) @@ -391,11 +334,11 @@ class GatedDeltaNetBlock(FFNBlock): 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, recurrent_state = llm_amd.gated_delta_prefill(q * self.head_k_dim**-0.5, k, v, beta, log_alpha.exp(), initial_state) + core = 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 - 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 + conv_state = conv_window[:, state_pos:state_pos+self.ssm_conv_kernel-1].cast(self.conv_state.dtype).contiguous() + return out, conv_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 [] @@ -421,19 +364,19 @@ class Transformer: # we specialize the JIT for prefill and rollout self.prefill_jit = TinyJit(self.forward) self.rollout_jit = TinyJit(self.forward) - self.recurrent_rollout_jit = TinyJit(functools.partial(self.forward_recurrent_decode, decode_len=self.max_context)) + self.recurrent_rollout_jit = TinyJit(self.forward_recurrent_decode) - def forward(self, tokens:Tensor, start_pos:int|UOp, temperature:Tensor, kv_len:int|UOp|None=None, valid_len:int|UOp|None=None) -> Tensor: + def forward(self, tokens:Tensor, start_pos:int|UOp, temperature:Tensor, 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, kv_len, valid_len) + for block in self.blk: x = block(x, start_pos, 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: - return tokens.assign(self.forward(tokens, start_pos, temperature, kv_len=decode_len, valid_len=valid_len)) + def forward_recurrent_decode(self, tokens:Tensor, start_pos:int|UOp, temperature:Tensor, valid_len:int|UOp|None=None) -> Tensor: + return tokens.assign(self.forward(tokens, start_pos, temperature, valid_len)) def __call__(self, tokens:Tensor, start_pos:int|UOp, temperature:Tensor, valid_len:int|UOp|None=None) -> Tensor: if not self.has_recurrent_block: @@ -458,6 +401,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) @@ -511,25 +457,17 @@ class Transformer: routed_scaling_factor=kv.get(f'{arch}.expert_weights_scale', 1.0), attn_output_gate=arch in ('qwen35', 'qwen35moe'), ssm=ssm, ssm_layers=ssm_layers, qkv_bias='blk.0.attn_q.bias' in state_dict, - expert_bias=f"blk.{kv.get(f'{arch}.leading_dense_block_count', 0)}.exp_probs_b.bias" in state_dict) - 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)) + expert_bias=f"blk.{kv.get(f'{arch}.leading_dense_block_count', 0)}.exp_probs_b.bias" in state_dict) + model = Transformer(config) + if getenv("HALF", 1): state_dict = {name:weight.cast('float16') for name,weight in state_dict.items()} 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 (offset:=owner.set_quantized(weight)) is not None: + 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_layers.append((owner, offset)) - state_dict[name] = owner.weight - elif getenv("HALF", 1): state_dict[name] = weight.cast('float16') + state_dict[key] = owner.weight nn.state.load_state_dict(model, state_dict, verbose=False, consume=True, realize=False) # NOTE: rope_freqs.weight (32,) is unused - for block in model.blk: - 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 @@ -540,23 +478,21 @@ class Transformer: def get_start_pos(self, tokens:list[int]) -> int: prefix_len = sum(1 for _ in itertools.takewhile(lambda ab: ab[0] == ab[1], zip(tokens[:-1], self._cached_tokens))) - if self.has_recurrent_block: return prefix_len if prefix_len == len(self._cached_tokens) else 0 - return min(block._reusable_prefix_len(prefix_len, len(self._cached_tokens)) for block in self.blk) + return 0 if self.has_recurrent_block and prefix_len != len(self._cached_tokens) else prefix_len def warmup(self, chunk_size:int=256): assert self.has_recurrent_block - warm_len = min(chunk_size, 256, self.max_context - 1) 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) + warm = self.generate([0] * min(chunk_size, 256, self.max_context-1), 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) + if resets := [r for block in self.blk if hasattr(block, "_state_reset_ops") for r in block._state_reset_ops()]: Tensor.realize(*resets) self._cached_tokens = [] def generate(self, tokens:list[int], chunk_size:int|None=None, temperature:float=0.0): @@ -568,21 +504,18 @@ class Transformer: # TODO: use UOp.variable for temperature once float variables are supported device = self.token_embd.weight.device temp = Tensor([temperature], device=device) - t = None if self.has_recurrent_block else \ - Tensor(tokens + [0] * (self.max_context + chunk_size - len(tokens)), dtype="int32", device=device).reshape(1, self.max_context + chunk_size) - if start_pos < len(self._cached_tokens) and (resets := [r for b in self.blk for r in b._state_reset_ops()]): Tensor.realize(*resets) + t = Tensor(tokens + [0] * (self.max_context + chunk_size - len(tokens)), dtype="int32", device=device).reshape(1, self.max_context + chunk_size) + if start_pos < len(self._cached_tokens) and \ + (resets := [r for b in self.blk if hasattr(b, "_state_reset_ops") for r in b._state_reset_ops()]): Tensor.realize(*resets) 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 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) - if self.has_recurrent_block and (start_pos < prompt_len or out is None): - assert isinstance(nt, int) - inp = out.assign(Tensor([[tokens[start_pos]]], dtype="int32", device=device)).realize() if decode_resume and out is not None else \ - Tensor(tokens[start_pos:start_pos+actual_nt] + [0] * (nt-actual_nt), dtype="int32", device=device).reshape(1, nt) + if decode_resume and start_pos < prompt_len and out is not None: + inp = out.assign(Tensor([[tokens[start_pos]]], dtype="int32", device=device)).realize() elif start_pos < prompt_len or out is None: - assert t is not None inp = t[:, sp:sp+nt] else: inp = out valid_len = v_toks.bind(actual_nt) if recurrent_prefill else None diff --git a/tinygrad/renderer/cstyle.py b/tinygrad/renderer/cstyle.py index e25ab0c605..68178ea831 100644 --- a/tinygrad/renderer/cstyle.py +++ b/tinygrad/renderer/cstyle.py @@ -162,6 +162,7 @@ class CStyleLanguage(Renderer): def render_index(self, x:UOp, buf:UOp, idx:UOp): if buf.addrspace == AddrSpace.ALU: + if buf.max_numel() == 1: return self[buf] # this is lane access in C if idx.op is not Ops.CONST: return f"({self[buf]})[{self[idx]}]" return self[buf]+(f"[{idx.val}]" if buf.max_numel() > self.gep_arr_threshold else f".{'xyzwabcd'[idx.val]}") diff --git a/tinygrad/schedule/rangeify.py b/tinygrad/schedule/rangeify.py index 64c5805c7c..3c4e3750aa 100644 --- a/tinygrad/schedule/rangeify.py +++ b/tinygrad/schedule/rangeify.py @@ -412,16 +412,7 @@ 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): - 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)) - 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), (UPat(Ops.STAGE, src=(UPat(), UPat(name="idx")), name="x"), lambda ctx,x,idx: bufferize_to_store(ctx, x, idx, allow_locals=False)), # INDEX of a buffer through the weak cast added above: index the buffer directly and cast the loaded value instead.