From 59df317b12e805e344f11fefa2475b3f2050c5d4 Mon Sep 17 00:00:00 2001 From: chenyu Date: Sun, 2 Aug 2026 13:31:55 -0400 Subject: [PATCH 1/3] use UOp.const to create new consts [PR] (#17368) replace arg won't work with ConstArg --- test/null/test_graph_rewrite.py | 18 +++++++++--------- test/null/test_uop_graph.py | 2 +- test/null/test_uops.py | 8 +++++++- test/null/test_uops_stats.py | 12 ++++++------ test/null/test_viz.py | 8 ++++---- tinygrad/codegen/decomp/dtype.py | 2 ++ tinygrad/dtype.py | 2 +- tinygrad/uop/ops.py | 2 +- 8 files changed, 31 insertions(+), 23 deletions(-) diff --git a/test/null/test_graph_rewrite.py b/test/null/test_graph_rewrite.py index e9d26d0113..a10d13ea8b 100644 --- a/test/null/test_graph_rewrite.py +++ b/test/null/test_graph_rewrite.py @@ -307,8 +307,8 @@ class TestRecurse(unittest.TestCase): def test_inf_loop(self): a = UOp.const(3) pm = PatternMatcher([ - (UPat(Ops.CONST, arg=3, name="x"), lambda x: x.replace(arg=4)), - (UPat(Ops.CONST, arg=4, name="x"), lambda x: x.replace(arg=3)), + (UPat(Ops.CONST, arg=3, name="x"), lambda x: UOp.const(4, x.dtype)), + (UPat(Ops.CONST, arg=4, name="x"), lambda x: UOp.const(3, x.dtype)), ]) with self.assertRaises(RuntimeError): graph_rewrite(a, pm) @@ -316,8 +316,8 @@ class TestRecurse(unittest.TestCase): def test_inf_loop_bottom_up(self): a = UOp.const(3) pm = PatternMatcher([ - (UPat(Ops.CONST, arg=3, name="x"), lambda x: x.replace(arg=4)), - (UPat(Ops.CONST, arg=4, name="x"), lambda x: x.replace(arg=3)), + (UPat(Ops.CONST, arg=3, name="x"), lambda x: UOp.const(4, x.dtype)), + (UPat(Ops.CONST, arg=4, name="x"), lambda x: UOp.const(3, x.dtype)), ]) with self.assertRaises(RuntimeError): graph_rewrite(a, pm, bottom_up=True) @@ -378,8 +378,8 @@ class TestWalkRewrite(unittest.TestCase): """A bouncing pattern applies once and stops instead of looping.""" a = UOp.const(3) pm = PatternMatcher([ - (UPat(Ops.CONST, arg=3, name="x"), lambda x: x.replace(arg=4)), - (UPat(Ops.CONST, arg=4, name="x"), lambda x: x.replace(arg=3)), + (UPat(Ops.CONST, arg=3, name="x"), lambda x: UOp.const(4, x.dtype)), + (UPat(Ops.CONST, arg=4, name="x"), lambda x: UOp.const(3, x.dtype)), ]) with self.assertRaises(RuntimeError): graph_rewrite(a, pm, bottom_up=True) @@ -456,8 +456,8 @@ class TestWalkRewrite(unittest.TestCase): """Bottom-up walk also applies once per node, no fixed-point iteration.""" a = UOp.const(3) pm = PatternMatcher([ - (UPat(Ops.CONST, arg=3, name="x"), lambda x: x.replace(arg=4)), - (UPat(Ops.CONST, arg=4, name="x"), lambda x: x.replace(arg=3)), + (UPat(Ops.CONST, arg=3, name="x"), lambda x: UOp.const(4, x.dtype)), + (UPat(Ops.CONST, arg=4, name="x"), lambda x: UOp.const(3, x.dtype)), ]) ret = graph_rewrite(a, pm, bottom_up=True, walk=True) self.assertIs(ret, UOp.const(4)) @@ -511,7 +511,7 @@ class TestWalkRewrite(unittest.TestCase): def bpm_match(ctx, x): ctx.append((x.val if x.op is Ops.CONST else x.op, "bpm")) # rewrite const(1) -> const(10), short-circuiting its subtree - if x.op is Ops.CONST and x.val == 1: return x.replace(arg=10) + if x.op is Ops.CONST and x.val == 1: return UOp.const(10, x.dtype) return None def pm_match(ctx, x): ctx.append((x.val if x.op is Ops.CONST else x.op, "pm")) diff --git a/test/null/test_uop_graph.py b/test/null/test_uop_graph.py index e4ce98c0db..e889df8611 100644 --- a/test/null/test_uop_graph.py +++ b/test/null/test_uop_graph.py @@ -593,7 +593,7 @@ class TestUOpTags(unittest.TestCase): def test_inc_by_one(self): g = UOp.const(1) + UOp.const(1) assert g.ssimplify() == 2 - pm_plus_1 = PatternMatcher([(UPat(Ops.CONST, name="x"), lambda x: x.replace(arg=x.val+1, tag=1) if x.tag is None else None)]) + pm_plus_1 = PatternMatcher([(UPat(Ops.CONST, name="x"), lambda x: UOp.const(x.val+1, x.dtype).rtag(1) if x.tag is None else None)]) pm_strip_tags = PatternMatcher([(UPat(GroupOp.All, name="x"), lambda x: x.replace(tag=None) if x.tag is not None else None)]) g = graph_rewrite(g, pm_plus_1) assert g.ssimplify() == 4 diff --git a/test/null/test_uops.py b/test/null/test_uops.py index 09f6e9a55a..247c611250 100644 --- a/test/null/test_uops.py +++ b/test/null/test_uops.py @@ -126,6 +126,12 @@ class TestConstFloatEq(unittest.TestCase): self.assertFalse(nan == Invalid) self.assertTrue(nan != Invalid) # __ne__ must defer to the reflected eq, not swallow NotImplemented + def test_invalid_eq_defers_to_reflected(self): + class HoldsInvalid: # a carrier that knows it holds Invalid. returning False for foreign types would silence its eq + def __eq__(self, other): return other is Invalid + self.assertTrue(Invalid == HoldsInvalid()) + self.assertFalse(Invalid != HoldsInvalid()) + def test_matchers_agree_on_nan(self): n = UOp.const(math.nan, dtypes.float32) for compiled in (False, True): @@ -447,7 +453,7 @@ class TestUPatHelpers(unittest.TestCase): class TestUopsObject(unittest.TestCase): def test_timing(self): - with Timing("create 10k uops:"): ret = [UOp(Ops.CONST, dtypes.int, arg=10000000+i) for i in range(10000)] + with Timing("create 10k uops:"): ret = [UOp.const(10000000+i, dtypes.int) for i in range(10000)] assert len(ret) == 10000 def test_nested(self): diff --git a/test/null/test_uops_stats.py b/test/null/test_uops_stats.py index 7bef6acb1b..edab948072 100644 --- a/test/null/test_uops_stats.py +++ b/test/null/test_uops_stats.py @@ -147,21 +147,21 @@ class TestUOpsStats(unittest.TestCase): #MULACC should have the same stats as MUL + ADD def test_mulacc(self): globl = UOp.param(0, dtypes.int, (3,)) - o1 = UOp(Ops.CONST, dtypes.int, tuple(), 1) - o2 = UOp(Ops.CONST, dtypes.int, tuple(), 2) + o1 = UOp.const(1, dtypes.int) + o2 = UOp.const(2, dtypes.int) u1 = globl.index(o1) u2 = globl.index(o2) - u3 = UOp(Ops.CONST, dtypes.int, tuple(), 3) + u3 = UOp.const(3, dtypes.int) u4 = UOp(Ops.MUL, src=(u1,u2)) u5 = UOp(Ops.ADD, src=(u4,u3)) uops = tuple(u5.toposort()) globl = UOp.param(0, dtypes.int, (3,)) - o1 = UOp(Ops.CONST, dtypes.int, tuple(), 1) - o2 = UOp(Ops.CONST, dtypes.int, tuple(), 2) + o1 = UOp.const(1, dtypes.int) + o2 = UOp.const(2, dtypes.int) u1 = globl.index(o1) u2 = globl.index(o2) - u3 = UOp(Ops.CONST, dtypes.int, tuple(), 3) + u3 = UOp.const(3, dtypes.int) u4 = UOp(Ops.MULACC, src=(u1,u2,u3)) uops_fma = tuple(u4.toposort()) diff --git a/test/null/test_viz.py b/test/null/test_viz.py index fad98fc4a4..954a4f2252 100644 --- a/test/null/test_viz.py +++ b/test/null/test_viz.py @@ -97,7 +97,7 @@ class TestViz(unittest.TestCase): # VIZ tracks rewrites up to and including the error def count_3(x:UOp): assert x.val <= 3 - return x.replace(arg=x.val+1) + return UOp.const(x.val+1, x.dtype) err_pm = PatternMatcher([(UPat.cvar("x"), count_3),]) a = UOp.const(1) with save_viz() as viz: @@ -202,8 +202,8 @@ class TestViz(unittest.TestCase): a = UOp.const(3) b = UOp.const(4) pm = PatternMatcher([ - (UPat(Ops.CONST, arg=3, name="x"), lambda x: x.replace(arg=4)), - (UPat(Ops.CONST, arg=4, name="x"), lambda x: x.replace(arg=3)), + (UPat(Ops.CONST, arg=3, name="x"), lambda x: UOp.const(4, x.dtype)), + (UPat(Ops.CONST, arg=4, name="x"), lambda x: UOp.const(3, x.dtype)), ]) with save_viz() as viz: # use smaller stack limit for faster test (default is 250000) @@ -224,7 +224,7 @@ class TestViz(unittest.TestCase): list(viz.get_details(0, 0)) def test_enter_calls_rewrite(self): - pm = PatternMatcher([(UPat(Ops.CONST, arg=3, name="x"), lambda x: x.replace(arg=4))]) + pm = PatternMatcher([(UPat(Ops.CONST, arg=3, name="x"), lambda x: UOp.const(4, x.dtype))]) with save_viz() as viz: inner = UOp.const(3) call = UOp(Ops.CALL, src=(UOp(Ops.SINK, src=(inner,)),)) diff --git a/tinygrad/codegen/decomp/dtype.py b/tinygrad/codegen/decomp/dtype.py index d682aad8f5..8e25f7965d 100644 --- a/tinygrad/codegen/decomp/dtype.py +++ b/tinygrad/codegen/decomp/dtype.py @@ -178,6 +178,8 @@ pm_float_decomp = PatternMatcher([ f2f(x.bitcast(f2f_dt[ctx[0]]), ctx[0], ctx[1]) if bc.dtype == ctx[0] else None), (UPat(Ops.CAST, dtypes.floats, src=(UPat.var("val"),), name="x"), lambda ctx,x,val: f2f_clamp(val.cast(ctx[1]), ctx[0]) if x.dtype == ctx[0] else None), + # a CONST has no srcs to cast, it restates its value at the emulating dtype + (UPat(Ops.CONST, dtypes.floats, name="x"), lambda ctx,x: UOp.const(x.val, ctx[1]) if x.dtype == ctx[0] else None), (UPat(GroupOp.All-{Ops.BITCAST}, dtypes.floats, name="x"), lambda ctx,x: x.replace(dtype=ctx[1], src=tuple(s.cast(ctx[1]) if s.dtype == ctx[0] else s for s in x.src)) if x.dtype == ctx[0] else None), diff --git a/tinygrad/dtype.py b/tinygrad/dtype.py index 9272b7e617..4558aeaa46 100644 --- a/tinygrad/dtype.py +++ b/tinygrad/dtype.py @@ -27,7 +27,7 @@ class InvalidType: def __new__(cls): if cls._instance is None: cls._instance = object.__new__(cls) return cls._instance - def __eq__(self, other): return self is other + def __eq__(self, other): return self is other if isinstance(other, InvalidType) else NotImplemented # foreign types get the reflected eq def __hash__(self): return id(self) def __repr__(self): return "Invalid" def __reduce__(self): return (InvalidType, ()) # unpickle returns the singleton diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 67129a675d..0b20d84b24 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -1764,7 +1764,7 @@ def lower_weak_node(u:UOp) -> UOp|None: else unwrap(dtype_from_uop(u.op, src, u.arg))) return u.replace(dtype=None, src=src[:start]+tuple(s.cast(dt) for s in src[start:])).cast(u.dtype) pm_lower_weak = PatternMatcher([ - (UPat(Ops.CONST, dtype=dtypes.weaks, name="u"), lambda u: u.replace(dtype=select_dtype(u)).cast(u.dtype)), + (UPat(Ops.CONST, dtype=dtypes.weaks, name="u"), lambda u: UOp.const(u.val, select_dtype(u)).cast(u.dtype)), # two stacked weak casts are a weakint value used as weakfloat (or vice versa): resolve the inner one at the outer kind's default. # a SINGLE weak cast is never rewritten here, each consumer absorbs it on its own edge (see lower_weak_srcs) (UPat(Ops.CAST, dtype=dtypes.weaks, src=(UPat(Ops.CAST, dtype=dtypes.weaks, src=(UPat.var("x"),)),), name="u"), From 09dabfe05ef22854d55a3ec38d7394e38b91cb7b Mon Sep 17 00:00:00 2001 From: chenyu Date: Sun, 2 Aug 2026 15:17:33 -0400 Subject: [PATCH 2/3] minor pm_float_decomp cleanup [PR] (#17371) make the rule order independently correct --- tinygrad/codegen/decomp/dtype.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tinygrad/codegen/decomp/dtype.py b/tinygrad/codegen/decomp/dtype.py index 8e25f7965d..cbb324ace0 100644 --- a/tinygrad/codegen/decomp/dtype.py +++ b/tinygrad/codegen/decomp/dtype.py @@ -180,7 +180,7 @@ pm_float_decomp = PatternMatcher([ f2f_clamp(val.cast(ctx[1]), ctx[0]) if x.dtype == ctx[0] else None), # a CONST has no srcs to cast, it restates its value at the emulating dtype (UPat(Ops.CONST, dtypes.floats, name="x"), lambda ctx,x: UOp.const(x.val, ctx[1]) if x.dtype == ctx[0] else None), - (UPat(GroupOp.All-{Ops.BITCAST}, dtypes.floats, name="x"), lambda ctx,x: + (UPat(GroupOp.All-GroupOp.Defines-{Ops.CAST, Ops.BITCAST, Ops.CONST}, dtypes.floats, name="x"), lambda ctx,x: x.replace(dtype=ctx[1], src=tuple(s.cast(ctx[1]) if s.dtype == ctx[0] else s for s in x.src)) if x.dtype == ctx[0] else None), (UPat(Ops.STORE, src=(UPat.var("idx"), UPat(Ops.BITCAST, dtypes.floats, name="val")), name='st'), lambda ctx,st,idx,val: From 05bc7c69940a69ff8ef2ba3e3abea32329557489 Mon Sep 17 00:00:00 2001 From: George Hotz <72895+geohot@users.noreply.github.com> Date: Sun, 2 Aug 2026 12:31:52 -0700 Subject: [PATCH 3/3] fix llm reasoning and Linear import (#17372) --- test/null/test_llm_server.py | 22 ++++++++++ test/unit/test_llm_server.py | 12 ++++++ tinygrad/llm/model.py | 51 +++++++++++----------- tinygrad/llm/serve.py | 82 +++++++++++++++++++++--------------- tinygrad/viz/serve.py | 2 +- 5 files changed, 108 insertions(+), 61 deletions(-) diff --git a/test/null/test_llm_server.py b/test/null/test_llm_server.py index 2493258bd5..9af50b0296 100644 --- a/test/null/test_llm_server.py +++ b/test/null/test_llm_server.py @@ -109,6 +109,28 @@ class TestLLMServer(unittest.TestCase): self.assertGreater(len(contents), 0) + def test_interrupted_stream_logs_tokens(self): + with patch.object(self.mock_model, "generate", side_effect=lambda ids, **kwargs: iter([300, 301, 999])), \ + patch("tinygrad.llm.serve.stderr_log") as log, patch("tinygrad.llm.serve.colored", side_effect=lambda text, color: text) as color: + stream = self.server.RequestHandlerClass.run_model(Mock(server=self.server), [200, 201, 202], "test") + next(stream) + next(stream) + stream.close() + interrupt = log.call_args.args[0] + self.assertFalse(interrupt.startswith("\n")) + self.assertTrue(interrupt.endswith("\n")) + self.assertIn("gen:", interrupt) + self.assertIn("out: 1", interrupt) + 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 + source, handler = Mock(), Mock() + source.__iter__ = Mock(return_value=iter([{}])) + handler.wfile.write.side_effect = BrokenPipeError + HTTPRequestHandler.stream_json(handler, source) + source.close.assert_called_once() + def test_non_streaming(self): resp = self.client.chat.completions.create( model="test-model", diff --git a/test/unit/test_llm_server.py b/test/unit/test_llm_server.py index bbc8c90907..934d47e61c 100644 --- a/test/unit/test_llm_server.py +++ b/test/unit/test_llm_server.py @@ -3,6 +3,7 @@ from unittest.mock import patch from tinygrad import Tensor, UOp from tinygrad.schedule import schedule_cache from tinygrad.llm.model import Transformer, TransformerConfig +from tinygrad.llm.serve import StreamRouter TEST_CONFIG = TransformerConfig(num_blocks=1, dim=64, hidden_dim=128, n_heads=2, n_kv_heads=2, norm_eps=1e-5, vocab_size=100, head_dim=32, rope_theta=10000.0, rope_dim=32, v_head_dim=32, max_context=32) @@ -10,6 +11,17 @@ V_START_POS = UOp.variable("start_pos", 0, TEST_CONFIG.max_context-1) V_TOKS = UOp.variable("toks", 1, 32) # 32 is the default chunk_size in generate class TestTransformerGenerate(unittest.TestCase): + def test_first_recurrent_generate_before_state_init(self): + model = Transformer(TEST_CONFIG) + model.has_recurrent_block = True + with patch.object(Transformer, '__call__', return_value=Tensor([[42]])): + self.assertEqual(next(model.generate([0])), 42) + + def test_template_starts_reasoning(self): + router = StreamRouter(reasoning=True) + self.assertEqual(list(router.route("reasoninganswer")), + [("reasoning_content", "reasoning"), ("content", "answer")]) + def test_kv_cache_reuse(self): """Test that generate reuses the KV cache when tokens extend the cached prefix.""" model = Transformer(TEST_CONFIG) diff --git a/tinygrad/llm/model.py b/tinygrad/llm/model.py index e802c7b394..3fab067602 100644 --- a/tinygrad/llm/model.py +++ b/tinygrad/llm/model.py @@ -2,6 +2,7 @@ from __future__ import annotations import functools, itertools, pathlib from dataclasses import dataclass, replace from tinygrad import Tensor, nn, UOp, TinyJit, getenv, function +from tinygrad.nn import Linear from tinygrad.llm.gguf import gguf_load from tinygrad.uop.ops import resolve @@ -12,7 +13,7 @@ def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0, device:str| return freqs.cos().cat(freqs.sin(), dim=-1).clone(device) class ExpertWeights: - """Like nn.Linear but with num_experts dimension. Weight shape: (num_experts, out_features, in_features).""" + """Like Linear but with num_experts dimension. Weight shape: (num_experts, out_features, in_features).""" def __init__(self, num_experts:int, in_features:int, out_features:int): self.weight = Tensor.zeros(num_experts, out_features, in_features) def __call__(self, sel:Tensor, x:Tensor) -> Tensor: @@ -83,20 +84,20 @@ class FFNBlock: # --- feed-forward (MoE or dense) ------------------------------------- if config.num_experts > 0: - self.ffn_gate_inp = nn.Linear(config.dim, config.num_experts, bias=False) # router + self.ffn_gate_inp = Linear(config.dim, config.num_experts, bias=False) # router if config.expert_bias: self.exp_probs_b = {"bias": Tensor.zeros(config.num_experts)} self.ffn_gate_exps = ExpertWeights(config.num_experts, config.dim, config.hidden_dim) self.ffn_up_exps = ExpertWeights(config.num_experts, config.dim, config.hidden_dim) self.ffn_down_exps = ExpertWeights(config.num_experts, config.hidden_dim, config.dim) if config.shared_expert_dim > 0: - self.ffn_gate_shexp = nn.Linear(config.dim, config.shared_expert_dim, bias=False) - self.ffn_up_shexp = nn.Linear(config.dim, config.shared_expert_dim, bias=False) - self.ffn_down_shexp = nn.Linear(config.shared_expert_dim, config.dim, bias=False) + self.ffn_gate_shexp = Linear(config.dim, config.shared_expert_dim, bias=False) + self.ffn_up_shexp = Linear(config.dim, config.shared_expert_dim, bias=False) + self.ffn_down_shexp = Linear(config.shared_expert_dim, config.dim, bias=False) if config.shared_expert_gate: self.ffn_gate_inp_shexp = {"weight": Tensor.zeros(config.dim)} else: - self.ffn_gate = nn.Linear(config.dim, config.hidden_dim, bias=False) - self.ffn_up = nn.Linear(config.dim, config.hidden_dim, bias=False) - self.ffn_down = nn.Linear(config.hidden_dim, config.dim, bias=False) + self.ffn_gate = Linear(config.dim, config.hidden_dim, bias=False) + self.ffn_up = Linear(config.dim, config.hidden_dim, bias=False) + self.ffn_down = Linear(config.hidden_dim, config.dim, bias=False) def _feed_forward(self, x:Tensor) -> Tensor: if hasattr(self, 'ffn_gate_exps'): @@ -145,10 +146,10 @@ class TransformerBlock(FFNBlock): # --- attention projections (all linear, bias-free) ------------------ q_proj_out = config.head_dim * config.n_heads * (2 if config.attn_output_gate else 1) kv_proj_out = config.head_dim * config.n_kv_heads - self.attn_q = nn.Linear(config.dim, q_proj_out, bias=config.qkv_bias) - self.attn_k = nn.Linear(config.dim, kv_proj_out, bias=config.qkv_bias) - self.attn_v = nn.Linear(config.dim, kv_proj_out, bias=config.qkv_bias) - self.attn_output = nn.Linear(config.head_dim * config.n_heads, config.dim, bias=False) + self.attn_q = Linear(config.dim, q_proj_out, bias=config.qkv_bias) + self.attn_k = Linear(config.dim, kv_proj_out, bias=config.qkv_bias) + self.attn_v = Linear(config.dim, kv_proj_out, bias=config.qkv_bias) + self.attn_output = Linear(config.head_dim * config.n_heads, config.dim, bias=False) if config.qk_norm: self.attn_q_norm, self.attn_k_norm = nn.RMSNorm(config.qk_norm, config.norm_eps), nn.RMSNorm(config.qk_norm, config.norm_eps) def _attention(self, x:Tensor, start_pos:int|UOp) -> Tensor: @@ -195,16 +196,16 @@ class MLATransformerBlock(FFNBlock): super().__init__(config) qk_nope_head_dim = config.head_dim - config.rope_dim if config.q_lora_rank > 0: - self.attn_q_a = nn.Linear(config.dim, config.q_lora_rank, bias=False) + self.attn_q_a = Linear(config.dim, config.q_lora_rank, bias=False) self.attn_q_a_norm = nn.RMSNorm(config.q_lora_rank, config.norm_eps) - self.attn_q_b = nn.Linear(config.q_lora_rank, config.n_heads * config.head_dim, bias=False) + self.attn_q_b = Linear(config.q_lora_rank, config.n_heads * config.head_dim, bias=False) else: - self.attn_q = nn.Linear(config.dim, config.n_heads * config.head_dim, bias=False) - self.attn_kv_a_mqa = nn.Linear(config.dim, config.kv_lora_rank + config.rope_dim, bias=False) + self.attn_q = Linear(config.dim, config.n_heads * config.head_dim, bias=False) + self.attn_kv_a_mqa = Linear(config.dim, config.kv_lora_rank + config.rope_dim, bias=False) self.attn_kv_a_norm = nn.RMSNorm(config.kv_lora_rank, config.norm_eps) self.attn_k_b = {"weight": Tensor.zeros(config.n_heads, config.kv_lora_rank, qk_nope_head_dim)} self.attn_v_b = {"weight": Tensor.zeros(config.n_heads, config.v_head_dim, config.kv_lora_rank)} - self.attn_output = nn.Linear(config.n_heads * config.v_head_dim, config.dim, bias=False) + self.attn_output = Linear(config.n_heads * config.v_head_dim, config.dim, bias=False) def _attention(self, x:Tensor, start_pos:int|UOp) -> Tensor: B, T, _ = x.shape @@ -244,18 +245,18 @@ class GatedDeltaNetBlock(FFNBlock): assert self.num_v_heads % self.num_k_heads == 0 self.head_v_dim, self.ssm_conv_kernel = ssm.inner_size // ssm.time_step_rank, ssm.conv_kernel self.conv_channels, self.q_dim = ssm.inner_size + 2*ssm.group_count*ssm.state_size, ssm.state_size*ssm.group_count - self.attn_qkv = nn.Linear(config.dim, self.conv_channels, bias=False) + self.attn_qkv = Linear(config.dim, self.conv_channels, bias=False) if ssm.kda: - self.ssm_g_a, self.ssm_g_b = nn.Linear(config.dim, self.head_v_dim, bias=False), nn.Linear(self.head_v_dim, ssm.inner_size, bias=False) - self.ssm_f_a, self.ssm_f_b = nn.Linear(config.dim, self.head_k_dim, bias=False), nn.Linear(self.head_k_dim, ssm.inner_size, bias=False) + self.ssm_g_a, self.ssm_g_b = Linear(config.dim, self.head_v_dim, bias=False), Linear(self.head_v_dim, ssm.inner_size, bias=False) + self.ssm_f_a, self.ssm_f_b = Linear(config.dim, self.head_k_dim, bias=False), Linear(self.head_k_dim, ssm.inner_size, bias=False) else: - self.attn_gate = nn.Linear(config.dim, ssm.inner_size, bias=False) - self.ssm_alpha = nn.Linear(config.dim, self.num_v_heads, bias=False) - self.ssm_beta = nn.Linear(config.dim, self.num_v_heads, bias=False) + self.attn_gate = Linear(config.dim, ssm.inner_size, bias=False) + self.ssm_alpha = Linear(config.dim, self.num_v_heads, bias=False) + self.ssm_beta = Linear(config.dim, self.num_v_heads, bias=False) self.ssm_conv1d = {"weight": Tensor.zeros(self.conv_channels, self.ssm_conv_kernel)} self.ssm_dt = {"bias": Tensor.zeros(ssm.inner_size if ssm.kda else self.num_v_heads)} self.ssm_a = Tensor.zeros(self.num_v_heads, 1) if ssm.kda else Tensor.zeros(self.num_v_heads) - self.ssm_norm, self.ssm_out = nn.RMSNorm(self.head_v_dim, config.norm_eps), nn.Linear(ssm.inner_size, config.dim, bias=False) + self.ssm_norm, self.ssm_out = nn.RMSNorm(self.head_v_dim, config.norm_eps), Linear(ssm.inner_size, config.dim, bias=False) def _attention(self, x:Tensor, start_pos:int|UOp) -> Tensor: B, T, _ = x.shape @@ -314,7 +315,7 @@ class Transformer: block_cls(dense_config if i < config.leading_dense_blocks else config) for i in range(config.num_blocks)] self.token_embd = nn.Embedding(config.vocab_size, config.dim) self.output_norm = nn.RMSNorm(config.dim, config.norm_eps) - self.output = nn.Linear(config.dim, config.vocab_size, bias=False) + self.output = Linear(config.dim, config.vocab_size, bias=False) self.max_context = config.max_context self.has_recurrent_block = any(isinstance(b, GatedDeltaNetBlock) for b in self.blk) self._cached_tokens: list[int] = [] diff --git a/tinygrad/llm/serve.py b/tinygrad/llm/serve.py index 15dd813acc..77b01f2735 100644 --- a/tinygrad/llm/serve.py +++ b/tinygrad/llm/serve.py @@ -34,9 +34,9 @@ def normalize_messages(messages:list[dict]) -> None: class StreamRouter: # routes streamed output text to (field, text) deltas, keeping tool_call regions in .buf for the final parse - def __init__(self): + def __init__(self, reasoning:bool=False): self.buf = "" - self.mode = "undecided" # output inside a think block is sent as reasoning_content + self.mode = "reasoning" if reasoning else "undecided" # output inside a think block is sent as reasoning_content def split(self, tag:str, final:bool) -> tuple[str, bool]: # split buf on the first full tag, holding back a partial tag at the end unless final if tag in self.buf: @@ -66,47 +66,58 @@ class Handler(HTTPRequestHandler): def do_GET(self): if self.path == "/v1/models": self.send_data(json.dumps({"object":"list","data":[{"id":self.server.model_name,"object":"model"}]}).encode()) else: self.send_data((pathlib.Path(__file__).parent / "chat.html").read_bytes(), content_type="text/html") - def run_model(self, ids:list[int], model_name:str, include_usage=False, max_tokens:int|None=None, temperature:float=0.0): + def run_model(self, ids:list[int], model_name:str, include_usage=False, max_tokens:int|None=None, temperature:float=0.0, + reasoning:bool=False): model, tok = self.server.model, self.server.tok prompt_tokens = len(ids) cache_start_pos = model.get_start_pos(ids) stderr_log(f"in:{colored(f'{cache_start_pos:5d}', 'green')} +{len(ids)-cache_start_pos:5d} {colored('--', 'BLACK')} ") tmpl = {"id":f"chatcmpl-{uuid.uuid4().hex[:24]}", "object":"chat.completion.chunk", "created":int(time.time()), "model":model_name} def chunk(d:dict): return {"choices": [{"index":0, "delta":d, "finish_reason":None}], **tmpl} - yield chunk({"role":"assistant", "content":""}) out: list[int] = [] finish_reason = "stop" - st = time.perf_counter() + st = pt = time.perf_counter() dec = tok.stream_decoder() - router = StreamRouter() - for next_id in model.generate(ids, temperature=temperature): - if len(out) == 0: stderr_log(f"prefill:{(prompt_tokens-cache_start_pos)/((pt:=time.perf_counter())-st):4.0f} tok/s {colored('--', 'BLACK')} ") - if tok.is_end(next_id): break - out.append(next_id) - for field, delta in router.route(dec(next_id)): yield chunk({field:delta}) - if max_tokens is not None and len(out) >= max_tokens: - finish_reason = "length" - break - for field, delta in router.route(dec(), final=True): yield chunk({field:delta}) - tool_calls: list[dict] = [] - for m in re.finditer(r"\s*(.*?)\s*(?:|$)", router.buf, re.DOTALL): - if (parsed := parse_tool_call(m.group(1))) is None: - stderr_log(f"failed to parse tool call: {m.group(1)[:200]}") - yield chunk({"content":m.group(0)}) # don't silently drop output the client can't use - else: - name, args = parsed - tool_calls.append({"index":len(tool_calls), "id":f"call_{uuid.uuid4().hex[:24]}", "type":"function", - "function":{"name":name, "arguments":args if isinstance(args, str) else json.dumps(args)}}) - if tool_calls: - yield chunk({"tool_calls":tool_calls}) - if finish_reason == "stop": finish_reason = "tool_calls" - yield {"choices": [{"index":0, "delta":{},"finish_reason":finish_reason}], **tmpl} - if include_usage: - yield {"choices": [], "usage": {"prompt_tokens": prompt_tokens, "completion_tokens": len(out), - "total_tokens": prompt_tokens + len(out)}, **tmpl} - et = time.perf_counter() - stderr_log(f"gen:{len(out)/(et-pt) if len(out) > 1 else 0:4.0f} tok/s {colored('--', 'BLACK')} " - f"out:{len(out):5d} {colored('--', 'BLACK')} total:{et-st:6.2f}s\n") + router = StreamRouter(reasoning) + def log_stats(interrupted:bool=False): + et = time.perf_counter() + total = f"total:{et-st:6.2f}s" + stderr_log(f"gen:{len(out)/(et-pt) if len(out) > 1 else 0:4.0f} tok/s {colored('--', 'BLACK')} " + f"out:{len(out):5d} {colored('--', 'BLACK')} {colored(total, 'red') if interrupted else total}\n") + completed = False + try: + yield chunk({"role":"assistant", "content":""}) + for next_id in model.generate(ids, temperature=temperature): + if len(out) == 0: + stderr_log(f"prefill:{(prompt_tokens-cache_start_pos)/((pt:=time.perf_counter())-st):4.0f} tok/s {colored('--', 'BLACK')} ") + if tok.is_end(next_id): break + out.append(next_id) + for field, delta in router.route(dec(next_id)): yield chunk({field:delta}) + if max_tokens is not None and len(out) >= max_tokens: + finish_reason = "length" + break + for field, delta in router.route(dec(), final=True): yield chunk({field:delta}) + tool_calls: list[dict] = [] + for m in re.finditer(r"\s*(.*?)\s*(?:|$)", router.buf, re.DOTALL): + if (parsed := parse_tool_call(m.group(1))) is None: + stderr_log(f"failed to parse tool call: {m.group(1)[:200]}") + yield chunk({"content":m.group(0)}) # don't silently drop output the client can't use + else: + name, args = parsed + tool_calls.append({"index":len(tool_calls), "id":f"call_{uuid.uuid4().hex[:24]}", "type":"function", + "function":{"name":name, "arguments":args if isinstance(args, str) else json.dumps(args)}}) + if tool_calls: + yield chunk({"tool_calls":tool_calls}) + if finish_reason == "stop": finish_reason = "tool_calls" + completed = True + yield {"choices": [{"index":0, "delta":{},"finish_reason":finish_reason}], **tmpl} + if include_usage: + yield {"choices": [], "usage": {"prompt_tokens": prompt_tokens, "completion_tokens": len(out), + "total_tokens": prompt_tokens + len(out)}, **tmpl} + log_stats() + except GeneratorExit: + if not completed: log_stats(interrupted=True) + raise def do_POST(self): request_st = time.perf_counter() @@ -129,7 +140,8 @@ class Handler(HTTPRequestHandler): # reply max_tokens = body.get("max_completion_tokens") or body.get("max_tokens") chunks = self.run_model(ids, body["model"], not body.get("stream") or body.get("stream_options",{}).get("include_usage", False), - max_tokens=max_tokens, temperature=float(body.get("temperature", 0.0))) + max_tokens=max_tokens, temperature=float(body.get("temperature", 0.0)), + reasoning=rendered.rstrip().endswith("")) if body.get("stream"): self.stream_json(chunks) else: out, reasoning, tool_calls, finish_reason = [], [], [], "stop" diff --git a/tinygrad/viz/serve.py b/tinygrad/viz/serve.py index 8574bb1a71..0548f8a3c8 100755 --- a/tinygrad/viz/serve.py +++ b/tinygrad/viz/serve.py @@ -37,7 +37,7 @@ class HTTPRequestHandler(BaseHTTPRequestHandler): self.wfile.flush() self.wfile.write("data: [DONE]\n\n".encode("utf-8")) # pass if client closed connection - except (BrokenPipeError, ConnectionResetError): return + except (BrokenPipeError, ConnectionResetError): source.close() from tinygrad.uop.ops import TrackedGraphRewrite, RewriteTrace, UOp, Ops, GroupOp, srender, sint, sym_infer, range_str, range_start, multirange_str from tinygrad.uop.ops import KernelInfo