forked from tinygrad/tinygrad
add flash attention to llm (codex slop)
This commit is contained in:
@@ -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()
|
||||
@@ -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)
|
||||
|
||||
+56
-23
@@ -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()))
|
||||
|
||||
Reference in New Issue
Block a user