wip: test change, do not merge

- checkbranch: fetch master from the base repo clone_url instead of the
  PR head remote (origin), which on GitHub is the fork for fork PRs
- replace the curl comment snippets with sticky_comment.py: stdlib-only,
  works with the REST API on both platforms, and keeps the
  skip_unchanged/ignore_empty behavior of the old action

.github/workflows/sticky_comment.py

@@ -0,0 +1,34 @@
+#!/usr/bin/env python3
+# Sticky PR comment via the REST API: find an existing comment containing MARKER and PATCH it, or POST a new one.
+# Works on GitHub and Gitea (stdlib only, replaces marocchino/sticky-pull-request-comment which needs GraphQL).
+# Env vars: GITHUB_TOKEN, GITHUB_API_URL, GITHUB_REPOSITORY (set by the runner), PR_NUMBER, MARKER, and BODY_FILE or MESSAGE.
+import json, os, sys, urllib.request
+
+api, repo = os.environ["GITHUB_API_URL"], os.environ["GITHUB_REPOSITORY"]
+pr, marker = os.environ["PR_NUMBER"], os.environ["MARKER"]
+body = open(os.environ["BODY_FILE"]).read() if os.environ.get("BODY_FILE") else os.environ["MESSAGE"]
+
+if not body.strip():
+  print("comment body is empty, not posting")
+  sys.exit(0)
+
+def req(url, method="GET", payload=None):
+  r = urllib.request.Request(url, data=None if payload is None else json.dumps(payload).encode(), method=method,
+    headers={"Authorization": f"token {os.environ['GITHUB_TOKEN']}", "Accept": "application/json", "Content-Type": "application/json"})
+  return json.load(urllib.request.urlopen(r))
+
+# find the latest sticky comment (paginate, 100 comments per page)
+existing, page = None, 1
+while True:
+  comments = req(f"{api}/repos/{repo}/issues/{pr}/comments?per_page=100&page={page}")
+  stickies = [c for c in comments if marker in (c.get("body") or "")]
+  if stickies: existing = stickies[-1]
+  if not comments or len(comments) < 100: break
+  page += 1
+
+if existing is not None and existing["body"] == body:
+  print("comment is already up to date")
+  sys.exit(0)
+url = f"{api}/repos/{repo}/issues/comments/{existing['id']}" if existing is not None else f"{api}/repos/{repo}/issues/{pr}/comments"
+resp = req(url, 'PATCH' if existing is not None else 'POST', {'body': body})
+print(f"{'updated' if existing is not None else 'created'} comment {resp['id']}")

.github/workflows/szdiff.yml

@@ -26,14 +26,14 @@ jobs:
       - name: Check whether branch is up-to-date
         id: brstat
         run: |
-          git remote add tinygrad https://github.com/tinygrad/tinygrad
-          git fetch tinygrad master
+          # fetch master from the base repo (tinygrad/tinygrad on GitHub, the mirror on Gitea), not the PR head remote
+          git fetch "${{ github.event.pull_request.base.repo.clone_url }}" master
           echo "${{ github.event.pull_request.head.sha }}"
-          git rev-list --left-right --count  tinygrad/master...${{ github.event.pull_request.head.sha }} | awk '{print "Behind "$1" - Ahead "$2""}'
-          count=$(git rev-list --left-right --count  tinygrad/master...${{ github.event.pull_request.head.sha }} | awk '{print $1}')
+          git rev-list --left-right --count FETCH_HEAD...${{ github.event.pull_request.head.sha }} | awk '{print "Behind "$1" - Ahead "$2""}'
+          count=$(git rev-list --left-right --count FETCH_HEAD...${{ github.event.pull_request.head.sha }} | awk '{print $1}')
           if [ $count -gt 0 ]
           then
-            echo "Current branch is behind tinygrad master branch!"
+            echo "Current branch is behind ${{ github.event.pull_request.base.repo.full_name }} master branch!"
             echo "stat=true" >> "$GITHUB_OUTPUT"
           else
             echo "stat=false" >> "$GITHUB_OUTPUT"
@@ -75,13 +75,13 @@ jobs:
           python sz.py "$BASE" "$PR" > loc_content.txt
       - name: Comment Code Line Diff
         continue-on-error: false
-        uses: marocchino/sticky-pull-request-comment@v3
-        with:
+        env:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
-          ignore_empty: true
-          skip_unchanged: true
-          recreate: true
-          path: loc_content.txt
+          PR_NUMBER: ${{ github.event.pull_request.number }}
+          MARKER: "### Changes"
+          BODY_FILE: loc_content.txt
+        # note: run the script from the base checkout, never from the PR checkout
+        run: python3 "$GITHUB_WORKSPACE/base/.github/workflows/sticky_comment.py"
 
   rebase:
     name: Core Library Line Difference
@@ -91,12 +91,14 @@ jobs:
     needs: checkbranch
     if: needs.checkbranch.outputs.branchstat == 'true'
     steps:
+      # pull_request_target: a plain checkout gets the base repo, so no PR code is executed
+      - uses: actions/checkout@v6
       - name: Comment Rebase
         continue-on-error: false
-        uses: marocchino/sticky-pull-request-comment@v3
-        with:
+        env:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
-          skip_unchanged: true
-          recreate: true
-          message: |
-            This branch currently is behind tinygrad/master. The line count difference bot is disabled.
+          PR_NUMBER: ${{ github.event.pull_request.number }}
+          MARKER: "line count difference bot is disabled"
+          MESSAGE: |
+            This branch currently is behind ${{ github.event.pull_request.base.repo.full_name }} master. The line count difference bot is disabled.
+        run: python3 .github/workflows/sticky_comment.py

