From f26128e8d2bc45161bdec87f7ae55ca9d34cb33b Mon Sep 17 00:00:00 2001 From: George Hotz Date: Mon, 3 Aug 2026 16:28:20 +0000 Subject: [PATCH] 400 lines --- extra/benchmark_qwen_amd.py | 15 ++- test/unit/test_attention.py | 9 +- test/unit/test_function.py | 5 +- tinygrad/llm/cli.py | 15 ++- tinygrad/llm/kernels/amd.py | 186 +++++++++++++--------------------- tinygrad/llm/model.py | 166 +++++++++++------------------- tinygrad/llm/serve.py | 3 +- tinygrad/schedule/rangeify.py | 8 +- 8 files changed, 157 insertions(+), 250 deletions(-) diff --git a/extra/benchmark_qwen_amd.py b/extra/benchmark_qwen_amd.py index 01cae37c9b..c7c070de55 100644 --- a/extra/benchmark_qwen_amd.py +++ b/extra/benchmark_qwen_amd.py @@ -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 diff --git a/test/unit/test_attention.py b/test/unit/test_attention.py index a73f15e008..2731f3277a 100644 --- a/test/unit/test_attention.py +++ b/test/unit/test_attention.py @@ -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]: diff --git a/test/unit/test_function.py b/test/unit/test_function.py index fa007124dd..3db67a9fa1 100644 --- a/test/unit/test_function.py +++ b/test/unit/test_function.py @@ -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]) diff --git a/tinygrad/llm/cli.py b/tinygrad/llm/cli.py index 5b62dd5a61..4dd352e6a4 100644 --- a/tinygrad/llm/cli.py +++ b/tinygrad/llm/cli.py @@ -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 diff --git a/tinygrad/llm/kernels/amd.py b/tinygrad/llm/kernels/amd.py index 53eeea1bc4..7cf0e8587d 100644 --- a/tinygrad/llm/kernels/amd.py +++ b/tinygrad/llm/kernels/amd.py @@ -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: diff --git a/tinygrad/llm/model.py b/tinygrad/llm/model.py index 6978eb6dd8..7faa66ea74 100644 --- a/tinygrad/llm/model.py +++ b/tinygrad/llm/model.py @@ -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 diff --git a/tinygrad/llm/serve.py b/tinygrad/llm/serve.py index bc1f2fe41d..34e414548f 100644 --- a/tinygrad/llm/serve.py +++ b/tinygrad/llm/serve.py @@ -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) diff --git a/tinygrad/schedule/rangeify.py b/tinygrad/schedule/rangeify.py index e74c4a888b..64c5805c7c 100644 --- a/tinygrad/schedule/rangeify.py +++ b/tinygrad/schedule/rangeify.py @@ -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),