sirhcm/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests 2 Actions 8 Activity

remove benchmark_qwen_amd, use extra/benchmark_llm.py

Browse Source
This commit is contained in:
geohot
2026-08-05 16:36:34 -07:00
parent db7e683d91
commit b09e879e95
2 changed files with 1 additions and 54 deletions
Download Patch File Download Diff File Expand all files Collapse all files
extra/benchmark_llm.py
+1
View File
@@ -1,3 +1,4 @@
# DEV=USB+AMD python extra/benchmark_llm.py --model /raid/models/Qwen3.6-27B-IQ4_XS.gguf --max-context 131072 --prompt-tokens 3072
import argparse, time
from tinygrad.llm.model import Transformer
extra/benchmark_qwen_amd.py
-54
View File
@@ -1,54 +0,0 @@
import argparse, time
from tinygrad import Context
from tinygrad.llm.model import Transformer
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--model", default="/raid/models/Qwen3.6-27B-IQ4_XS.gguf")
parser.add_argument("--max-context", type=int, default=131072)
parser.add_argument("--prompt-tokens", type=int, default=3072)
parser.add_argument("--decode-tokens", type=int, default=16)
parser.add_argument("--chunk-size", type=int, default=32)
parser.add_argument("--skip-resume-check", action="store_true")
args = parser.parse_args()
startup_st = 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()
with Context(BEAM=0): model.warmup(args.chunk_size)
print(f"warm {time.perf_counter()-st:.3f}s", flush=True)
assert time.perf_counter()-startup_st < 120
states = [getattr(block, name) for block in model.blk for name in ("cache_kv", "cache_kv_scale", "conv_state", "recurrent_state")
if hasattr(block, name)]
assert all(str(state.device).startswith("AMD") and state.uop.is_realized for state in states)
assert all(block.cache_kv.shape[3] >= args.max_context for block in model.blk if hasattr(block, "cache_kv"))
assert model.prefill_jit.cnt >= 2 and model.rollout_jit.cnt >= 2
print(f"preallocated {sum(state.nbytes() for state in states)/2**30:.3f} GiB state on AMD", flush=True)
prompt = [257] + [1000+i%1000 for i in range(args.prompt_tokens-1)]
gen, st = model.generate(prompt, chunk_size=args.chunk_size), time.perf_counter()
output = [next(gen)]
pt = time.perf_counter()
prefill = args.prompt_tokens/(pt-st)
print(f"prefill {prefill:.3f} tok/s", flush=True)
for _ in range(args.decode_tokens): output.append(next(gen))
et = time.perf_counter()
decode = args.decode_tokens/(et-pt)
print(f"decode {decode:.3f} tok/s output {output}", flush=True)
# prefill is sequential-chunk-latency bound: 412/515/759/808 tok/s at chunk 32/64/128/256, decode is flat ~41
if args.prompt_tokens % args.chunk_size == 0: assert prefill > 400 and decode > 40
if args.prompt_tokens == 3000 and args.decode_tokens == 16:
assert output == [13, 271, 248068, 198, 8160, 579, 264, 7047, 1817, 25, 271, 16, 13, 220, 2972, 2014, 53983]
if not args.skip_resume_check:
follow = model._cached_tokens + [1234+i for i in range(8)]
full_prompt = list(follow)
resume_pos, gen, st = model.get_start_pos(follow), model.generate(follow, chunk_size=args.chunk_size), time.perf_counter()
resumed_token = next(gen)
print(f"resume {len(follow)-1} tokens from {resume_pos} in {time.perf_counter()-st:.3f}s token {resumed_token}", flush=True)
model._cached_tokens = [-1]
st = time.perf_counter()
full_token = next(model.generate(full_prompt, chunk_size=args.chunk_size))
print(f"full {time.perf_counter()-st:.3f}s token {full_token} match {resumed_token == full_token}", flush=True)
assert resumed_token == full_token