tinygrad/codegen/__init__.py

@@ -330,7 +330,10 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
   sink = graph_rewrite(sink, symbolic_simple+pm_expand_broadcast+pm_add_loads, name="*** expand broadcast / add loads")
 
   # devectorize
-  sink = graph_rewrite(sink, symbolic_simple+devectorizer2+indexing_simplify, ctx=ren, name="devectorize2")
+  sink = graph_rewrite(sink, symbolic_simple+devectorizer2, ctx=ren, name="devectorize2")
+
+  # simplify indexing
+  sink = graph_rewrite(sink, indexing_simplify, name="simplify load/store indexing")
 
   # some coalescing misses without this
   sink = graph_rewrite(sink, sym, name="early symbolic")

tinygrad/helpers.py

@@ -620,3 +620,5 @@ class count:
     cur = self.n
     self.n += self.step
     return cur
+
+# test change for the szdiff bot

tinygrad/llm/serve.py

@@ -1,5 +1,5 @@
 from __future__ import annotations
-import json, pathlib, queue, re, threading, time, traceback, typing, uuid
+import json, pathlib, re, time, typing, uuid
 from typing import TYPE_CHECKING
 from tinygrad.helpers import DEBUG, colored, stderr_log
 from tinygrad.viz.serve import TCPServerWithReuse, HTTPRequestHandler
@@ -60,11 +60,6 @@ class StreamRouter:
     if emit: yield "content", emit
     if found: self.mode, self.buf = "tool", "<tool_call>" + self.buf
 
-class Request:
-  def __init__(self, ids:list[int], model_name:str, include_usage:bool, max_tokens:int|None, temperature:float):
-    self.ids, self.model_name, self.include_usage, self.max_tokens, self.temperature = ids, model_name, include_usage, max_tokens, temperature
-    self.out: queue.Queue[dict|None] = queue.Queue()
-
 class Handler(HTTPRequestHandler):
   server: LLMServer
   def log_request(self, code='-', size='-'): pass
@@ -72,10 +67,46 @@ class Handler(HTTPRequestHandler):
     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):
-    # hand off to the scheduler thread (the only thread that touches the model), drain our private stream
-    req = Request(ids, model_name, include_usage, max_tokens, temperature)
-    self.server.waiting.put(req)
-    while (c := req.out.get()) is not None: yield c
+    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()
+    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"<tool_call>\s*(.*?)\s*(?:</tool_call>|$)", 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")
 
   def do_POST(self):
     request_st = time.perf_counter()
@@ -121,58 +152,4 @@ class Handler(HTTPRequestHandler):
 class LLMServer(TCPServerWithReuse):
   def __init__(self, server_address:tuple, model:Transformer, model_name:str, tok:SimpleTokenizer, template:typing.Any):
     self.model, self.model_name, self.tok, self.template = model, model_name, tok, template
-    self.waiting: queue.Queue[Request] = queue.Queue()
-    threading.Thread(target=self.scheduler, daemon=True, name="LLMScheduler").start()
     super().__init__(server_address, Handler)
-
-  def scheduler(self):
-    # single thread owns the model; handler threads enqueue requests and drain req.out
-    while True:
-      req = self.waiting.get()                                   # admit FCFS
-      try:
-        for c in self._generate(req): req.out.put(c)             # TODO: batch decodes across requests (batched_step)
-      except Exception: traceback.print_exc()
-      req.out.put(None)                                          # retire: None ends the client's stream
-
-  def _generate(self, req:Request):
-    ids, model_name, include_usage, max_tokens, temperature = req.ids, req.model_name, req.include_usage, req.max_tokens, req.temperature
-    model, tok = self.model, self.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()
-    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"<tool_call>\s*(.*?)\s*(?:</tool_call>|$)", 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")

tinygrad/viz/serve.py

@@ -13,10 +13,8 @@ from tinygrad.renderer.amd.dsl import Inst
 from tinygrad.renderer.amd import detect_format
 
 # NOTE: using HTTPServer forces a potentially slow socket.getfqdn
-# NOTE: ThreadingMixIn means handlers must be thread-safe: LLMServer funnels model access through its scheduler thread
-class TCPServerWithReuse(socketserver.ThreadingMixIn, socketserver.TCPServer):
+class TCPServerWithReuse(socketserver.TCPServer):
   allow_reuse_address = True
-  daemon_threads = True
   def __init__(self, server_address, RequestHandlerClass):
     print(f"*** started server on http://127.0.0.1:{server_address[1]} at {time.perf_counter()-START_TIME:.2f} s")
     super().__init__(server_address, RequestHandlerClass)
@@ -524,12 +522,10 @@ def pma_timeline(blob:bytes, sm_version:int) -> list[ProfileEvent]:
 
 # ** Assembly static analyzers
 
-# redirect_stdout mutates global sys.stdout, serialize capture across request threads
-stdout_lock = threading.Lock()
 def get_stdout(f: Callable) -> str:
   buf = io.StringIO()
   try:
-    with stdout_lock, redirect_stdout(buf), redirect_stderr(buf): f()
+    with redirect_stdout(buf), redirect_stderr(buf): f()
   except Exception: traceback.print_exc(file=buf)
   return buf.getvalue()