llm kernels: adapt to Ops.BIND removal

Variables are 0-d ALU BUFFERs in the tensor graph and take the ALU PARAM form
inside kernels (UOp.variable(param=True)). Add kernel_var helper for the
conversion, and keep start_pos in bound form at the graph level so function
implicit-input collection and the schedule's binds rename-back line up.

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

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)