From 8a5fdf1e67dcbad4698fd587410fd582fda6481e Mon Sep 17 00:00:00 2001 From: George Hotz Date: Fri, 17 Jul 2026 18:46:34 +0000 Subject: [PATCH] add flash attention to llm (codex slop) --- extra/benchmark_llm.py | 104 ++++++++++++++++++++++++++++++ extra/gemm/amd_flash_attention.py | 101 +++++++++++++++++++++-------- tinygrad/llm/model.py | 79 ++++++++++++++++------- 3 files changed, 234 insertions(+), 50 deletions(-) create mode 100644 extra/benchmark_llm.py diff --git a/extra/benchmark_llm.py b/extra/benchmark_llm.py new file mode 100644 index 0000000000..dc18d1334a --- /dev/null +++ b/extra/benchmark_llm.py @@ -0,0 +1,104 @@ +#!/usr/bin/env python3 +"""Benchmark tinygrad LLM prefill and decode independently. + +Examples: + python -m extra.benchmark_llm --model qwen3:0.6b --max-context 32768 + python -m extra.benchmark_llm --model /path/to/model.gguf --prompt-tokens 8192 --realize +""" +from __future__ import annotations + +import argparse, json, statistics, time +from dataclasses import asdict, dataclass + +from tinygrad.helpers import fetch +from tinygrad.llm.cli import models +from tinygrad.llm.model import Transformer + + +@dataclass +class Result: + prompt_tokens: int + decode_tokens: int + time_to_first_token_s: float + prefill_tokens_per_s: float + decode_jit_s: float + decode_tokens_per_s: float + decode_p50_ms: float + decode_p95_ms: float + + +def percentile(values:list[float], percentile:float) -> float: + ordered = sorted(values) + return ordered[round((len(ordered) - 1) * percentile)] + + +def synthetic_prompt(length:int, vocab_size:int, salt:int) -> list[int]: + # Avoid tokenizer and chat-template work while exercising the same embedding/model path. + # Changing token zero guarantees that Transformer.get_start_pos cannot reuse an earlier KV cache. + assert length > 0 and vocab_size > 256 + return [256 + salt % (vocab_size - 256)] + [256 + (i * 7919) % (vocab_size - 256) for i in range(1, length)] + + +def benchmark(model:Transformer, prompt:list[int], decode_tokens:int, chunk_size:int) -> Result: + gen = model.generate(prompt.copy(), chunk_size=chunk_size) + begin = time.perf_counter() + next(gen) + ttft = time.perf_counter() - begin + + # A bucketed rollout JIT captures on its first two calls at this context size. + begin = time.perf_counter() + next(gen) + next(gen) + decode_jit = time.perf_counter() - begin + decode_times: list[float] = [] + for _ in range(decode_tokens): + begin = time.perf_counter() + next(gen) + decode_times.append(time.perf_counter() - begin) + + return Result(len(prompt), decode_tokens, ttft, len(prompt) / ttft, decode_jit, decode_tokens / sum(decode_times), + statistics.median(decode_times) * 1e3, percentile(decode_times, 0.95) * 1e3) + + +def main() -> None: + parser = argparse.ArgumentParser(description="Measure LLM prefill and steady-state decode speed") + parser.add_argument("--model", default="qwen3:0.6b", help="Model preset or local GGUF path") + parser.add_argument("--max-context", type=int, default=32768) + parser.add_argument("--prompt-tokens", type=int, nargs="+", default=[128, 2048, 8192]) + parser.add_argument("--decode-tokens", type=int, default=32) + parser.add_argument("--chunk-size", type=int, default=256) + parser.add_argument("--realize", action="store_true", help="Unpack model weights once at load time") + parser.add_argument("--json", action="store_true", help="Print machine-readable results") + args = parser.parse_args() + + if args.decode_tokens < 1: parser.error("--decode-tokens must be positive") + if args.chunk_size < 1: parser.error("--chunk-size must be positive") + if max(args.prompt_tokens) + args.decode_tokens >= args.max_context: + parser.error("prompt plus decode tokens must fit within --max-context") + + path = fetch(models.get(args.model, args.model)) + model, kv = Transformer.from_gguf(path, args.max_context, realize=args.realize) + vocab_size = len(kv["tokenizer.ggml.tokens"]) + + # TinyJit captures on its first two calls. Prime prefill and rollout without leaving a reusable prefix. + warm_length = min(max(args.chunk_size * 2, 2), args.max_context - 4) + for salt in (0, 10_000): + warmup = model.generate(synthetic_prompt(warm_length, vocab_size, salt), chunk_size=args.chunk_size) + for _ in range(4): next(warmup) + + results = [benchmark(model, synthetic_prompt(n, vocab_size, salt=i+1), args.decode_tokens, args.chunk_size) + for i, n in enumerate(args.prompt_tokens)] + if args.json: + print(json.dumps({"model": args.model, "max_context": args.max_context, "chunk_size": args.chunk_size, + "realize": args.realize, "results": [asdict(x) for x in results]}, indent=2)) + return + + print(f"model={args.model} max_context={args.max_context} chunk_size={args.chunk_size} realize={args.realize}") + print(f"{'prompt':>8} {'TTFT':>10} {'prefill':>14} {'decode JIT':>12} {'decode':>14} {'decode p50':>12} {'decode p95':>12}") + for result in results: + print(f"{result.prompt_tokens:8d} {result.time_to_first_token_s:9.3f}s {result.prefill_tokens_per_s:11.1f} t/s " + f"{result.decode_jit_s:10.3f}s " + f"{result.decode_tokens_per_s:11.1f} t/s {result.decode_p50_ms:9.2f} ms {result.decode_p95_ms:9.2f} ms") + + +if __name__ == "__main__": main() diff --git a/extra/gemm/amd_flash_attention.py b/extra/gemm/amd_flash_attention.py index 73740d8165..cbbd4f92a9 100644 --- a/extra/gemm/amd_flash_attention.py +++ b/extra/gemm/amd_flash_attention.py @@ -1,7 +1,7 @@ from tinygrad import Tensor, UOp, getenv from tinygrad.uop.ops import AxisType, KernelInfo, Ops from tinygrad.dtype import AddrSpace, dtypes -from tinygrad.helpers import DEBUG, GlobalCounters, Context +from tinygrad.helpers import GlobalCounters, Context import math BLOCK_M, BLOCK_N = 64, 64 @@ -35,11 +35,18 @@ def warp_reduce_sum(val, lane): val = val + warp_shfl_xor(val, offset, lane) return val -def amd_flash_attention(o:UOp, q:UOp, k:UOp, v:UOp) -> UOp: - # inputs are (B*H, N, D) - BH, N, D = q.shape - assert N % BLOCK_M == 0 and N % BLOCK_N == 0, f"N={N} must be divisible by BLOCK_M={BLOCK_M} and BLOCK_N={BLOCK_N}" - assert D % WMMA_K == 0 and D % LANES_PER_WAVE_N == 0, f"D={D} must be divisible by WMMA_K={WMMA_K} and LANES_PER_WAVE_N={LANES_PER_WAVE_N}" +def _amd_flash_attention(o:UOp, q:UOp, k:UOp, v:UOp, causal:bool, valid_kv_len:int|UOp|None=None) -> UOp: + # inputs are q=(B*H, M, D), k/v=(B*H, N, D). For causal attention q is the final M tokens of k/v. + BH, M, D = q.shape + physical_n = k.shape[1] + N = physical_n if valid_kv_len is None else valid_kv_len + 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}" + assert isinstance(D, int) and D % WMMA_K == 0 and D % LANES_PER_WAVE_N == 0, \ + f"D={D} must be divisible by WMMA_K={WMMA_K} and LANES_PER_WAVE_N={LANES_PER_WAVE_N}" assert BLOCK_M % (WAVES_M * WMMA_M) == 0 and BLOCK_N % LANES_PER_WAVE_N == 0 TM = BLOCK_M // (WAVES_M * LANES_PER_WAVE_M) TN = BLOCK_N // (WAVES_N * LANES_PER_WAVE_N) @@ -47,12 +54,11 @@ def amd_flash_attention(o:UOp, q:UOp, k:UOp, v:UOp) -> UOp: SCALE = 1.0 / math.sqrt(D) block_bh = UOp.range(BH, 0, AxisType.GLOBAL) - block_m = UOp.range(N // BLOCK_M, 1, AxisType.GLOBAL) + block_m = UOp.range(M // BLOCK_M, 1, AxisType.GLOBAL) - q = q.reshape(BH, N//BLOCK_M, BLOCK_M, D)[block_bh, block_m] - k = k.reshape(BH, N//BLOCK_N, BLOCK_N, D)[block_bh] - v = v.reshape(BH, N//BLOCK_N, BLOCK_N, D)[block_bh] - o = o.reshape(BH, N//BLOCK_M, BLOCK_M, D)[block_bh, block_m] + q = q.reshape(BH, M//BLOCK_M, BLOCK_M, D)[block_bh, block_m] + k, v = k[block_bh // gqa_group], v[block_bh // gqa_group] + o = o.reshape(BH, M//BLOCK_M, BLOCK_M, D)[block_bh, block_m] wave_m = UOp.range(WAVES_M, 2, AxisType.LOCAL) wave_n = UOp.range(WAVES_N, 3, AxisType.LOCAL) @@ -76,14 +82,18 @@ def amd_flash_attention(o:UOp, q:UOp, k:UOp, v:UOp) -> UOp: l_i = l_i.after(l_i.store(l_i.const_like(0))) # ====== KV tile loop ====== - n_tile = UOp.range(N // BLOCK_N, 100, AxisType.REDUCE) + # Causal blocks never need KV tiles strictly to their right. Besides saving work, this avoids an all + # -inf tile, whose online-softmax update would otherwise contain -inf - -inf. + n_tiles = (N - M + (block_m + 1) * BLOCK_M + BLOCK_N - 1) // BLOCK_N if causal else N // BLOCK_N + n_tile = UOp.range(n_tiles, 100, AxisType.REDUCE) # load Q + K into LDS (Q reloaded each iteration since P overwrites slot 0) Q_lds = QP_lds[:, :D] Q_store = Q_lds.after(n_tile).reshape(THREADS_PER_BLOCK, ELEMS_PER_THREAD)[tid].store( q.reshape(THREADS_PER_BLOCK, ELEMS_PER_THREAD)[tid]) - K_store = KV_lds.reshape(THREADS_PER_BLOCK, ELEMS_PER_THREAD)[tid].store( - k[n_tile].reshape(THREADS_PER_BLOCK, ELEMS_PER_THREAD)[tid]) + load_k = UOp.range(ELEMS_PER_THREAD, 90, AxisType.LOOP) + K_store = KV_lds.reshape(THREADS_PER_BLOCK, ELEMS_PER_THREAD)[tid, load_k].store( + k.reshape(physical_n*D)[n_tile*BLOCK_N*D + tid*ELEMS_PER_THREAD + load_k]).end(load_k) qk_load_barrier = UOp.barrier(UOp.group(Q_store, K_store)) Q_lds = Q_lds.after(qk_load_barrier) KV_lds_k = KV_lds.after(qk_load_barrier) @@ -104,6 +114,16 @@ def amd_flash_attention(o:UOp, q:UOp, k:UOp, v:UOp) -> UOp: # -- softmax in registers with warp shuffles -- S_reg = S_reg.after(S_reg.store(S_reg * SCALE)) + if causal: + # WMMA accumulator ownership: each lane owns an 8x4 fragment of the 64x64 score tile. + # q is aligned to the right of k, matching PyTorch's causal_lower_right mask. + rm = UOp.range(TM, 250, AxisType.LOOP) + rn = UOp.range(TN, 251, AxisType.LOOP) + q_idx = N - 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 + masked = (k_idx <= q_idx).where(S_reg[rm, rn], S_reg[rm, rn].const_like(-math.inf)) + S_reg = S_reg.after(S_reg[rm, rn].store(masked).end(rm, rn)) + # per-thread local row max over TN=4 elements, then warp reduce across 16 lanes m_ij = UOp.placeholder((TM,), dtypes.float, slot=7, addrspace=AddrSpace.REG) m_ij = m_ij.after(m_ij.after(n_tile).store(m_ij.const_like(-math.inf))) @@ -125,11 +145,12 @@ def amd_flash_attention(o:UOp, q:UOp, k:UOp, v:UOp) -> UOp: # write P = exp(S - m_ij) to P_lds (reuses slot 0, Q no longer needed) P_lds = QP_lds[:, :BLOCK_N] - P_write = P_lds.reshape(WAVES_M, TM // WMMA_ACC, WMMA_ACC, LANES_PER_WAVE_M, WAVES_N, TN, LANES_PER_WAVE_N) - P_write = P_write.permute((0, 4, 3, 6, 1, 2, 5)).reshape(THREADS_PER_BLOCK, TM, TN) + P_write = P_lds.reshape(WAVES_M, TM, LANES_PER_WAVE_M, 1, WAVES_N, TN, LANES_PER_WAVE_N, 1) + P_write = P_write.permute((0, 4, 2, 6, 1, 3, 5, 7)).reshape(THREADS_PER_BLOCK, TM, TN) P_store = P_write[tid].store(S_reg.cast(dtypes.half)) # -- online softmax correction -- + beta_i = UOp.placeholder((TM,), dtypes.float, slot=9, addrspace=AddrSpace.REG) ri4 = UOp.range(TM, 330, AxisType.LOOP) m_new_val = m_i[ri4].maximum(m_ij[ri4]) alpha_val = ((m_i[ri4] - m_new_val) * LOG2E).exp2() @@ -139,29 +160,43 @@ def amd_flash_attention(o:UOp, q:UOp, k:UOp, v:UOp) -> UOp: acc[ri4, rj4].store(alpha_val * acc[ri4, rj4]).end(rj4), l_i[ri4].store(alpha_val * l_i[ri4] + beta_val * p_sum[ri4]), m_i[ri4].store(m_new_val), + beta_i[ri4].store(beta_val), ).end(ri4) acc = acc.after(correction) l_i = l_i.after(correction) m_i = m_i.after(correction) + beta_i = beta_i.after(correction) - # load V into KV_lds (must wait for QK WMMA to finish reading K from KV_lds) - V_store = KV_lds.after(qk_done).reshape(THREADS_PER_BLOCK, ELEMS_PER_THREAD)[tid].store( - v[n_tile].reshape(THREADS_PER_BLOCK, ELEMS_PER_THREAD)[tid]) + # Load V transposed into LDS: PV's B operand is logically (D, BLOCK_N), while global V is (BLOCK_N, D). + # It reuses K's slot and must wait for QK WMMA to finish reading that slot. + V_lds = UOp.placeholder((D, BLOCK_N + LDS_PAD), dtypes.half, slot=1, addrspace=AddrSpace.LOCAL)[:, :BLOCK_N] + V_copy = V_lds.after(qk_done).permute(1, 0) + load_v = UOp.range(ELEMS_PER_THREAD, 390, AxisType.LOOP) + V_store = V_copy.reshape(THREADS_PER_BLOCK, ELEMS_PER_THREAD)[tid, load_v].store( + v.reshape(physical_n*D)[n_tile*BLOCK_N*D + tid*ELEMS_PER_THREAD + load_v]).end(load_v) pv_barrier = UOp.barrier(UOp.group(P_store, V_store)) P_lds = P_lds.after(pv_barrier) - KV_lds_v = KV_lds.after(pv_barrier) + V_lds = V_lds.after(pv_barrier) - # -- acc += P @ V via WMMA -- + # -- acc += beta * (P @ V) via WMMA -- + pv_acc = UOp.placeholder((TM, TD), dtypes.float, slot=10, addrspace=AddrSpace.REG) + pv_acc = pv_acc.after(pv_acc.after(n_tile).store(pv_acc.const_like(0))).after(pv_barrier) k_pv = UOp.range(BLOCK_N // WMMA_K, 400, AxisType.REDUCE) tm2 = UOp.range(TM // WMMA_ACC, 401, AxisType.LOOP) tn2 = UOp.range(TD, 402, AxisType.LOOP) - acc_frag = acc.reshape(TM // WMMA_ACC, WMMA_ACC, TD).permute(0, 2, 1)[tm2, tn2] + pv_frag = pv_acc.reshape(TM // WMMA_ACC, WMMA_ACC, TD).permute(0, 2, 1)[tm2, tn2] p_frag = P_lds.reshape(WAVES_M, TM // WMMA_ACC, WMMA_M, BLOCK_N // WMMA_K, WMMA_K)[wave_m, tm2, lane_n, k_pv] - v_frag = KV_lds_v.reshape(WAVES_N, TD, WMMA_N, BLOCK_N // WMMA_K, WMMA_K)[wave_n, tn2, lane_n, k_pv] - pv = UOp.wmma(p_frag, v_frag, acc_frag.after(k_pv), *WMMA_ARG) + v_frag = V_lds.reshape(WAVES_N, TD, WMMA_N, BLOCK_N // WMMA_K, WMMA_K)[wave_n, tn2, lane_n, k_pv] + pv = UOp.wmma(p_frag, v_frag, pv_frag.after(k_pv), *WMMA_ARG) + pv_done = pv_frag.store(pv).end(tm2, tn2).end(k_pv) + pv_acc = pv_acc.after(pv_done) + + ri5 = UOp.range(TM, 410, AxisType.LOOP) + rj5 = UOp.range(TD, 411, AxisType.LOOP) + accumulate = acc[ri5, rj5].store(acc[ri5, rj5] + beta_i[ri5] * pv_acc[ri5, rj5]).end(ri5, rj5) # end KV tile loop - n_tile_end = acc_frag.store(pv).end(tm2, tn2).end(k_pv).barrier().end(n_tile) + n_tile_end = accumulate.barrier().end(n_tile) acc = acc.after(n_tile_end) l_i = l_i.after(n_tile_end) m_i = m_i.after(n_tile_end) @@ -170,10 +205,22 @@ def amd_flash_attention(o:UOp, q:UOp, k:UOp, v:UOp) -> UOp: acc = acc.after(acc.store(acc * (1 / l_i).reshape(TM, 1).expand(TM, TD))) # store output - o = o.reshape(WAVES_M, TM // WMMA_ACC, WMMA_ACC, LANES_PER_WAVE_M, WAVES_N, TD, LANES_PER_WAVE_N) - o = o.permute((0, 4, 3, 6, 1, 2, 5)).reshape(THREADS_PER_BLOCK, TM, TD) + o = o.reshape(WAVES_M, TM, LANES_PER_WAVE_M, 1, WAVES_N, TD, LANES_PER_WAVE_N, 1) + 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 amd_flash_attention(o:UOp, q:UOp, k:UOp, v:UOp) -> UOp: + return _amd_flash_attention(o, q, k, v, causal=False) + +def amd_flash_attention_causal(o:UOp, q:UOp, k:UOp, v:UOp) -> UOp: + return _amd_flash_attention(o, q, k, v, causal=True) + +def amd_flash_attention_causal_cached(o:UOp, q:UOp, cache_kv:UOp, *, valid_kv_len:int|UOp) -> UOp: + _, B, H_KV, N, D = cache_kv.shape + k = cache_kv[0].reshape(B*H_KV, N, D) + v = cache_kv[1].reshape(B*H_KV, N, D) + return _amd_flash_attention(o, q, k, v, causal=True, valid_kv_len=valid_kv_len) + if __name__ == "__main__": B, H, N, D = getenv("B", 1), getenv("H", 32), getenv("N", 1024), getenv("D", 64) q = Tensor.rand(B, H, N, D).cast(dtypes.half) diff --git a/tinygrad/llm/model.py b/tinygrad/llm/model.py index dbeeebdc16..a065096fea 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 tinygrad import Tensor, nn, UOp, TinyJit, getenv, function +from tinygrad import Device, Tensor, nn, UOp, TinyJit, getenv, function from tinygrad.llm.gguf import gguf_load from tinygrad.uop.ops import resolve @@ -125,14 +125,14 @@ 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) -> Tensor: raise NotImplementedError + def _attention(self, x:Tensor, start_pos:int|UOp, use_flash:bool=False, kv_len:int|None=None) -> Tensor: raise NotImplementedError - def __call__(self, x: Tensor, start_pos: int|UOp): + def __call__(self, x: Tensor, start_pos: int|UOp, use_flash:bool=False, kv_len:int|None=None): self._init_state(x) # we pass in the weights implicitly so we unpack the GGUF on the fly @function(precompile=True, allow_implicit=True) def _run(x:Tensor, start_pos:int|UOp): - h = x + self._attention(self.attn_norm(x), start_pos) + h = x + self._attention(self.attn_norm(x), start_pos, use_flash, kv_len) return (h + self._feed_forward(self.ffn_norm(h))).contiguous() return _run(x, start_pos) @@ -150,7 +150,7 @@ class TransformerBlock(FFNBlock): self.attn_output = nn.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) -> Tensor: + def _attention(self, x:Tensor, start_pos:int|UOp, use_flash:bool=False, kv_len:int|None=None) -> Tensor: 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) @@ -167,9 +167,11 @@ 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 - assigned_kv = Tensor(self.cache_kv.uop.after(self.cache_kv[:, :, :, start_pos:start_pos+T, :].uop.store(Tensor.stack(k, v).uop))) - k = assigned_kv[0, :, :, 0:start_pos+T, :] - v = assigned_kv[1, :, :, 0:start_pos+T, :] + assigned_kv = Tensor(self.cache_kv.uop.after( + self.cache_kv[:, :, :, start_pos:start_pos+T, :].uop.store(Tensor.stack(k, v).cast(self.cache_kv.dtype).uop))) + cache_len = start_pos + T if kv_len is None else kv_len + k = assigned_kv[0, :, :, 0:cache_len, :] + v = assigned_kv[1, :, :, 0:cache_len, :] #self.cache_kv[:, :, :, start_pos:start_pos+T, :].assign(Tensor.stack(k, v)) #k = self.cache_kv[0, :, :, 0:start_pos+T, :] @@ -177,16 +179,30 @@ 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 - mask = Tensor.full((1, 1, T, start_pos+T), float("-inf"), dtype=x.dtype, buffer=False).triu(start_pos+1) \ - if resolve(T != 1) else None - attn = q.scaled_dot_product_attention(k, v, attn_mask=mask, enable_gqa=True) # (B,H,T,Hd) + if use_flash: + from extra.gemm.amd_flash_attention import amd_flash_attention_causal_cached + flash_start_pos = start_pos.unbind()[0] if isinstance(start_pos, UOp) else start_pos + valid_kv_len = flash_start_pos + T + q_flat = q.half().reshape(B*self.config.n_heads, T, self.config.head_dim) + out = Tensor.empty(B*self.config.n_heads, T, self.config.head_dim, dtype="float32", device=x.device) + attn = Tensor.custom_kernel(out, q_flat, assigned_kv, + fxn=functools.partial(amd_flash_attention_causal_cached, valid_kv_len=valid_kv_len))[0].reshape(B, self.config.n_heads, T, -1) + else: + mask:Tensor|None + if kv_len is not None: + mask = Tensor.full((1, 1, 1, kv_len), float("-inf"), dtype=x.dtype, buffer=False).triu(start_pos+1) + else: + mask = Tensor.full((1, 1, T, start_pos+T), float("-inf"), dtype=x.dtype, buffer=False).triu(start_pos+1) \ + if resolve(T != 1) else None + attn = q.half().scaled_dot_product_attention(k, v, 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())) def _init_state(self, x:Tensor): if not hasattr(self, "cache_kv"): # TODO: how is the dtype of this determined? - self.cache_kv = Tensor.empty(2, x.shape[0], self.config.n_kv_heads, self.config.max_context, self.config.head_dim, device=x.device) + self.cache_kv = Tensor.empty(2, x.shape[0], self.config.n_kv_heads, self.config.max_context, self.config.head_dim, + dtype="float16", device=x.device) self.freqs_cis = precompute_freqs_cis(self.config.rope_dim, self.config.max_context, self.config.rope_theta, device=x.device) class MLATransformerBlock(FFNBlock): @@ -205,7 +221,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 = nn.Linear(config.n_heads * config.v_head_dim, config.dim, bias=False) - def _attention(self, x:Tensor, start_pos:int|UOp) -> Tensor: + def _attention(self, x:Tensor, start_pos:int|UOp, use_flash:bool=False, kv_len:int|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) @@ -250,7 +266,7 @@ class GatedDeltaNetBlock(FFNBlock): self.ssm_a = Tensor.zeros(self.num_v_heads) self.ssm_norm, self.ssm_out = nn.RMSNorm(self.head_v_dim, config.norm_eps), nn.Linear(ssm.inner_size, config.dim, bias=False) - def _attention(self, x:Tensor, start_pos:int|UOp) -> Tensor: + def _attention(self, x:Tensor, start_pos:int|UOp, use_flash:bool=False, kv_len:int|None=None) -> Tensor: B, T, _ = x.shape assert T == 1, "GatedDeltaNetBlock currently only supports T=1" @@ -308,17 +324,26 @@ class Transformer: self._cached_tokens: list[int] = [] # 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.rollout_jits:dict[int, TinyJit] = {} - def forward(self, tokens:Tensor, start_pos:int|UOp, temperature:Tensor) -> Tensor: + def forward(self, tokens:Tensor, start_pos:int|UOp, temperature:Tensor, use_flash:bool=False, kv_len:int|None=None) -> Tensor: x = self.token_embd(tokens).float() # (B, T, D) - for block in self.blk: x = block(x, start_pos) + for block in self.blk: x = block(x, start_pos, use_flash, kv_len) logits = self.output(self.output_norm(x))[:, -1, :] # Gumbel-max trick: argmax(logits/temp - log(-log(uniform))) is equivalent to sampling from softmax(logits/temp) return (logits / temperature.maximum(1e-12) - (Tensor.rand_like(logits).maximum(1e-12).log().neg()).log()).argmax(-1, keepdim=True) - def __call__(self, tokens:Tensor, start_pos:int|UOp, temperature:Tensor) -> Tensor: - return (self.prefill_jit if resolve(tokens.shape[1] != 1) else self.rollout_jit)(tokens.contiguous(), start_pos, temperature) + def __call__(self, tokens:Tensor, start_pos:int|UOp, temperature:Tensor, use_flash:bool=False) -> Tensor: + if resolve(tokens.shape[1] == 1): + pos = start_pos.unbind()[1] if isinstance(start_pos, UOp) else start_pos + min_bucket = max(1, getenv("DECODE_BUCKET", 256)) + kv_len = min(self.max_context, max(min_bucket, 1 << pos.bit_length())) + if kv_len not in self.rollout_jits: self.rollout_jits[kv_len] = TinyJit(functools.partial(self.forward, kv_len=kv_len)) + jit = self.rollout_jits[kv_len] + else: + jit = self.flash_prefill_jit if use_flash else self.prefill_jit + return jit(tokens.contiguous(), start_pos, temperature) @staticmethod def from_gguf(gguf:Tensor|str|pathlib.Path, max_context:int|None=None, @@ -393,7 +418,7 @@ class Transformer: prefix_len = sum(1 for _ in itertools.takewhile(lambda ab: ab[0] == ab[1], zip(tokens[:-1], self._cached_tokens))) return min(block._reusable_prefix_len(prefix_len, len(self._cached_tokens)) for block in self.blk) - def generate(self, tokens:list[int], chunk_size:int=32, temperature:float=0.0): + def generate(self, tokens:list[int], chunk_size:int=256, temperature:float=0.0): if self.has_recurrent_block: chunk_size = 1 v_start_pos = UOp.variable("start_pos", 0, self.max_context-1) v_toks = UOp.variable("toks", 1, chunk_size) @@ -406,9 +431,17 @@ class Transformer: 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) out, prompt_len = None, len(tokens) while len(tokens) < self.max_context: - sp, nt = v_start_pos.bind(start_pos), v_toks.bind(min(chunk_size, len(tokens) - start_pos)) - out = self(t[:, sp:sp+nt] if start_pos < prompt_len or out is None else out, sp, temp).realize() - start_pos += nt.val + remaining = len(tokens) - start_pos + use_flash = bool(getenv("AMD_FLASH_ATTENTION", 1)) and start_pos > 0 and remaining >= chunk_size and chunk_size % 64 == 0 and \ + not self.has_recurrent_block + if use_flash: + device = str(getattr(self.blk[0], "cache_kv").device) + use_flash = device.startswith("AMD") and Device[device].renderer.target.arch.startswith("gfx11") + sp = v_start_pos.bind(start_pos) + nt = chunk_size if use_flash else v_toks.bind(min(chunk_size, remaining)) + inp = t[:, sp:sp+nt] if start_pos < prompt_len or out is None else out + out = (self(inp, sp, temp, use_flash=True) if use_flash else self(inp, sp, temp)).realize() + start_pos += nt if isinstance(nt, int) else nt.val # chunked prefill: keep processing until all prompt tokens are consumed if start_pos < len(tokens): continue tokens.append(int(out.item()))