llm: parallelize long decode reduction

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))

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]

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()