diff --git a/tinygrad/llm/kernels/__init__.py b/tinygrad/llm/kernels/__init__.py index 974bb0afdf..1470f928a0 100644 --- a/tinygrad/llm/kernels/__init__.py +++ b/tinygrad/llm/kernels/__init__.py @@ -6,6 +6,11 @@ from tinygrad.dtype import AddrSpace from tinygrad.helpers import prod from tinygrad.uop.ops import AxisType, KernelInfo, Ops +def kernel_var(x:UOp) -> UOp: + # a Variable is a 0-d ALU BUFFER in the tensor graph; inside kernels it takes the ALU PARAM form (same name keeps the value binding) + return x.substitute({v: UOp.variable(v.expr, v.vmin, v.vmax, dtype=v.dtype, multiple_of=v.arg.multiple_of, param=True) + for v in x.toposort() if v.is_variable}) + def amd_custom_kernels_supported(device:str|tuple[str, ...]|None) -> bool: # the custom kernels are tuned for RDNA3 (gfx11): the WMMA register layouts don't match gfx12 (RDNA4) # or CDNA (MFMA-only, wave64), and the dp4a builtins and 32-lane wave ops aren't portable either. @@ -83,5 +88,5 @@ def gated_delta_prefill(q:Tensor, k:Tensor, v:Tensor, beta:Tensor, alpha:Tensor, contig = tuple(x.uop if x.uop.op is Ops.AFTER else x.uop.contiguous() for x in srcs) params = tuple(UOp.placeholder_like(x, slot=i) for i,x in enumerate(contig)) assert start_pos.uop.is_bound_var - call = kernel(*params, start_pos.uop.src[0]).call(*contig, start_pos.uop) + call = kernel(*params, kernel_var(start_pos.uop.src[0])).call(*contig, start_pos.uop) return Tensor(contig[0].after(call)) diff --git a/tinygrad/llm/kernels/amd.py b/tinygrad/llm/kernels/amd.py index fd550051af..0686f47e28 100644 --- a/tinygrad/llm/kernels/amd.py +++ b/tinygrad/llm/kernels/amd.py @@ -2,7 +2,7 @@ from __future__ import annotations import functools, math from typing import Callable, cast from tinygrad import Tensor, UOp -from tinygrad.llm.kernels import Linear +from tinygrad.llm.kernels import Linear, kernel_var from tinygrad.uop.ops import AxisType, KernelInfo, Ops, resolve from tinygrad.dtype import AddrSpace, dtypes @@ -27,6 +27,7 @@ def _reg(shape:tuple[int, ...], slot:int, value:float, dep:UOp|None=None) -> UOp @functools.cache def _amd_flash_attention_decode_partial(out, stats, q, cache_kv, cache_scale, valid_kv_len, max_kv_len, block_n): + if isinstance(valid_kv_len, UOp): valid_kv_len = kernel_var(valid_kv_len.unbind_all()[0]) _, 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 @@ -84,6 +85,7 @@ def amd_flash_attention_decode(q:Tensor, cache_kv:Tensor, valid_kv_len:int|UOp, @functools.cache def _amd_flash_attention(o:UOp, q:UOp, cache:UOp, kv_scale:UOp, valid_kv_len:int|UOp) -> UOp: + if isinstance(valid_kv_len, UOp): valid_kv_len = kernel_var(valid_kv_len.unbind_all()[0]) BH, M, D = q.shape _, 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) @@ -188,7 +190,8 @@ def quantized_attention(q:Tensor, stacked_kv:Tensor, cache_kv:Tensor, cache_scal # each store goes on its own buffer's AFTER: sharing both stores across both AFTERs leaves # un-ended stores with open ranges in the kernel graph assigned_kv, assigned_scale = Tensor(cache_kv.uop.after(store_kv)), Tensor(cache_scale.uop.after(store_scale)) - valid_end = start_pos.unbind_all()[0]+T if isinstance(start_pos, UOp) else start_pos+T + # keep start_pos in its bound form at the graph level, the kernel builders unbind it to the kernel-side PARAM form + valid_end = start_pos+T return amd_flash_attention_decode(q.half(), assigned_kv, valid_end, assigned_scale, cast(int, cache_kv.shape[3])) if resolve(T == 1) else \ flash_attention_causal_cached(q.half(), assigned_kv, valid_end, assigned_scale)