Compare commits

...
Author SHA1 Message Date
geohot eec82b0db4 revert wrong custom kernel fix 2026-08-05 18:27:40 -07:00
George HotzandGitHub 470c032a5e fix slice + non contig kernels (#17423)
* movement: resolve negative int slice bounds against symbolic sizes

negative int bounds in a slice against a symbolic dim were passed through
unresolved, giving wrong views. resolve them against the (possibly
symbolic) size, like slice.indices does for int dims

* schedule: realize custom kernel inputs that don't resolve to a buffer state

rangeify assigns ranges backward from consumers and CALL contributes none,
so the subgraph above a custom kernel input gets no ranges unless something
in it is realized, and reduce conversion crashes with a KeyError. realize
call inputs that don't resolve to a buffer state.

only view-only movement ops preserve the underlying buffer: anything
computed (ALU, REDUCE, ...) must be realized even if one of its sources
resolves to a buffer, since the whole subgraph above the call has no
ranges. unwrapping src[0] unconditionally missed const branches hanging
off non-src[0] children and silently resolved REDUCEs to their source
buffer. includes regression tests for pure const, mixed buffer+const, and
view-over-buffer inputs
2026-08-05 16:18:17 -07:00
geohot a8a8030bc9 add benchmark_llm script 2026-08-05 15:59:51 -07:00
4 changed files with 46 additions and 1 deletions
+31
View File
@@ -0,0 +1,31 @@
import argparse, time
from tinygrad.llm.model import Transformer
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--model", required=True, help="path to gguf model")
parser.add_argument("--max-context", type=int, default=8192, help="max context length (default: %(default)s)")
parser.add_argument("--prompt-tokens", type=int, default=1024, help="number of prompt tokens (default: %(default)s)")
parser.add_argument("--decode-tokens", type=int, default=16, help="number of tokens to decode (default: %(default)s)")
parser.add_argument("--chunk-size", type=int, default=32, help="chunk size for prefill (default: %(default)s)")
args = parser.parse_args()
st = time.perf_counter()
model, _ = Transformer.from_gguf(args.model, args.max_context)
print(f"load {time.perf_counter()-st:.3f}s", flush=True)
st = time.perf_counter()
model.warmup()
print(f"warm {time.perf_counter()-st:.3f}s", flush=True)
prompt = [257] + [1000+i%1000 for i in range(args.prompt_tokens-1)]
gen = model.generate(prompt, chunk_size=args.chunk_size)
st = time.perf_counter()
# first token is time-to-first-token; counted as part of prefill
output = [next(gen)]
pt = time.perf_counter()
print(f"prefill {args.prompt_tokens/(pt-st):.3f} tok/s", flush=True)
for _ in range(args.decode_tokens): output.append(next(gen))
et = time.perf_counter()
print(f"decode {args.decode_tokens/(et-pt):.3f} tok/s output {output}", flush=True)
+9
View File
@@ -6,6 +6,15 @@ from examples.gpt2 import Attention
import numpy as np
class TestSymbolicOps(unittest.TestCase):
def test_negative_slice(self):
a = Tensor.rand(3, 10, 4)
for i in range(3, 10):
vi = Variable("i", 1, 10).bind(i)
# negative int bounds against a symbolic dim must resolve against the size, like slice.indices
np.testing.assert_allclose(a[:, :vi][:, -3:-1].numpy(), a[:, :i][:, -3:-1].numpy(), atol=1e-6, rtol=1e-6)
np.testing.assert_allclose(a[:, :vi][:, -1:].numpy(), a[:, :i][:, -1:].numpy(), atol=1e-6, rtol=1e-6)
np.testing.assert_allclose(a[:, :vi][:, -1].numpy(), a[:, :i][:, -1].numpy(), atol=1e-6, rtol=1e-6)
def test_plus1(self):
def f(a): return (a+1).realize()
a = Tensor.rand(3, 10)
+4 -1
View File
@@ -90,8 +90,11 @@ class MovementMixin:
if resolve(index.step == 0, False): raise ValueError(f"{index=} cannot have 0 as step")
start, stop = 0 if index.start is None else index.start, size if index.stop is None else index.stop
step = 1 if index.step is None else index.step
# resolve negative int bounds against the (possibly symbolic) size, like slice.indices
if isinstance(start, int) and start < 0: start = start + size
if isinstance(stop, int) and stop < 0: stop = stop + size
if all_int((start, stop, step)):
# handle int slicing (resolve negative bounds, clamp, stride)
# handle int slicing (clamp, stride)
*bound, stride = index.indices(int(size.vmax) if isinstance(size, UOp) else size)
bound = [0, 0] if stride * (bound[1] - bound[0]) < 0 else ([bound[1]+1, bound[0]+1] if stride < 0 else bound)
return {"size":ceildiv(bound[1]-bound[0], abs(stride)), "boundary":tuple(bound), "stride":stride, "collapse_dim":False}
+2
View File
@@ -25,6 +25,8 @@ def realize_store_after_src(ctx:dict[UOp, None], dest:UOp, src:UOp):
# you don't usually have to do this for assign unless there's a WAR hazard like TestAssign.test_assign_double_diamond_reduce
if dest.base in src.backward_slice_with_self: ctx[src] = None
BUFFER_STATE_OPS: set[Ops] = {Ops.AFTER, Ops.BUFFER, Ops.PARAM, Ops.MSELECT, Ops.MSTACK, Ops.BIND}
pm_generate_realize_map = PatternMatcher([
# always realize
(UPat({Ops.CONTIGUOUS, Ops.STORE}, name="tr"), realize),