From e69df9e8b6eb3821f8d5809da2bdfd877cd31448 Mon Sep 17 00:00:00 2001 From: George Hotz Date: Tue, 21 Jul 2026 15:59:46 +0000 Subject: [PATCH] llm: parallelize long decode reduction --- extra/benchmark_llm.py | 24 ++++++++++++-- extra/gemm/amd_flash_attention.py | 40 ++++++++++++++++++++++++ test/backend/test_amd_flash_attention.py | 12 ++++--- 3 files changed, 69 insertions(+), 7 deletions(-) diff --git a/extra/benchmark_llm.py b/extra/benchmark_llm.py index ea451a07a4..260fc96b3b 100644 --- a/extra/benchmark_llm.py +++ b/extra/benchmark_llm.py @@ -10,6 +10,7 @@ from __future__ import annotations import argparse, gc, json, statistics, time from dataclasses import asdict, dataclass +from tinygrad import Device, Tensor, Variable from tinygrad.helpers import fetch from tinygrad.llm.cli import models from tinygrad.llm.model import Transformer @@ -54,12 +55,26 @@ def benchmark(model:Transformer, prompt:list[int], decode_tokens:int, chunk_size statistics.median(decode_times) * 1e3, percentile(decode_times, 0.95) * 1e3) +def benchmark_decode_position(model:Transformer, position:int, decode_tokens:int) -> Result: + # Shape-only tail benchmark: warmed recurrent state and zero KV entries isolate decode cost without a full O(context) prefill. + token = Tensor([[0]], dtype="int32", device=Device.DEFAULT).realize() + temperature = Tensor([0.0], device=Device.DEFAULT).realize() + decode_times = [] + for pos in range(position, position + decode_tokens): + begin = time.perf_counter() + model(token, Variable("start_pos", 0, model.max_context-1).bind(pos), temperature).realize().item() + decode_times.append(time.perf_counter() - begin) + return Result(position, decode_tokens, 0.0, 0.0, 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("--decode-position", type=int, help="Benchmark decode at this position without prefilling preceding KV entries") 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") @@ -67,8 +82,10 @@ def main() -> None: 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: + if args.decode_position is None and max(args.prompt_tokens) + args.decode_tokens >= args.max_context: parser.error("prompt plus decode tokens must fit within --max-context") + if args.decode_position is not None and args.decode_position + args.decode_tokens >= args.max_context: + parser.error("decode position 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) @@ -77,8 +94,9 @@ def main() -> None: model.warmup(args.chunk_size) gc.freeze() - 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)] + results = [benchmark_decode_position(model, args.decode_position, args.decode_tokens)] if args.decode_position is not None else \ + [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)) diff --git a/extra/gemm/amd_flash_attention.py b/extra/gemm/amd_flash_attention.py index 2829acb6c2..e4f736318b 100644 --- a/extra/gemm/amd_flash_attention.py +++ b/extra/gemm/amd_flash_attention.py @@ -134,6 +134,38 @@ def _amd_flash_attention_decode_reduce(out:UOp, partial:UOp, stats:UOp, valid_ch stores = [out[b, head, 0, d].store(numerator[i]/denominator[0]) for i,d in enumerate(dims)] return UOp.group(*stores).end(lane, block_bh).sink(arg=KernelInfo(name="flash_decode_reduce", opts_to_apply=())) +@functools.cache +def _amd_flash_attention_decode_reduce_partial(out:UOp, out_stats:UOp, partial:UOp, stats:UOp, + valid_chunks:int|UOp, group_size:int) -> UOp: + B, H, _, D = partial.shape + DV = D // WARP_SIZE + block = UOp.range(B*H*((valid_chunks+group_size-1)//group_size), 0, AxisType.GLOBAL) + lane = UOp.range(WARP_SIZE, 1, AxisType.LOCAL) + group, bh = block % ((valid_chunks+group_size-1)//group_size), block // ((valid_chunks+group_size-1)//group_size) + b, head = bh // H, bh % H + dims = tuple(lane + i*WARP_SIZE for i in range(DV)) + start, count = group*group_size, (valid_chunks-group*group_size).minimum(group_size) + + row_max = UOp.placeholder((1,), dtypes.float, slot=0, addrspace=AddrSpace.REG) + row_max = row_max.after(row_max.store(row_max.const_like(-math.inf))) + max_chunk = UOp.range(count, 100, AxisType.REDUCE) + max_done = row_max.store(row_max.after(max_chunk).maximum(stats[b, head, start+max_chunk, 0])).end(max_chunk) + row_max = row_max.after(max_done) + + numerator = UOp.placeholder((DV,), dtypes.float, slot=1, addrspace=AddrSpace.REG) + denominator = UOp.placeholder((1,), dtypes.float, slot=2, addrspace=AddrSpace.REG) + init = UOp.group(numerator.store(numerator.const_like(0)), denominator.store(denominator.const_like(0))) + numerator, denominator = numerator.after(init), denominator.after(init) + chunk = UOp.range(count, 101, AxisType.REDUCE) + scale = ((stats[b, head, start+chunk, 0]-row_max[0])*LOG2E).exp2() + update = UOp.group(numerator.store(numerator.after(chunk) + UOp.stack(*(partial[b, head, start+chunk, d] for d in dims))*scale), + denominator.store(denominator.after(chunk) + stats[b, head, start+chunk, 1]*scale)).end(chunk) + numerator, denominator = numerator.after(update), denominator.after(update) + stores = [out[b, head, group, d].store(numerator[i]) for i,d in enumerate(dims)] + \ + [out_stats[b, head.valid(lane.eq(0)), group, 0].store(row_max[0]), + out_stats[b, head.valid(lane.eq(0)), group, 1].store(denominator[0])] + return UOp.group(*stores).end(lane, block).sink(arg=KernelInfo(name="flash_decode_reduce_partial", opts_to_apply=())) + def amd_flash_attention_decode(q:Tensor, cache_kv:Tensor, valid_kv_len:int|UOp, max_kv_len:int|None=None) -> Tensor: _, B, H_KV, N, D = cache_kv.shape _, H, M, _ = q.shape @@ -146,6 +178,14 @@ def amd_flash_attention_decode(q:Tensor, cache_kv:Tensor, valid_kv_len:int|UOp, partial, stats = Tensor.custom_kernel(partial, stats, q, cache_kv, fxn=functools.partial(_amd_flash_attention_decode_partial, valid_kv_len=valid_kv_len, max_kv_len=max_kv_len, block_n=block_n))[:2] live_chunks = (valid_kv_len+block_n-1)//block_n + if max_kv_len > 8192: + reduce_group = 8 + reduced_chunks = (chunks+reduce_group-1)//reduce_group + reduced = Tensor.empty(B, H, reduced_chunks, D, dtype="float32", device=q.device) + reduced_stats = Tensor.empty(B, H, reduced_chunks, 2, dtype="float32", device=q.device) + partial, stats = Tensor.custom_kernel(reduced, reduced_stats, partial, stats, + fxn=functools.partial(_amd_flash_attention_decode_reduce_partial, valid_chunks=live_chunks, group_size=reduce_group))[:2] + live_chunks = (live_chunks+reduce_group-1)//reduce_group out = Tensor.empty(B, H, 1, D, dtype="float32", device=q.device) return Tensor.custom_kernel(out, partial, stats, fxn=functools.partial(_amd_flash_attention_decode_reduce, valid_chunks=live_chunks))[0] diff --git a/test/backend/test_amd_flash_attention.py b/test/backend/test_amd_flash_attention.py index 6ff720b077..df19de8009 100644 --- a/test/backend/test_amd_flash_attention.py +++ b/test/backend/test_amd_flash_attention.py @@ -7,20 +7,20 @@ from extra.gemm.amd_flash_attention import amd_flash_attention_decode @unittest.skipUnless(Device.DEFAULT.startswith("AMD"), "AMD flash attention required") class TestAMDFlashAttention(unittest.TestCase): - def test_short_decode_is_finite_and_matches_reference(self): + def _test_decode(self, max_kv_len:int, valid_kv_len:int): rng = np.random.default_rng(1) q_np = rng.standard_normal((1, 16, 1, 256)).astype(np.float16) - kv_np = rng.standard_normal((2, 1, 2, 8192, 256)).astype(np.float16) + kv_np = rng.standard_normal((2, 1, 2, max_kv_len, 256)).astype(np.float16) q, kv = Tensor(q_np).realize(), Tensor(kv_np).realize() @TinyJit - def decode(q:Tensor, kv:Tensor): return amd_flash_attention_decode(q, kv, 25, 8192).realize() + def decode(q:Tensor, kv:Tensor): return amd_flash_attention_decode(q, kv, valid_kv_len, max_kv_len).realize() out = None for _ in range(3): out = decode(q, kv).numpy() assert out is not None q_ref = q_np[0, :, 0].astype(np.float32) - k_ref, v_ref = kv_np[:, 0, :, :25].astype(np.float32) + k_ref, v_ref = kv_np[:, 0, :, :valid_kv_len].astype(np.float32) expected = np.empty((16, 256), dtype=np.float32) for head in range(16): scores = q_ref[head] @ k_ref[head // 8].T / np.sqrt(256) @@ -30,5 +30,9 @@ class TestAMDFlashAttention(unittest.TestCase): self.assertTrue(np.isfinite(out).all()) np.testing.assert_allclose(out[0, :, 0], expected, rtol=2e-3, atol=2e-3) + def test_short_decode_is_finite_and_matches_reference(self): self._test_decode(8192, 25) + + def test_hierarchical_decode_matches_reference(self): self._test_decode(16384, 4097) + if __name__ == "__main__": unittest.main()