From 104ee90ccfefe44944e0525b94133830a42e3f1f Mon Sep 17 00:00:00 2001 From: YassineYousfi Date: Mon, 3 Aug 2026 14:19:05 -0700 Subject: [PATCH 1/3] usb: wait for PCIe link after power on (#17380) --- tinygrad/runtime/support/usb.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/tinygrad/runtime/support/usb.py b/tinygrad/runtime/support/usb.py index b7e8b8186e..7c735c0056 100644 --- a/tinygrad/runtime/support/usb.py +++ b/tinygrad/runtime/support/usb.py @@ -90,10 +90,11 @@ class CustomASM24Controller: def __init__(self, usb:USB3): self.usb = usb - # Custom firmware now boots with PCIe off. Power it on before probing the link. - ltssm = self.read(0xB450, 1)[0] - if ltssm != 0x78: self.set_pcie_power(True) - ltssm = self.read(0xB450, 1)[0] + # Custom firmware now boots with PCIe off. Power it on before probing the link with a 5s grace period. + if (ltssm:=self.read(0xB450, 1)[0]) != 0x78: + self.set_pcie_power(True) + grace_period = time.monotonic() + 5. + while time.monotonic() < grace_period and (ltssm:=self.read(0xB450, 1)[0]) != 0x78: time.sleep(0.1) if ltssm != 0x78: raise RuntimeError(f"PCIe link not up (LTSSM=0x{ltssm:02X}), custom firmware not ready") def set_pcie_power(self, enabled:bool, timeout:int=10000): self.usb.control_write(0xF3, value=int(enabled), timeout=timeout) From 33755a34657d25920914badbe32a9d70489669c7 Mon Sep 17 00:00:00 2001 From: chenyu Date: Mon, 3 Aug 2026 17:24:25 -0400 Subject: [PATCH 2/3] improve threefry codegen [pr] (#17379) decomp uint64 can handle part of it --- tinygrad/codegen/decomp/op.py | 8 ++++---- tinygrad/mixin/rand.py | 2 +- tinygrad/uop/symbolic.py | 12 ++++-------- 3 files changed, 9 insertions(+), 13 deletions(-) diff --git a/tinygrad/codegen/decomp/op.py b/tinygrad/codegen/decomp/op.py index 6b0526b498..6a48cdca53 100644 --- a/tinygrad/codegen/decomp/op.py +++ b/tinygrad/codegen/decomp/op.py @@ -47,17 +47,17 @@ def fast_idiv(ren: Renderer, x: UOp, d: int, dont_cast=False) -> UOp|None: def threefry2x32(x: UOp, key: UOp): # split x and key from uint64 to two uint32 - x0, x1 = (x & 0xffffffff).cast(dtypes.uint32), ((x // 2**32) & 0xffffffff).cast(dtypes.uint32) - key0, key1 = (key & 0xffffffff).cast(dtypes.uint32), ((key // 2**32) & 0xffffffff).cast(dtypes.uint32) + x0, x1 = x.cast(dtypes.uint32), (x >> 32).cast(dtypes.uint32) + key0, key1 = key.cast(dtypes.uint32), (key >> 32).cast(dtypes.uint32) rotations = [[13, 15, 26, 6], [17, 29, 16, 24]] ks = [key1, key0 ^ key1 ^ 0x1BD11BDA, key0] xr:list[UOp] = [x0 + ks[-1], x1 + ks[0]] for i in range(5): - for r in rotations[i % 2]: xr[0], xr[1] = (x0 := xr[0] + xr[1]), x0 ^ ((xr[1] * 2**r) + (xr[1] // 2**(32 - r))) + for r in rotations[i % 2]: xr[0], xr[1] = (x0 := xr[0] + xr[1]), x0 ^ ((xr[1] << r) + (xr[1] >> (32 - r))) xr = [(xr[0] + ks[i % 3]), (xr[1] + ks[(i + 1) % 3] + i + 1)] - return xr[1].cast(dtypes.uint64) * 2**32 | xr[0].cast(dtypes.uint64) + return (xr[1].cast(dtypes.uint64) << 32) | xr[0].cast(dtypes.uint64) # ***** decomposition patterns ***** diff --git a/tinygrad/mixin/rand.py b/tinygrad/mixin/rand.py index 122f4e1f41..e1e09fd94d 100644 --- a/tinygrad/mixin/rand.py +++ b/tinygrad/mixin/rand.py @@ -12,7 +12,7 @@ class RandMixin(OpMixin): def _threefry_random_bits(key, counts0, counts1): x = (counts1.cast(dtypes.uint64) << 32) | counts0.cast(dtypes.uint64) x = x.threefry((key[1].cast(dtypes.uint64) << 32) | key[0].cast(dtypes.uint64)) - return (x & 0xffffffff).cast(dtypes.uint32).cat(((x >> 32) & 0xffffffff).cast(dtypes.uint32)) + return x.cast(dtypes.uint32).cat((x >> 32).cast(dtypes.uint32)) @classmethod def random_bits(cls, key:Self, counter:Self, num:int) -> Self: diff --git a/tinygrad/uop/symbolic.py b/tinygrad/uop/symbolic.py index 8830abcd70..76c600f90d 100644 --- a/tinygrad/uop/symbolic.py +++ b/tinygrad/uop/symbolic.py @@ -127,7 +127,6 @@ symbolic_simple = pm_data_invalid + PatternMatcher([ (UPat.var("x") ^ UPat.var("x"), lambda x: x.const_like(0)), # x^x -> 0 (UPat.var("x") & 0, lambda x: x.const_like(0)), # x&0 -> 0 # (x&mask)>>k -> x>>k when mask only clears bits below k - # TODO: combine this with "# rules for threefry" below ((UPat.var("x") & UPat.cvar("mask")) >> UPat.cvar("k"), lambda x,mask,k: x >> k.val if mask.val | ((1 << k.val) - 1) == -1 else None), ((UPat.var("x") & UPat.cvar("mask")) // UPat.cvar("c"), @@ -168,13 +167,10 @@ symbolic_simple = pm_data_invalid + PatternMatcher([ (UPat.var("x").alu(Ops.POW, UPat.cvar("c")), simplify_pow), # positive const ** x (UPat.cvar("c").alu(Ops.POW, UPat.var("x")), lambda c,x: c if c.val == 1 else (x*math.log2(c.val)).exp2() if c.val > 0 else None), - # rules for threefry - ((UPat.var('x', dtypes.uint64)&0xFFFFFFFF).cast(dtypes.uint32), lambda x: x.cast(dtypes.uint32)), - (((UPat.var(None, dtypes.uint64)*(1<<32)) | UPat.var('y', dtypes.uint32).cast(dtypes.uint64)).cast(dtypes.uint32), lambda y: y), - (((UPat.var('x', dtypes.uint64)*(1<<32)) | UPat.var(None, dtypes.uint32).cast(dtypes.uint64))//(1<<32), lambda x: x), - (((UPat.var(None, dtypes.uint64)<<32) | UPat.var('y', dtypes.uint32).cast(dtypes.uint64)).cast(dtypes.uint32), lambda y: y), - (((UPat.var('x', dtypes.uint64)<<32) | UPat.var(None, dtypes.uint32).cast(dtypes.uint64))//(1<<32), lambda x: x), - (((UPat.var('x', dtypes.uint64)<<32) | UPat.var(None, dtypes.uint32).cast(dtypes.uint64))>>32, lambda x: x), + # unpack a uint64 packed from two uint32 (threefry) + (((UPat.var(None, dtypes.uint64)<<32) | UPat.var('y', dtypes.uint32).cast(dtypes.uint64)).cast(dtypes.uint32), lambda y: y), + (((UPat.var('x', dtypes.uint32).cast(dtypes.uint64)<<32) | UPat.var(None, dtypes.uint32).cast(dtypes.uint64))>>32, + lambda x: x.cast(dtypes.uint64)), # ** simple where folding ** # a conditional with the same results either way is a noop, also fold const conditionals (UPat.var().where(UPat.var("val"), UPat.var("val")), lambda val: val), From c2625c78cbe6ca16195432f8d6d087060e037b4b Mon Sep 17 00:00:00 2001 From: George Hotz <72895+geohot@users.noreply.github.com> Date: Mon, 3 Aug 2026 15:10:21 -0700 Subject: [PATCH 3/3] scalar ALU index fix + llm: preserve_thinking (#17381) * cstyle: scalar ALU index fix, serve: preserve_thinking, test: fix Handler import - cstyle.py: return scalar directly when ALU buffer has 1 element - cli.py: add preserve_thinking param to FallbackTemplate.render - serve.py: pass preserve_thinking=True when rendering chat completions - test_llm_server.py: fix import to use Handler from llm.serve * real fix --- test/backend/test_wait_loop.py | 14 ++++++++++++++ test/null/test_llm_server.py | 4 ++-- tinygrad/codegen/__init__.py | 3 +++ tinygrad/llm/cli.py | 2 +- tinygrad/llm/serve.py | 2 +- 5 files changed, 21 insertions(+), 4 deletions(-) diff --git a/test/backend/test_wait_loop.py b/test/backend/test_wait_loop.py index 83e0543e61..4ce681c631 100644 --- a/test/backend/test_wait_loop.py +++ b/test/backend/test_wait_loop.py @@ -81,6 +81,20 @@ def loop_in_loop_kernel(C:UOp) -> UOp: return C[0].store(i[0].load()).sink(arg=KernelInfo(name="loop_in_loop", opts_to_apply=())) +def scalar_alu_index_kernel(out: UOp) -> UOp: + val = UOp.param(1, dtypes.int, (1,), vmin_vmax=(0, 100), addrspace=AddrSpace.ALU, name="val") + idx = UOp(Ops.INDEX, dtypes.int, (val, UOp.const(0, dtypes.int))) + return out[0].store(idx + 1).sink(arg=KernelInfo(name="scalar_alu_index")) + +class TestScalarALUIndex(unittest.TestCase): + def test_scalar_alu_index(self): + # regression: indexing a scalar (1,)-shaped ALU param must render the variable + # directly, not as data.x (which is a compile error on a scalar int) + out = Tensor.zeros(1, dtype=dtypes.int).contiguous().realize() + result = Tensor.custom_kernel(out, fxn=scalar_alu_index_kernel)[0] + run_linear(result.schedule_linear(), var_vals={"val": 42}) + self.assertEqual(result.item(), 43) + class TestWaitLoop(unittest.TestCase): def test_wait_loop(self): c = Tensor.empty(1, dtype=dtypes.int) diff --git a/test/null/test_llm_server.py b/test/null/test_llm_server.py index 9af50b0296..c4a3d39b6e 100644 --- a/test/null/test_llm_server.py +++ b/test/null/test_llm_server.py @@ -124,11 +124,11 @@ class TestLLMServer(unittest.TestCase): self.assertTrue(any(args[0].startswith("total:") and args[1] == "red" for args, _ in color.call_args_list)) def test_stream_disconnect_closes_source(self): - from tinygrad.viz.serve import HTTPRequestHandler + from tinygrad.llm.serve import Handler source, handler = Mock(), Mock() source.__iter__ = Mock(return_value=iter([{}])) handler.wfile.write.side_effect = BrokenPipeError - HTTPRequestHandler.stream_json(handler, source) + Handler.stream_json(handler, source) source.close.assert_called_once() def test_non_streaming(self): diff --git a/tinygrad/codegen/__init__.py b/tinygrad/codegen/__init__.py index 3f37a4c1dd..e0b8730c16 100644 --- a/tinygrad/codegen/__init__.py +++ b/tinygrad/codegen/__init__.py @@ -150,6 +150,9 @@ devectorizer2 = mop_cleanup+pm_mops+PatternMatcher([ (UPat(GroupOp.Elementwise|{Ops.LOAD,Ops.STORE}, name="b"), do_devectorize), # INDEX without src is nothing (TODO: this should be in mop_cleanup) (UPat(Ops.INDEX, src=(UPat.var('x'),)), lambda x: x), + # INDEX of a scalar ALU param is the param itself + (UPat(Ops.INDEX, src=(UPat(Ops.PARAM, name="v"),), allow_any_len=True), + lambda v: v if v.addrspace == AddrSpace.ALU and v.max_numel() == 1 else None), # unpack WMMA (UPat(Ops.WMMA, name="u"), do_stack_wmma), # stacked INDEX is many INDEX diff --git a/tinygrad/llm/cli.py b/tinygrad/llm/cli.py index a8c7d75ce4..cba67a4a1c 100644 --- a/tinygrad/llm/cli.py +++ b/tinygrad/llm/cli.py @@ -113,7 +113,7 @@ class FallbackTemplate: if self.tok.preset == 'glm4': return "" if self.tok.preset == 'tekken': return "[/INST]" return self.tok.decode([self.tok.eos_id]) - def render(self, messages:list[dict], tools=None, add_generation_prompt:bool=True) -> str: + def render(self, messages:list[dict], tools=None, add_generation_prompt:bool=True, preserve_thinking:bool=False) -> str: out = self.tok.decode([] if self.tok.bos_id is None else [self.tok.bos_id]) + ("" if self.tok.preset == 'glm4' else "") for msg in messages: out += self.role(msg["role"]) diff --git a/tinygrad/llm/serve.py b/tinygrad/llm/serve.py index 77b01f2735..34e414548f 100644 --- a/tinygrad/llm/serve.py +++ b/tinygrad/llm/serve.py @@ -128,7 +128,7 @@ class Handler(HTTPRequestHandler): if self.path == "/v1/chat/completions": # render and tokenize normalize_messages(body["messages"]) - rendered = self.server.template.render(messages=body["messages"], tools=body.get("tools"), add_generation_prompt=True) + rendered = self.server.template.render(messages=body["messages"], tools=body.get("tools"), add_generation_prompt=True, preserve_thinking=True) ids: list[int] = self.server.tok.encode(rendered) stderr_log(f"prep:{(time.perf_counter()-request_st)*1e3:5.0f} ms {colored('--', 'BLACK')} ") if len(ids) >= self.server.model.max_context: