mirror of
https://github.com/tinygrad/tinygrad.git
synced 2026-08-25 20:46:07 +00:00
Compare commits
3
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
eec82b0db4 | ||
|
|
470c032a5e | ||
|
|
a8a8030bc9 |
@@ -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)
|
||||
@@ -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)
|
||||
|
||||
@@ -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}
|
||||
|
||||
@@ -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),
|
||||
|
||||
Reference in New Issue
Block a user