ELMer FUD

2026-09-03 22:53:44 +08:00 · 2026-09-02 20:38:12 -04:00
320 changed files with 8599 additions and 14642 deletions
@@ -159,21 +159,7 @@ All four files must be updated together.

 ## Manifest

-The current test branch uses manifest v25 and requests v25 only. Seed the new
-manifest from the previous catalog, then replace entries as artifacts are
-rebuilt with the pinned runtime:
-
-```bash
-cp /path/to/model_names_v24.json /path/to/model_names_v25.json
-```
-
-The current tinygrad pin is `f6fc4e3f2c3db5fae1e19cbfbc3ad9fc579a12ae`, from
-`openpilot` `origin/master` (`bump tg + TC_MIN_GLOBALS`). StarPilot's
-multi-model `modeld` remains in place; do not replace it with upstream's
-single-model `modeld`.
-
-For the older namespace migration workflow, generate the base manifest after
-compilation and namespace the release artifacts as v23:
+Generate the base manifest after compilation, then namespace the release artifacts as v23:

 ```bash
 python3 scripts/model_rebuild_pipeline.py manifest \
@@ -189,9 +175,9 @@ python3 scripts/namespace_model_artifacts.py \

 The namespace command changes IDs such as `tr1422` to `tr14223`, renames the
 compiled and upload-ready files, and writes an ID map. It preserves display
-names and behavioral versions. The current model manager requests v25 only; the
-manifest is fetched from `Models/model_names_v25.json`. Devices still running
-the prior branch continue to request their existing manifest version.
+names and behavioral versions. The current model manager requests v23 only;
+the manifest is fetched from `Models/model_names_v23.json`, while v22 remains
+available for devices that have not updated yet.

 After importing newly compiled sources, normalize the release namespace before
 copying files into either resource repository:
@@ -219,4 +205,4 @@ Compilation validates JIT capture/replay, pickle round-trip, finite outputs, met
 4. Confirm `driverStateV2` on both supported camera resolutions.
 5. Test download, selection, deletion, randomization, migration, and fallback in both device UIs and Galaxy.

-The built-in RDF artifact is `selfdrive/modeld/models/driving_tinygrad.pkl`. If migration cannot download the selected v25 artifact, StarPilot switches to that built-in model.
+The built-in RDF artifact is `selfdrive/modeld/models/driving_tinygrad.pkl`. If migration cannot download the selected v23 artifact, StarPilot switches to that built-in model.
@@ -56,13 +56,14 @@ def build_compile_env(*, supercombo: bool = False) -> dict[str, str]:
  existing_pythonpath = env.get("PYTHONPATH", "")
  env["PYTHONPATH"] = f"{REPO_ROOT}{os.pathsep}{existing_pythonpath}" if existing_pythonpath else str(REPO_ROOT)
  defaults = {
-    "DEBUG": "0",
    "FLOAT16": "1",
    "IMAGE": "1" if supercombo else "2",
    "JIT_BATCH_SIZE": "0",
    "NOLOCALS": "1",
    "OPENPILOT_HACKS": "1",
-  }
+  } | ({} if supercombo else {
+    "DEBUG": "0",
+  })
  for key, default in defaults.items():
    try:
      int(str(env.get(key)), 0)
@@ -31,7 +31,7 @@ OPENPILOT_REPO = "commaai/openpilot"
 RESOURCES_REPO = os.environ.get("STARPILOT_RESOURCES_REPO", "firestar5683/StarPilot-Resources")
 HF_BUCKET = os.environ.get("STARPILOT_HF_BUCKET", "StarPilot-Driving/StarPilot-Resources")
 RESOURCE_BRANCH = "Models"
-MANIFEST_VERSION = "v25"
+MANIFEST_VERSION = "v24"
 DEFAULT_BEHAVIOR_VERSION = "v16"
 DEVICE_ROOT = "/data/openpilot"
 REPOSITORY_FILE_LIMIT = 100_000_000
@@ -79,14 +79,14 @@ def test_runtime_scan_excludes_model_weights_but_flags_runtime_code():


 def test_update_manifest_replaces_one_entry(tmp_path: Path):
-  manifest = tmp_path / "model_names_v25.json"
+  manifest = tmp_path / "model_names_v24.json"
  manifest.write_text(json.dumps({"models": [{"id": "old"}]}) + "\n")
  info = parse_pasted_release(RELEASE_TEXT, "bmrlnapv4", "v16")
  path = update_manifest(
    tmp_path,
    info,
    {"size": 123, "sha256": "a" * 64},
-    "v25",
+    "v24",
  )
  payload = json.loads(path.read_text())
  assert len(payload["models"]) == 2
@@ -7,10 +7,6 @@ import struct
 from openpilot.system.hardware import HARDWARE, TICI
 os.environ['GMMU'] = '0'
 os.environ['DEV'] = 'QCOM' if TICI else 'LLVM'
-try:
-  int(os.getenv('DEBUG', '0'), 0)
-except ValueError:
-  os.environ['DEBUG'] = '0'
 from tinygrad.device import Device
 from tinygrad.tensor import Tensor
 import time
@@ -1 +1 @@
-31902b114b7fb8455af694d83333a86a44112b83be662e064ffbd67e8daafe72  driving_tinygrad.pkl
+a77db33c2e2d6a7570dc2a4a70c2b877429ee8bd9ca5dfeda74b5a41231aaff9  driving_tinygrad.pkl
@@ -24,7 +24,7 @@ from openpilot.starpilot.common.starpilot_utilities import delete_file
 from openpilot.starpilot.common.starpilot_variables import MODELS_PATH
 from openpilot.system.hardware.usb import chestnut_firmware_ready

-MANIFEST_CANDIDATES = ("v25",)
+MANIFEST_CANDIDATES = ("v24",)
 MODEL_NAMESPACE_SUFFIX = "3"
 DEFAULT_MODEL_KEY = "rdf43"
 LOCAL_MODEL_PREFIX = "local-"
@@ -15,12 +15,12 @@ from openpilot.starpilot.assets.model_manager import MANIFEST_CANDIDATES, ModelM
 from openpilot.starpilot.common.model_versions import UNIFIED_ARTIFACT_FORMAT


-def test_v25_is_the_only_manifest_candidate():
-  assert MANIFEST_CANDIDATES == ("v25",)
+def test_v24_is_the_only_manifest_candidate():
+  assert MANIFEST_CANDIDATES == ("v24",)


-def test_v25_manifest_is_loaded_from_models_checkout():
-  assert ModelManager._manifest_paths("v25") == ("Models/model_names_v25.json",)
+def test_v24_manifest_is_loaded_from_models_checkout():
+  assert ModelManager._manifest_paths("v24") == ("Models/model_names_v24.json",)


 def test_resource_sources_prefer_huggingface_then_github(monkeypatch):
@@ -33,9 +33,9 @@ def test_resource_sources_prefer_huggingface_then_github(monkeypatch):


 def test_huggingface_manifest_has_root_and_manifests_fallbacks():
-  assert ModelManager._hf_manifest_paths("v25") == (
-    "model_names_v25.json",
-    "manifests/model_names_v25.json",
+  assert ModelManager._hf_manifest_paths("v24") == (
+    "model_names_v24.json",
+    "manifests/model_names_v24.json",
  )


@@ -214,7 +214,7 @@ class BlueZClient:
        continue
      props = interfaces[DEVICE_IFACE]
      uuids = [str(value).lower() for value in props.get("UUIDs", [])]
-      audio, controller = device_capabilities(uuids, int(props.get("Class", 0)), str(props.get("Icon", "")))
+      audio, controller, serial = device_capabilities(uuids, int(props.get("Class", 0)), str(props.get("Icon", "")))
      device = {
        "path": path,
        "address": str(props.get("Address", "")),
@@ -227,9 +227,10 @@ class BlueZClient:
        "uuids": uuids,
        "audio": audio,
        "controller": controller,
+        "serial": serial,
      }
      if include_hidden or show_pairing_device(device["address"], device["name"], device["paired"], device["trusted"], device["connected"],
-                                               device["blocked"], audio, controller, include_discovering):
+                                               device["blocked"], audio, controller, serial, include_discovering):
        devices.append(device)
    return sorted(devices, key=lambda device: (not device["connected"], not device["paired"], -(device["rssi"] or -127), device["name"].lower()))

@@ -9,11 +9,13 @@ from typing import Any
 from openpilot.common.params import Params
 from openpilot.common.swaglog import cloudlog
 from openpilot.starpilot.system.bluetooth.bluez import BlueZClient
+from openpilot.starpilot.system.bluetooth.elm327 import ELM327Session
 from openpilot.starpilot.system.bluetooth.protocol import BLUETOOTH_SOCKET_PATH
 from openpilot.starpilot.system.bluetooth.radio import BluetoothRadio


-OFFROAD_COMMANDS = {"set_power", "start_scan", "stop_scan", "pair", "forget", "test_audio", "pairing_response"}
+OFFROAD_COMMANDS = {"set_power", "start_scan", "stop_scan", "pair", "forget", "test_audio", "pairing_response",
+                    "elm_open", "elm_command", "elm_read_dtcs"}
 SCAN_DURATION = 20.0
 AUDIO_TEST_START_DELAY = 3.0
 AUDIO_TEST_HOLD_TIME = 3.0
@@ -21,13 +23,15 @@ AUDIO_TEST_HOLD_TIME = 3.0

 class BluetoothController:
  def __init__(self, params: Params | None = None, bluez_factory=BlueZClient, radio: BluetoothRadio | None = None,
-               params_memory: Params | None = None, sleep=time.sleep):
+               params_memory: Params | None = None, sleep=time.sleep, elm_factory=ELM327Session):
    self.params = params or Params()
    self.params_memory = params_memory or Params(memory=True)
    self._bluez_factory = bluez_factory
+    self._elm_factory = elm_factory
    self._radio = radio or BluetoothRadio()
    self._lock = threading.RLock()
    self._bluez: BlueZClient | None = None
+    self._elm: ELM327Session | None = None
    self._pairing_address = ""
    self._pairing_error = ""
    self._last_reconnect = 0.0
@@ -41,6 +45,7 @@ class BluetoothController:
    self.params.remove("BluetoothAudioTestActive")
    self.params_memory.remove("TestAlert")
    with self._lock:
+      self._close_elm()
      if self._bluez is not None:
        self._bluez.close()
        self._bluez = None
@@ -75,6 +80,7 @@ class BluetoothController:

  def _reset_client(self) -> None:
    with self._lock:
+      self._close_elm()
      if self._bluez is not None:
        try:
          self._bluez.close()
@@ -82,6 +88,32 @@ class BluetoothController:
          pass
        self._bluez = None

+  def _close_elm(self) -> None:
+    with self._lock:
+      session = self._elm
+      self._elm = None
+      if session is None:
+        return
+      try:
+        session.close()
+      except Exception:
+        cloudlog.warning("ELM327 session close failed")
+
+  def _invalidate_elm(self, session: ELM327Session) -> None:
+    with self._lock:
+      if self._elm is not session:
+        return
+      self._close_elm()
+
+  def _active_elm(self, address: str) -> ELM327Session:
+    with self._lock:
+      session = self._elm
+      if session is None:
+        raise RuntimeError("ELM327 session is not open")
+      if str(session.address).upper() != address.upper():
+        raise RuntimeError("ELM327 session is open for another device")
+      return session
+
  def _offroad(self) -> bool:
    return self.params.get_bool("IsOffroad")

@@ -178,6 +210,7 @@ class BluetoothController:
              pass
            raise
        else:
+          self._close_elm()
          try:
            client = self._bluez
            if client is not None:
@@ -209,6 +242,9 @@ class BluetoothController:
    elif command == "disconnect":
      self._client().disconnect(address)
    elif command == "forget":
+      with self._lock:
+        if self._elm is not None and str(self._elm.address).upper() == address.upper():
+          self._close_elm()
      self._client().remove(address)
      if (self.params.get("BluetoothAudioAddress", encoding="utf-8") or "").upper() == address.upper():
        self.params.remove("BluetoothAudioAddress")
@@ -234,6 +270,50 @@ class BluetoothController:
      self._audio_test_deadline = deadline
      threading.Thread(target=self._test_audio_worker, args=(address, deadline), daemon=True).start()
      return {"audio_test_delay_ms": max(0, round((deadline - time.monotonic()) * 1000))}
+    elif command == "elm_open":
+      if not address:
+        raise RuntimeError("Bluetooth device address is required")
+      if not self.params.get_bool("BluetoothEnabled"):
+        raise RuntimeError("Bluetooth is disabled")
+      with self._lock:
+        if self._elm is not None and str(self._elm.address).upper() == address.upper():
+          return {"adapter": self._elm.adapter_name}
+        device = self._client().device_for_address(address)
+        if not device.get("paired"):
+          raise RuntimeError("Pair the Bluetooth device before opening ELM327")
+        if not device.get("serial"):
+          raise RuntimeError("Bluetooth device does not advertise Serial Port Profile")
+        self._close_elm()
+        session = self._elm_factory(address)
+        try:
+          adapter = str(session.open())
+        except Exception:
+          try:
+            session.close()
+          except Exception:
+            pass
+          raise
+        session.adapter_name = adapter
+        self._elm = session
+        return {"adapter": adapter}
+    elif command == "elm_close":
+      with self._lock:
+        if self._elm is not None and str(self._elm.address).upper() == address.upper():
+          self._close_elm()
+    elif command == "elm_command":
+      session = self._active_elm(address)
+      try:
+        return {"response": session.command(str(request.get("value", "")))}
+      except Exception:
+        self._invalidate_elm(session)
+        raise
+    elif command == "elm_read_dtcs":
+      session = self._active_elm(address)
+      try:
+        return session.read_dtcs()
+      except Exception:
+        self._invalidate_elm(session)
+        raise
    elif command == "pairing_response":
      if not self._client().agent.respond(str(request.get("prompt_id", "")), bool(request.get("accepted", False)), str(request.get("value", ""))):
        raise RuntimeError("Pairing request is no longer active")
@@ -252,10 +332,14 @@ class BluetoothController:
    while True:
      time.sleep(2)
      if not self.params.get_bool("BluetoothEnabled"):
+        self._close_elm()
        continue
      try:
        status = self.status()
+        if self._elm is not None and not status["offroad"]:
+          self._close_elm()
        if not status["available"] or not status["powered"]:
+          self._close_elm()
          continue
        now = time.monotonic()
        self._maintain_scan(status, now)
@@ -0,0 +1,220 @@
+from __future__ import annotations
+
+import re
+import socket
+import threading
+
+
+DEFAULT_CHANNEL = 1
+OPEN_TIMEOUT = 10.0
+DEFAULT_COMMAND_TIMEOUT = 10.0
+DTC_COMMAND_TIMEOUT = 25.0
+MAX_COMMAND_LENGTH = 256
+MAX_RESPONSE_SIZE = 64 * 1024
+RECV_SIZE = 4096
+
+
+class DTCParseError(ValueError):
+  def __init__(self, message: str, raw: str):
+    super().__init__(message)
+    self.raw = raw
+
+
+def decode_dtc(first: int, second: int) -> str:
+  prefixes = "PCBU"
+  prefix = prefixes[(first >> 6) & 0x03]
+  return f"{prefix}{(first >> 4) & 0x03:X}{first & 0x0F:X}{second >> 4:X}{second & 0x0F:X}"
+
+
+_HEX_BYTE = re.compile(r"(?i)(?<![0-9a-f])([0-9a-f]{2})(?![0-9a-f])")
+_FORMATTED_LENGTH = re.compile(r"(?i)^[0-9a-f]{3}$")
+_FORMATTED_FRAME = re.compile(r"(?i)^([0-9a-f]+):\s*(.*)$")
+_HEX_BYTE_TOKEN = re.compile(r"(?i)^[0-9a-f]{2}$")
+
+
+def _reassemble_formatted_responses(raw: str) -> list[str]:
+  lines = []
+  expected_length = None
+  expected_index = 0
+  assembled = bytearray()
+  for raw_line in raw.splitlines():
+    line = raw_line.strip()
+    if expected_length is None:
+      if _FORMATTED_LENGTH.fullmatch(line):
+        expected_length = int(line, 16)
+        if expected_length == 0:
+          raise DTCParseError("Formatted ELM response has an invalid length", raw)
+        expected_index = 0
+        assembled.clear()
+      else:
+        lines.append(line)
+      continue
+
+    if _FORMATTED_LENGTH.fullmatch(line):
+      raise DTCParseError("Formatted ELM response started before the previous block completed", raw)
+    match = _FORMATTED_FRAME.fullmatch(line)
+    if match is None:
+      raise DTCParseError("Formatted ELM response has a malformed continuation", raw)
+    index, byte_text = match.groups()
+    if int(index, 16) != expected_index:
+      raise DTCParseError("Formatted ELM response has a missing continuation", raw)
+    byte_tokens = byte_text.split()
+    if not byte_tokens or any(_HEX_BYTE_TOKEN.fullmatch(token) is None for token in byte_tokens):
+      raise DTCParseError("Formatted ELM response has invalid hex bytes", raw)
+    assembled.extend(int(token, 16) for token in byte_tokens)
+    expected_index += 1
+    if len(assembled) >= expected_length:
+      lines.append(" ".join(f"{value:02X}" for value in assembled[:expected_length]))
+      expected_length = None
+      assembled.clear()
+
+  if expected_length is not None:
+    raise DTCParseError("Formatted ELM response is incomplete", raw)
+  return lines
+
+
+def parse_dtcs(raw: str) -> list[str]:
+  response_lines = [line.strip().upper() for line in raw.splitlines() if line.strip()]
+  if response_lines and all(line == "NO DATA" for line in response_lines):
+    return []
+
+  codes = []
+  seen = set()
+  valid_response = False
+  for line in _reassemble_formatted_responses(raw):
+    values = [int(match, 16) for match in _HEX_BYTE.findall(line)]
+    try:
+      response_index = values.index(0x43)
+    except ValueError:
+      continue
+
+    payload = values[response_index + 1:]
+    if not payload:
+      raise DTCParseError("Mode 03 response has no payload", raw)
+    if len(payload) % 2:
+      count = payload[0]
+      expected_length = count * 2
+      if len(payload) - 1 < expected_length:
+        raise DTCParseError(f"Mode 03 response claims {expected_length} DTC bytes, received {len(payload) - 1}", raw)
+      payload = payload[1:1 + expected_length]
+
+    valid_response = True
+    for index in range(0, len(payload) - 1, 2):
+      first, second = payload[index:index + 2]
+      if first == 0 and second == 0:
+        continue
+      code = decode_dtc(first, second)
+      if code not in seen:
+        seen.add(code)
+        codes.append(code)
+  if not valid_response:
+    raise DTCParseError("No valid Mode 03 response", raw)
+  return codes
+
+
+class ELM327Session:
+  def __init__(self, address: str, channel: int = DEFAULT_CHANNEL):
+    self.address = address
+    self.channel = channel
+    self.socket: socket.socket | None = None
+    self.lock = threading.RLock()
+    self.adapter_name = ""
+
+  def _close_unlocked(self) -> None:
+    client_socket = self.socket
+    self.socket = None
+    self.adapter_name = ""
+    if client_socket is not None:
+      try:
+        client_socket.close()
+      except Exception:
+        pass
+
+  def close(self) -> None:
+    with self.lock:
+      self._close_unlocked()
+
+  def _receive_until_prompt_unlocked(self, timeout: float) -> bytes:
+    if self.socket is None:
+      raise RuntimeError("ELM327 session is not open")
+    self.socket.settimeout(timeout)
+    response = bytearray()
+    while True:
+      chunk = self.socket.recv(RECV_SIZE)
+      if not chunk:
+        raise RuntimeError("ELM327 connection closed")
+      response.extend(chunk)
+      if len(response) > MAX_RESPONSE_SIZE:
+        raise RuntimeError("ELM327 response exceeded 64 KiB")
+      if b">" in response:
+        return bytes(response)
+
+  @staticmethod
+  def _clean_response(raw: bytes, command: str) -> str:
+    response = raw.split(b">", 1)[0].decode("ascii", errors="replace")
+    response = response.replace("\r\n", "\n").replace("\r", "\n")
+    lines = response.split("\n")
+    while lines and not lines[0].strip():
+      lines.pop(0)
+    if lines and lines[0].strip() == command:
+      lines.pop(0)
+    return "\n".join(lines).strip()
+
+  def _exchange_unlocked(self, command: str, timeout: float) -> str:
+    if self.socket is None:
+      raise RuntimeError("ELM327 session is not open")
+    try:
+      self.socket.settimeout(timeout)
+      self.socket.sendall(command.encode("ascii") + b"\r")
+      return self._clean_response(self._receive_until_prompt_unlocked(timeout), command)
+    except Exception as error:
+      self._close_unlocked()
+      raise RuntimeError(f"ELM327 transport failed: {error}") from error
+
+  def open(self) -> str:
+    with self.lock:
+      if self.socket is not None:
+        return self.adapter_name
+      try:
+        self.socket = socket.socket(socket.AF_BLUETOOTH, socket.SOCK_STREAM, socket.BTPROTO_RFCOMM)
+        self.socket.settimeout(OPEN_TIMEOUT)
+        self.socket.connect((self.address, self.channel))
+        adapter_name = self._exchange_unlocked("ATI", OPEN_TIMEOUT)
+        if not adapter_name or adapter_name.strip().upper() in {"?", "ERROR", "COMMAND UNKNOWN", "UNKNOWN COMMAND", "NO DATA"}:
+          raise RuntimeError("ELM327 adapter rejected ATI")
+        self.adapter_name = adapter_name
+        for setup_command in ("ATE0", "ATL0", "ATH0"):
+          self._exchange_unlocked(setup_command, OPEN_TIMEOUT)
+        return self.adapter_name
+      except Exception as error:
+        self._close_unlocked()
+        if isinstance(error, RuntimeError) and str(error).startswith("ELM327 open failed:"):
+          raise
+        raise RuntimeError(f"ELM327 open failed: {error}") from error
+
+  def command(self, command: str, timeout: float = DEFAULT_COMMAND_TIMEOUT) -> str:
+    if not isinstance(command, str):
+      raise ValueError("ELM327 command must be text")
+    if "\r" in command or "\n" in command:
+      raise ValueError("ELM327 command cannot contain carriage returns or newlines")
+    command = command.strip()
+    if not command:
+      raise ValueError("ELM327 command cannot be empty")
+    if len(command) > MAX_COMMAND_LENGTH:
+      raise ValueError("ELM327 command is too long")
+    try:
+      command.encode("ascii")
+    except UnicodeEncodeError as error:
+      raise ValueError("ELM327 command must contain ASCII characters") from error
+    if timeout <= 0:
+      raise ValueError("ELM327 command timeout must be positive")
+
+    with self.lock:
+      return self._exchange_unlocked(command, timeout)
+
+  def read_dtcs(self) -> dict[str, str | list[str]]:
+    with self.lock:
+      for setup_command in ("ATD", "ATE0", "ATL0", "ATS1", "ATH0", "ATCAF1", "ATSP0"):
+        self.command(setup_command)
+      raw = self.command("03", timeout=DTC_COMMAND_TIMEOUT)
+    return {"codes": parse_dtcs(raw), "raw": raw}
@@ -16,6 +16,7 @@ BLUETOOTH_RADIO_HELPER = "/usr/comma/bluetooth-radio"
 A2DP_SINK_UUID = "0000110b-0000-1000-8000-00805f9b34fb"
 HID_UUID = "00001124-0000-1000-8000-00805f9b34fb"
 HOG_UUID = "00001812-0000-1000-8000-00805f9b34fb"
+SPP_UUID = "00001101-0000-1000-8000-00805f9b34fb"
 COMMAND_TIMEOUTS = {
  "set_power": 90.0,
  "start_scan": 20.0,
@@ -24,6 +25,10 @@ COMMAND_TIMEOUTS = {
  "disconnect": 20.0,
  "forget": 20.0,
  "test_audio": 10.0,
+  "elm_open": 15.0,
+  "elm_close": 5.0,
+  "elm_command": 20.0,
+  "elm_read_dtcs": 30.0,
 }
 TRUE_VALUES = {"1", "true", "yes", "on"}

@@ -40,6 +45,7 @@ class BluetoothDevice:
  uuids: tuple[str, ...] = ()
  audio: bool = False
  controller: bool = False
+  serial: bool = False

  @classmethod
  def from_dict(cls, value: dict[str, Any]) -> "BluetoothDevice":
@@ -54,6 +60,7 @@ class BluetoothDevice:
      uuids=tuple(str(uuid).lower() for uuid in value.get("uuids", ())),
      audio=bool(value.get("audio", False)),
      controller=bool(value.get("controller", False)),
+      serial=bool(value.get("serial", False)),
    )


@@ -86,21 +93,22 @@ class BluetoothStatus:
    )


-def device_capabilities(uuids: list[str] | tuple[str, ...], bluetooth_class: int = 0, icon: str = "") -> tuple[bool, bool]:
+def device_capabilities(uuids: list[str] | tuple[str, ...], bluetooth_class: int = 0, icon: str = "") -> tuple[bool, bool, bool]:
  normalized = {str(uuid).lower() for uuid in uuids}
  major_class = (int(bluetooth_class) >> 8) & 0x1F
  audio = A2DP_SINK_UUID in normalized or major_class == 0x04 or icon in {"audio-card", "audio-headphones", "audio-headset"}
  controller = HID_UUID in normalized or HOG_UUID in normalized or major_class == 0x05 or icon in {"input-gaming", "input-mouse", "input-keyboard"}
-  return audio, controller
+  serial = SPP_UUID in normalized
+  return audio, controller, serial


 def show_pairing_device(address: str, name: str, paired: bool, trusted: bool, connected: bool, blocked: bool,
-                        audio: bool, controller: bool, discovering: bool = False) -> bool:
+                        audio: bool, controller: bool, serial: bool = False, discovering: bool = False) -> bool:
  known = paired or trusted or connected
  normalized_address = "".join(character for character in address.upper() if character.isalnum())
  normalized_name = "".join(character for character in name.upper() if character.isalnum())
  named = bool(name) and name != "Unknown device" and normalized_name != normalized_address
-  return known or (named and not blocked and (audio or controller))
+  return known or (named and not blocked and (audio or controller or serial))


 class _DesktopFakeBluetooth:
@@ -304,5 +312,17 @@ class BluetoothClient:
    result = self.call("test_audio", address=address)
    return max(0.0, float(result.get("audio_test_delay_ms", 0)) / 1000.0)

+  def elm_open(self, address: str) -> dict[str, Any]:
+    return self.call("elm_open", address=address)
+
+  def elm_close(self, address: str) -> dict[str, Any]:
+    return self.call("elm_close", address=address)
+
+  def elm_command(self, address: str, value: str) -> dict[str, Any]:
+    return self.call("elm_command", address=address, value=value)
+
+  def elm_read_dtcs(self, address: str) -> dict[str, Any]:
+    return self.call("elm_read_dtcs", address=address)
+
  def respond(self, prompt_id: str, accepted: bool, value: str = "") -> None:
    self.call("pairing_response", prompt_id=prompt_id, accepted=accepted, value=value)
@@ -6,10 +6,11 @@ import numpy as np
 import pytest

 from openpilot.starpilot.system.bluetooth.audio import BluetoothAudioSink
+import openpilot.starpilot.system.bluetooth.daemon as bluetooth_daemon
 from openpilot.starpilot.system.bluetooth.bluez import PairingAgent
 from openpilot.starpilot.system.bluetooth.daemon import BluetoothController
 from openpilot.starpilot.system.bluetooth.protocol import (A2DP_SINK_UUID, HID_UUID, BluetoothClient, BluetoothDevice, BluetoothStatus,
-                                                           device_capabilities, show_pairing_device)
+                                                           SPP_UUID, device_capabilities, show_pairing_device)
 from openpilot.system import hardware
 from openpilot.system.ui.lib.bluetooth_manager import BluetoothManager

@@ -64,6 +65,7 @@ class FakeBlueZ:
      "connected": False,
      "audio": True,
      "controller": False,
+      "serial": False,
    }

  def close(self):
@@ -163,9 +165,38 @@ class FakeProcess:
    self.stopped = True


+class FakeELM:
+  instances = []
+
+  def __init__(self, address):
+    self.address = address
+    self.adapter_name = "Fake ELM327"
+    self.closed = False
+    self.commands = []
+    self.opened = False
+    FakeELM.instances.append(self)
+
+  def open(self):
+    self.opened = True
+    return self.adapter_name
+
+  def close(self):
+    self.closed = True
+
+  def command(self, value):
+    self.commands.append(value)
+    if value == "fail":
+      raise RuntimeError("transport failed")
+    return f"response for {value}"
+
+  def read_dtcs(self):
+    return {"codes": ["P0133"], "raw": "43 01 33"}
+
+
 def test_protocol_round_trip_and_capabilities():
-  audio, controller = device_capabilities([A2DP_SINK_UUID, HID_UUID])
+  audio, controller, serial = device_capabilities([A2DP_SINK_UUID, HID_UUID])
  assert audio and controller
+  assert not serial
  status = BluetoothStatus.from_dict({
    "available": True,
    "enabled": True,
@@ -174,12 +205,155 @@ def test_protocol_round_trip_and_capabilities():
  assert status.devices == (BluetoothDevice("00:11:22:33:44:55", "Combo", uuids=(A2DP_SINK_UUID, HID_UUID), audio=True, controller=True),)


+def test_serial_capability_round_trips_and_is_discoverable():
+  audio, controller, serial = device_capabilities([SPP_UUID.upper()])
+  assert not audio and not controller and serial
+
+  status = BluetoothStatus.from_dict({
+    "devices": [{"address": "00:11:22:33:44:55", "name": "OBDII", "serial": True}],
+  })
+  assert status.devices[0].serial
+  assert show_pairing_device("00:11:22:33:44:55", "OBDII", False, False, False, False,
+                             audio=False, controller=False, serial=True)
+
+
+def test_serial_device_is_not_auto_reconnected_but_audio_device_is(monkeypatch):
+  class StopMaintenance(Exception):
+    pass
+
+  params = FakeParams(IsOffroad=True, BluetoothEnabled=True)
+  client = FakeBlueZ()
+  client.powered = True
+  client.device.update(paired=True, trusted=True, connected=False, audio=False, controller=False, serial=True)
+  sleeps = 0
+
+  def sleep(_delay):
+    nonlocal sleeps
+    sleeps += 1
+    if sleeps > 1:
+      raise StopMaintenance
+
+  monkeypatch.setattr(bluetooth_daemon.time, "sleep", sleep)
+  controller = BluetoothController(params, lambda: client, FakeRadio(), sleep=sleep, elm_factory=FakeELM)
+  controller._bluez = client
+  controller._last_reconnect = -100.0
+  with pytest.raises(StopMaintenance):
+    controller.maintain_connections()
+  assert client.actions == []
+
+  client.device.update(audio=True, serial=False)
+  sleeps = 0
+  controller._last_reconnect = -100.0
+  with pytest.raises(StopMaintenance):
+    controller.maintain_connections()
+  assert client.actions == [("connect", client.device["address"])]
+
+
+def make_elm_controller(elm_factory=FakeELM):
+  params = FakeParams(IsOffroad=True, BluetoothEnabled=True)
+  client = FakeBlueZ()
+  client.device.update(paired=True, trusted=True, serial=True)
+  controller = BluetoothController(params, lambda: client, FakeRadio(), elm_factory=elm_factory)
+  return controller, client, params
+
+
+def test_elm_is_lazy_and_requires_a_paired_serial_device():
+  FakeELM.instances = []
+  controller, client, params = make_elm_controller()
+  assert FakeELM.instances == []
+  controller.status()
+  assert FakeELM.instances == []
+
+  result = controller.handle({"command": "elm_open", "address": client.device["address"]})
+  assert result == {"adapter": "Fake ELM327"}
+  assert len(FakeELM.instances) == 1
+
+  params.values["IsOffroad"] = False
+  with pytest.raises(RuntimeError, match="offroad"):
+    controller.handle({"command": "elm_command", "address": client.device["address"], "value": "ATI"})
+
+  params.values["IsOffroad"] = True
+  client.device["paired"] = False
+  controller.handle({"command": "elm_close", "address": client.device["address"]})
+  with pytest.raises(RuntimeError, match="Pair"):
+    controller.handle({"command": "elm_open", "address": client.device["address"]})
+
+  client.device.update(paired=True, serial=False)
+  with pytest.raises(RuntimeError, match="Serial Port Profile"):
+    controller.handle({"command": "elm_open", "address": client.device["address"]})
+
+
+def test_elm_commands_use_one_session_and_close_on_transport_failure():
+  FakeELM.instances = []
+  controller, client, _ = make_elm_controller()
+  address = client.device["address"]
+  controller.handle({"command": "elm_open", "address": address})
+  assert controller.handle({"command": "elm_open", "address": address}) == {"adapter": "Fake ELM327"}
+  assert controller.handle({"command": "elm_command", "address": address, "value": "ATI"}) == {"response": "response for ATI"}
+  assert controller.handle({"command": "elm_read_dtcs", "address": address}) == {"codes": ["P0133"], "raw": "43 01 33"}
+  with pytest.raises(RuntimeError, match="transport"):
+    controller.handle({"command": "elm_command", "address": address, "value": "fail"})
+  assert controller._elm is None
+  assert FakeELM.instances[0].closed
+
+
+def test_elm_open_replaces_a_different_session_and_close_is_allowed_onroad():
+  FakeELM.instances = []
+  controller, client, params = make_elm_controller()
+  first = client.device["address"]
+  second = "AA:BB:CC:DD:EE:FF"
+  controller.handle({"command": "elm_open", "address": first})
+  controller.handle({"command": "elm_open", "address": second})
+  assert len(FakeELM.instances) == 2
+  assert FakeELM.instances[0].closed
+  assert not FakeELM.instances[1].closed
+
+  params.values["IsOffroad"] = False
+  controller.handle({"command": "elm_close", "address": second})
+  assert FakeELM.instances[1].closed and controller._elm is None
+
+
+def test_elm_cleanup_happens_before_poweroff_forget_shutdown_and_onroad(monkeypatch):
+  class StopMaintenance(Exception):
+    pass
+
+  for cleanup in ("power", "forget", "shutdown", "onroad"):
+    FakeELM.instances = []
+    controller, client, params = make_elm_controller()
+    address = client.device["address"]
+    controller.handle({"command": "elm_open", "address": address})
+    session = FakeELM.instances[0]
+
+    if cleanup == "power":
+      controller.handle({"command": "set_power", "enabled": False})
+    elif cleanup == "forget":
+      controller.handle({"command": "forget", "address": address})
+    elif cleanup == "shutdown":
+      controller.close()
+    else:
+      params.values["IsOffroad"] = False
+      sleeps = 0
+
+      def sleep(_delay):
+        nonlocal sleeps
+        sleeps += 1
+        if sleeps > 1:
+          raise StopMaintenance
+
+      monkeypatch.setattr(bluetooth_daemon.time, "sleep", sleep)
+      controller._sleep = sleep
+      with pytest.raises(StopMaintenance):
+        controller.maintain_connections()
+
+    assert session.closed and controller._elm is None
+
+
 def test_pairing_list_filters_anonymous_and_irrelevant_advertisements():
  assert not show_pairing_device("00:11:22:33:44:55", "00:11:22:33:44:55", False, False, False, False, False, False)
  assert not show_pairing_device("00:11:22:33:44:55", "Nearby sensor", False, False, False, False, False, False)
  assert show_pairing_device("00:11:22:33:44:55", "Media Remote", False, False, False, False, False, True)
-  assert show_pairing_device("00:11:22:33:44:55", "Media Remote", False, False, False, False, False, True, True)
-  assert not show_pairing_device("00:11:22:33:44:55", "Nearby sensor", False, False, False, False, False, False, True)
+  assert show_pairing_device("00:11:22:33:44:55", "Media Remote", False, False, False, False, False, True, discovering=True)
+  assert not show_pairing_device("00:11:22:33:44:55", "Nearby sensor", False, False, False, False, False, False, discovering=True)
  assert show_pairing_device("00:11:22:33:44:55", "Known device", True, True, False, False, False, False)


@@ -0,0 +1,262 @@
+from collections import deque
+import threading
+
+import pytest
+
+from openpilot.starpilot.system.bluetooth import elm327
+
+
+ADDRESS = "00:11:22:33:44:55"
+
+
+class FakeSocket:
+  def __init__(self, responses):
+    self.responses = deque(deque(response) for response in responses)
+    self.pending = deque()
+    self.sent = []
+    self.connected_to = None
+    self.timeouts = []
+    self.closed = False
+    self.close_calls = 0
+    self.recv_error = None
+    self.send_error = None
+
+  def settimeout(self, timeout):
+    self.timeouts.append(timeout)
+
+  def connect(self, address):
+    self.connected_to = address
+
+  def sendall(self, value):
+    if self.send_error is not None:
+      raise self.send_error
+    self.sent.append(value)
+    self.pending = self.responses.popleft() if self.responses else deque()
+
+  def recv(self, _size):
+    if self.recv_error is not None:
+      raise self.recv_error
+    return self.pending.popleft() if self.pending else b""
+
+  def close(self):
+    self.close_calls += 1
+    self.closed = True
+
+
+def startup_responses(identity=b"ELM327 v1.5"):
+  return [
+    [b"ATI\r\n", identity + b"\r\n>"],
+    [b"ATE0\r\nOK\r\n>"],
+    [b"ATL0\r\nOK\r\n>"],
+    [b"ATH0\r\nOK\r\n>"],
+  ]
+
+
+def make_session(monkeypatch, responses):
+  fake = FakeSocket(responses)
+  monkeypatch.setattr(elm327.socket, "AF_BLUETOOTH", 31, raising=False)
+  monkeypatch.setattr(elm327.socket, "BTPROTO_RFCOMM", 3, raising=False)
+  monkeypatch.setattr(elm327.socket, "socket", lambda *args: fake)
+  return elm327.ELM327Session(ADDRESS), fake
+
+
+def test_open_uses_rfccomm_channel_one_and_validates_ati(monkeypatch):
+  session, fake = make_session(monkeypatch, startup_responses())
+
+  assert session.open() == "ELM327 v1.5"
+  assert fake.connected_to == (ADDRESS, 1)
+  assert fake.sent == [b"ATI\r", b"ATE0\r", b"ATL0\r", b"ATH0\r"]
+  assert session.adapter_name == "ELM327 v1.5"
+
+
+def test_open_rejects_empty_or_obviously_rejected_ati(monkeypatch):
+  for identity in (b"", b"?", b"ERROR"):
+    session, fake = make_session(monkeypatch, startup_responses(identity))
+    with pytest.raises(RuntimeError, match="ATI"):
+      session.open()
+    assert session.socket is None and fake.closed
+
+
+def test_command_removes_exact_echo_and_reads_split_prompt_response(monkeypatch):
+  responses = startup_responses() + [[b"ATR", b"V\r\n12.4V\r", b"\n>"]]
+  session, fake = make_session(monkeypatch, responses)
+  session.open()
+
+  assert session.command(" ATRV ") == "12.4V"
+  assert fake.sent[-1] == b"ATRV\r"
+
+
+def test_malformed_response_bytes_decode_with_replacement(monkeypatch):
+  responses = startup_responses() + [[b"ATI\r\n\xffOK\r\n>"]]
+  session, _ = make_session(monkeypatch, responses)
+  session.open()
+
+  assert session.command("ATI") == "�OK"
+
+
+@pytest.mark.parametrize("command", ["", "   ", "ATI\r", "ATI\n", "AT\r\nI", "A" * (elm327.MAX_COMMAND_LENGTH + 1)])
+def test_command_rejects_invalid_input(monkeypatch, command):
+  session, _ = make_session(monkeypatch, [])
+  with pytest.raises(ValueError):
+    session.command(command)
+
+
+def test_command_rejects_non_ascii_input(monkeypatch):
+  session, _ = make_session(monkeypatch, [])
+  with pytest.raises(ValueError, match="ASCII"):
+    session.command("ATé")
+
+
+def test_response_size_limit_closes_session(monkeypatch):
+  responses = startup_responses() + [[b"x" * (elm327.MAX_RESPONSE_SIZE + 1)]]
+  session, fake = make_session(monkeypatch, responses)
+  session.open()
+
+  with pytest.raises(RuntimeError, match="64 KiB"):
+    session.command("ATI")
+  assert session.socket is None and fake.closed
+
+
+@pytest.mark.parametrize("error", [TimeoutError("timed out"), OSError("disconnected")])
+def test_timeout_or_eof_closes_session(monkeypatch, error):
+  responses = startup_responses() + [[]]
+  session, fake = make_session(monkeypatch, responses)
+  session.open()
+  fake.recv_error = error if isinstance(error, TimeoutError) else None
+
+  if isinstance(error, TimeoutError):
+    with pytest.raises(RuntimeError, match="transport"):
+      session.command("ATI")
+  else:
+    with pytest.raises(RuntimeError, match="connection closed"):
+      session.command("ATI")
+  assert session.socket is None and fake.closed
+
+
+def test_send_failure_closes_session(monkeypatch):
+  responses = startup_responses() + [[b"OK>"]]
+  session, fake = make_session(monkeypatch, responses)
+  session.open()
+  fake.send_error = OSError("send failed")
+
+  with pytest.raises(RuntimeError, match="transport"):
+    session.command("ATI")
+  assert session.socket is None and fake.closed
+
+
+def test_close_is_idempotent(monkeypatch):
+  session, fake = make_session(monkeypatch, startup_responses())
+  session.open()
+  session.close()
+  session.close()
+
+  assert fake.close_calls == 1
+  assert session.socket is None
+
+
+def test_simultaneous_commands_are_serialized(monkeypatch):
+  class SerializedSocket(FakeSocket):
+    def __init__(self, responses):
+      super().__init__(responses)
+      self.command_started = threading.Event()
+      self.release_command = threading.Event()
+      self._command_sends = 0
+
+    def sendall(self, value):
+      super().sendall(value)
+      self._command_sends += 1
+      if self._command_sends == 5:
+        self.command_started.set()
+        assert self.release_command.wait(timeout=1.0)
+
+  fake = SerializedSocket(startup_responses() + [[b"VALUE1>"], [b"VALUE2>"]])
+  monkeypatch.setattr(elm327.socket, "AF_BLUETOOTH", 31, raising=False)
+  monkeypatch.setattr(elm327.socket, "BTPROTO_RFCOMM", 3, raising=False)
+  monkeypatch.setattr(elm327.socket, "socket", lambda *args: fake)
+  session = elm327.ELM327Session(ADDRESS)
+  session.open()
+
+  results = []
+  first = threading.Thread(target=lambda: results.append(session.command("ONE")))
+  second = threading.Thread(target=lambda: results.append(session.command("TWO")))
+  first.start()
+  assert fake.command_started.wait(timeout=1.0)
+  second.start()
+  assert len(fake.sent) == 5
+  fake.release_command.set()
+  first.join(timeout=1.0)
+  second.join(timeout=1.0)
+
+  assert sorted(results) == ["VALUE1", "VALUE2"]
+  assert len(fake.sent) == 6
+
+
+def test_read_dtcs_runs_known_setup_and_returns_mode_three_result(monkeypatch):
+  responses = startup_responses() + [[b"OK>"] for _ in range(7)] + [[b"43 01 33 04 20 00 00>"]]
+  session, fake = make_session(monkeypatch, responses)
+  session.open()
+
+  assert session.read_dtcs() == {"codes": ["P0133", "P0420"], "raw": "43 01 33 04 20 00 00"}
+  assert fake.sent[-8:] == [b"ATD\r", b"ATE0\r", b"ATL0\r", b"ATS1\r", b"ATH0\r", b"ATCAF1\r", b"ATSP0\r", b"03\r"]
+
+
+@pytest.mark.parametrize("raw_command", ["ATS0", "ATCAF0"])
+def test_read_dtcs_restores_parser_state_after_raw_command(monkeypatch, raw_command):
+  responses = startup_responses() + [[b"OK>"]] + [[b"OK>"] for _ in range(7)] + [[b"43 01 33 00 00 00 00>"]]
+  session, fake = make_session(monkeypatch, responses)
+  session.open()
+  session.command(raw_command)
+
+  assert session.read_dtcs() == {"codes": ["P0133"], "raw": "43 01 33 00 00 00 00"}
+  assert fake.sent[-8:] == [b"ATD\r", b"ATE0\r", b"ATL0\r", b"ATS1\r", b"ATH0\r", b"ATCAF1\r", b"ATSP0\r", b"03\r"]
+
+
+def test_non_can_dtc_is_decoded_and_padding_ignored():
+  assert elm327.parse_dtcs("43 01 33 00 00 00 00") == ["P0133"]
+
+
+def test_can_dtc_count_byte_is_skipped():
+  assert elm327.parse_dtcs("43 02 01 33 04 20") == ["P0133", "P0420"]
+
+
+def test_no_data_returns_no_codes():
+  assert elm327.parse_dtcs("NO DATA") == []
+
+
+def test_can_dtc_multiframe_response_is_reassembled():
+  raw = "008\n0: 43 03 00 59 01 54\n1: 01 55"
+  assert elm327.parse_dtcs(raw) == ["P0059", "P0154", "P0155"]
+
+
+def test_larger_numbered_multiframe_response_is_reassembled():
+  raw = "00E\n0: 43 06 01 33 04 20\n1: 00 59 01 54\n2: 01 55 01 56"
+  assert elm327.parse_dtcs(raw) == ["P0133", "P0420", "P0059", "P0154", "P0155", "P0156"]
+
+
+@pytest.mark.parametrize("raw", [
+  "008\n0: 43 03 00 59 01 54",
+  "008\n1: 43 03 00 59 01 54\n1: 01 55",
+  "008\n0: 43 03 00 59 GG 54\n1: 01 55",
+  "008\n0: 43 03 00 59 01 54\n00E",
+])
+def test_invalid_numbered_multiframe_response_raises(raw):
+  with pytest.raises(elm327.DTCParseError):
+    elm327.parse_dtcs(raw)
+
+
+@pytest.mark.parametrize("raw", ["UNABLE TO CONNECT", "STOPPED", "?", "BUS ERROR", "7F 03 11", "43"])
+def test_non_mode_three_response_raises(raw):
+  with pytest.raises(elm327.DTCParseError):
+    elm327.parse_dtcs(raw)
+
+
+def test_multiple_ecu_lines_are_ordered_and_deduplicated():
+  raw = "43 01 33 00 00\n43 04 20 00 00\n43 01 33 00 00"
+  assert elm327.parse_dtcs(raw) == ["P0133", "P0420"]
+
+
+def test_malformed_can_count_raises_with_raw_response():
+  raw = "43 02 01 33"
+  with pytest.raises(elm327.DTCParseError) as error:
+    elm327.parse_dtcs(raw)
+  assert error.value.raw == raw
@@ -398,6 +398,116 @@
  animation-delay: 0.28s;
 }

+.bluetoothElmPanel {
+  background: var(--secondary-bg);
+  border: 1px solid rgba(169, 140, 229, 0.45);
+  border-radius: var(--border-radius-lg);
+  box-shadow: var(--shadow-sm);
+  padding: 18px;
+}
+
+.bluetoothElmHeader,
+.bluetoothElmActions,
+.bluetoothElmCommand > div {
+  align-items: center;
+  display: flex;
+  gap: 10px;
+}
+
+.bluetoothElmHeader {
+  justify-content: space-between;
+  gap: 18px;
+}
+
+.bluetoothElmHeader h3,
+.bluetoothElmHeader p,
+.bluetoothElmCodes p,
+.bluetoothElmResponse pre,
+.bluetoothElmCommand label {
+  margin: 0;
+}
+
+.bluetoothElmHeader p,
+.bluetoothElmCodes p {
+  color: var(--text-muted);
+  margin-top: 4px;
+}
+
+.bluetoothElmHeader strong {
+  color: #cbb2fa;
+  overflow: hidden;
+  text-overflow: ellipsis;
+  white-space: nowrap;
+}
+
+.bluetoothElmActions {
+  margin-top: 16px;
+}
+
+.bluetoothElmActions button,
+.bluetoothElmCommand button {
+  background: linear-gradient(135deg, #765bb6, #9474ce);
+  border: 0;
+  border-radius: var(--border-radius-md);
+  color: #fff;
+  cursor: pointer;
+  font-weight: 700;
+  padding: 10px 14px;
+}
+
+.bluetoothElmActions .bluetoothSecondaryButton {
+  background: var(--input-bg);
+  border: 1px solid var(--sidebar-border-color);
+  color: var(--text-color);
+}
+
+.bluetoothElmActions button:disabled,
+.bluetoothElmCommand button:disabled,
+.bluetoothElmCommand input:disabled {
+  cursor: not-allowed;
+  opacity: 0.45;
+}
+
+.bluetoothElmCodes,
+.bluetoothElmCommand,
+.bluetoothElmResponse {
+  margin-top: 16px;
+}
+
+.bluetoothElmCommand label {
+  color: var(--text-muted);
+  display: block;
+  font-size: 0.86rem;
+  margin-bottom: 6px;
+}
+
+.bluetoothElmCommand > div {
+  align-items: stretch;
+}
+
+.bluetoothElmCommand input {
+  background: var(--input-bg);
+  border: 1px solid var(--sidebar-border-color);
+  border-radius: var(--border-radius-md);
+  color: var(--text-color);
+  flex: 1;
+  font: inherit;
+  min-width: 0;
+  padding: 10px 12px;
+}
+
+.bluetoothElmResponse pre {
+  background: var(--input-bg);
+  border: 1px solid var(--sidebar-border-color);
+  border-radius: var(--border-radius-md);
+  color: var(--text-color);
+  margin-top: 6px;
+  max-height: 220px;
+  overflow: auto;
+  padding: 12px;
+  white-space: pre-wrap;
+}
+
@keyframes bluetoothSpin {
  to { transform: rotate(360deg); }
 }
@@ -20,6 +20,11 @@ const state = reactive({
  prompt: null,
  audioTestAddress: "",
  audioTestLabel: "",
+  elmAddress: "",
+  elmName: "",
+  elmAdapter: "",
+  elmResponse: "",
+  elmCodes: null,
  error: "",
 })

@@ -46,7 +51,9 @@ function schedulePoll(delay = pollDelay()) {
  pollTimer = setTimeout(async () => {
    pollTimer = null
    try {
-      if (bluetoothPageActive() && document.visibilityState !== "hidden" && state.busy !== "power") {
+      if (state.elmAddress && (document.visibilityState === "hidden" || !bluetoothPageActive())) {
+        closeElm()
+      } else if (bluetoothPageActive() && document.visibilityState !== "hidden" && state.busy !== "power") {
        await refresh()
      }
    } finally {
@@ -97,8 +104,10 @@ async function request(operation, body = {}) {
    }
    state.error = ""
    await refresh()
+    return payload
  } catch (error) {
    state.error = error?.message || "Bluetooth operation failed"
+    return null
  } finally {
    state.busy = ""
    if (operation === "power") state.powerTarget = null
@@ -106,6 +115,54 @@ async function request(operation, body = {}) {
  }
 }

+function clearElmState() {
+  state.elmAddress = ""
+  state.elmName = ""
+  state.elmAdapter = ""
+  state.elmResponse = ""
+  state.elmCodes = null
+}
+
+function closeElm() {
+  const address = state.elmAddress
+  if (!address) return
+  clearElmState()
+  request("elm_close", { address })
+}
+
+async function openElm(address) {
+  const device = state.devices.find((item) => normalizedAddress(item) === String(address || "").toUpperCase())
+  const payload = await request("elm_open", { address })
+  if (!payload || !device) return
+  state.elmAddress = address
+  state.elmName = device.name || address
+  state.elmAdapter = String(payload.adapter || "")
+  state.elmResponse = ""
+  state.elmCodes = null
+}
+
+async function readElmCodes() {
+  const address = state.elmAddress
+  if (!address) return
+  const payload = await request("elm_read_dtcs", { address })
+  if (!payload || state.elmAddress !== address) return
+  state.elmCodes = Array.isArray(payload.codes) ? payload.codes.map(String) : []
+  state.elmResponse = String(payload.raw || "")
+}
+
+async function sendElmCommand() {
+  const address = state.elmAddress
+  const input = document.getElementById("bluetoothElmCommand")
+  const command = input?.value.trim() || ""
+  if (!address || !command) {
+    state.error = "Enter an ELM327 command."
+    return
+  }
+  const payload = await request("elm_command", { address, value: command })
+  if (!payload || state.elmAddress !== address) return
+  state.elmResponse = `${command}\n${String(payload.response || "")}`.trim()
+}
+
 async function refreshOnce() {
  const statusUrl = `${galaxyPath("/api/bluetooth/status")}?_=${Date.now()}`
  const response = await fetch(statusUrl, { cache: "no-store" })
@@ -126,6 +183,10 @@ async function refreshOnce() {
    devices,
  })
  state.devices = devices
+  if (state.elmAddress && (!state.enabled || !state.offroad ||
+      !devices.some((device) => normalizedAddress(device) === state.elmAddress.toUpperCase() && device.paired))) {
+    closeElm()
+  }
  if (state.deviceSignature !== deviceSignature) {
    state.deviceSignature = deviceSignature
    state.revision++
@@ -223,7 +284,11 @@ function initialize() {
  window.addEventListener("focus", refresh)
  window.addEventListener("pageshow", refresh)
  document.addEventListener("visibilitychange", () => {
-    if (document.visibilityState !== "hidden" && bluetoothPageActive()) refresh()
+    if (document.visibilityState === "hidden" || !bluetoothPageActive()) {
+      closeElm()
+    } else {
+      refresh()
+    }
  })
  refresh()
  schedulePoll(0)
@@ -248,6 +313,7 @@ function deviceCapabilities(device) {
  const capabilities = []
  if (device.audio) capabilities.push("Audio")
  if (device.controller) capabilities.push("Controller")
+  if (device.serial) capabilities.push("Serial")
  return capabilities.join(" · ") || "Bluetooth device"
 }

@@ -353,7 +419,14 @@ function renderDeviceActions(device) {
    actions.push("<button data-bluetooth-operation=\"pair\" data-address=\"" + address + "\"" +
      renderDisabledAttribute(!state.offroad || !!state.busy || pairing) + ">" + (pairing ? "Pairing…" : "Pair") + "</button>")
  }
-  if (device.paired || device.connected) {
+  if (device.paired && device.serial) {
+    actions.push("<button data-bluetooth-operation=\"elm_open\" data-address=\"" + address + "\"" +
+      renderDisabledAttribute(!state.offroad || !!state.busy) + ">ELM327</button>")
+    actions.push("<button class=\"bluetoothIconButton bluetoothForgetButton\" data-bluetooth-operation=\"forget\" data-address=\"" +
+      address + "\" data-device-name=\"" + name + "\" title=\"Forget device\" aria-label=\"Forget " + name + "\"" +
+      renderDisabledAttribute(!state.offroad || !!state.busy) + "><i class=\"bi bi-trash3\" aria-hidden=\"true\"></i></button>")
+  }
+  if ((device.paired || device.connected) && !device.serial) {
    const operation = device.connected ? "disconnect" : "connect"
    actions.push("<button data-bluetooth-operation=\"" + operation + "\" data-address=\"" + address + "\"" +
      renderDisabledAttribute(!!state.busy) + ">" + (device.connected ? "Disconnect" : "Connect") + "</button>")
@@ -407,7 +480,48 @@ function handleDeviceListClick(event) {
  const operation = button.dataset.bluetoothOperation
  const address = button.dataset.address || ""
  if (operation === "forget" && !window.confirm("Forget " + (button.dataset.deviceName || "this device") + "?")) return
-  request(operation, { address })
+  if (operation === "elm_open") {
+    openElm(address)
+  } else {
+    request(operation, { address })
+  }
+}
+
+function renderElmPanel() {
+  if (!state.elmAddress) return ""
+  return html`
+    <section class="bluetoothElmPanel">
+      <div class="bluetoothElmHeader">
+        <div>
+          <h3>ELM327</h3>
+          <p>${() => state.elmName || state.elmAddress}</p>
+        </div>
+        <strong>${() => state.elmAdapter || "Connecting…"}</strong>
+      </div>
+      <div class="bluetoothElmActions">
+        <button disabled="${() => !state.offroad || !!state.busy}" @click="${readElmCodes}">Read Codes</button>
+        <button class="bluetoothSecondaryButton" disabled="${() => !!state.busy}" @click="${closeElm}">Close</button>
+      </div>
+      ${() => state.elmCodes !== null ? html`
+        <div class="bluetoothElmCodes">
+          <strong>Stored Codes</strong>
+          <p>${() => state.elmCodes.length ? state.elmCodes.join(" · ") : "No stored codes reported."}</p>
+        </div>
+      ` : ""}
+      <div class="bluetoothElmCommand">
+        <label for="bluetoothElmCommand">Command</label>
+        <div>
+          <input id="bluetoothElmCommand" type="text" autocomplete="off" placeholder="ATI"
+                 disabled="${() => !state.offroad || !!state.busy}" />
+          <button disabled="${() => !state.offroad || !!state.busy}" @click="${sendElmCommand}">Send</button>
+        </div>
+      </div>
+      <div class="bluetoothElmResponse">
+        <strong>Response</strong>
+        <pre>${() => state.elmResponse || "—"}</pre>
+      </div>
+    </section>
+  `
 }

 export function Bluetooth() {
@@ -440,6 +554,7 @@ export function Bluetooth() {
          <span>The test sound is sent at NOW. The audible gap is Bluetooth latency.</span>
        </div>
      ` : ""}
+      ${() => renderElmPanel()}

      <div class="bluetoothToolbar">
        <button disabled="${() => !state.offroad || !state.enabled || !!state.busy}"
@@ -132,7 +132,15 @@ class FakeBluetoothClient:

  def call(self, command, **payload):
    self.calls.append((command, payload))
-    return {"audio_test_delay_ms": 3000} if command == "test_audio" else {}
+    if command == "test_audio":
+      return {"audio_test_delay_ms": 3000}
+    if command == "elm_open":
+      return {"adapter": "Fake ELM327"}
+    if command == "elm_command":
+      return {"response": "OK"}
+    if command == "elm_read_dtcs":
+      return {"codes": ["P0133"], "raw": "43 01 33"}
+    return {}


 def test_bluetooth_status_api(monkeypatch):
@@ -203,6 +211,26 @@ def test_bluetooth_api_dispatches_operations(monkeypatch):
  ]


+def test_bluetooth_api_dispatches_elm_payload_and_allows_close_onroad(monkeypatch):
+  FakeBluetoothClient.calls = []
+  client, fake_params = _params_client(monkeypatch, {"IsOffroad": True}, "mici")
+  monkeypatch.setattr(the_galaxy, "BluetoothClient", FakeBluetoothClient)
+  address = "00:11:22:33:44:55"
+
+  assert client.post("/api/bluetooth/elm_open", json={"address": address}).get_json()["adapter"] == "Fake ELM327"
+  assert client.post("/api/bluetooth/elm_command", json={"address": address, "value": "ATI"}).get_json()["response"] == "OK"
+  assert client.post("/api/bluetooth/elm_read_dtcs", json={"address": address}).get_json()["codes"] == ["P0133"]
+  fake_params.values["IsOffroad"] = False
+  assert client.post("/api/bluetooth/elm_close", json={"address": address}).status_code == 200
+
+  assert FakeBluetoothClient.calls == [
+    ("elm_open", {"address": address}),
+    ("elm_command", {"address": address, "value": "ATI"}),
+    ("elm_read_dtcs", {"address": address}),
+    ("elm_close", {"address": address}),
+  ]
+
+
 def test_wheel_controls_status_includes_favorite_slots(monkeypatch):
  client, _ = _params_client(monkeypatch, {"IsOffroad": True, "FavoriteSlots": []}, "mici")
  monkeypatch.setattr(the_galaxy, "wheel_control_status", lambda *_args: {"available": True, "mappings": [], "devices": []})
@@ -4994,11 +4994,16 @@ def setup(app):
      "select_audio": "select_audio",
      "test_audio": "test_audio",
      "pairing_response": "pairing_response",
+      "elm_open": "elm_open",
+      "elm_close": "elm_close",
+      "elm_command": "elm_command",
+      "elm_read_dtcs": "elm_read_dtcs",
    }
    command = commands.get(operation)
    if command is None:
      return jsonify({"error": "Unknown Bluetooth operation."}), 404
-    offroad_only = {"power", "scan", "stop_scan", "pair", "forget", "test_audio", "pairing_response"}
+    offroad_only = {"power", "scan", "stop_scan", "pair", "forget", "test_audio", "pairing_response",
+                    "elm_open", "elm_command", "elm_read_dtcs"}
    if operation in offroad_only and not params.get_bool("IsOffroad"):
      return jsonify({"error": "Bluetooth settings can only be changed offroad."}), 409

@@ -5016,6 +5021,8 @@ def setup(app):
      payload["address"] = str(data.get("address", ""))
      if not payload["address"] and command != "select_audio":
        return jsonify({"error": "Bluetooth device address is required."}), 400
+      if command == "elm_command":
+        payload["value"] = str(data.get("value", ""))
    try:
      client = BluetoothClient(timeout=10.0)
      if command == "set_power":
@@ -97,6 +97,32 @@ class BluetoothManager:
              self._operations.pop(normalized_address, None)
    threading.Thread(target=worker, daemon=True).start()

+  def _run_result(self, fn, *args, operation: str = "", address: str = "", callback=None) -> None:
+    normalized_address = address.upper()
+    if normalized_address:
+      with self._lock:
+        self._operations[normalized_address] = operation
+
+    def worker():
+      result = None
+      error = None
+      try:
+        result = fn(*args)
+      except Exception as exception:
+        error = str(exception)
+      finally:
+        if normalized_address:
+          with self._lock:
+            if self._operations.get(normalized_address) == operation:
+              self._operations.pop(normalized_address, None)
+      if callback is not None:
+        callback(result, error)
+      elif error:
+        with self._lock:
+          self._operation_error = error
+
+    threading.Thread(target=worker, daemon=True).start()
+
  def set_power(self, enabled: bool) -> None:
    with self._lock:
      if self._power_pending:
@@ -146,5 +172,17 @@ class BluetoothManager:
          self._audio_test_deadline = 0.0
    threading.Thread(target=worker, daemon=True).start()

+  def elm_open(self, address: str, callback=None) -> None:
+    self._run_result(self._client.elm_open, address, operation="elm_open", address=address, callback=callback)
+
+  def elm_close(self, address: str, callback=None) -> None:
+    self._run_result(self._client.elm_close, address, operation="elm_close", address=address, callback=callback)
+
+  def elm_command(self, address: str, value: str, callback=None) -> None:
+    self._run_result(self._client.elm_command, address, value, operation="elm_command", address=address, callback=callback)
+
+  def elm_read_dtcs(self, address: str, callback=None) -> None:
+    self._run_result(self._client.elm_read_dtcs, address, operation="elm_read_dtcs", address=address, callback=callback)
+
  def respond(self, prompt_id: str, accepted: bool, value: str = "") -> None:
    self._run(self._client.respond, prompt_id, accepted, value)
@@ -2,6 +2,7 @@ from __future__ import annotations

 from dataclasses import dataclass
 from functools import partial
+import threading

 import pyray as rl

@@ -14,7 +15,7 @@ from openpilot.system.ui.widgets import DialogResult, Widget
 from openpilot.system.ui.widgets.button import Button, ButtonStyle
 from openpilot.system.ui.widgets.confirm_dialog import ConfirmDialog, alert_dialog
 from openpilot.system.ui.widgets.keyboard import Keyboard
-from openpilot.system.ui.widgets.label import gui_label
+from openpilot.system.ui.widgets.label import gui_label, gui_text_box
 from openpilot.system.ui.widgets.option_dialog import MultiOptionDialog
 from openpilot.system.ui.widgets.toggle import Toggle

@@ -45,11 +46,15 @@ def device_status_text(device: BluetoothDevice, operation: str, selected_audio:
    capabilities.append(tr("audio output") if selected_audio.upper() == device.address.upper() else tr("audio"))
  if device.controller:
    capabilities.append(tr("controller"))
+  if device.serial:
+    capabilities.append(tr("serial"))
  capability_text = " / ".join(capabilities)

  if device.connected:
    return tr("Connected") + (f" / {capability_text}" if capability_text else "")
  if device.paired:
+    if device.serial:
+      return tr("Paired - tap to use ELM327") + (f" / {capability_text}" if capability_text else "")
    return tr("Paired - tap to connect")
  return tr("Tap to pair") + (f" / {capability_text}" if capability_text else "")

@@ -58,6 +63,8 @@ def device_action_allowed(device: BluetoothDevice, operation: str, offroad: bool
  """Mirror the daemon's operation policy before a row can receive a tap."""
  if operation:
    return False
+  if device.serial and not offroad:
+    return False
  if not offroad and not device.paired:
    return False
  return True
@@ -175,6 +182,139 @@ class BluetoothAudioTestDialog(Widget):
    self._done_button.render(button_rect)


+class ELM327Dialog(Widget):
+  """Ephemeral, offroad-only ELM327 controls opened from the Bluetooth panel."""
+  def __init__(self, manager: BluetoothManager, device: BluetoothDevice):
+    super().__init__()
+    self._manager = manager
+    self._address = device.address
+    self._name = device.name
+    self._state_lock = threading.Lock()
+    self._closed = False
+    self._connecting = True
+    self._adapter = ""
+    self._response = ""
+    self._codes: list[str] | None = None
+    self._error = ""
+    self._keyboard = Keyboard(max_text_size=256, min_text_size=1, password_mode=False)
+    self._read_button = Button(tr("Read Codes"), self._read_codes, button_style=ButtonStyle.PRIMARY, font_size=42)
+    self._command_button = Button(tr("Send Command"), self._send_command, button_style=ButtonStyle.NORMAL, font_size=42)
+    self._done_button = Button(tr("Done"), gui_app.pop_widget, button_style=ButtonStyle.NORMAL, font_size=42)
+
+  def show_event(self):
+    super().show_event()
+    self._manager.elm_open(self._address, callback=self._on_open_result)
+
+  def hide_event(self):
+    with self._state_lock:
+      self._closed = True
+    self._manager.elm_close(self._address)
+    super().hide_event()
+
+  def _on_open_result(self, result, error):
+    should_close = False
+    with self._state_lock:
+      self._connecting = False
+      if error:
+        self._error = error
+      else:
+        self._adapter = str((result or {}).get("adapter", ""))
+        self._response = ""
+        self._codes = None
+        should_close = self._closed
+    if should_close:
+      self._manager.elm_close(self._address)
+
+  def _on_command_result(self, command: str, result, error):
+    with self._state_lock:
+      if error:
+        self._error = error
+      else:
+        response = str((result or {}).get("response", ""))
+        self._response = f"{command}\n{response}".strip()
+
+  def _on_read_result(self, result, error):
+    with self._state_lock:
+      if error:
+        self._error = error
+      else:
+        self._codes = [str(code) for code in (result or {}).get("codes", [])]
+        self._response = str((result or {}).get("raw", ""))
+
+  def _send_command(self):
+    with self._state_lock:
+      if self._connecting or self._error or self._closed:
+        return
+    self._keyboard.reset(min_text_size=1)
+    self._keyboard.set_title(tr("Send ELM327 command"), tr("Raw commands are available offroad only."))
+    self._keyboard.set_callback(self._on_command_entered)
+    gui_app.push_widget(self._keyboard)
+
+  def _on_command_entered(self, result: DialogResult):
+    command = self._keyboard.text.strip()
+    self._keyboard.clear()
+    if result != DialogResult.CONFIRM or not command:
+      return
+    with self._state_lock:
+      self._response = f"{command}\nSending..."
+      self._error = ""
+    self._manager.elm_command(self._address, command, callback=partial(self._on_command_result, command))
+
+  def _read_codes(self):
+    with self._state_lock:
+      self._codes = None
+      self._error = ""
+    self._manager.elm_read_dtcs(self._address, callback=self._on_read_result)
+
+  def _update_state(self):
+    with self._state_lock:
+      ready = bool(self._adapter) and not self._connecting and not self._error and not self._closed
+    status = self._manager.status
+    operation = self._manager.operation_for(self._address)
+    enabled = ready and status.offroad and not operation
+    self._read_button.set_enabled(enabled)
+    self._command_button.set_enabled(enabled)
+    self._done_button.set_enabled(True)
+
+  def _render(self, rect: rl.Rectangle):
+    dialog_rect = rl.Rectangle(rect.x + 90, rect.y + 60, rect.width - 180, rect.height - 120)
+    rl.draw_rectangle_rounded(dialog_rect, 0.03, 20, DIALOG_BACKGROUND)
+
+    with self._state_lock:
+      connecting = self._connecting
+      adapter = self._adapter
+      response = self._response
+      codes = self._codes
+      error = self._error
+    gui_label(rl.Rectangle(dialog_rect.x + 45, dialog_rect.y + 35, dialog_rect.width - 90, 70), tr("ELM327"),
+              font_size=62, font_weight=FontWeight.BOLD)
+    gui_label(rl.Rectangle(dialog_rect.x + 45, dialog_rect.y + 105, dialog_rect.width - 90, 52), self._name,
+              font_size=42, color=TEXT_SECONDARY)
+    identity = tr("Connecting...") if connecting else adapter
+    gui_label(rl.Rectangle(dialog_rect.x + 45, dialog_rect.y + 157, dialog_rect.width - 90, 52), identity,
+              font_size=42, color=TEXT_CONNECTED if adapter else TEXT_SECONDARY)
+
+    button_y = dialog_rect.y + 225
+    button_width = (dialog_rect.width - 135) / 3
+    self._read_button.render(rl.Rectangle(dialog_rect.x + 45, button_y, button_width, 90))
+    self._command_button.render(rl.Rectangle(dialog_rect.x + 60 + button_width, button_y, button_width, 90))
+    self._done_button.render(rl.Rectangle(dialog_rect.x + 75 + button_width * 2, button_y, button_width, 90))
+
+    if error:
+      gui_text_box(rl.Rectangle(dialog_rect.x + 45, dialog_rect.y + 330, dialog_rect.width - 90, 80), error,
+                   font_size=38, color=rl.Color(255, 150, 150, 255))
+    if codes is not None:
+      gui_label(rl.Rectangle(dialog_rect.x + 45, dialog_rect.y + 410, 300, 45), tr("Stored Codes"),
+                font_size=38, font_weight=FontWeight.BOLD)
+      code_text = "\n".join(codes) if codes else tr("No stored codes reported.")
+      gui_text_box(rl.Rectangle(dialog_rect.x + 45, dialog_rect.y + 455, dialog_rect.width - 90, 95), code_text,
+                   font_size=38, color=TEXT_SECONDARY)
+    gui_label(rl.Rectangle(dialog_rect.x + 45, dialog_rect.y + 560, 300, 45), tr("Response"),
+              font_size=38, font_weight=FontWeight.BOLD)
+    gui_text_box(rl.Rectangle(dialog_rect.x + 45, dialog_rect.y + 605, dialog_rect.width - 90, dialog_rect.height - 650),
+                 response or "—", font_size=36, color=TEXT_SECONDARY)
+
+
 class BluetoothManagerUI(Widget):
  """Big UI Bluetooth settings panel backed by the existing Bluetooth manager daemon."""
  def __init__(self, manager: BluetoothManager):
@@ -264,6 +404,8 @@ class BluetoothManagerUI(Widget):
      return
    if not device.paired:
      self._manager.pair(device.address)
+    elif device.serial:
+      gui_app.push_widget(ELM327Dialog(self._manager, device))
    elif not device.connected:
      self._manager.connect(device.address)
    else:
@@ -4,7 +4,7 @@ inputs:
  python-version:
    description: 'Python version to use'
    required: false
-    default: '3.14'
+    default: '' # if you don't set a version, the native python version will be used
  key:
    description: 'Key for the python cache'
    required: false
@@ -41,12 +41,12 @@ inputs:
    description: "Install LLVM?"
    required: false
    default: 'false'
-  qemu:
-    description: "Install qemu?"
+  tinydreno:
+    description: "Install tinydreno"
    required: false
    default: 'false'
-  ninja:
-    description: "Install ninja?"
+  qemu:
+    description: "Install qemu"
    required: false
    default: 'false'
 runs:
@@ -59,18 +59,18 @@ runs:
        echo "OMP_NUM_THREADS=1" >> "$GITHUB_ENV"
        # no buffers should be over 300MB in CI
        echo "MAX_BUFFER_SIZE=300000000" >> "$GITHUB_ENV"
-        if [[ "$RUNNER_OS" == "Linux" ]]; then
-          echo "VIRTUAL_ENV=/opt/venv/${{ inputs.python-version }}" >> "$GITHUB_ENV"
-          echo "UV_PYTHON_INSTALL_DIR=/opt/python" >> "$GITHUB_ENV"
-        else
-          echo "VIRTUAL_ENV=${{ github.workspace }}/.venv" >> "$GITHUB_ENV"
-        fi

    - name: Set up uv
      uses: astral-sh/setup-uv@08807647e7069bb48b6ef5acd8ec9567f424441b
      with:
        enable-cache: 'false' # see below for manual caching

+    - name: Set up Python ${{ inputs.python-version }}
+      uses: actions/setup-python@v6
+      if: inputs.python-version != ''
+      with:
+        python-version: ${{ inputs.python-version }}
+
    # **** Caching packages ****

    - name: Cache Python packages (PR)
@@ -109,15 +109,15 @@ runs:
      if: inputs.deps != ''
      shell: bash
      run: |
-        uv venv --allow-existing --python ${{ inputs.python-version }} "$VIRTUAL_ENV"
+        uv venv .venv
        DEPS="${{ inputs.deps }}"
-        uv pip install --python "$VIRTUAL_ENV" -e ".[${DEPS// /,}]" ${{ inputs.pydeps }} --torch-backend cpu --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/Triton-Nightly/pypi/simple/
+        uv pip install --python .venv -e ".[${DEPS// /,}]" ${{ inputs.pydeps }} --torch-backend cpu --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/Triton-Nightly/pypi/simple/
    - name: Install dependencies in venv (without extra)
      if: inputs.deps == ''
      shell: bash
      run: |
-        uv venv --allow-existing --python ${{ inputs.python-version }} "$VIRTUAL_ENV"
-        uv pip install --python "$VIRTUAL_ENV" -e . ${{ inputs.pydeps }}
+        uv venv .venv
+        uv pip install --python .venv -e . ${{ inputs.pydeps }}
    - name: Prune uv cache
      if: github.event_name != 'pull_request'
      shell: bash
@@ -125,15 +125,16 @@ runs:
    - name: Configure venv
      shell: bash
      run: |
+        echo "VIRTUAL_ENV=${{ github.workspace }}/.venv" >> "$GITHUB_ENV"
        if [[ "$RUNNER_OS" == "Windows" ]]; then
-          echo "$VIRTUAL_ENV/Scripts" >> "$GITHUB_PATH"
+          echo "${{ github.workspace }}/.venv/Scripts" >> "$GITHUB_PATH"
        else
-          echo "$VIRTUAL_ENV/bin" >> "$GITHUB_PATH"
+          echo "${{ github.workspace }}/.venv/bin" >> "$GITHUB_PATH"
        fi

    # ******************* apt *******************
    - name: Setup apt
-      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true')
+      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true')
      shell: bash
      run: |
        sudo mkdir -p /var/cache/apt/archives
@@ -161,7 +162,7 @@ runs:
        echo "deb http://apt.llvm.org/$(lsb_release -cs)/ llvm-toolchain-$(lsb_release -cs)-20 main" | sudo tee /etc/apt/sources.list.d/llvm.list

    - name: Compute Package List + Hash
-      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true')
+      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true')
      id: apt-pkgs
      shell: bash
      run: |
@@ -186,37 +187,25 @@ runs:
        if [[ "${{ inputs.qemu }}" == "true" ]]; then
          pkgs+=" qemu-user-static"
        fi
-        # **** ninja ****
-        if [[ "${{ inputs.ninja }}" == "true" ]]; then
-          pkgs+=" ninja-build"
-        fi

        echo "pkgs=$pkgs" >> "$GITHUB_OUTPUT"
        echo "hash=$(echo -n "$pkgs" | sha256sum | cut -d' ' -f1)" >> "$GITHUB_OUTPUT"

-        installed=true
-        for pkg in $pkgs; do
-          info=$(dpkg-query -W -f='${db:Status-Abbrev} ${Version}' "$pkg" 2> /dev/null || true)
-          echo "${pkg}: ${info:-not in dpkg database}"
-          [[ "$info" == ii* ]] || installed=false
-        done
-        echo "installed=$installed" >> "$GITHUB_OUTPUT"
-
    - name: Cache apt (PR)
-      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true') && github.event_name == 'pull_request' && steps.apt-pkgs.outputs.installed == 'false'
+      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true') && github.event_name == 'pull_request'
      uses: actions/cache/restore@v5
      with:
        path: /var/cache/apt/archives/
        key: ${{ runner.os }}-${{ runner.arch }}-apt-${{ steps.apt-pkgs.outputs.hash }}-${{ env.CACHE_VERSION }}
    - name: Cache apt
-      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true') && github.event_name != 'pull_request' && steps.apt-pkgs.outputs.installed == 'false'
+      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true') && github.event_name != 'pull_request'
      uses: actions/cache@v5
      with:
        path: /var/cache/apt/archives/
        key: ${{ runner.os }}-${{ runner.arch }}-apt-${{ steps.apt-pkgs.outputs.hash }}-${{ env.CACHE_VERSION }}

    - name: Run apt Update + Install
-      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true') && steps.apt-pkgs.outputs.installed == 'false'
+      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true')
      shell: bash
      run: |
        sudo apt -qq update || true
@@ -288,6 +277,12 @@ runs:
      shell: bash
      run: brew install llvm@20

+    # *** tinydreno ***
+    - name: Install tinydreno (linux)
+      if: inputs.tinydreno == 'true' && runner.os == 'Linux'
+      shell: bash
+      run: sudo curl -fL https://github.com/sirhcm/tinydreno/raw/refs/heads/master/libllvm-qcom.so -o /usr/lib/libllvm-qcom.so
+
    # *** OpenCL ***
    - name: Install rusticl
      if: inputs.opencl == 'true'
@@ -35,15 +35,15 @@ jobs:
        key: 'autogen'
        amd: 'true'
        llvm: 'true'
-        deps: 'autogen'
+        pydeps: 'pyyaml mako'
    - name: Install autogen support packages
      run: sudo apt-get install -y --no-install-recommends libclang-20-dev llvm-20-dev hip-dev libusb-1.0-0-dev libdrm-dev liburing-dev
    - name: Regenerate autogen files
      run: |
-        find tinygrad/runtime/autogen -type f -name "*.py" -not -path "*/amd/*" -not -name "__init__.py" -not -name "metal.py" -not -name "iokit.py" -not -name "corefoundation.py" -not -name "libclang.py" -delete
+        find tinygrad/runtime/autogen -type f -name "*.py" -not -path "*/amd/*" -not -name "__init__.py" -not -name "comgr.py" -not -name "metal.py" -not -name "iokit.py" -not -name "corefoundation.py" -not -name "libclang.py" -delete
        python3 -c "from tinygrad.runtime.autogen import opencl"
        python3 -c "from tinygrad.runtime.autogen import cuda, nvrtc, nvjitlink, nv_570, nv_580, nv_610, nv"
-        python3 -c "from tinygrad.runtime.autogen import comgr, comgr_3, hsa, hip, amd_gpu, sqtt, rocprof, amdgpu_kd, amdgpu_drm"
+        python3 -c "from tinygrad.runtime.autogen import comgr_3, hsa, hip, amd_gpu, sqtt, rocprof, amdgpu_kd, amdgpu_drm"
        python3 -c "from tinygrad.runtime.autogen.am import *"
        python3 -c "from tinygrad.runtime.autogen.nv_regs import *"
        python3 -c "from tinygrad.runtime.autogen import libc, kfd, io_uring, pci, vfio"
@@ -54,7 +54,7 @@ jobs:
        python3 -c "from tinygrad.runtime.autogen import mesa"
        python3 -c "from tinygrad.runtime.autogen import avcodec"
        python3 -c "from tinygrad.runtime.autogen import llvm_qcom"
-        python3 -c "from tinygrad.runtime.autogen import mlx5, bnxt"
+        python3 -c "from tinygrad.runtime.autogen import mlx5"
        python3 -c "from tinygrad.runtime.autogen import ggml_common"
        REGEN=1 python3 -c "from tinygrad.runtime.autogen import libclang"
    - name: Check for differences
@@ -102,3 +102,42 @@ jobs:
      with:
        name: autogen-macos-patch
        path: autogen-macos.patch
+
+  autogen-comgr-2:
+    name: In-tree Autogen (comgr 2)
+    runs-on: ubuntu-24.04
+    timeout-minutes: 15
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v6
+    - name: Setup Environment
+      uses: ./.github/actions/setup-tinygrad
+      with:
+        key: 'autogen-comgr'
+    - name: Install autogen support packages
+      run: |
+        wget https://repo.radeon.com/rocm/rocm.gpg.key -O - | gpg --dearmor | sudo tee /etc/apt/keyrings/rocm.gpg > /dev/null
+        sudo tee /etc/apt/sources.list.d/rocm.list <<EOF
+        deb [arch=amd64 signed-by=/etc/apt/keyrings/rocm.gpg] https://repo.radeon.com/rocm/apt/6.2 $(lsb_release -cs) main
+        EOF
+        echo -e 'Package: *\nPin: release o=repo.radeon.com\nPin-Priority: 600' | sudo tee /etc/apt/preferences.d/rocm-pin-600
+        sudo apt -qq update || true
+        sudo apt-get install -y --no-install-recommends libclang-20-dev comgr
+    - name: Regenerate autogen files
+      run: |
+        rm tinygrad/runtime/autogen/comgr.py
+        python3 -c "from tinygrad.runtime.autogen import comgr"
+    - name: Check for differences
+      run: |
+        if ! git diff --quiet; then
+          git diff
+          git diff > autogen-comgr2.patch
+          echo "Autogen mismatch detected. Patch available at: ${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}#artifacts"
+          exit 1
+        fi
+    - name: Upload patch artifact
+      if: failure()
+      uses: actions/upload-artifact@v7
+      with:
+        name: autogen-comgr2-patch
+        path: autogen-comgr2.patch
@@ -88,13 +88,13 @@ jobs:
      fail-fast: false
      matrix:
        dev: ['METAL', 'AMD', 'NV']
-    timeout-minutes: 30
+    timeout-minutes: 60
    defaults:
      run:
        shell: bash -e -o pipefail {0}
    env:
      DEV: ${{ matrix.dev }}
-      HCQ2: '0'
+      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
    if: github.repository_owner == 'tinygrad'
    steps:
    - name: Checkout Code
@@ -102,11 +102,16 @@ jobs:
    - name: Setup (AMD)
      if: ${{ matrix.dev == 'AMD' }}
      run: |
-        ./extra/hcq/hcq_smi.py amd rmmod --expect
-        ./extra/hcq/hcq_smi.py amd kill_pids --sudoless
+        ./extra/amdpci/setup_python_cap.sh
+        ./extra/hcq/hcq_smi.py amd rmmod
+        ./extra/hcq/hcq_smi.py amd kill_pids
    - name: Setup (NV)
      if: ${{ matrix.dev == 'NV' }}
-      run: lsof -tQ /dev/nvidia* | { xargs -r kill -9 || true; }
+      run: sudo lsof -tQ /dev/nvidia* | { xargs -r sudo kill -9 || true; }
+    - name: Symlink models and datasets
+      run: |
+        mkdir -p weights
+        ln -s /raid/weights/LLaMA-3 weights/LLaMA-3
    - name: setup staging db
      if: github.ref == 'refs/heads/update_benchmark_staging'
      run: |
@@ -116,14 +121,18 @@ jobs:
      run: python3 test/external/process_replay/reset.py
    - name: Run llama3.2
      run: BENCHMARK_LOG=llama32_3b-f16 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 -m tinygrad.llm -m llama3.2:3b-f16 --benchmark --warmup
-    - name: Run qwen3.8
-      # qwen3.8:27b doesn't fit on mac
+    - name: Run qwen3.6
+      # qwen3.6:35b-a3b doesn't fit on mac
      if: ${{ matrix.dev != 'METAL' }}
-      run: BENCHMARK_LOG=qwen38_27b JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 -m tinygrad.llm -m qwen3.8:27b --benchmark --warmup
+      run: BENCHMARK_LOG=qwen36_35b-a3b JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 -m tinygrad.llm -m qwen3.6:35b-a3b --benchmark --warmup
    - name: Run olmoe
      # just metal for now
      if: ${{ matrix.dev == 'METAL' }}
      run: BENCHMARK_LOG=olmoe JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 -m tinygrad.llm -m olmoe --benchmark --warmup
+    - name: Run LLaMA-3 8B on 4 GPUs with BEAM
+      # only run on machines with multiple gpus
+      if: ${{ matrix.dev != 'METAL' }}
+      run: BENCHMARK_LOG=llama3_beam_4gpu JITBEAM=2 IGNORE_BEAM_CACHE=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 4 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0
    - name: Run process replay tests
      uses: ./.github/actions/process-replay

@@ -134,13 +143,13 @@ jobs:
      fail-fast: false
      matrix:
        dev: ['METAL', 'AMD', 'NV']
-    timeout-minutes: 10
+    timeout-minutes: 60
    defaults:
      run:
        shell: bash -e -o pipefail {0}
    env:
      DEV: ${{ matrix.dev }}
-      HCQ2: '0'
+      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
    if: github.repository_owner == 'tinygrad'
    steps:
    - name: Checkout Code
@@ -148,11 +157,12 @@ jobs:
    - name: Setup (AMD)
      if: ${{ matrix.dev == 'AMD' }}
      run: |
-        ./extra/hcq/hcq_smi.py amd rmmod --expect
-        ./extra/hcq/hcq_smi.py amd kill_pids --sudoless
+        ./extra/amdpci/setup_python_cap.sh
+        ./extra/hcq/hcq_smi.py amd rmmod
+        ./extra/hcq/hcq_smi.py amd kill_pids
    - name: Setup (NV)
      if: ${{ matrix.dev == 'NV' }}
-      run: lsof -tQ /dev/nvidia* | { xargs -r kill -9 || true; }
+      run: sudo lsof -tQ /dev/nvidia* | { xargs -r sudo kill -9 || true; }
    - name: setup staging db
      if: github.ref == 'refs/heads/update_benchmark_staging'
      run: |
@@ -172,6 +182,10 @@ jobs:
      # slow on metal
      if: ${{ matrix.dev != 'METAL' }}
      run: time BENCHMARK_LOG=cifar DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
+    - name: Run full CIFAR training steps w 6 GPUS
+      # only run on machines with multiple gpus
+      if: ${{ matrix.dev != 'METAL' }}
+      run: time BENCHMARK_LOG=cifar_6gpu CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
    - name: Run process replay tests
      uses: ./.github/actions/process-replay

@@ -182,13 +196,13 @@ jobs:
      fail-fast: false
      matrix:
        dev: ['AMD', 'NV']
-    timeout-minutes: 5
+    timeout-minutes: 60
    defaults:
      run:
        shell: bash -e -o pipefail {0}
    env:
      DEV: ${{ matrix.dev }}
-      HCQ2: '0'
+      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
    if: github.repository_owner == 'tinygrad'
    steps:
    - name: Checkout Code
@@ -196,11 +210,12 @@ jobs:
    - name: Setup (AMD)
      if: ${{ matrix.dev == 'AMD' }}
      run: |
-        ./extra/hcq/hcq_smi.py amd rmmod --expect
-        ./extra/hcq/hcq_smi.py amd kill_pids --sudoless
+        ./extra/amdpci/setup_python_cap.sh
+        ./extra/hcq/hcq_smi.py amd rmmod
+        ./extra/hcq/hcq_smi.py amd kill_pids
    - name: Setup (NV)
      if: ${{ matrix.dev == 'NV' }}
-      run: lsof -tQ /dev/nvidia* | { xargs -r kill -9 || true; }
+      run: sudo lsof -tQ /dev/nvidia* | { xargs -r sudo kill -9 || true; }
    - name: Symlink models and datasets
      run: |
        mkdir -p extra/datasets
@@ -212,8 +227,15 @@ jobs:
        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
    - name: reset process replay
      run: test/external/process_replay/reset.py
+    - name: Run MLPerf resnet eval on training data
+      run: time BENCHMARK_LOG=resnet_eval MODEL=resnet python3 examples/mlperf/model_eval.py
    - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
      run: BENCHMARK_LOG=resnet_10steps DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py
+    - name: Run 10 MLPerf ResNet50 training steps (6 gpu)
+      run: BENCHMARK_LOG=resnet_10steps_6gpu CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py
+    - name: Run 10 MLPerf Bert training steps (6 gpu)
+      # TODO: remove BERT_LAYERS once scheduler is fast
+      run: BENCHMARK_LOG=bert_10steps_6gpu CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=72 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
    - name: Run process replay tests
      uses: ./.github/actions/process-replay

@@ -224,13 +246,13 @@ jobs:
      fail-fast: false
      matrix:
        dev: ['METAL', 'AMD', 'NV']
-    timeout-minutes: 15
+    timeout-minutes: 60
    defaults:
      run:
        shell: bash -e -o pipefail {0}
    env:
      DEV: ${{ matrix.dev }}
-      HCQ2: '0'
+      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
    if: github.repository_owner == 'tinygrad'
    steps:
    - name: Checkout Code
@@ -238,11 +260,12 @@ jobs:
    - name: Setup (AMD)
      if: ${{ matrix.dev == 'AMD' }}
      run: |
-        ./extra/hcq/hcq_smi.py amd rmmod --expect
-        ./extra/hcq/hcq_smi.py amd kill_pids --sudoless
+        ./extra/amdpci/setup_python_cap.sh
+        ./extra/hcq/hcq_smi.py amd rmmod
+        ./extra/hcq/hcq_smi.py amd kill_pids
    - name: Setup (NV)
      if: ${{ matrix.dev == 'NV' }}
-      run: lsof -tQ /dev/nvidia* | { xargs -r kill -9 || true; }
+      run: sudo lsof -tQ /dev/nvidia* | { xargs -r sudo kill -9 || true; }
    - name: setup staging db
      if: github.ref == 'refs/heads/update_benchmark_staging'
      run: |
@@ -262,59 +285,6 @@ jobs:
    - name: Run process replay tests
      uses: ./.github/actions/process-replay

-  multigpubenchmark:
-    name: Multi-GPU Benchmarks (DEV=${{ matrix.dev }})
-    runs-on: [self-hosted, "${{ matrix.dev == 'AMD' && 'tinybox' || 'tinyboxgreen' }}"]
-    strategy:
-      fail-fast: false
-      matrix:
-        dev: ['AMD', 'NV']
-    timeout-minutes: 20
-    defaults:
-      run:
-        shell: bash -e -o pipefail {0}
-    env:
-      DEV: ${{ matrix.dev }}
-      HCQ2: '0'
-    if: github.repository_owner == 'tinygrad'
-    steps:
-    - name: Checkout Code
-      uses: actions/checkout@v6
-    - name: Setup (AMD)
-      if: ${{ matrix.dev == 'AMD' }}
-      run: |
-        ./extra/hcq/hcq_smi.py amd rmmod --expect
-        ./extra/hcq/hcq_smi.py amd kill_pids --sudoless
-    - name: Setup (NV)
-      if: ${{ matrix.dev == 'NV' }}
-      run: lsof -tQ /dev/nvidia* | { xargs -r kill -9 || true; }
-    - name: Symlink models and datasets
-      run: |
-        mkdir -p weights
-        mkdir -p extra/datasets
-        ln -s /raid/weights/LLaMA-3 weights/LLaMA-3
-        ln -s /raid/datasets/imagenet extra/datasets/imagenet
-    - name: setup staging db
-      if: github.ref == 'refs/heads/update_benchmark_staging'
-      run: |
-        echo "CACHEDB=/tmp/staging.db" >> $GITHUB_ENV
-        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
-    - name: reset process replay
-      run: python3 test/external/process_replay/reset.py
-    - name: Run LLaMA-3 8B on 4 GPUs with BEAM
-      run: BENCHMARK_LOG=llama3_beam_4gpu JITBEAM=2 IGNORE_BEAM_CACHE=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 4 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0
-    - name: Run full CIFAR training steps w 6 GPUS
-      run: time BENCHMARK_LOG=cifar_6gpu CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
-    - name: Run MLPerf resnet eval on training data
-      run: time BENCHMARK_LOG=resnet_eval MODEL=resnet python3 examples/mlperf/model_eval.py
-    - name: Run 10 MLPerf ResNet50 training steps (6 gpu)
-      run: BENCHMARK_LOG=resnet_10steps_6gpu CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py
-    - name: Run 10 MLPerf Bert training steps (6 gpu)
-      # TODO: remove BERT_LAYERS once scheduler is fast
-      run: BENCHMARK_LOG=bert_10steps_6gpu CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=72 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
-    - name: Run process replay tests
-      uses: ./.github/actions/process-replay
-
  tests:
    name: Tests (DEV=${{ matrix.dev }})
    runs-on: [self-hosted, "${{ matrix.dev == 'METAL' && 'macOS' || matrix.dev == 'AMD' && 'tinybox' || 'tinyboxgreen' }}"]
@@ -322,7 +292,7 @@ jobs:
      fail-fast: false
      matrix:
        dev: ['METAL', 'AMD', 'NV']
-    timeout-minutes: 10
+    timeout-minutes: 60
    defaults:
      run:
        shell: bash -e -o pipefail {0}
@@ -335,11 +305,12 @@ jobs:
    - name: Setup (AMD)
      if: ${{ matrix.dev == 'AMD' }}
      run: |
-        ./extra/hcq/hcq_smi.py amd rmmod --expect
-        ./extra/hcq/hcq_smi.py amd kill_pids --sudoless
+        ./extra/amdpci/setup_python_cap.sh
+        ./extra/hcq/hcq_smi.py amd rmmod
+        ./extra/hcq/hcq_smi.py amd kill_pids
    - name: Setup (NV)
      if: ${{ matrix.dev == 'NV' }}
-      run: lsof -tQ /dev/nvidia* | { xargs -r kill -9 || true; }
+      run: sudo lsof -tQ /dev/nvidia* | { xargs -r sudo kill -9 || true; }
    - name: setup staging db
      if: github.ref == 'refs/heads/update_benchmark_staging'
      run: |
@@ -403,7 +374,7 @@ jobs:
      run: python test/external/external_benchmark_multitensor_allreduce.py
    - name: HEVC Decode Benchmark
      if: ${{ matrix.dev == 'NV' }}
-      run: IGNORE_BEAM_CACHE=1 VALIDATE=1 MAX_FRAMES=100 ASSERT_FPS=1400 JITBEAM=1 PYTHONPATH=. python3 extra/hevc/decode.py
+      run: VALIDATE=1 MAX_FRAMES=100 ASSERT_FPS=1400 JITBEAM=1 PYTHONPATH=. python3 extra/hevc/decode.py
    - uses: actions/upload-artifact@v7
      if: ${{ matrix.dev != 'AMD' }}
      with:
@@ -416,7 +387,7 @@ jobs:
  testusbgpu:
    name: UsbGPU Benchmark
    runs-on: [self-hosted, macOS]
-    timeout-minutes: 3
+    timeout-minutes: 10
    defaults:
      run:
        shell: bash -e -o pipefail {0}
@@ -431,66 +402,32 @@ jobs:
        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
    - name: Kill stale pids
      run: |
-        ./extra/hcq/hcq_smi.py amd kill_pids --sudoless
-        ./extra/hcq/hcq_smi.py nv kill_pids --sudoless
+        PYTHONPATH=. ./extra/hcq/hcq_smi.py amd kill_pids
+        PYTHONPATH=. ./extra/hcq/hcq_smi.py nv kill_pids
+    # since sudo is required for usbgpu on macos, do not write bytecode, as some of the files are owned by root
    - name: UsbGPU boot time
-      run: GMMU=0 DEBUG=2 AM_RESET=1 DEV=USB+AMD time python3.11 test/test_tiny.py TestTiny.test_plus
+      run: sudo -E PYTHONDONTWRITEBYTECODE=1 PYTHONPATH=. GMMU=0 DEBUG=2 AM_RESET=1 DEV=USB+AMD time python3.11 test/test_tiny.py TestTiny.test_plus
    - name: UsbGPU tiny tests
-      run: GMMU=0 DEV=USB+AMD python3.11 test/test_tiny.py
+      run: sudo -E PYTHONDONTWRITEBYTECODE=1 PYTHONPATH=. GMMU=0 DEV=USB+AMD python3.11 test/test_tiny.py
    - name: UsbGPU copy speeds
-      run: SIZE=64000000 PYTHONPATH=. GMMU=0 DEV=USB+AMD python3.11 test/external/external_test_usb_asm24.py TestDevCopySpeeds
+      run: sudo -E PYTHONDONTWRITEBYTECODE=1 PYTHONPATH=. GMMU=0 DEV=USB+AMD python3.11 test/external/external_test_usb_asm24.py TestDevCopySpeeds
+    #- name: UsbGPU openpilot test
+    #  run: sudo -E PYTHONPATH=. GMMU=0 DEV=USB+AMD GRAPH_ONE_KERNEL=1 python3.11 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
    - name: UsbGPU (USB4/TB) install script
-      run: sh extra/setup_tinygpu_osx.sh
+      run: PYTHONPATH=. sh extra/setup_tinygpu_osx.sh
    - name: UsbGPU (USB4/TB) boot time
-      run: DEBUG=3 DEV=PCI+NV:NAK time python3.11 test/test_tiny.py TestTiny.test_plus
+      run: PYTHONPATH=. DEBUG=3 DEV=PCI+NV:NAK time python3.11 test/test_tiny.py TestTiny.test_plus
    - name: UsbGPU (USB4/TB) tiny tests
-      run: DEV=PCI+NV:NAK python3.11 test/test_tiny.py
+      run: PYTHONPATH=. DEV=PCI+NV:NAK python3.11 test/test_tiny.py

-  testcomma:
-    strategy:
-      matrix:
-        dev: ['QCOM', 'QCOM:IR3']
-        version: ['0.11.0', '0.11.2']
-        model: ['vision', 'policy', 'supercombo', 'dmonitoring']
-        # exclude non-existent models
-        exclude: [{ version: '0.11.0', model: supercombo }, { version: '0.11.2', model: vision }, { version: '0.11.2', model: policy }]
-        include:
-          - version: '0.11.0'
-            model: vision
-            url: https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/driving_vision.onnx
-            timing: 18
-          - version: '0.11.0'
-            model: policy
-            url: https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/driving_policy.onnx
-            timing: 3.4
-          - version: '0.11.0'
-            model: dmonitoring
-            url: https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/dmonitoring_model.onnx
-            timing: 13
-          - version: '0.11.2'
-            model: supercombo
-            url: https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/433f85f956837606ad1f1cbee4aa7e2158ad23c768dea914b20436c97232741b
-            timing: 28
-          - dev: QCOM:IR3
-            version: '0.11.2'
-            model: supercombo
-            timing: 29
-          - version: '0.11.2'
-            model: dmonitoring
-            url: https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/3e7b31dfbc0a5234f1baf196513b77fc6af12204b8a8ffe8ee0417e48352f316
-            timing: 12.5
-      fail-fast: false
-    name: openpilot ${{ matrix.version }} compile3 ${{ matrix.model }} (DEV=${{ matrix.dev }})
-    runs-on: [self-hosted, Linux, comma4]
-    timeout-minutes: 5
+  testcommalatest:
+    name: comma Benchmark (0.11.2)
+    runs-on: [self-hosted, Linux, comma]
+    timeout-minutes: 12
    defaults:
      run:
        shell: bash -e -o pipefail {0}
    if: github.repository_owner == 'tinygrad'
-    env:
-      DEV: ${{ matrix.dev }}
-      ASSERT_MIN_STEP_TIME: ${{ matrix.timing }}
-      BENCHMARK_LOG: ${{ matrix.dev == 'QCOM:IR3' && 'ir3_' || '' }}openpilot_${{ matrix.version }}_${{ matrix.model }}
    steps:
    - name: Checkout Code
      uses: actions/checkout@v6
@@ -501,10 +438,45 @@ jobs:
        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
    - name: reset process replay
      run: test/external/process_replay/reset.py
-    - name: compile
-      run: FLOAT16=1 IMAGE=1 taskset -c 4-7 python3 examples/openpilot/compile3.py ${{ matrix.url }} openpilot.pkl
-    - name: run pickle
-      run: BENCHMARK_LOG="${BENCHMARK_LOG}_run_pickle" RUN_PICKLE=1 taskset -c 4-7 python3 examples/openpilot/compile3.py - openpilot.pkl
+    - name: openpilot compile3 0.11.2 supercombo
+      run: BENCHMARK_LOG=openpilot_0_11_2_supercombo PYTHONPATH="." ASSERT_MIN_STEP_TIME=26 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/433f85f956837606ad1f1cbee4aa7e2158ad23c768dea914b20436c97232741b
+    - name: openpilot compile3 0.11.2 supercombo (from pickle)
+      run: BENCHMARK_LOG=openpilot_0_11_2_supercombo_run_pickle RUN_PICKLE=1 PYTHONPATH="." ASSERT_MIN_STEP_TIME=26 DEV=QCOM taskset -c 4-7 python3 examples/openpilot/compile3.py
+    - name: IR3 openpilot compile3 0.11.2 supercombo
+      run: BENCHMARK_LOG=ir3_openpilot_0_11_2_supercombo PYTHONPATH="." ASSERT_MIN_STEP_TIME=41 DEV=QCOM:IR3 FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/433f85f956837606ad1f1cbee4aa7e2158ad23c768dea914b20436c97232741b
+    - name: openpilot compile3 0.11.2 dmonitoring
+      run: BENCHMARK_LOG=openpilot_0_11_2_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=11 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/3e7b31dfbc0a5234f1baf196513b77fc6af12204b8a8ffe8ee0417e48352f316
+    - name: Run process replay tests
+      uses: ./.github/actions/process-replay
+
+  testcommaold:
+    name: comma Benchmark (0.11.0)
+    runs-on: [self-hosted, Linux, comma]
+    timeout-minutes: 12
+    defaults:
+      run:
+        shell: bash -e -o pipefail {0}
+    if: github.repository_owner == 'tinygrad'
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v6
+    - name: setup staging db
+      if: github.ref == 'refs/heads/update_benchmark_staging'
+      run: |
+        echo "CACHEDB=/tmp/staging.db" >> $GITHUB_ENV
+        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
+    - name: reset process replay
+      run: test/external/process_replay/reset.py
+    - name: openpilot compile3 0.11.0 driving_vision
+      run: BENCHMARK_LOG=openpilot_0_11_0_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=17 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/driving_vision.onnx
+    - name: openpilot compile3 0.11.0 driving_vision (from pickle)
+      run: BENCHMARK_LOG=openpilot_0_11_0_vision_run_pickle RUN_PICKLE=1 PYTHONPATH="." ASSERT_MIN_STEP_TIME=17 DEV=QCOM taskset -c 4-7 python3 examples/openpilot/compile3.py
+    - name: IR3 openpilot compile3 0.11.0 driving_vision
+      run: BENCHMARK_LOG=ir3_openpilot_0_11_0_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=18 DEV=QCOM:IR3 FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/driving_vision.onnx
+    - name: openpilot compile3 0.11.0 driving_policy
+      run: BENCHMARK_LOG=openpilot_0_11_0_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=3.2 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/driving_policy.onnx
+    - name: openpilot compile3 0.11.0 dmonitoring
+      run: BENCHMARK_LOG=openpilot_0_11_0_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=11 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/dmonitoring_model.onnx
    - name: Run process replay tests
      uses: ./.github/actions/process-replay

@@ -517,6 +489,15 @@ jobs:
        shell: bash -e -o pipefail {0}
    if: github.repository_owner == 'tinygrad'
    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v6
+    - name: setup staging db
+      if: github.ref == 'refs/heads/update_benchmark_staging'
+      run: |
+        echo "CACHEDB=/tmp/staging.db" >> $GITHUB_ENV
+        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
+    - name: reset process replay
+      run: test/external/process_replay/reset.py
    - name: Checkout Code
      uses: actions/checkout@v6
    - name: setup staging db
@@ -529,8 +510,8 @@ jobs:
    - name: benchmark MobileNetV2 on DSP
      run: |
        # generate quantized weights
-        ln -s ~/tinygrad/extra/datasets/imagenet extra/datasets/imagenet
-        ln -s ~/tinygrad/testsig-*.so .
+        ln -s /data/home/tiny/tinygrad/extra/datasets/imagenet extra/datasets/imagenet
+        ln -s /data/home/tiny/tinygrad/testsig-*.so .
        PYTHONPATH=. DEV=CPU QUANT=1 CNT=0 python3 examples/test_onnx_imagenet.py https://github.com/xamcat/mobcat-samples/raw/refs/heads/master/onnx_runtime/InferencingSample/InferencingSample/mobilenetv2-7.onnx /tmp/model.quant.onnx
        # benchmark on DSP with NOOPT=1, the devectorizer has issues
        PYTHONPATH=. DEV=DSP NOOPT=1 CNT=2 DEBUG=2 python3 examples/test_onnx_imagenet.py /tmp/model.quant.onnx
@@ -540,7 +521,7 @@ jobs:
  testcommausbgpubenchmark:
    name: UsbGPU Benchmark (comma)
    runs-on: [self-hosted, Linux, comma4]
-    timeout-minutes: 10
+    timeout-minutes: 20
    defaults:
      run:
        shell: bash -e -o pipefail {0}
@@ -560,7 +541,7 @@ jobs:
    - name: openpilot run_pickle big_driving_supercombo
      run: BENCHMARK_LOG=usbgpu_openpilot_big_driving_supercombo_run_pickle RUN_PICKLE=1 PICKLE_OOB=1 PYTHONPATH="." GMMU=0 DEV=USB+AMD ASSERT_MIN_STEP_TIME=50 python3 examples/openpilot/compile3.py - openpilot.pkl
    - name: Test copy speeds
-      run: SIZE=64000000 PYTHONPATH=. GMMU=0 DEV=USB+AMD python3 test/external/external_test_usb_asm24.py TestDevCopySpeeds
+      run: SIZE=64e6 PYTHONPATH=. GMMU=0 DEV=USB+AMD python3 test/external/external_test_usb_asm24.py TestDevCopySpeeds

  driverbenchmarks:
    name: PCI Driver Benchmark (DEV=${{ matrix.dev }})
@@ -569,7 +550,7 @@ jobs:
      fail-fast: false
      matrix:
        dev: ['AMD', 'NV']
-    timeout-minutes: 5
+    timeout-minutes: 20
    defaults:
      run:
        shell: bash -e -o pipefail {0}
@@ -581,8 +562,9 @@ jobs:
      uses: actions/checkout@v6
    - name: Setup
      run: |
-        ./extra/hcq/hcq_smi.py ${{ matrix.dev }} rmmod --expect
-        ./extra/hcq/hcq_smi.py ${{ matrix.dev }} kill_pids --sudoless
+        ./extra/amdpci/setup_python_cap.sh
+        ./extra/hcq/hcq_smi.py ${{ matrix.dev == 'AMD' && 'amd' || 'nv' }} rmmod
+        ./extra/hcq/hcq_smi.py ${{ matrix.dev == 'AMD' && 'amd' || 'nv' }} kill_pids
        mkdir -p extra/datasets
        ln -s /raid/datasets/imagenet extra/datasets/imagenet
    - name: setup staging db
@@ -617,9 +599,6 @@ jobs:
      run: |
        GRAPH_ONE_KERNEL=1 NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyDefaulttoCPUJit
        GRAPH_ONE_KERNEL=1 NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyCPUtoDefaultJit
-    - name: HEVC Decode Benchmark
-      if: ${{ matrix.dev == 'NV' }}
-      run: IGNORE_BEAM_CACHE=1 VALIDATE=1 MAX_FRAMES=100 ASSERT_FPS=1400 JITBEAM=1 PYTHONPATH=. python3 extra/hevc/decode.py
    - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
      if: ${{ matrix.dev == 'NV' }}
      run: BENCHMARK_LOG=resnet_10steps MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py
@@ -642,12 +621,12 @@ jobs:
  llvmspeed:
    name: LLVM Speed
    runs-on: [self-hosted, Linux, tinyboxrandom]
-    timeout-minutes: 10
+    timeout-minutes: 20
    if: github.repository_owner == 'tinygrad'
    steps:
    - name: Checkout Code
      uses: actions/checkout@v6
    - name: Speed Test
-      run: DEV=CPU:LLVM python3 test/speed/external_test_speed_v_torch.py
+      run: DEV=CPU:LLVM THREADS=0 python3 test/speed/external_test_speed_v_torch.py
    - name: Speed Test (BEAM=2)
-      run: IGNORE_BEAM_CACHE=1 BEAM=2 DEV=CPU:LLVM python3 test/speed/external_test_speed_v_torch.py
+      run: BEAM=2 DEV=CPU:LLVM THREADS=0 python3 test/speed/external_test_speed_v_torch.py
@@ -8,7 +8,7 @@ permissions:
  contents: write
 jobs:
  deploy:
-    runs-on: ubuntu-24.04
+    runs-on: ubuntu-latest
    steps:
      - uses: actions/checkout@v6
      - name: Configure Git Credentials
@@ -166,7 +166,7 @@ jobs:
        uses: ./.github/actions/setup-tinygrad
        with:
          key: windows-${{ matrix.dev }}-minimal
-          deps: testing_minimal
+          deps: testing_unit
          pydeps: ${{ matrix.dev == 'WEBGPU' && 'dawn-python' || '' }}
      - name: Set env
        shell: bash
@@ -179,3 +179,35 @@ jobs:
      - name: Run test_tiny
        shell: bash
        run: python -m pytest -n=auto test/test_tiny.py --durations=20
+
+
+  qcomclcompiletests:
+    name: Compile-only (QCOM CL)
+    runs-on: ubuntu-24.04-arm
+    timeout-minutes: 15
+    steps:
+      - name: Checkout Code
+        uses: actions/checkout@v6
+      - name: Setup Environment
+        uses: ./.github/actions/setup-tinygrad
+        with:
+          key: compile-qcomcl
+          deps: testing_unit
+          tinydreno: 'true'
+      - name: Set env
+        shell: bash
+        run: printf "DEV=NULL:QCOMCL:a630\nNULL_ALLOW_COPYOUT=1" >> $GITHUB_ENV
+      - name: Run test_ops
+        shell: bash
+        run: |
+          python -c "from tinygrad import Device; assert Device.DEFAULT == 'NULL'"
+          DEBUG=4 python3 test/backend/test_ops.py TestOps.test_add
+          python -m pytest -n=auto test/backend/test_ops.py --durations=20
+      - name: Run test_ops (IMAGE)
+        shell: bash
+        env:
+          IMAGE: 1
+          DEV: "NULL:QCOMCL:a630,IMAGE_PITCH_ALIGNMENT=64"
+        run: |
+          DEBUG=4 python test/backend/test_ops.py TestOps.test_gemm | grep read_imagef
+          python -m pytest -n=auto test/backend/test_ops.py --durations=20
@@ -10,7 +10,7 @@ on:

 jobs:
  deploy:
-    runs-on: ubuntu-24.04
+    runs-on: ubuntu-latest
    steps:
    - uses: actions/checkout@v6
    - name: Set up Python
@@ -10,7 +10,7 @@ concurrency:
 jobs:
  checkbranch:
    name: Check PR Branch status
-    runs-on: ubuntu-24.04
+    runs-on: ubuntu-latest
    outputs:
      branchstat: ${{ steps.brstat.outputs.stat}}
    steps:
@@ -44,7 +44,7 @@ jobs:
    permissions:
      contents: read
      pull-requests: write
-    runs-on: ubuntu-24.04
+    runs-on: ubuntu-latest
    needs: checkbranch
    if: needs.checkbranch.outputs.branchstat == 'false'
    steps:
@@ -87,7 +87,7 @@ jobs:
    name: Core Library Line Difference
    permissions:
      pull-requests: write
-    runs-on: ubuntu-24.04
+    runs-on: ubuntu-latest
    needs: checkbranch
    if: needs.checkbranch.outputs.branchstat == 'true'
    steps:
@@ -21,7 +21,7 @@ concurrency:
 jobs:
  docs:
    name: Docs
-    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+    runs-on: &linux ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
    timeout-minutes: 10
    env:
      CHECK_OOB: 0
@@ -31,7 +31,8 @@ jobs:
    - name: Setup Environment
      uses: ./.github/actions/setup-tinygrad
      with:
-        deps: "docs testing_minimal"
+        deps: docs
+        pydeps: "capstone torch"
    - name: Build wheel and show size
      run: |
        uv build --wheel
@@ -60,7 +61,7 @@ jobs:

  torchbackend:
    name: Torch Backend Tests
-    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+    runs-on: *linux
    timeout-minutes: 15
    steps:
    - name: Checkout Code
@@ -72,7 +73,10 @@ jobs:
        deps: testing_unit
        pydeps: "pillow torchvision expecttest"
        llvm: 'true'
-        ninja: 'true'
+    - name: Install ninja
+      run: |
+        sudo apt update || true
+        sudo apt install -y --no-install-recommends ninja-build
    - name: Test ResNet-18
      run: DEBUG=2 python3 extra/torch_backend/example.py
    - name: Test one op in torch tests
@@ -82,26 +86,9 @@ jobs:
    - name: Custom tests
      run: DEV=CPU:LLVM GPUS=4 TINY_BACKEND=1 python3 -m pytest -nauto extra/torch_backend/test.py extra/torch_backend/test_inplace.py extra/torch_backend/test_multigpu.py extra/torch_backend/test_kernel_fusion.py --durations=20

-  torchbackendtrain:
-    name: Torch Backend Training
-    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
-    timeout-minutes: 15
-    steps:
-    - name: Checkout Code
-      uses: actions/checkout@v6
-    - name: Setup Environment
-      uses: ./.github/actions/setup-tinygrad
-      with:
-        key: torch-backend-pillow-torchvision-et-pt
-        deps: testing_unit
-        llvm: 'true'
-        ninja: 'true'
-    - name: Test beautiful_mnist in torch with TINY_BACKEND
-      run: STEPS=20 DEV=CPU TARGET_EVAL_ACC_PCT=90.0 MAX_BUFFER_SIZE=0 TINY_BACKEND=1 python3 examples/other_mnist/beautiful_mnist_torch.py
-
  bepython:
    name: Python Backend
-    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+    runs-on: *linux
    timeout-minutes: 15
    steps:
    - name: Checkout Code
@@ -139,7 +126,7 @@ jobs:

  linter:
    name: Linters
-    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+    runs-on: *linux
    timeout-minutes: 10

    steps:
@@ -170,7 +157,7 @@ jobs:

  nulltest:
    name: Null Tests
-    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+    runs-on: *linux
    timeout-minutes: 15

    steps:
@@ -204,7 +191,7 @@ jobs:

  unittest:
    name: Unit Tests
-    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+    runs-on: *linux
    timeout-minutes: 15

    steps:
@@ -233,7 +220,7 @@ jobs:
    - name: Run process replay tests
      uses: ./.github/actions/process-replay
    - name: Repo line count <= 26000 lines
-      run: MAX_LINE_COUNT=26500 python sz.py
+      run: MAX_LINE_COUNT=26000 python sz.py

  spec:
    strategy:
@@ -241,7 +228,7 @@ jobs:
      matrix:
        group: [1, 2]
    name: SPEC=2 (${{ matrix.group }})
-    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+    runs-on: *linux
    timeout-minutes: 15
    steps:
    - name: Checkout Code
@@ -257,7 +244,7 @@ jobs:

  fuzzing:
    name: Fuzzing
-    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+    runs-on: *linux
    timeout-minutes: 10
    steps:
    - name: Checkout Code
@@ -273,7 +260,7 @@ jobs:

  testopenclimage:
    name: CL IMAGE Tests
-    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+    runs-on: *linux
    timeout-minutes: 15
    steps:
      - name: Checkout Code
@@ -293,7 +280,7 @@ jobs:

  testopenpilot:
    name: openpilot Compile Tests
-    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+    runs-on: *linux
    timeout-minutes: 15
    steps:
      - name: Checkout Code
@@ -322,7 +309,7 @@ jobs:

  testonnxcpu:
    name: ONNX (CPU) Tests
-    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+    runs-on: *linux
    timeout-minutes: 20

    steps:
@@ -341,7 +328,7 @@ jobs:

  testoptim:
    name: Optimization Tests
-    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+    runs-on: *linux
    timeout-minutes: 20
    steps:
      - name: Checkout Code
@@ -373,7 +360,7 @@ jobs:

  testllm:
    name: Test LLM
-    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+    runs-on: *linux
    timeout-minutes: 15
    env:
      CHECK_OOB: 0
@@ -390,17 +377,17 @@ jobs:
        run: |
          parallel --link --tagstring '[{1}]' '{2}' \
            ::: llama 'llama q4' qwen3.5 qwen \
-            ::: $'echo "What\'s a male chicken called? Answer with only one word." | python3 -m tinygrad.llm --no_chat_template --model llama3.2:1b | tee /dev/stderr | grep -i rooster' \
-                $'echo "What\'s a male chicken called? Answer with only one word." | python3 -m tinygrad.llm --no_chat_template --model llama3.2:1b-q4 | tee /dev/stderr | grep -i rooster' \
-                $'echo "What\'s a male chicken called? Answer with only one word." | python3 -m tinygrad.llm --no_chat_template --model qwen3.5:0.8b | tee /dev/stderr | grep -i rooster' \
-                $'echo "What\'s a female chicken called? Answer with only one word." | python3 -m tinygrad.llm --no_chat_template --model qwen3:0.6b | tee /dev/stderr | grep -i hen'
+            ::: $'echo "What\'s a male chicken called? Answer with only one word." | python3 -m tinygrad.llm --model llama3.2:1b | tee /dev/stderr | grep -i rooster' \
+                $'echo "What\'s a male chicken called? Answer with only one word." | python3 -m tinygrad.llm --model llama3.2:1b-q4 | tee /dev/stderr | grep -i rooster' \
+                $'echo "What\'s a male chicken called? Answer with only one word." | python3 -m tinygrad.llm --model qwen3.5:0.8b | tee /dev/stderr | grep -i rooster' \
+                $'echo "What\'s a female chicken called? Answer with only one word." | python3 -m tinygrad.llm --model qwen3:0.6b | tee /dev/stderr | grep -i hen'
                # NOTE: qwen is dumb and only knows about female chickens

 # ****** Models Tests ******

  testmodels:
    name: Models
-    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+    runs-on: *linux
    timeout-minutes: 15
    steps:
      - name: Checkout Code
@@ -420,7 +407,7 @@ jobs:

  testdsp:
    name: Linux (DSP)
-    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+    runs-on: *linux
    timeout-minutes: 15
    steps:
    - name: Checkout Code
@@ -448,7 +435,7 @@ jobs:
          - 'WEBGPU'

    name: Linux (DEV=${{ matrix.dev }})
-    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+    runs-on: *linux
    timeout-minutes: 20
    steps:
    - name: Checkout Code
@@ -474,7 +461,7 @@ jobs:

  testamdasm:
    name: AMD ASM IDE
-    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+    runs-on: *linux
    timeout-minutes: 20
    env:
      DEV: MOCKKFD+AMD
@@ -504,7 +491,7 @@ jobs:
      - name: Run AMD renderer tests (AMD:LLVM)
        run: DEV=MOCKKFD+AMD:LLVM python -m pytest -n=auto test/amd/ --durations 20
      - name: Run SQTT profiling tests
-        run: VIZ=-2 python3 -m pytest -n=auto test/amd/test_sqtt_profiler.py
+        run: PROFILE=1 SQTT=1 python3 -m pytest -n=auto test/amd/test_sqtt_profiler.py
      - name: Run AMD emulated tests on NULL backend
        env:
          AMD: 0
@@ -520,7 +507,7 @@ jobs:

  hcq2:
    name: hcq2
-    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+    runs-on: *linux
    timeout-minutes: 5
    steps:
      - name: Checkout Code
@@ -534,15 +521,16 @@ jobs:
      - name: Run HCQ2 tests
        run: HCQ_RUNTIME_DEV=PYTHON HCQ2=1 DEV=MOCKKFD+AMD FORWARD_ONLY=1 PYTHONPATH=. python test/test_tiny.py
      - name: Run HCQ2 multi-device tests
-        run: HCQ_RUNTIME_DEV=PYTHON HCQ2=1 DEV=MOCKKFD+AMD FORWARD_ONLY=1 PYTHONPATH=. python -m pytest -n=auto test/backend/test_multitensor.py
+        run: |
+          HCQ_RUNTIME_DEV=PYTHON HCQ2=1 DEV=MOCKKFD+AMD FORWARD_ONLY=1 PYTHONPATH=. python test/unit/test_multitensor.py \
+            TestMultiTensor.test_simple_add TestMultiTensor.test_shard_reduce \
+            TestMultiTensor.test_backward_sum TestMultiTensor.test_matmul_shard_0_0
      - name: Run HCQ2 JIT tests
        run: HCQ_RUNTIME_DEV=PYTHON HCQ2=1 DEV=MOCKKFD+AMD FORWARD_ONLY=1 PYTHONPATH=. python test/unit/test_jit.py
-      - name: Run HCQ2 unit tests
-        run: HCQ_RUNTIME_DEV=PYTHON HCQ2=1 DEV=MOCKKFD+AMD FORWARD_ONLY=1 PYTHONPATH=. python -m pytest test/device/test_hcq2.py

  testmockam:
    name: Linux (am)
-    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+    runs-on: *linux
    timeout-minutes: 15
    env:
      DEV: MOCKPCI+AMD
@@ -578,7 +566,7 @@ jobs:
        arch: [gfx1100, gfx1201, gfx950]

    name: Linux (${{ matrix.backend }} ${{ matrix.arch }})
-    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+    runs-on: *linux
    timeout-minutes: 15
    env:
      DEV: MOCKKFD+AMD:${{ matrix.backend == 'amdllvm' && 'LLVM' || '' }}:${{ matrix.arch }}
@@ -601,7 +589,7 @@ jobs:
        if: ${{ matrix.backend == 'amd' && matrix.arch == 'gfx950' }}
        run: PYTHONPATH=. DEV=NULL:HIP:gfx950 MXFP4=1 LLAMA_LAYERS=2 BENCHMARK=3 NULL_ALLOW_COPYOUT=1 NO_HIPCC=1 ROCM_PATH=/opt/rocm JITBEAM=0 examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/profile.sh
      - name: Run pytest (amd)
-        run: python -m pytest -n=auto test/backend/test_ops.py test/backend/test_dtype.py test/backend/test_dtype_alu.py test/backend/test_linearizer.py test/backend/test_randomness.py test/backend/test_jit.py test/backend/test_graph.py test/backend/test_multitensor.py test/device/test_hcq.py test/external/external_test_am.py test/backend/test_asm_gemm.py::TestAsmGEMM --durations=20
+        run: python -m pytest -n=auto test/backend/test_ops.py test/backend/test_dtype.py test/backend/test_dtype_alu.py test/backend/test_linearizer.py test/backend/test_randomness.py test/backend/test_jit.py test/backend/test_graph.py test/backend/test_multitensor.py test/device/test_hcq.py test/external/external_test_am.py test/backend/test_asm_gemm.py::TestAsmGEMM test/opt/test_tensor_cores.py --durations=20
      - name: Run disk copy tests
        run: python -m pytest test/unit/test_disk_tensor.py -k test_copy_from_disk
      - name: Run TRANSCENDENTAL math
@@ -616,7 +604,7 @@ jobs:
        backend: [ptx, nv]

    name: Linux (${{ matrix.backend }})
-    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+    runs-on: *linux
    timeout-minutes: 20
    env:
      FORWARD_ONLY: 1
@@ -650,17 +638,10 @@ jobs:
    strategy:
      fail-fast: false
      matrix:
-        dev:
-          - 'NULL:IR3:a630'
-          - 'NULL:QCOMCL:a630'
-          - 'NULL:NAK:sm_120'
-    name: Compile-only (DEV=${{ matrix.dev }})
-    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+        backend: [ir3, nak]
+    name: Compile-only (${{ matrix.backend }})
+    runs-on: *linux
    timeout-minutes: 15
-    env:
-      NULL_ALLOW_COPYOUT: 1
-      DEV: ${{ matrix.dev }}${{ contains(matrix.dev, 'a630') && ',IMAGE_PITCH_ALIGNMENT=64' || '' }}
-      IMAGE: ${{ contains(matrix.dev, 'a630') && '1' || '0' }}
    steps:
      - name: Checkout Code
        uses: actions/checkout@v6
@@ -669,15 +650,21 @@ jobs:
        with:
          key: compile-${{ matrix.backend }}
          deps: "testing_unit mesa"
-          qemu: ${{ contains(matrix.dev, 'QCOMCL') }}
-      - name: Test IMAGE
+      - name: Set env
        shell: bash
-        if: contains(matrix.dev, 'a630')
-        run: DEBUG=7 python3 test/backend/test_ops.py TestOps.test_gemm | grep isam
+        run: printf "NULL_ALLOW_COPYOUT=1\n${{ matrix.backend == 'ir3' && 'DEV=NULL:IR3:a630' || matrix.backend == 'nak' && 'DEV=NULL:NAK:sm_120' }}" >> $GITHUB_ENV
      - name: Run test_ops
        shell: bash
        run: |
          python -c "from tinygrad import Device; assert Device.DEFAULT == 'NULL'"
          DEBUG=4 python3 test/backend/test_ops.py TestOps.test_add
-          # QCOMCL compiles in qemu, too slow for parallel workers
-          ${{ contains(matrix.dev, 'QCOMCL') && 'PARALLEL=0' || '' }} python -m pytest -n=auto test/backend/test_ops.py --durations=20
+          python -m pytest -n=auto test/backend/test_ops.py --durations=20
+      - name: Run test_ops (IMAGE)
+        if: matrix.backend == 'ir3'
+        shell: bash
+        env:
+          IMAGE: 1
+          DEV: "NULL:IR3:a630,IMAGE_PITCH_ALIGNMENT=64"
+        run: |
+          DEBUG=4 python3 test/backend/test_ops.py TestOps.test_gemm | grep image_load
+          python -m pytest -n=auto test/backend/test_ops.py --durations=20
@@ -69,4 +69,3 @@ mutants
 dagre/
 graphlib/
 uv.lock
-pi_session_window0.jsonl
@@ -28,7 +28,7 @@ repos:
        pass_filenames: false
      - id: tests
        name: comprehensive test suite
-        entry: env OMP_NUM_THREADS=1 SKIP_SLOW_TEST=1 PYTHONPATH="." python3 -m pytest -n=6 test/backend/test_ops.py test/backend/test_schedule.py test/backend/test_assign.py test/backend/test_tensor.py test/backend/test_jit.py test/unit/test_schedule_cache.py test/null/test_pattern_matcher.py test/null/test_uop_symbolic.py test/unit/test_helpers.py
+        entry: env OMP_NUM_THREADS=1 SKIP_SLOW_TEST=1 PYTHONPATH="." python3 -m pytest -n=6 test/backend/test_ops.py test/backend/test_schedule.py test/unit/test_assign.py test/backend/test_tensor.py test/backend/test_jit.py test/unit/test_schedule_cache.py test/null/test_pattern_matcher.py test/null/test_uop_symbolic.py test/unit/test_helpers.py
        language: system
        always_run: true
        pass_filenames: false
@@ -4,4 +4,3 @@
 - Run `python -m mypy tinygrad/` to typecheck
 - Run `python -m ruff check .` to lint
 - Read `./tinygrad/viz/README.md` for profiling and debugging rewrite rules
- Do not do amend commits. Always do a new commit if a force push to origin would be required.
@@ -140,7 +140,7 @@ Documentation along with a quick start guide can be found on the [docs website](
 ```python
 from tinygrad import Tensor

-x = Tensor.eye(3).clone()  # clone to make it a buffer
+x = Tensor.eye(3)
 y = Tensor([[2.0,0,-2.0]])
 z = y.matmul(x).sum()
 z.backward()
@@ -1 +1 @@
-f6fc4e3f2c3db5fae1e19cbfbc3ad9fc579a12ae
+8611fe22a7fcc7d1928bbde19ded66277cb12f3e
@@ -2,7 +2,7 @@ import os, pytest, signal, threading

@pytest.hookimpl(wrapper=True)
 def pytest_runtest_call(item):
-  t = threading.Timer(int(os.getenv("TEST_TIMEOUT", 90)), os.kill, args=(os.getpid(), signal.SIGABRT))
+  t = threading.Timer(int(os.getenv("TEST_TIMEOUT", 300)), os.kill, args=(os.getpid(), signal.SIGABRT))
  t.start()
  try: yield
  finally:
@@ -122,7 +122,7 @@ def example_5_custom_assembly(a:Tensor, correct):
        offset_dwords = (self.labels[inst._target] - inst._pos - inst.size()) // 4
        if not -32768 <= offset_dwords <= 32767: raise ValueError(f"branch to '{inst._target}' offset {offset_dwords} exceeds simm16 range")
        inst.simm16 = offset_dwords
-      return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=(x, dtypes.void)) for x in self.instructions]))))
+      return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in self.instructions]))))

  CU_COUNT = 32
  LANES = 64
@@ -52,7 +52,7 @@ In `kernel.py` we have a set of `OptOps`, these control the parameters of the sp

 The main bottleneck in most kernels is accessing memory. In a freshman algorithms class, you'll learn about cache aware matrix multiplication, and this is all forms of that. While the same math is run, the order in which you run it can have large impacts on the speed depending on if the data you are loading. OptOps will change this order.

-Memory, even cache, is often much slower than accessing the register file. The amount of times data is used in math is called the "arithmetic intensity". For operations like BS=1 GEMV, the arithmetic intensity is 1, but for GEMMs and convs it can be much higher. Splitting an axis into UPCAST can increase this, but be careful of making them too large, as if there's too much register pressure on the GPU the warp scheduler may not be able to fit many warps, or even worse, it could be spilling to local memory.
+Memory, even cache, is often much slower than accessing the register file. The amount of times data is used in math is called the "arithmetic intensity". For operations like BS=1 GEMV, the arithmetic intensity is 1, but for GEMMs and convs it can be much higher. OptOps like UPCAST and UNROLL can increase this, but be careful of making them too large, as if there's too much register pressure on the GPU the warp scheduler may not be able to fit many warps, or even worse, it could be spilling to local memory.

 4090s have 1 TB/s of ram bandwidth and ~160 TFLOPS of compute, so you need to use each loaded value ~100 times. The L1 cache has around 40 TB/s of bandwidth, so in order to get full compute utilization you need to use each value ~4 times.

@@ -5,7 +5,7 @@ from multiprocessing import Queue, Process, shared_memory, connection, Lock

 import numpy as np
 from tinygrad import dtypes, Tensor
-from tinygrad.helpers import getenv, prod, Context, round_up, tqdm, OSX, CPU_COUNT
+from tinygrad.helpers import getenv, prod, Context, round_up, tqdm, OSX, NUM_CPU_THREADS
 from tinygrad.nn.state import TensorIO

 ### ResNet
@@ -131,7 +131,7 @@ def batch_load_resnet(batch_size=64, val=False, shuffle=True, seed=None, pad_fir
    else: X = Tensor.empty(*sz, dtype=dtypes.uint8, device=f"disk:/dev/shm/{shm_name}")
    Y = [None] * (batch_size*BATCH_COUNT)

-    for _ in range(CPU_COUNT):
+    for _ in range(NUM_CPU_THREADS.value):
      p = Process(target=loader_process, args=(q_in, q_out, X, seed))
      p.daemon = True
      p.start()
@@ -212,7 +212,7 @@ def batch_load_train_bert(BS:int, seed:int|None=None):
    rng.shuffle(fs)
    train_files.append(fs.pop(0))

-  cycle_length = min(CPU_COUNT, len(train_files))
+  cycle_length = min(NUM_CPU_THREADS.value, len(train_files))
  assert cycle_length > 0, "cycle_length must be greater than 0"

  dataset = InterleavedDataset(train_files, cycle_length)
@@ -301,7 +301,7 @@ def batch_load_unet3d(preprocessed_dataset_dir:Path, batch_size:int=6, val:bool=
    X = Tensor.empty(*sz, dtype=dtypes.float32, device=f"disk:/dev/shm/{shm_name_x}")
    Y = Tensor.empty(*sz, dtype=dtypes.uint8, device=f"disk:/dev/shm/{shm_name_y}")

-    for _ in range(CPU_COUNT):
+    for _ in range(NUM_CPU_THREADS.value):
      proc = Process(target=load_unet3d_data, args=(preprocessed_dataset_dir, seed, queue_in, queue_out, X, Y))
      proc.daemon = True
      proc.start()
@@ -437,7 +437,7 @@ def batch_load_retinanet(dataset, val:bool, base_dir:Path, batch_size:int=32, sh
  dataset_iter = iter(image_ids)

  try:
-    for _ in range(CPU_COUNT):
+    for _ in range(NUM_CPU_THREADS.value):
      proc = Process(
        target=load_retinanet_data,
        args=(base_dir, val, queue_in, queue_out, imgs, boxes, labels),
@@ -1667,14 +1667,15 @@ def train_llama3():
 def train_gptoss():
  from examples.mlperf.models.gpt_oss import GPTOSS, GPT_OSS_20B, apply_grad, FP8_DTYPE
  from examples.mlperf.lr_schedulers import CosineAnnealingLRWithWarmup
-  from examples.mlperf.optim import GradAccClipAdamW, GradAccClipAdamWGroup, fclip_grads
+  from examples.mlperf.optim import GradAccClipAdamW, GradAccClipAdamWGroup, clip_grads

  BENCHMARK = getenv("BENCHMARK")

  config = {}
  BASEDIR            = config["BASEDIR"]                = Path(getenv("BASEDIR", "/raid/datasets/c4-8b/"))
  BS                 = config["BS"]                     = getenv("BS", 16)
-  GBS                = config["GLOBAL_BATCH_SIZE"]      = BS
+  grad_acc           = config["GRADIENT_ACC_STEPS"]     = getenv("GRADIENT_ACC_STEPS", 1)
+  GBS                = config["GLOBAL_BATCH_SIZE"]      = BS * grad_acc
  SEED               = config["SEED"]                   = getenv("SEED", 5760)
  DATA_SEED          = config["DATA_SEED"]              = getenv("DATA_SEED", SEED)
  SEQLEN             = config["SEQLEN"]                 = getenv("SEQLEN", 8192)
@@ -1736,13 +1737,13 @@ def train_gptoss():
  params_wd = [p for p in params if p.ndim >= 3]
  params_no_wd = [p for p in params if p.ndim < 3]
  optim = GradAccClipAdamWGroup(
-    GradAccClipAdamW(params_wd, lr=0.0, b1=opt_adamw_beta_1, b2=opt_adamw_beta_2, eps=opt_adamw_epsilon, weight_decay=opt_adamw_weight_decay, grad_acc=1, device=optim_device),
-    GradAccClipAdamW(params_no_wd, lr=0.0, b1=opt_adamw_beta_1, b2=opt_adamw_beta_2, eps=opt_adamw_epsilon, weight_decay=0.0, grad_acc=1, device=optim_device),
+    GradAccClipAdamW(params_wd, lr=0.0, b1=opt_adamw_beta_1, b2=opt_adamw_beta_2, eps=opt_adamw_epsilon, weight_decay=opt_adamw_weight_decay, grad_acc=grad_acc, device=optim_device),
+    GradAccClipAdamW(params_no_wd, lr=0.0, b1=opt_adamw_beta_1, b2=opt_adamw_beta_2, eps=opt_adamw_epsilon, weight_decay=0.0, grad_acc=grad_acc, device=optim_device),
  )

  for p in optim.params:
-    p.grad = p.zeros_like(dtype=dtypes.bfloat16 if p.dtype == FP8_DTYPE else p.dtype).contiguous()
-    if getattr(p, "_zero2", False): p.grad = optim.optimizers[0]._zero_shard(p.grad)
+    grad_dtype = dtypes.bfloat16 if p.dtype == FP8_DTYPE else p.dtype
+    p.grad = p.zeros_like(dtype=grad_dtype).contiguous()
  grads = [p.grad for p in optim.params]

  from extra.gemm.cdna_asm_gemm import _mx_block_scale
@@ -1769,32 +1770,31 @@ def train_gptoss():

  @TinyJit
  @Context(TRAINING=1)
-  def step(tokens:Tensor):
+  def minibatch(tokens:Tensor):
    if is_dp: tokens = tokens.to(None).shard(device, 0)
    if not is_sharding: tokens = tokens.to(None)
-
    logits:Tensor = model(tokens[:, :-1], save=True)
-    if getenv("FUSED_CE", 0):
-      from extra.llama_kernels.fused_ce import fused_ce_loss
-      loss = fused_ce_loss(logits.cast(dtypes.bfloat16), tokens[:, 1:], label_smoothing=0.0)
-    else:
-      loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
+    loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])

    for g, new_g in zip(grads, loss.gradient(*optim.params)):
      apply_grad(g, new_g.uop)

-    Tensor.realize(loss, *grads)
+    loss_cpu = loss.flatten().float().to("CPU")
+    return loss_cpu.realize(*grads)

-    clipped_grads, grad_norm = fclip_grads(grads, 1.0)
-    optim.fstep(clipped_grads, grad_norm)
+  @TinyJit
+  def optim_step():
+    grad_norm = clip_grads(grads, grad_acc, 1.0)
+    optim.fstep(grads, grad_norm)
    scheduler.step()

-    loss_cpu = loss.flatten().float().to("CPU")
+    for g in grads: g.assign(0)
+
    lr_cpu = optim.lr.float().to("CPU")
    grad_norm_cpu = grad_norm.float().to("CPU")
-    Tensor.realize(loss_cpu, lr_cpu, grad_norm_cpu, *grads, *fp8_inv_scales)
+    Tensor.realize(lr_cpu, grad_norm_cpu, *grads, *fp8_inv_scales)

-    return loss_cpu, lr_cpu, grad_norm_cpu
+    return lr_cpu, grad_norm_cpu

  @TinyJit
  @Context(TRAINING=0)
@@ -1843,20 +1843,30 @@ def train_gptoss():
      profile_marker(f"train @ {i}")
      st = time.perf_counter()

-      ist = time.perf_counter()
+      stopped = False
+      losses, data_time, dev_time = [], 0, 0
+      for _ in range(grad_acc if i >= 2 else 1):
+        ist = time.perf_counter()
+        try: tokens = next(train_iter)
+        except StopIteration:
+          stopped = True
+          break
+        mst = time.perf_counter()
+        data_time += mst - ist
+        losses.append(minibatch(tokens).item())
+        dev_time += time.perf_counter() - mst
+      if stopped: break

-      try: tokens = next(train_iter)
-      except StopIteration: break
-      mst = time.perf_counter()
-      data_time = mst - ist
-
-      ret = step(tokens)
-      dev_time = time.perf_counter() - mst
-
-      loss, lr, grad_norm = ret[0].item(), ret[1].item(), ret[2].item()
+      gt = time.perf_counter()
+      ret = optim_step()
+      lr, grad_norm = ret[0].item(), ret[1].item()
      et = time.perf_counter()

+      loss = sum(losses) / len(losses)
+      optim_time = et - gt
+      dev_time += optim_time
      step_time = et - st
+      gbs_time = gt - st
      if BENCHMARK: step_times.append(step_time)

      i += 1
@@ -1866,7 +1876,7 @@ def train_gptoss():
      gflops = GlobalCounters.global_ops / 1e9 / dev_time
      mfu = ((6 * num_params * SEQLEN * GBS) / (dev_time * device_count * 4.6e15)) * 100
      tqdm.write(
-          f"{i:5} {step_time:.3f} s step, {dev_time:.3f} s dev, {data_time:.3f} s data, {loss:.4f} loss, " \
+          f"{i:5} {step_time:.3f} s step, {gbs_time:.3f} s gbs, {optim_time:.3f} s optim, {data_time:.3f} s data, {loss:.4f} loss, " \
          f"{lr:.12f} LR, {grad_norm:.6f} grad_norm, {mem_gb:.2f} GB used, {gflops:9.2f} GFLOPS, {mfu:5.2f}% MFU")
      if DEBUG >= 1: tqdm.write("  mem per device: " + ', '.join(f"{dev}: {mem/1e9:.2f} GB" for dev, mem in sorted(GlobalCounters.mem_used_per_device.items())))

@@ -1876,6 +1886,8 @@ def train_gptoss():
          "train/lr": lr,
          "train/grad_norm": grad_norm,
          "train/step_time": step_time,
+          "train/gbs_time": gbs_time,
+          "train/optim_time": optim_time,
          "train/dev_time": dev_time,
          "train/data_time": data_time,
          "train/mem": mem_gb,
@@ -12,7 +12,7 @@ from tinygrad.helpers import Timing, colored, GlobalCounters, profile_marker
 from tinygrad.uop.ops import Ops, UOp
 from extra.models.llama import apply_rotary_emb
 from extra.llama_kernels.rmsnorm import rmsnorm
-from extra.gemm.cdna_asm_gemm import _mx_block_scale, _mx_block_scale_3d, quantize_mxfp8, asm_gemm, can_use_asm_gemm
+from extra.gemm.cdna_asm_gemm import _mx_block_scale, _mx_block_scale_3d, quantize_mxfp8
 from extra.gemm.moe_gemm import grouped_mx_gemm
 from extra.gemm.moe_routing import route, dispatch, combine

@@ -61,25 +61,7 @@ def dequant_weight(w_q:Tensor, w_scale:Tensor) -> Tensor:
  call = UOp.maketuple(fxn.uop).call(w_q.uop, w_scale.uop, grad_fxn=_dequant_bwd)
  return Tensor(call.gettuple(0))

-def matmul_mx(x:Tensor|tuple[Tensor, Tensor], w_q:Tensor, w_scale:Tensor) -> Tensor:
-  if isinstance(x, tuple):
-    assert ASM_GEMM, "pre-quantized MXFP8 input requires ASM_GEMM"
-    from extra.gemm.cdna_asm_gemm import asm_gemm, can_use_asm_gemm, mx_pack
-    x_q, x_e8 = x
-    l_shape, padded = x_q.shape[:-1], x_q.shape[-1]
-    x_q, x_e8 = x_q.reshape(-1, padded), x_e8.reshape(-1, padded // 32)
-    K, N = w_q.shape[1], w_q.shape[0]
-    assert padded >= K and (padded - K) % 32 == 0 and x_e8.shape[-1] == padded // 32
-    wq, ws = w_q, w_scale
-    if (pad := padded - K):
-      wq = wq.pad(((0, 0), (0, pad)))
-      ws = ws.pad(((0, 0), (0, pad // 32)), value=127).cast(dtypes.uint8)
-    if (npad := (-N) % 256):
-      wq = wq.pad(((0, npad), (0, 0)))
-      ws = ws.pad(((0, npad), (0, 0)), value=127).cast(dtypes.uint8)
-    assert can_use_asm_gemm(x_q, wq.T)
-    out = asm_gemm(x_q, wq.T, mx=True, mx_scales=(mx_pack(x_e8), x_e8, mx_pack(ws), ws), mx_w_stored=True)
-    return (out[:, :N] if npad else out).reshape(*l_shape, N).cast(dtypes.bfloat16)
+def matmul_mx(x:Tensor, w_q:Tensor, w_scale:Tensor) -> Tensor:
  l_shape = x.shape[:-1]
  if ASM_GEMM:
    from extra.gemm.cdna_asm_gemm import asm_gemm, can_use_asm_gemm, mx_pack
@@ -164,7 +146,6 @@ class GPTOSS:
      return w_q, w_e8.is_param_(False)
    if moe:
      qs = [_one(*shape[1:]) for _ in range(shape[0])]
-      for q in qs: q[0]._zero2 = True  # grad arrives sharded on the expert axis under ZeRO-2 (moe_gemm)
      return [q[0] for q in qs], [q[1] for q in qs]
    return _one(*shape)

@@ -193,30 +174,20 @@ class GPTOSS:
  def attention(self, x:Tensor, freqs_cis:Tensor, mask:Tensor, sliding:bool, *, attention_norm:Tensor, wqkv:Tensor,
                wqkv_scale:Tensor, wqkv_bias:Tensor, wo:Tensor, wo_scale:Tensor, wo_bias:Tensor, sinks:Tensor):
    bsz, seqlen, _ = x.shape
-
-    if getenv("FUSED_RMSNORM_MUL", 0):
-      from extra.gptoss_kernels.rmsnorm import rmsnorm_mul
-      x_normed, rrms = rmsnorm_mul(x, attention_norm, self.norm_eps)
-      norm_saves = [x_normed, rrms]
-    else:
-      x_normed, rrms = rmsnorm(x, self.norm_eps)
-      qkv = matmul_mx(x_normed * attention_norm, wqkv, wqkv_scale) + wqkv_bias
-      norm_saves = [x_normed, rrms]
-
+    x_normed, rrms = rmsnorm(x, self.norm_eps)
+    qkv = matmul_mx(x_normed * attention_norm, wqkv, wqkv_scale) + wqkv_bias
    qkv = qkv.reshape(bsz, seqlen, self.n_kv_heads, self.n_rep + 2, self.head_dim)
    xq = qkv[:, :, :, :self.n_rep].reshape(bsz, seqlen, self.n_heads, self.head_dim)
    xk, xv = qkv[:, :, :, self.n_rep], qkv[:, :, :, self.n_rep + 1]
    xq, xk = apply_rotary_emb(xq, xk, freqs_cis)
    xq, xk, xv = xq.cast(dtypes.bfloat16), xk.cast(dtypes.bfloat16), xv.cast(dtypes.bfloat16)  # (B,N,H,D)/(B,N,KV,D)

-    fa_saves = []
-    if getenv("HK_FLASH_ATTENTION"):
-      from extra.thunder.amd.fa import flash_attention
-      attn, _, l_vec = flash_attention(xq, xk, xv, is_causal=True, write_flat=True, sinks=sinks, window=self.sliding_window if sliding else 0)
-      attn = attn.reshape(bsz, seqlen, self.n_heads * self.head_dim)
-      fa_saves = [xq, xk, xv, l_vec]
-    elif sliding:
+    if sliding:
      attn = self._sliding_attention(xq, xk, xv, sinks)
+    elif getenv("HK_FLASH_ATTENTION"):
+      from extra.thunder.amd.fa import flash_attention
+      attn, *_ = flash_attention(xq, xk, xv, is_causal=True, write_flat=True, sinks=sinks)
+      attn = attn.reshape(bsz, seqlen, self.n_heads * self.head_dim)
    else:
      xqm = xq.reshape(bsz, seqlen, self.n_kv_heads, self.n_rep, self.head_dim).permute(0, 2, 3, 1, 4)
      xkm, xvm = xk.permute(0, 2, 1, 3).unsqueeze(2), xv.permute(0, 2, 1, 3).unsqueeze(2)
@@ -228,19 +199,13 @@ class GPTOSS:
      attn = (w @ xvm).permute(0, 3, 1, 2, 4).reshape(bsz, seqlen, self.n_heads * self.head_dim)

    out = matmul_mx(attn, wo, wo_scale) + wo_bias
-    return out, [attn] + norm_saves + fa_saves
+    return out, [x_normed, rrms, attn]

  def feed_forward(self, x:Tensor, *, ffn_norm:Tensor, gate:Tensor, gate_bias:Tensor,
                   w_gate_up:Tensor, w_gate_up_scale:Tensor, w_gate_up_bias:Tensor,
                   w_down:Tensor, w_down_scale:Tensor, w_down_bias:Tensor):
-    if getenv("FUSED_RMSNORM_MUL", 0):
-      from extra.gptoss_kernels.rmsnorm import rmsnorm_mul
-      x_normed, rrms = rmsnorm_mul(x, ffn_norm, self.norm_eps)
-      inp = x_normed
-    else:
-      x_normed, rrms = rmsnorm(x, self.norm_eps)
-      inp = x_normed * ffn_norm
-
+    x_normed, rrms = rmsnorm(x, self.norm_eps)
+    inp = x_normed * ffn_norm
    logits = inp.float() @ gate.float().T + gate_bias.float()
    dim, inter = self.dim, self.intermediate_size

@@ -255,7 +220,6 @@ class GPTOSS:
      z = grouped_mx_gemm(_pad_cols(y.cast(dtypes.bfloat16)), (w_down, w_down_scale), r.off)[:, :dim] \
          + (onehot @ w_down_bias.float()).cast(dtypes.bfloat16)
      out = combine(z, r, inp.shape[0], self.experts_per_tok).reshape(bsz, seqlen, dim)
-      return out, [x_normed, rrms, xg, h, y, z, r.weights, r.dest_row, r.off]
    else:
      thresh = logits.topk(self.experts_per_tok)[0][..., -1:]
      weights = (logits >= thresh).where(logits, -float("inf")).softmax(-1)
@@ -299,11 +263,7 @@ class GPTOSS:
                        w_down=self.w_down[i], w_down_scale=self.w_down_scale[i], w_down_bias=self.w_down_bias[i])
      h, *_ = self.run_layer(h, freqs_cis, mask_full, i % 2 == 0, attn_kwargs, ffn_kwargs, save=save)

-    h_normed = self.norm(h)
-    pad = (-self.dim) % 256
-    h_padded, w_padded = h_normed.pad((None, None, (0, pad))), self.output.pad(((0, 0), (0, pad)))
-    if ASM_GEMM and can_use_asm_gemm(h_padded, w_padded.T): logits = asm_gemm(h_padded, w_padded.T)
-    else: logits = h_normed @ self.output.T
+    logits = self.norm(h) @ self.output.T
    return logits

 def _get_pads(uop:UOp) -> list[UOp]:
@@ -314,14 +274,14 @@ def apply_grad(grad_buf:Tensor, new_grad:UOp):
  pads = _get_pads(new_grad)
  if len(pads) <= 1:
    new_grad = new_grad.cast(grad_buf.dtype)
-    grad_buf.uop = grad_buf.uop.after(grad_buf.uop.store(new_grad))
+    grad_buf.uop = grad_buf.uop.after(grad_buf.uop.store(grad_buf.uop + new_grad))
    return
  cur = grad_buf.uop
  for pad in sorted(pads, key=lambda p: p.marg[0][0] if p.op == Ops.PAD else 0, reverse=True):
    if pad.op == Ops.PAD:
-      grad_shrink = tuple((p[0], s+p[0]) for s,p in zip(pad.src[0].shape, pad.marg))
+      grad_shrink = tuple([(p[0], s+p[0]) for s,p in zip(pad.src[0].shape, pad.marg)])
      buf_slice = cur.shrink(grad_shrink)
-      cur = cur.after(buf_slice.store(pad.src[0].cast(cur.dtype)))
+      cur = cur.after(buf_slice.store(buf_slice + pad.src[0].cast(cur.dtype)))
    else:
      cur = cur.after(cur.store(cur + pad.cast(cur.dtype)))
  grad_buf.uop = cur
@@ -15,7 +15,7 @@ def stochastic_round_bf16(x:Tensor) -> Tensor:
  bits = x.bitcast(dtypes.uint32)
  if isinstance(x.device, tuple):
    shape = x.uop.shard_shape if x.uop.axis is not None else x.shape
-    noise = Tensor(UOp(Ops.MSTACK, src=tuple(Tensor.rand(*shape, device=d).uop for d in x.device)))
+    noise = Tensor(UOp(Ops.MSTACK, dtypes.default_float, tuple(Tensor.rand(*shape, device=d).uop for d in x.device)))
  else:
    noise = x.rand_like()
  noise = (noise * 0xFFFF).cast(dtypes.uint32)
@@ -27,11 +27,6 @@ def clip_grads(grads:list[Tensor], grad_acc, clip_norm) -> Tensor:
  for g in grads: g.assign((g * (clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)).cast(g.dtype))
  return total_norm

-def fclip_grads(grads:list[Tensor], clip_norm) -> Tensor:
-  total_norm = Tensor.stack(*[g.float().square().sum() for g in grads]).sum().sqrt().contiguous()
-  scale = (clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)
-  return [(g * scale).cast(g.dtype) for g in grads], total_norm
-
 class GradAccClipAdamW(Optimizer):
  def __init__(self, params:list[Tensor], lr=0.001, b1=0.9, b2=0.999, eps=1e-6, weight_decay=0.0, grad_acc=1, clip_norm=1.0, device=None, fused=FUSE_OPTIM):
    super().__init__(params, lr, device, fused)
@@ -1,13 +0,0 @@
-from pathlib import Path
-
-from examples.mlperf.dataloader import get_llama3_dataset
-from tinygrad.helpers import getenv
-
-BASEDIR = Path(getenv("BASEDIR", "/raid/datasets/c4-8b/"))
-SAMPLES = getenv("SAMPLES", 1_200_000 * 32)
-EVAL_SAMPLES = getenv("EVAL_SAMPLES", 1024)
-SEQLEN = getenv("SEQLEN", 8192)
-DATA_SEED = getenv("DATA_SEED", 5760)
-
-get_llama3_dataset(SAMPLES, SEQLEN, BASEDIR, seed=DATA_SEED, val=False, small=True)
-get_llama3_dataset(EVAL_SAMPLES, SEQLEN, BASEDIR, seed=0, val=True, small=True)
@@ -26,7 +26,7 @@ export FUSED_SILU_W13=${FUSED_SILU_W13:-1}
 export SPLIT_W13=${SPLIT_W13:-0}
 export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-0}

-export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="float32"
+export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
 export DP=${DP:-8} MP=${MP:-1} BS=${BS:-16} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2}
 export GBS=$((BS * GRADIENT_ACC_STEPS))

@@ -35,7 +35,7 @@ export BASEDIR="/raid/datasets/c4-8b/"
 export SMALL=1
 export LLAMA3_SIZE=${LLAMA3_SIZE:-"8B"}
 export EVAL_TARGET=3.3 EVAL_FREQ=12288
-export LR="1e-3" END_LR="1e-4" WARMUP_SAMPLES=2048 MAX_STEPS=1200000
+export LR="1e-3" END_LR="1e-4" WARMUP_SAMPLES=4096 MAX_STEPS=1200000
 export WARMUP_STEPS=$((WARMUP_SAMPLES / GBS))
 export SAMPLES=$((MAX_STEPS * GBS))
 export SEQLEN=${SEQLEN:-8192}
@@ -44,7 +44,7 @@ export SEED=${SEED:-5760}
 export DATA_SEED=${DATA_SEED:-5760}

 export JITBEAM=${JITBEAM:-3}
-export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
+export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=1

 export FAKEDATA=${FAKEDATA:-1} BENCHMARK=${BENCHMARK:-10}
 if [ -z "$FULL_LAYERS" ]; then
@@ -1,8 +1,8 @@
 #!/usr/bin/env bash

 export PYTHONPATH="."
-export ROCM_PATH=${ROCM_PATH:-/opt/rocm-7.1.1}
-export PATH="$ROCM_PATH/bin:$PATH"
+export PATH="/opt/rocm-7.1.1/bin:$PATH"
+export ROCM_PATH="/opt/rocm-7.1.1"
 export DEV=${DEV:-AMD}
 export CHECK_OOB=0
 export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
@@ -26,7 +26,7 @@ export FUSED_SILU_W13=${FUSED_SILU_W13:-1}
 export SPLIT_W13=${SPLIT_W13:-0}
 export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-0}

-export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="float32"
+export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
 export DP=${DP:-8} MP=${MP:-1} BS=${BS:-16} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2}
 export GBS=$((BS * GRADIENT_ACC_STEPS))

@@ -35,7 +35,7 @@ export BASEDIR="/raid/datasets/c4-8b/"
 export SMALL=1
 export LLAMA3_SIZE=${LLAMA3_SIZE:-"8B"}
 export EVAL_TARGET=3.3 EVAL_FREQ=12288
-export LR="1e-3" END_LR="1e-4" WARMUP_SAMPLES=2048 MAX_STEPS=1200000
+export LR="1e-3" END_LR="1e-4" WARMUP_SAMPLES=4096 MAX_STEPS=1200000
 export WARMUP_STEPS=$((WARMUP_SAMPLES / GBS))
 export SAMPLES=$((MAX_STEPS * GBS))
 export SEQLEN=${SEQLEN:-8192}
@@ -44,6 +44,6 @@ export SEED=${SEED:-$RANDOM}
 export DATA_SEED=${DATA_SEED:-5760}

 export JITBEAM=${JITBEAM:-3}
-export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
+export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=1

 python3 examples/mlperf/model_train.py
@@ -17,7 +17,7 @@ export USE_ATOMICS=1
 export ASM_GEMM=1
 export WQKV=1
 export MASTER_WEIGHTS=1
-export MXFP4=1
+export FP8=1
 export ALLREDUCE_CAST=1
 export FAST_CE=1
 export FUSED_INPUT_QUANTIZE=1
@@ -26,7 +26,7 @@ export FUSED_ADD_NORM_MUL_QUANTIZE=1
 export FUSED_SILU_W13=1
 export SPLIT_W13=0

-export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="float32"
+export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
 export DP=8 MP=1 BS=16 EVAL_BS=8 GRADIENT_ACC_STEPS=2
 export GBS=$((BS * GRADIENT_ACC_STEPS))

@@ -44,7 +44,7 @@ export SEED=$RANDOM
 export DATA_SEED=$SEED

 export JITBEAM=3
-export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
+export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=1

 export LOGMLPERF=1

@@ -11,7 +11,6 @@ export DEVICE_IN_FUNCTION_BUG=1
 export DEBUG=${DEBUG:-2}
 export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
 export ASM_GEMM=${ASM_GEMM:-1}
-export GROUPED_MOE=${GROUPED_MOE:-1}
 export ALL2ALL=${ALL2ALL:-1}
 export LATE_ALLREDUCE=${LATE_ALLREDUCE:-0}
 export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
@@ -11,7 +11,6 @@ export DEVICE_IN_FUNCTION_BUG=1
 export DEBUG=${DEBUG:-0}
 export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
 export ASM_GEMM=${ASM_GEMM:-1}
-export GROUPED_MOE=${GROUPED_MOE:-1}
 export ALL2ALL=${ALL2ALL:-1}
 export LATE_ALLREDUCE=${LATE_ALLREDUCE:-0}
 export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
@@ -107,21 +107,14 @@ def compile(onnx_file):
  return inputs, test_val

 def test_vs_compile(run, inputs, test_val=None):
-  if (log:=bool(getenv("BENCHMARK_LOG", ""))): from extra.bench_log import WallTimeEvent, BenchEvent

  # run 20 times
  step_times = []
  for _ in range(20):
    st = time.perf_counter()
-    if log:
-      with WallTimeEvent(BenchEvent.STEP):
-        out = run(**inputs)
-        mt = time.perf_counter()
-        val = out.numpy()
-    else:
-      out = run(**inputs)
-      mt = time.perf_counter()
-      val = out.numpy()
+    out = run(**inputs)
+    mt = time.perf_counter()
+    val = out.numpy()
    et = time.perf_counter()
    step_times.append((et-st)*1e3)
    print(f"enqueue {(mt-st)*1e3:6.2f} ms -- total run {step_times[-1]:6.2f} ms")
@@ -167,6 +160,12 @@ def test_vs_onnx(new_inputs, test_val, onnx_file, tol):
  print("test vs onnx passed")
  return timings

+def bench(run, inputs):
+  from extra.bench_log import WallTimeEvent, BenchEvent
+  for _ in range(10):
+    with WallTimeEvent(BenchEvent.STEP):
+      run(**inputs).numpy()
+
 if __name__ == "__main__":
  if getenv("RUN_PICKLE"):
    with open(OUTPUT, "rb") as f: pickle_loaded = load_pickle(f)
@@ -182,3 +181,6 @@ if __name__ == "__main__":
    test_vs_compile(pickle_loaded, inputs, outputs)
    if getenv("SELFTEST"):
      test_vs_onnx(inputs, outputs, onnx_file, 1e-4)
+
+  if getenv("BENCHMARK_LOG", ""):
+    bench(pickle_loaded, inputs)
@@ -84,8 +84,7 @@ class AMSMI(AMDev):
    with open(f"/sys/bus/pci/devices/{self.pcibus}/power_state", "r") as f: return f.read().strip().rstrip()

 class SMICtx:
-  def __init__(self, dev_filter=None):
-    self.dev_filter = dev_filter
+  def __init__(self):
    self.devs = []
    self.opened_pcidevs = []
    self.opened_pci_resources = {}
@@ -136,7 +135,6 @@ class SMICtx:
    pattern = os.path.join('/tmp', 'am_*.lock')
    for d in [f[8:-5] for f in glob.glob(pattern)]:
      if d.startswith("usb"): continue
-      if self.dev_filter is not None and d != self.dev_filter: continue
      if d not in self.opened_pcidevs:
        self._open_am_device(d)

@@ -408,7 +406,7 @@ if __name__ == "__main__":

  try:
    if not args.list: os.system('clear')
-    smi_ctx = SMICtx(args.dev)
+    smi_ctx = SMICtx()
    while True:
      smi_ctx.rescan_devs()
      smi_ctx.draw(args.list)
@@ -35,7 +35,7 @@ class WallTimeEvent:
    return self
  def __exit__(self, *_):
    self.time = time.monotonic() - self.start
-    _events[self.event]["wall"].append((self.time, BENCHMARK_LOG.value))
+    _events[self.event]["wall"].append(self.time)
    return False

 class KernelTimeEvent:
@@ -47,19 +47,19 @@ class KernelTimeEvent:
    self.start = GlobalCounters.time_sum_s
    return self
  def __exit__(self, *_):
-    _events[self.event]["kernel"].append((GlobalCounters.time_sum_s - self.start, BENCHMARK_LOG.value))
+    _events[self.event]["kernel"].append(GlobalCounters.time_sum_s - self.start)
    return False

 def log_event_instant(event:InstantBenchEvent, value:float):
-  _events[event].append((value, BENCHMARK_LOG.value))
+  _events[event].append(value)

 if BENCHMARK_LOG:
  INFLUXDB_HOST = getenv("INFLUXDB_HOST", "")
  INFLUXDB_ORG = getenv("INFLUXDB_ORG", "tiny")
  INFLUXDB_TOKEN = getenv("INFLUXDB_TOKEN", "")

-  def _create_point(run_id, i, attempt, ref, commit, name, value, log_name, run):
-    point = Point(log_name.replace(':', '_').replace('.', '_')).tag("id", run_id).tag("index", i)
+  def _create_point(run_id, i, attempt, ref, commit, name, value, run):
+    point = Point(BENCHMARK_LOG.value).tag("id", run_id).tag("index", i)
    point = point.tag("device", Device.DEFAULT)
    point = point.tag("attempt", attempt).tag("ref", ref).tag("commit", commit)
    point = point.field(name, value).field("x", run)
@@ -91,12 +91,12 @@ if BENCHMARK_LOG:
      run_id = str(uuid.uuid4())
      if isinstance(event, BenchEvent):
        for event_type, values in _events[event].items():
-          for i, (value, log_name) in enumerate(values):
-            point = _create_point(run_id, i, attempt, ref, commit, f"{event.value}_{event_type}", value, log_name, run)
+          for i, value in enumerate(values):
+            point = _create_point(run_id, i, attempt, ref, commit, f"{event.value}_{event_type}", value, run)
            points.append(point)
      else:
-        for i, (value, log_name) in enumerate(_events[event]):
-          point = _create_point(run_id, i, attempt, ref, commit, event.value, value, log_name, run)
+        for i, value in enumerate(_events[event]):
+          point = _create_point(run_id, i, attempt, ref, commit, event.value, value, run)
          points.append(point)

    write_options = WriteOptions(write_type=WriteType.synchronous, retry_interval=5000, max_retries=5, max_retry_delay=30000, exponential_base=2)
@@ -1,238 +0,0 @@
-import ctypes, struct
-from tinygrad.helpers import ceildiv, getenv, wait_cond, DEBUG
-from tinygrad.runtime.autogen import bnxt, pci
-from tinygrad.runtime.support.system import PCIDevice, System, ipv4_to_gid
-
-BNXT_DEBUG = getenv("BNXT_DEBUG", 0)
-BNXT_ACCESS, BNXT_INIT_MASK, BNXT_RTR_MASK, BNXT_RTS_MASK = 3, 0xd, 0x41515ad, 0xae005
-BNXT_CHIMP_COMM, BNXT_CHIMP_COMM_TRIGGER = 0x0, 0x100
-BNXT_BACKING_STORE = ((0, 2), (1, 0), (2, 2), (3, 0), (4, 2), (5, 0), (6, 0), (14, 2), (15, 0))
-
-def db_value(xid, typ, index, epoch):
-  return (xid & bnxt.DBC_DBC_XID_MASK | bnxt.DBC_DBC_PATH_ROCE | typ | bnxt.BNXT_QPLIB_DBR_VALID) << 32 | \
-         index & bnxt.DBC_DBC_INDEX_MASK | epoch << bnxt.BNXT_QPLIB_DBR_EPOCH_SHIFT
-
-def _pbl(dev, paddrs, queue=False):
-  if len(paddrs) == 1: return 0, paddrs[0]
-  values = [p | bnxt.PTU_PTE_VALID for p in paddrs]
-  if queue:
-    values[-1] |= bnxt.PTU_PTE_LAST
-    if len(values) > 1: values[-2] |= bnxt.PTU_PTE_NEXT_TO_LAST
-  table, table_paddrs = dev.pci_dev.alloc_sysmem(ceildiv(len(values), 512) * 0x1000)
-  table[:len(values) * 8] = struct.pack(f"<{len(values)}Q", *values)
-  if len(table_paddrs) == 1: return 1, table_paddrs[0]
-  top, top_paddrs = dev.pci_dev.alloc_sysmem(0x1000)
-  top[:len(table_paddrs) * 8] = struct.pack(f"<{len(table_paddrs)}Q", *(p | bnxt.PTU_PTE_VALID for p in table_paddrs))
-  return 2, top_paddrs[0]
-
-def _queue(dev, stride:int=16, aux=False):
-  mem, paddrs = dev.pci_dev.alloc_sysmem(0x1000 + aux * 0x400)
-  level, base = _pbl(dev, paddrs, queue=True)
-  return {"mem":mem, "paddrs":paddrs, "stride":stride, "prod":0, "cons":0, "level":level, "base":base}
-
-def _qread(q, i):
-  off = (i & 15) * q["stride"]
-  return q["mem"][off:off + q["stride"]]
-
-def _qwrite(q, i, data, aux=False):
-  off = 0x1000 + i % 128 * 8 if aux else (i & 15) * q["stride"]
-  q["mem"][off:off + len(data)] = data
-
-class BNXTDev:
-  def __init__(self, pci_dev:PCIDevice, ip:str=getenv("BNXT_IP", "10.0.0.1")):
-    self.pci_dev, self.devfmt = pci_dev, pci_dev.pcibus
-    self.bar0, self.db = pci_dev.map_bar(0, fmt='I'), pci_dev.map_bar(2, fmt='Q')
-    pci_dev.write_config(pci.PCI_COMMAND, pci_dev.read_config(pci.PCI_COMMAND, 2) | pci.PCI_COMMAND_MASTER, 2)
-    self.resp, self.resp_pa = pci_dev.alloc_sysmem(0x1000)
-    self.seq = 0
-
-    ver = self.hwrm("ver_get")
-    if DEBUG >= 2: print(f"bnxt {self.devfmt}: firmware {ver.hwrm_fw_maj_8b}.{ver.hwrm_fw_min_8b}.{ver.hwrm_fw_bld_8b}")
-    self.hwrm("func_reset", timeout_ms=40000)
-    caps = self.hwrm("func_qcaps", fid=0xffff)
-    self.mac, self.port_id = int.from_bytes(bytes(caps.mac_address), 'big'), caps.port_id
-    self.hwrm("func_drv_rgtr")
-    self.db_off = self.hwrm("func_qcfg", fid=0xffff).legacy_l2_db_size_kb * 1024
-
-    self.setup_backing_store()
-    self._open_rcfw()
-    self._open_l2()
-    self.local_gid = ipv4_to_gid(ip)
-    gids, mac = (ctypes.c_uint32 * 4)(*(int.from_bytes(self.local_gid[i:i + 4], 'big') for i in (12, 8, 4, 0))), self.mac.to_bytes(6, 'big')
-    smac = (ctypes.c_uint16 * 3)(*(int.from_bytes(mac[i:i + 2], 'big') for i in (0, 2, 4)))
-    self.gid_id = self.rcfw("add_gid", gid=gids, src_mac=smac).xid
-
-    if DEBUG >= 2: print(f"bnxt {self.devfmt}: booted mac={self.mac.to_bytes(6, 'big').hex(':')} gid={self.local_gid.hex()}")
-
-  def hwrm(self, name, timeout_ms=10000, **fields):
-    inp, out = getattr(bnxt, f"struct_hwrm_{name}_input"), getattr(bnxt, f"struct_hwrm_{name}_output")
-    opcode = getattr(bnxt, f"HWRM_{name.upper()}")
-    self.seq = (self.seq + 1) & 0xffff
-    data = bytes(inp(req_type=opcode, cmpl_ring=bnxt.BNXT_HWRM_NO_CMPL_RING, seq_id=self.seq, target_id=bnxt.BNXT_HWRM_TARGET,
-                     resp_addr=self.resp_pa[0], **fields))
-    self.resp[:] = bytes(len(self.resp))
-    System.memory_barrier()
-    for i, w in enumerate(memoryview(bytearray(data.ljust(bnxt.HWRM_MAX_REQ_LEN, b'\0'))).cast('I')):
-      self.bar0[BNXT_CHIMP_COMM // 4 + i] = w
-    self.bar0[BNXT_CHIMP_COMM_TRIGGER // 4] = 1
-    def hdr(): return bnxt.struct_hwrm_resp_hdr.from_buffer_copy(bytes(self.resp[:8]))
-    wait_cond(lambda: (n := hdr().resp_len) and hdr().seq_id == self.seq and self.resp[n - 1], timeout_ms=timeout_ms, msg=f"HWRM {name}")
-    ret = out.from_buffer_copy(bytes(self.resp[:ctypes.sizeof(out)]))
-    assert ret.error_code == 0, f"HWRM {name}: {ret.error_code}"
-    return ret
-
-  def setup_backing_store(self):
-    counts: dict[int, int] = {}
-    for typ, extra in BNXT_BACKING_STORE:
-      caps = self.hwrm("func_backing_store_qcaps_v2", type=typ)
-      size, splits = caps.entry_size, tuple(getattr(caps, f"split_entry_{j}") for j in range(caps.subtype_valid_cnt))
-      counts[typ] = n = counts[0] if typ == 15 else max(caps.min_num_entries, sum(splits) + extra)
-      # a zero bitmap means the type has a single instance 0
-      for instance in [i for i in range(8) if caps.instance_bit_map >> i & 1] or [0]:
-        mem, paddrs = self.pci_dev.alloc_sysmem(ceildiv(n * size, 0x1000) * 0x1000)
-        if caps.ctx_init_value:
-          for off in range(caps.ctx_init_offset, len(mem), size): mem[off] = caps.ctx_init_value
-        lvl, base = _pbl(self, paddrs)
-        self.hwrm("func_backing_store_cfg_v2", type=typ, instance=instance, entry_size=size, num_entries=n, page_dir=base,
-          page_size_pbl_level=lvl, subtype_valid_cnt=len(splits),
-          flags=bnxt.FUNC_BACKING_STORE_CFG_V2_REQ_FLAGS_BS_CFG_ALL_DONE if typ == 15 else 0,
-          **{f"split_entry_{j}": v for j, v in enumerate(splits)})
-
-  def _open_rcfw(self):
-    self.rcfw_first = True
-
-    self.creq = _queue(self)
-    self.creq_id = self.hwrm("ring_alloc", ring_type=bnxt.RING_ALLOC_REQ_RING_TYPE_NQ, page_tbl_addr=self.creq["base"],
-      page_size=12, page_tbl_depth=self.creq["level"], length=16, int_mode=bnxt.RING_ALLOC_REQ_INT_MODE_MSIX).ring_id
-
-    self.cmdq = _queue(self)
-    self.doorbell(self.creq_id, bnxt.DBC_DBC_TYPE_NQ_ARM, 0, 0)
-    init = bnxt.struct_cmdq_init(cmdq_pbl=self.cmdq["base"], creq_ring_id=self.creq_id,
-                                 cmdq_size_cmdq_lvl=16 << bnxt.CMDQ_INIT_CMDQ_SIZE_SFT)
-
-    System.memory_barrier()
-    for i, w in enumerate(memoryview(bytearray(bytes(init))).cast('I')): self.bar0[bnxt.RCFW_COMM_BASE_OFFSET // 4 + i] = w
-
-    _, p = self.pci_dev.alloc_sysmem(0x1000)
-    self.rcfw("initialize_fw", stat_ctx_id=self.hwrm("stat_ctx_alloc", stats_dma_addr=p[0], stats_dma_length=176).stat_ctx_id,
-      flags=bnxt.CMDQ_INITIALIZE_FW_FLAGS_HW_REQUESTER_RETX_SUPPORTED)
-
-    # RoCE notification ring: never armed or serviced, but CQ and L2 ring allocation require one
-    nq = _queue(self)
-    self.nq_id = self.hwrm("ring_alloc", ring_type=bnxt.RING_ALLOC_REQ_RING_TYPE_NQ, page_tbl_addr=nq["base"],
-      page_size=12, page_tbl_depth=nq["level"], length=16, logical_id=1, int_mode=bnxt.RING_ALLOC_REQ_INT_MODE_MSIX).ring_id
-
-  def rcfw(self, name, timeout_ms=20000, **fields):
-    req_t, resp_t = getattr(bnxt, f"struct_cmdq_{name}"), getattr(bnxt, f"struct_creq_{name}_resp")
-    op = getattr(bnxt, f"CMDQ_BASE_OPCODE_{name.upper()}")
-    data = bytes(req_t(opcode=op, cmd_size=(slots := ceildiv(ctypes.sizeof(req_t), 16)), **fields)).ljust(slots * 16, b'\0')
-    for i in range(slots): _qwrite(self.cmdq, self.cmdq["prod"] + i, data[i * 16:(i + 1) * 16])
-
-    self.cmdq["prod"] += slots
-    prod = self.cmdq["prod"] & 0xffff
-    if self.rcfw_first: prod, self.rcfw_first = prod | 1 << bnxt.FIRMWARE_FIRST_FLAG, False
-
-    System.memory_barrier()
-
-    self.bar0[(bnxt.RCFW_COMM_BASE_OFFSET + bnxt.RCFW_PF_VF_COMM_PROD_OFFSET) // 4] = prod
-    self.bar0[(bnxt.RCFW_COMM_BASE_OFFSET + bnxt.RCFW_COMM_TRIG_OFFSET) // 4] = bnxt.RCFW_CMDQ_TRIG_VAL
-
-    def poll():
-      h = bnxt.struct_creq_base.from_buffer_copy(bytes(_qread(self.creq, self.creq["cons"])))
-      return bool(h.v & bnxt.CREQ_BASE_V) != bool((self.creq["cons"] // 16) & 1)
-    wait_cond(poll, timeout_ms=timeout_ms, msg=f"RCFW {name}")
-
-    ret = resp_t.from_buffer_copy(bytes(_qread(self.creq, self.creq["cons"])))
-    self.creq["cons"] += 1
-
-    # NQ_ARM also publishes the CREQ consumer index, which is what frees ring space for the next command
-    self.doorbell(self.creq_id, bnxt.DBC_DBC_TYPE_NQ_ARM, self.creq["cons"] & 15, (self.creq["cons"] // 16) & 1)
-    assert ret.status == 0, f"RCFW {name}: {ret.status}"
-
-    if BNXT_DEBUG >= 1: print(f"bnxt {self.devfmt}: rcfw {name} xid={getattr(ret, 'xid', 0):#x}")
-    return ret
-
-  def doorbell(self, xid, typ, index, epoch):
-    System.memory_barrier()
-    self.db[self.db_off // 8] = db_value(xid, typ, index, epoch)
-
-  # L2 receive path, required for RoCE ingress even though no ethernet receive buffers are posted
-  def _open_l2(self):
-    cq = _queue(self)
-    ci = self.hwrm("ring_alloc", enables=bnxt.RING_ALLOC_REQ_ENABLES_NQ_RING_ID_VALID, ring_type=bnxt.RING_ALLOC_REQ_RING_TYPE_L2_CMPL,
-      page_tbl_addr=cq["base"], page_size=12, page_tbl_depth=cq["level"], length=16, nq_ring_id=self.nq_id).ring_id
-    rx = _queue(self)
-    ri = self.hwrm("ring_alloc", enables=bnxt.RING_ALLOC_REQ_ENABLES_NQ_RING_ID_VALID |
-      bnxt.RING_ALLOC_REQ_ENABLES_RX_BUF_SIZE_VALID, ring_type=bnxt.RING_ALLOC_REQ_RING_TYPE_RX, page_tbl_addr=rx["base"],
-      page_size=12, page_tbl_depth=rx["level"], length=16, rx_buf_size=640, nq_ring_id=self.nq_id).ring_id
-    vi = self.hwrm("vnic_alloc").vnic_id
-    self.hwrm("vnic_cfg", enables=bnxt.VNIC_CFG_REQ_ENABLES_MRU | bnxt.VNIC_CFG_REQ_ENABLES_DEFAULT_RX_RING_ID |
-      bnxt.VNIC_CFG_REQ_ENABLES_DEFAULT_CMPL_RING_ID, vnic_id=vi, mru=9018,
-      default_rx_ring_id=ri, default_cmpl_ring_id=ci)
-    self.hwrm("cfa_l2_filter_alloc", flags=bnxt.CFA_L2_FILTER_ALLOC_REQ_FLAGS_PATH_RX,
-      enables=bnxt.CFA_L2_FILTER_ALLOC_REQ_ENABLES_L2_ADDR | bnxt.CFA_L2_FILTER_ALLOC_REQ_ENABLES_L2_ADDR_MASK |
-      bnxt.CFA_L2_FILTER_ALLOC_REQ_ENABLES_DST_ID, l2_addr=tuple(self.mac.to_bytes(6, 'big')), l2_addr_mask=(0xff,) * 6, dst_id=vi)
-
-  def register_mem(self, paddrs:list[int], size:int, log_page_size:int=12) -> int:
-    level, base = _pbl(self, paddrs[:ceildiv(size, 1 << log_page_size)])
-    return self.rcfw("register_mr", flags=bnxt.CMDQ_REGISTER_MR_FLAGS_ALLOC_MR,
-      log2_pg_size_lvl=level << bnxt.CMDQ_REGISTER_MR_LVL_SFT | log_page_size << bnxt.CMDQ_REGISTER_MR_LOG2_PG_SIZE_SFT,
-      access=bnxt.CMDQ_REGISTER_MR_ACCESS_LOCAL_WRITE | bnxt.CMDQ_REGISTER_MR_ACCESS_REMOTE_WRITE,
-      log2_pbl_pg_size=12, pbl=base, va=paddrs[0], mr_size=size).xid
-
-class BNXTQP:
-  def __init__(self, dev:BNXTDev):
-    self.dev, self.sq_psn, self.msn = dev, 0, 0
-
-    self.cqq = _queue(dev, ctypes.sizeof(bnxt.struct_cq_base))
-    self.cq_id = dev.rcfw("create_cq", cq_size=16, pbl=self.cqq["base"],
-                          pg_size_lvl=self.cqq["level"], cq_fco_cnq_id=dev.nq_id).xid
-
-    self.sq = _queue(dev, aux=True)
-    self.qpn = dev.rcfw("create_qp", type=bnxt.CMDQ_CREATE_QP_TYPE_RC,
-      sq_size=16, sq_fwo_sq_sge=1, scq_cid=self.cq_id, rcq_cid=self.cq_id,
-      sq_pbl=self.sq["base"], sq_pg_size_sq_lvl=self.sq["level"]).xid
-    self.qp_op(1, BNXT_INIT_MASK, access=BNXT_ACCESS, pkey=0xffff)
-
-  def qp_op(self, state, mask, network_type=0, **fields):
-    self.dev.rcfw("modify_qp", qp_cid=self.qpn, modify_mask=mask,
-      network_type_en_sqd_async_notify_new_state=state | network_type, **fields)
-
-  def connect(self, qpn:int, gid:bytes, mac:int):
-    network_type = bnxt.CMDQ_MODIFY_QP_NETWORK_TYPE_ROCEV2_IPV4
-    dgid = (ctypes.c_uint32 * 4)(*(int.from_bytes(gid[i:i + 4], 'little') for i in (0, 4, 8, 12)))
-    dmac = (ctypes.c_uint16 * 3)(*(int.from_bytes(mac.to_bytes(6, 'big')[i:i + 2], 'little') for i in (0, 2, 4)))
-
-    self.qp_op(2, BNXT_RTR_MASK, network_type=network_type, qp_type=bnxt.CMDQ_MODIFY_QP_QP_TYPE_RC, access=BNXT_ACCESS,
-      pkey=0xffff, dgid=dgid, sgid_index=self.dev.gid_id, hop_limit=64, dest_mac=dmac,
-      path_mtu_pingpong_push_enable=bnxt.CMDQ_MODIFY_QP_PATH_MTU_MTU_1024, max_dest_rd_atomic=4,
-      dest_qp_id=qpn)
-    self.qp_op(3, BNXT_RTS_MASK, network_type=network_type, qp_type=bnxt.CMDQ_MODIFY_QP_QP_TYPE_RC, access=BNXT_ACCESS,
-      max_rd_atomic=1)
-
-    if BNXT_DEBUG >= 1: print(f"bnxt: QP {self.qpn:#x} connected (remote={qpn:#x})")
-
-  def _poll(self, timeout):
-    def poll():
-      base = bnxt.struct_cq_base.from_buffer_copy(bytes(_qread(self.cqq, self.cqq["cons"])))
-      return bool(base.cqe_type_toggle & bnxt.CQ_BASE_TOGGLE) == (not bool((self.cqq["cons"] // 16) & 1))
-    wait_cond(poll, timeout_ms=timeout, msg="BNXT CQ")
-    raw = bytes(_qread(self.cqq, self.cqq["cons"]))
-    self.cqq["cons"] += 1
-    self.dev.doorbell(self.cq_id, bnxt.DBC_DBC_TYPE_CQ, self.cqq["cons"] & 15, (self.cqq["cons"] // 16) & 1)
-    return raw
-
-  def rdma_write(self, rva, rkey, lva, lkey, size, timeout_ms=20000):
-    start = self.sq["prod"] & 15
-    hdr = bytes(bnxt.struct_sq_rdma_hdr(wqe_type=bnxt.SQ_RDMA_HDR_WQE_TYPE_WRITE_WQE,
-                flags=bnxt.SQ_SEND_FLAGS_SIGNAL_COMP, wqe_size=3, length=size, remote_va=rva, remote_key=rkey))
-    for i, data in enumerate((hdr[:16], hdr[16:32], bytes(bnxt.struct_sq_sge(va_or_pa=lva, l_key=lkey, size=size)))):
-      _qwrite(self.sq, start + i, data)
-    nxt = (self.sq_psn + max(1, ceildiv(size, 1024))) & 0xffffff
-    value = start << bnxt.SQ_MSN_SEARCH_START_IDX_SFT | nxt << bnxt.SQ_MSN_SEARCH_NEXT_PSN_SFT | self.sq_psn
-    _qwrite(self.sq, self.msn, struct.pack("<Q", value), aux=True)
-
-    self.msn, self.sq_psn, self.sq["prod"] = (self.msn + 1) % 128, nxt, self.sq["prod"] + 3
-    self.dev.doorbell(self.qpn, bnxt.DBC_DBC_TYPE_SQ, self.sq["prod"] & 15, (self.sq["prod"] // 16) & 1)
-    cqe = bnxt.struct_cq_req.from_buffer_copy(self._poll(timeout_ms))
-    assert cqe.status == 0
@@ -1,118 +0,0 @@
-#!/usr/bin/env python3
-"""Send and validate one RDMA WRITE between two Broadcom BNXT hosts.
-
-This follows ``extra/mlx_driver/connect.py``: sync the driver, start the remote
-endpoint over SSH, exchange QP/GID/MAC/MR metadata, move both RC QPs to RTS,
-write bytes into the remote MR, and verify the bytes on the remote host.
-
-Both PCI functions must be unbound from bnxt_en/bnxt_re first.
-"""
-import json
-import os
-import subprocess
-import sys
-from typing import Any, IO
-
-TINYGRAD = os.path.normpath(os.path.join(os.path.dirname(os.path.abspath(__file__)), "../.."))
-sys.path.insert(0, TINYGRAD)
-
-from extra.bnxt_driver.bnxtdev import BNXTDev, BNXTQP
-from tinygrad.runtime.support.system import PCIDevice
-
-REMOTE_HOST = os.getenv("REMOTE_HOST", "192.168.52.213")
-REMOTE_USER = os.getenv("REMOTE_USER", "nimlgen")
-LOCAL_PCI = os.getenv("BNXT_PCI", "0000:41:00.0")
-REMOTE_PCI = os.getenv("REMOTE_PCI", "0000:41:00.0")
-LOCAL_IP = os.getenv("LOCAL_IP", "10.0.200.5")
-REMOTE_IP = os.getenv("REMOTE_IP", "10.0.200.6")
-MESSAGE = os.getenv("RDMA_MESSAGE", "Test message, rdma works!").encode()
-REMOTE = f"{REMOTE_USER}@{REMOTE_HOST}"
-SSH = ["ssh", "-o", "BatchMode=yes", "-o", "ConnectTimeout=10", "-o", "StrictHostKeyChecking=accept-new", REMOTE]
-SYNC_FILES = ("tinygrad/runtime/autogen/bnxt.py", "tinygrad/runtime/support/system.py",
-              "extra/bnxt_driver/bnxtdev.py", "extra/bnxt_driver/connect.py")
-
-def read_json(stream:IO[str], what:str) -> dict[str, Any]:
-  for line in iter(stream.readline, ""):
-    print(f"  [remote] {line}", end="")
-    try: value = json.loads(line)
-    except json.JSONDecodeError: continue
-    if isinstance(value, dict): return value
-  raise RuntimeError(f"remote exited before publishing {what}")
-
-def wait_line(stream:IO[str], text:str) -> str:
-  for line in iter(stream.readline, ""):
-    print(f"  [remote] {line}", end="")
-    if text in line: return line
-  raise RuntimeError(f"remote exited before reporting {text!r}")
-
-def send_line(stream:IO[str], value:str|dict[str, Any]):
-  stream.write((json.dumps(value) if isinstance(value, dict) else value) + "\n")
-  stream.flush()
-
-def qp_info(dev:BNXTDev, qp:BNXTQP) -> dict[str, Any]:
-  return {"qpn":qp.qpn, "mac":dev.mac.to_bytes(6, "big").hex(), "gid":dev.local_gid.hex()}
-
-def server():
-  dev = BNXTDev(PCIDevice("bnxt", os.getenv("BNXT_PCI", "0000:41:00.0")), ip=os.getenv("BNXT_IP", REMOTE_IP))
-  qp = BNXTQP(dev)
-  print(json.dumps(qp_info(dev, qp)), flush=True)
-
-  peer = json.loads(sys.stdin.readline())
-  qp.connect(peer["qpn"], bytes.fromhex(peer["gid"]), int(peer["mac"], 16))
-  print("connected", flush=True)
-
-  target, target_paddrs = dev.pci_dev.alloc_sysmem(0x1000)
-  target[:0x1000] = bytes(0x1000)
-  rkey = dev.register_mem(target_paddrs, 0x1000)
-  print(json.dumps({"target_addr":target_paddrs[0], "rkey":rkey}), flush=True)
-
-  assert sys.stdin.readline().strip() == "done"
-  received = bytes(target).rstrip(b"\0")
-  print(f"AS TEXT: {received.decode(errors='replace')!r}", flush=True)
-  print(json.dumps({"data":received.hex()}), flush=True)
-
-def sync_remote():
-  if os.getenv("SYNC", "1") == "0": return
-  print("syncing BNXT driver to remote")
-  subprocess.run(["rsync", "-azR", *SYNC_FILES, f"{REMOTE}:~/tinygrad/"], cwd=TINYGRAD, check=True)
-
-def start_remote() -> subprocess.Popen[str]:
-  print("booting remote")
-  command = (f"cd ~/tinygrad && sudo env PYTHONPATH=. PYTHONUNBUFFERED=1 BNXT_DEBUG={os.getenv('BNXT_DEBUG', '0')} "
-             f"BNXT_PCI={REMOTE_PCI} BNXT_IP={REMOTE_IP} python3 extra/bnxt_driver/connect.py --server")
-  return subprocess.Popen(SSH + [command], stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=sys.stderr, text=True)
-
-def client():
-  assert 0 < len(MESSAGE) <= 0x1000
-  sync_remote()
-  remote = start_remote()
-  assert remote.stdin is not None and remote.stdout is not None
-  remote_info = read_json(remote.stdout, "QP information")
-  print("booting local")
-  dev = BNXTDev(PCIDevice("bnxt", LOCAL_PCI), ip=LOCAL_IP)
-  qp = BNXTQP(dev)
-
-  send_line(remote.stdin, qp_info(dev, qp))
-  wait_line(remote.stdout, "connected")
-  qp.connect(remote_info["qpn"], bytes.fromhex(remote_info["gid"]), int(remote_info["mac"], 16))
-  print("both QPs in RTS")
-
-  remote_target = read_json(remote.stdout, "MR information")
-  source, source_paddrs = dev.pci_dev.alloc_sysmem(0x1000)
-  source[:len(MESSAGE)] = MESSAGE
-  lkey = dev.register_mem(source_paddrs, 0x1000)
-  print(f"RDMA WRITE {len(MESSAGE)}B to remote phys 0x{remote_target['target_addr']:x}")
-  qp.rdma_write(remote_target["target_addr"], remote_target["rkey"], source_paddrs[0], lkey, len(MESSAGE))
-
-  send_line(remote.stdin, "done")
-  wait_line(remote.stdout, "AS TEXT")
-  result = read_json(remote.stdout, "RDMA result")
-  assert bytes.fromhex(result["data"]) == MESSAGE
-  print("RDMA WRITE data verified")
-
-  remote.stdin.close()
-  assert remote.wait() == 0
-  print("RDMA WRITE test complete")
-
-if __name__ == "__main__":
-  server() if "--server" in sys.argv else client()
@@ -1,46 +0,0 @@
-#!/usr/bin/env python3
-"""Local BNXT RoCEv2 RDMA WRITE loopback using the firmware's PHY loopback mode.
-
-The kernel bnxt_en/bnxt_re modules must be unloaded first.
-
-  sudo PYTHONPATH=. BNXT_PCI=0000:41:00.0 BNXT_IP=10.0.200.5 python3 extra/bnxt_driver/loopback.py
-"""
-import os
-import sys
-import time
-
-sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "../.."))
-
-from extra.bnxt_driver.bnxtdev import BNXTDev, BNXTQP
-from tinygrad.runtime.autogen import bnxt
-from tinygrad.runtime.support.system import PCIDevice
-
-BUF_SIZE = 0x1000
-BNXT_PCI = os.getenv("BNXT_PCI", "0000:41:00.0")
-BNXT_IP = os.getenv("BNXT_IP", "10.0.200.5")
-
-if __name__ == "__main__":
-  print(f"[init] BNXT at {BNXT_PCI}")
-  dev = BNXTDev(PCIDevice("bnxt", BNXT_PCI), ip=BNXT_IP)
-  tx_qp, rx_qp = BNXTQP(dev), BNXTQP(dev)
-  print(f"[init] loopback-connect TX QP 0x{tx_qp.qpn:x} <-> RX QP 0x{rx_qp.qpn:x}")
-  tx_qp.connect(rx_qp.qpn, dev.local_gid, dev.mac)
-  rx_qp.connect(tx_qp.qpn, dev.local_gid, dev.mac)
-
-  src, src_paddrs = dev.pci_dev.alloc_sysmem(BUF_SIZE)
-  dst, dst_paddrs = dev.pci_dev.alloc_sysmem(BUF_SIZE)
-  message = b"Hello from BNXT RoCE PHY loopback!"
-  src[:BUF_SIZE], dst[:BUF_SIZE] = bytes(BUF_SIZE), bytes(BUF_SIZE)
-  src[:len(message)] = message
-  lkey = dev.register_mem(src_paddrs, BUF_SIZE)
-  rkey = dev.register_mem(dst_paddrs, BUF_SIZE)
-
-  print("[loopback] enabling local PHY loopback")
-  dev.hwrm("port_phy_cfg", port_id=dev.port_id, enables=bnxt.PORT_PHY_CFG_REQ_ENABLES_LPBK, lpbk=bnxt.PORT_PHY_CFG_REQ_LPBK_LOCAL)
-  time.sleep(1)
-  tx_qp.rdma_write(dst_paddrs[0], rkey, src_paddrs[0], lkey, len(message))
-  got = bytes(dst[:len(message)])
-  print(f"[result] {got!r}")
-  assert got == message
-  print("BNXT RoCE PHY loopback RDMA WRITE passed")
-  dev.hwrm("port_phy_cfg", port_id=dev.port_id, enables=bnxt.PORT_PHY_CFG_REQ_ENABLES_LPBK, lpbk=bnxt.PORT_PHY_CFG_REQ_LPBK_NONE)
@@ -35,7 +35,7 @@ def compile_net(linear:UOp, output_bufs:List[Buffer]) -> Tuple[Dict[str,str], Li
    return name

  for call in iter_kernel_calls(linear):
-    arg_uops = [b for b in call.src[1:] if not b.is_bound_var]
+    arg_uops = [b for b in call.src[1:] if b.op is not Ops.BIND]
    prg = to_program(call.src[0], Device[arg_uops[0].device].renderer)
    info = prg.arg
    functions[info.function_name] = prg.src[2].arg
@@ -241,7 +241,8 @@ export default {model_name};
 def export_model(model, target:str, *inputs, model_name: Optional[str] = "model", stream_weights=False):
  assert Device.DEFAULT in EXPORT_SUPPORTED_DEVICE, f"only {', '.join(EXPORT_SUPPORTED_DEVICE)} are supported"

-  with Context(JIT=2): linear, output_bufs = jit_model(model, *inputs)
+  # NOTE: NUM_CPU_THREADS=1, since export does not support threading
+  with Context(JIT=2, NUM_CPU_THREADS=1): linear, output_bufs = jit_model(model, *inputs)
  functions, statements, bufs, bufs_to_save = compile_net(linear, output_bufs)
  state = get_state_dict(model)
  weight_names = {(id(b), b.offset, b.size, b.dtype): name for name, x in state.items() if (b:=x.uop.base.realized) is not None}
@@ -462,7 +462,7 @@ def test_matmul():
    lds = UOp.placeholder((lds_size,), dtypes.uint8, 0, AddrSpace.LOCAL)
    sink = UOp.sink(A.base, B.base, C.base, lds, *gidxs, *lidxs, arg=KernelInfo(name=colored("kernel", "cyan"),
                                                                                  estimates=Estimates(ops=N*N*N*2, mem=N*N*4*3)))
-    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=(x, dtypes.void)) for x in insts]))))
+    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
  c = Tensor.custom_kernel(a, b, c, fxn=asm_kernel)[2]
  linear = c.schedule_linear()

@@ -122,10 +122,9 @@ def custom_mxfp4_gemm(C:UOp, A:UOp, B:UOp, scale_a:UOp, scale_b:UOp, *extra:UOp,
  groups_x, groups_y = UOp.special(ceildiv(N, tile_n), "gidx0"), UOp.special(ceildiv(M, tile_m), "gidx1")
  lds = UOp.placeholder((163840,), dtypes.uint8, 0, AddrSpace.LOCAL)
  sink = UOp.sink(C.base, A.base, B.base, scale_a.base, scale_b.base, *(x.base for x in extra), lds, threads, groups_x, groups_y,
-                  arg=KernelInfo(f"mxfp4_gemm_{M}_{N}_{K}",
-                                 estimates=Estimates(ops=2*M*N*K, mem=(M*half_k+N*half_k)*A.dtype.itemsize+M*N*C.dtype.itemsize)))
+                  arg=KernelInfo(f"custom_mxfp4_gemm_{M}_{N}_{K}", estimates=Estimates(ops=2*M*N*K)))
  insts = build_kernel(M, N, K, tile_m, tile_n)
-  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple(UOp(Ops.INS, arg=(x, dtypes.void)) for x in insts))))
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple(UOp(Ops.INS, arg=x) for x in insts))))

 def _mxfp4_gemm_quantized(a_q:Tensor, b_q:Tensor, scale_a:Tensor, scale_b:Tensor) -> Tensor:
  M, half_k = a_q.shape
@@ -215,7 +214,7 @@ def custom_uop_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
  k = UOp.range(K, 0, AxisType.REDUCE)
  mul = (A.flatten().index((m*UOp.const(K)+k))*
         B.flatten().index((k*UOp.const(N)+n))).cast(dtypes.float32)
-  red = mul.reduce(k, arg=Ops.ADD).cast(C.dtype)
+  red = mul.reduce(k, arg=Ops.ADD, dtype=dtypes.float32).cast(C.dtype)
  store = C.flatten().index((m*UOp.const(N)+n)).store(red).end(m, n)
  return store.sink(arg=KernelInfo(name=f'uop_gemm_{M}_{N}_{K}'))

@@ -113,6 +113,7 @@ if __name__ == "__main__":
      }
    elif GEMM_VARIATION == "hcopt" and M == N == K == 4096 and DTYPE_IN == dtypes.half and DTYPE_OUT == dtypes.half and DTYPE_ACC == dtypes.float:
      print("Using CUDA and generated hcopt")
+      # [Opt(op=OptOps.TC, axis=0, amt=0), Opt(op=OptOps.UPCAST, axis=0, amt=4), Opt(op=OptOps.UPCAST, axis=1, amt=4), Opt(op=OptOps.LOCAL, axis=1, amt=4)]
      prog = CUDAProgram(device, "wmma_example", compiler.compile(open(os.path.join(script_dir, 'max_kernels/nv.fp16_fp32_fp16.hcopt.cu')).read()))
      args = (c, a, b)
      kwargs = {
@@ -1,32 +1,10 @@
 import functools, pathlib
 from tinygrad import Tensor, dtypes
-from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType
-from tinygrad.helpers import getenv
+from tinygrad.uop.ops import UOp, Ops, KernelInfo
 from tinygrad.renderer import Estimates
 from tinygrad.runtime.support.compiler_amd import HIPCCCompiler
 from extra.gemm.cdna_asm_gemm import quantize_mxfp8, _mx_block_scale, _mx_block_scale_3d

-ZERO_OPTIM = getenv("ZERO_OPTIM", 0)
-
-def reduce_scatter_devaxis(out:Tensor, shard_axis:int=0) -> Tensor:
-  # out: sharded on the device axis, shape (ndev, *rest); return the device-axis sum left sharded on shard_axis.
-  u = out.uop
-  devs, rest = u.device, u.shape[1:]
-  assert rest[shard_axis] % len(devs) == 0, f"reduce_scatter needs even shards: {rest[shard_axis]} % {len(devs)}"
-  # reach the raw per-device buffer below the UNSHARD, keeping the AFTERs so reads stay ordered after the kernel writes
-  node, barriers = u, []
-  while node.op is not Ops.UNSHARD:
-    if node.op is Ops.AFTER: barriers += node.src[1:]
-    node = node.src[0]
-  mbuf = node.src[0].after(*barriers) if barriers else node.src[0]
-  sz = rest[shard_axis] // len(devs)
-  shards = []
-  for i in range(len(devs)):
-    bounds = tuple((0,s) if a != shard_axis else (i*sz,(i+1)*sz) for a,s in enumerate(rest))
-    contribs = [mbuf.mselect(j).reshape(rest).shrink(bounds).copy_to_device(devs[i]) for j in range(len(devs))]
-    shards.append(functools.reduce(lambda a,b: a.alu(Ops.ADD, b), contribs))
-  return Tensor(UOp.mstack(*shards).unshard(shard_axis, UOp.range(len(devs), -1, AxisType.DEVICE)), device=devs)
-
@functools.cache
 def custom_hk_grouped_mxfp8_gemm(C:UOp, A:UOp, B:UOp, scale_A:UOp, scale_B:UOp, *extra:UOp, dname:str, n_experts:int) -> UOp:
  M, K = A.shape
@@ -80,8 +58,7 @@ def grouped_mx_wgrad(g:Tensor, xg:Tensor, expert_off:Tensor, n_experts:int) -> T
  out = Tensor(inv.uop.unshard(0), device=g.device) if is_multi else inv
  out = Tensor.custom_kernel(out, gT, xT, g_si, x_si, expert_off,
                             fxn=functools.partial(custom_hk_grouped_mxfp8_wgrad, dname=dname, n_experts=n_experts))[0]
-  if is_multi and ZERO_OPTIM: out = reduce_scatter_devaxis(out, 0)
-  else: out = out.sum(0) if is_multi else out.squeeze(0)
+  out = out.sum(0) if is_multi else out.squeeze(0)
  return out.reshape(n_experts, N, K)

 def mx_pack_3d(e8:Tensor) -> Tensor:
@@ -53,7 +53,7 @@ def _ggather_bwd(gradient:UOp, kernel:UOp) -> tuple:
    g, m, j, jo, ji = _kv_ranges(Gk, M, Dk, _blk_for(Dk))
    row = idx.index(g, m).cast(dtypes.weakint)
    val = gout.index(g, m, j).load().cast(dtypes.float32)
-    atomic = UOp(Ops.CUSTOM, src=(gtab.index(g, row, j), val), arg=(atomic_str, dtypes.void))
+    atomic = UOp(Ops.CUSTOM, dtypes.void, (gtab.index(g, row, j), val), arg=atomic_str)
    return atomic.end(g, m, jo, ji).sink(arg=KernelInfo(name=f"ggather_bwd_{M}_{Dk}", opts_to_apply=()))
  grad_table = Tensor.custom_kernel(gt, go, Tensor(idx_u, device=dev), fxn=_bwd_kernel)[0]
  return (None, grad_table.cast(table_u.dtype).uop, None)
@@ -223,7 +223,7 @@ def test_matmul():
    lds = UOp.placeholder((lds_size,), dtypes.uint8, 0, AddrSpace.LOCAL)
    sink = UOp.sink(A.base, B.base, C.base, lds, *gidxs, *lidxs,
                    arg=KernelInfo(name=colored("kernel","cyan"), estimates=Estimates(ops=N*N*N*2, mem=N*N*2*3)))
-    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=(x, dtypes.void)) for x in insts]))))
+    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))

  c = Tensor.custom_kernel(a, b, c, fxn=asm_kernel)[2]
  linear = c.schedule_linear()
@@ -1,7 +1,6 @@
 from tinygrad import Tensor, dtypes, Context
 from tinygrad.helpers import getenv
 from tinygrad.codegen.opt import Opt, OptOps
-from tinygrad.uop.ops import AxisType
 from tinygrad.engine.realize import run_linear
 from dataclasses import replace

@@ -14,17 +13,17 @@ if __name__ == "__main__":
  C = A.matmul(B)
  if getenv("GEMV"):
    opts = [
-      Opt(op=OptOps.SPLIT, axis=1, arg=(8, AxisType.UNROLL)),
-      Opt(op=OptOps.SPLIT, axis=1, arg=(32, AxisType.GROUP_REDUCE)),
+      Opt(op=OptOps.UNROLL, axis=0, amt=8),
+      Opt(op=OptOps.GROUP, axis=0, amt=32),
    ]
  else:
    opts = [
      Opt(op=OptOps.TC, axis=0, amt=0),
-      Opt(op=OptOps.SPLIT, axis=0, arg=(4, AxisType.UPCAST)),
-      Opt(op=OptOps.SPLIT, axis=1, arg=(8, AxisType.UPCAST)),
-      Opt(op=OptOps.SPLIT, axis=0, arg=(2, AxisType.LOCAL)),
-      Opt(op=OptOps.SPLIT, axis=1, arg=(2, AxisType.LOCAL)),
-      Opt(op=OptOps.SPLIT, axis=0, arg=(2, AxisType.LOCAL)),
+      Opt(op=OptOps.UPCAST, axis=0, amt=4),
+      Opt(op=OptOps.UPCAST, axis=1, amt=8),
+      Opt(op=OptOps.LOCAL, axis=0, amt=2),
+      Opt(op=OptOps.LOCAL, axis=1, amt=2),
+      Opt(op=OptOps.LOCAL, axis=0, amt=2),
    ]
  linear = C.schedule_linear()
  call = linear.src[-1]
@@ -79,7 +79,7 @@ if __name__ == "__main__":
  linear, var_vals = C.linear_with_vars()
  last_call = linear.src[-1]
  ast = last_call.src[0]
-  bufs = [s.buffer for s in last_call.src[1:] if not s.is_bound_var]
+  bufs = [s.buffer for s in last_call.src[1:] if s.op is not Ops.BIND]

  src = compiled.asm["ptx"]
  # specify the shared memory here so we don't need to do it dynamically
@@ -1,29 +0,0 @@
-from __future__ import annotations
-import functools
-from tinygrad import Tensor
-from tinygrad.uop.ops import UOp
-
-def rmsnorm_mul_fwd(x_in:Tensor, weight:Tensor, eps:float) -> tuple[Tensor, Tensor]:
-  x = x_in.float()
-  rrms = (x.square().mean(-1, keepdim=True) + eps).rsqrt()
-  return ((x * rrms) * weight.float()).cast(x_in.dtype), rrms
-
-@functools.cache
-def _rmsnorm_mul_fwd_fxn(x_in_p, w_p, eps, device):
-  return rmsnorm_mul_fwd(Tensor(x_in_p, device=device), Tensor(w_p, device=device), eps)
-
-def _rmsnorm_mul_bwd(grad:UOp, call:UOp) -> tuple:
-  x = Tensor(call.src[1]).float(); weight = Tensor(call.src[2]).float()
-  rrms = Tensor(call.gettuple(1))
-  x_normed = x * rrms                                  # recompute unweighted normed (x is call.src[1])
-  d_y = Tensor(grad).float()
-  dxn = d_y * weight                                   # d/d(x_normed)
-  d_x = rrms * (dxn - x_normed * (dxn * x_normed).mean(-1, keepdim=True))
-  dw = d_y * x_normed
-  d_weight = dw.sum(axis=tuple(range(dw.ndim - 1)))    # reduce batch/seq -> [dim]
-  return (d_x.cast(call.src[1].dtype).uop, d_weight.cast(call.src[2].dtype).uop)
-
-def rmsnorm_mul(x_in:Tensor, weight:Tensor, eps:float) -> tuple[Tensor, Tensor]:
-  fxn = _rmsnorm_mul_fwd_fxn(x_in.as_param(0).uop, weight.as_param(1).uop, eps, x_in.device)
-  call = UOp.maketuple(fxn[0].uop, fxn[1].uop).call(x_in.uop, weight.uop, grad_fxn=_rmsnorm_mul_bwd)
-  return Tensor(call.gettuple(0)), Tensor(call.gettuple(1))
@@ -16,12 +16,9 @@ def _do_reset_device(pci_bus): os.system(f"sudo sh -c 'echo 1 > /sys/bus/pci/dev
 def _is_module_loaded(name: str) -> bool: return os.path.isdir(f"/sys/module/{name}")

 def cmd_remove_module(args):
-  modules = ["nvidia_drm", "nvidia_modeset", "nvidia_uvm", "nvidia"] if args.backend == "nv" else ["amdgpu"]
+  modules = ["nvidia_drm", "nvidia_modeset", "nvidia_uvm", "nvidia", "ast"] if args.backend == "nv" else ["amdgpu"]
  to_unload = [m for m in modules if _is_module_loaded(m)]
  if not to_unload: print("Kernel modules are not loaded")
-  elif getattr(args, "expect", False):
-    print(f"Kernel modules are loaded: {to_unload}")
-    sys.exit(1)
  else:
    print("Removing kernel modules:", ", ".join(to_unload))
    try: subprocess.run(["sudo", "modprobe", "-r", *to_unload], check=True)
@@ -63,19 +60,17 @@ def cmd_show_pids(args):

 def cmd_kill_pids(args):
  devs = scan_devs_based_on_lock(prefix:={"amd":"am", "nv":"nv"}[args.backend], args)
-  use_sudo = not getattr(args, "sudoless", False)

  for dev in devs:
    for i in range(128):
      if i > 0: time.sleep(0.2)

      try:
-        try: pid = subprocess.check_output((['sudo'] if use_sudo else []) +
-                                           ['lsof', temp(f'{prefix}_{dev}.lock')]).decode('utf-8').strip().split('\n')[1].split()[1]
+        try: pid = subprocess.check_output(['sudo', 'lsof', temp(f'{prefix}_{dev}.lock')]).decode('utf-8').strip().split('\n')[1].split()[1]
        except subprocess.CalledProcessError: break

        print(f"Killing process {pid} (which uses {dev})")
-        subprocess.run((['sudo'] if use_sudo else []) + ['kill', '-9', pid], check=True)
+        subprocess.run(['sudo', 'kill', '-9', pid], check=True)
      except subprocess.CalledProcessError as e:
        print(f"Failed to kill process for device {dev}: {e}", file=sys.stderr)

@@ -84,7 +79,6 @@ def add_common_commands(parent_subparsers):
  p_insmod.set_defaults(func=cmd_insert_module)

  p_rmmod = parent_subparsers.add_parser("rmmod", help="Remove a kernel module")
-  p_rmmod.add_argument("--expect", action="store_true", help="Just assert that module is already unloaded")
  p_rmmod.set_defaults(func=cmd_remove_module)

  p_reset = parent_subparsers.add_parser("reset", help="Reset a device")
@@ -97,20 +91,17 @@ def add_common_commands(parent_subparsers):

  p_reset = parent_subparsers.add_parser("kill_pids", help="Kill pids of processes using the device")
  p_reset.add_argument("--pci_bus", default="", help="PCI bus ID of the device")
-  p_reset.add_argument("--sudoless", action="store_true", help="Do not use sudo when detecting or killing pids")
  p_reset.set_defaults(func=cmd_kill_pids)

 if __name__ == "__main__":
  parser = argparse.ArgumentParser()
  backend_subparsers = parser.add_subparsers(dest="backend", required=True, metavar="{nv,amd}", help="Hardware backend to target")

-  nv_parser = backend_subparsers.add_parser("nv", aliases=["NV"], help="NVIDIA GPUs")
-  nv_parser.set_defaults(backend="nv")
+  nv_parser = backend_subparsers.add_parser("nv", help="NVIDIA GPUs")
  nv_commands = nv_parser.add_subparsers(dest="command", required=True)
  add_common_commands(nv_commands)

-  amd_parser = backend_subparsers.add_parser("amd", aliases=["AMD"], help="AMD GPUs")
-  amd_parser.set_defaults(backend="amd")
+  amd_parser = backend_subparsers.add_parser("amd", help="AMD GPUs")
  amd_commands = amd_parser.add_subparsers(dest="command", required=True)
  add_common_commands(amd_commands)

@@ -1,32 +1,33 @@
 from __future__ import annotations
-from typing import cast
-import os, ctypes, struct, functools, importlib, mmap, errno, contextlib, sys, itertools, atexit
+from typing import cast, Any, Callable
+import os, ctypes, struct, hashlib, functools, importlib, mmap, errno, array, contextlib, sys, weakref, itertools, collections, atexit
 assert sys.platform != 'win32'
 from dataclasses import dataclass
-from tinygrad.runtime.support.hcq2 import HCQ2Compiled, HCQAllocator, make_buf, hcq_size_var
+from tinygrad.runtime.support.hcq2 import HCQ2Compiled, HCQAllocator, HCQ2Buffer, encode_kernargs_clike, make_cmdbuf
+from tinygrad.runtime.support.hcq2 import make_binary_patch
 from tinygrad.uop.ops import sint, UOp
-from tinygrad.device import BufferSpec, Buffer, Device
+from tinygrad.device import Compiled, BufferSpec, Buffer, Device
 from tinygrad.dtype import dtypes
-from tinygrad.helpers import getenv, round_up, data64_le, DEBUG, PROFILE, lo32, hi32
-from tinygrad.helpers import ceildiv, unwrap, pluralize, to_tuple
+from tinygrad.helpers import getenv, round_up, data64_le, DEBUG, PROFILE, ProfileEvent, lo32, hi32, colored, prod, ContextVar, TracingKey
+from tinygrad.helpers import VIZ, ceildiv, unwrap, pluralize, to_tuple
 from tinygrad.renderer.cstyle import HIPRenderer, HIPCCRenderer
 from tinygrad.renderer.llvmir import AMDLLVMRenderer
-from tinygrad.runtime.autogen import kfd, hsa, amdgpu_kd, amdgpu_drm
+from tinygrad.runtime.autogen import kfd, hsa, sqtt, amdgpu_kd, amdgpu_drm
 from tinygrad.runtime.autogen.am import am
 from tinygrad.runtime.support.elf import elf_loader
 from tinygrad.runtime.support.hcq import FileIOInterface, HCQBuffer, MMIOInterface, hcq_filter_visible_devices
 from tinygrad.runtime.support.am.amdev import AMDev, AMMemoryManager
 from tinygrad.runtime.support.amd import AMDReg, AMDIP, import_module, import_soc, import_pmc
 from tinygrad.runtime.support.system import PCIIfaceBase, PCIAllocationMeta, USBPCIDevice, MAP_FIXED, MAP_NORESERVE
-from tinygrad.runtime.support.usb import USB3, pm_usb_bufferize
+from tinygrad.runtime.support.usb import USB3
 from tinygrad.runtime.support.memory import AddrSpace, BumpAllocator
-from tinygrad.runtime.ops_amd import SQTT, PMC
-from tinygrad.runtime.ops_amd import EVENT_INDEX_PARTIAL_FLUSH, WAIT_REG_MEM_FUNCTION_GEQ
+from tinygrad.runtime.ops_amd import SQTT, SQTT_ITRACE_SE_MASK, SQTT_LIMIT_SE, SQTT_SIMD_SEL, SQTT_TOKEN_EXCLUDE, PMC
+from tinygrad.runtime.ops_amd import EVENT_INDEX_PARTIAL_FLUSH, WAIT_REG_MEM_FUNCTION_EQ, WAIT_REG_MEM_FUNCTION_NEQ, WAIT_REG_MEM_FUNCTION_GEQ
 if getenv("IOCTL"): import extra.hip_gpu_driver.hip_ioctl  # noqa: F401 # pylint: disable=unused-import

-from tinygrad.engine.realize import get_call_arg_uops, get_call_var_uops
+from tinygrad.engine.realize import get_runtime, pm_flatten_linear
 from tinygrad.uop import FastEnum, auto
-from tinygrad.uop.ops import Ops, UPat, PatternMatcher
+from tinygrad.uop.ops import Ops, UPat, PatternMatcher, graph_rewrite

 # *****************
 # PM4
@@ -35,11 +36,9 @@ class PM4Ops(FastEnum):
  SET_SH_REG = auto(); SET_UCONFIG_REG = auto(); WAIT_REG_MEM = auto(); ACQUIRE_MEM = auto()  # noqa: E702
  RELEASE_MEM = auto(); DISPATCH_DIRECT = auto(); EVENT_WRITE = auto()  # noqa: E702

-def _dw(vals) -> int: return sum(2 if isinstance(x, UOp) and x.dtype.itemsize == 8 else 1 for x in vals)
-
 def pkt3(ctx, op:PM4Ops, *vals):
-  return UOp(Ops.LINEAR, src=tuple(x if isinstance(x, UOp) else UOp.const(x, dtypes.uint32)
-    for x in (ctx.pm4.PACKET3(getattr(ctx.pm4, f"PACKET3_{op.name}"), _dw(vals) - 1), *vals)))
+  return UOp(Ops.INS, arg=op, src=tuple(UOp.const(x, dtypes.uint32)
+    for x in (ctx.pm4.PACKET3(getattr(ctx.pm4, f"PACKET3_{op.name}"), len(vals) - 1), *vals)))

 def wreg(ctx, reg:AMDReg, *args:sint, **kwargs:int):
  if bool(args) == bool(kwargs): raise RuntimeError('One (and only one) of *args or **kwargs must be specified')
@@ -53,7 +52,7 @@ def wreg(ctx, reg:AMDReg, *args:sint, **kwargs:int):
 def wait_reg_mem(ctx, value, mask=0xffffffff, mem=None, reg=None, reg_done=0, op=WAIT_REG_MEM_FUNCTION_GEQ):
  wrm_info_dw = ctx.pm4.WAIT_REG_MEM_MEM_SPACE(int(mem is not None)) | ctx.pm4.WAIT_REG_MEM_OPERATION(int(mem is None and reg_done > 0)) \
              | ctx.pm4.WAIT_REG_MEM_FUNCTION(op) | ctx.pm4.WAIT_REG_MEM_ENGINE(0)
-  return pkt3(ctx, PM4Ops.WAIT_REG_MEM, wrm_info_dw, *((mem,) if mem is not None else (reg, reg_done)), value, mask, 4)
+  return pkt3(ctx, PM4Ops.WAIT_REG_MEM, wrm_info_dw, *(data64_le(mem) if mem is not None else (reg, reg_done)), value, mask, 4)

 def acquire_mem(ctx, addr=0x0, sz=(1 << 64)-1, gli=1, glm=1, glk=1, glv=1, gl1=1, gl2=1):
  if ctx.target[0] != 9:
@@ -84,18 +83,16 @@ def release_mem(ctx, address=0x0, value=0, data_sel=0, int_sel=2, ctxid=0, cache
    event_dw = ctx.pm4.EVENT_TYPE(ctx.pm4.CACHE_FLUSH_AND_INV_TS_EVENT) | ctx.pm4.EVENT_INDEX(ctx.pm4.event_index__mec_release_mem__end_of_pipe)
    memsel_dw = ctx.pm4.DATA_SEL(data_sel) | ctx.pm4.INT_SEL(int_sel)
    ctxid = 0
-  addr_w = address if isinstance(address, UOp) else UOp.const(address, dtypes.uint64)
-  val_w = value.cast(dtypes.uint64) if isinstance(value, UOp) else UOp.const(value, dtypes.uint64)
-  return pkt3(ctx, PM4Ops.RELEASE_MEM, event_dw | cache_flags_dw, memsel_dw, addr_w, val_w, ctxid)
+  return pkt3(ctx, PM4Ops.RELEASE_MEM, event_dw | cache_flags_dw, memsel_dw, *data64_le(address), *data64_le(value), ctxid)

 def memory_barrier(ctx):
  pf = '' if ctx.nbio.version[0] == 2 else '0' if ctx.nbio.version[:2] != (7, 11) else '1'
-  return UOp(Ops.LINEAR, src=(
+  return UOp(Ops.LINEAR, dtypes.void, (
    wait_reg_mem(ctx, reg=getattr(ctx.nbio, f'regBIF_BX_PF{pf}_GPU_HDP_FLUSH_REQ').addr[0],
                 reg_done=getattr(ctx.nbio, f'regBIF_BX_PF{pf}_GPU_HDP_FLUSH_DONE').addr[0], value=0xffffffff),
    acquire_mem(ctx)))

-def pm4_wait(ctx, dst, val): return wait_reg_mem(ctx, val.cast(dtypes.uint32), mem=dst.getaddr(ctx.devs))
+def pm4_wait(ctx, dst, val): return wait_reg_mem(ctx, val, mem=dst.getaddr(ctx.devs))

 def pm4_barrier(ctx): return memory_barrier(ctx)

@@ -109,130 +106,155 @@ def pm4_timestamp(ctx, dst):
                     ctx.pm4.int_sel__mec_release_mem__none)

 def pm4_program(ctx, call, prg):
-  data, lib = amd_build_program(ctx.dev, prg)
-  info = prg.arg
-
-  # kernargs: a nested blob linear inside a getaddr, input addresses and variable values are filled per call through the input table
-  ka_words = [get_call_arg_uops(call)[gi].getaddr(ctx.devs) for gi in info.globals] + list(get_call_var_uops(call, prg))
-  pad = data.kernargs_alloc_size - sum(w.dtype.itemsize for w in ka_words)
-  assert pad >= 0 and pad % 4 == 0, f"bad kernargs padding {pad}"
-  ka = UOp(Ops.LINEAR, src=tuple(ka_words) + (UOp.const(0, dtypes.uint32),) * (pad // 4)).rtag("kernargs")
-
-  prog_addr = lib.getaddr(ctx.devs) + data.entry_point_offset
+  data, info = prg.arg
+  lib_gpu = prg.src[0]
+  args = encode_kernargs_clike(call, prg, ctx.devs)
+  prog_addr = lib_gpu.getaddr(ctx.devs) + data.entry_point_offset
  scratch_addr = UOp.placeholder((data.private_segment_size,), dtypes.uint8, 0, device=ctx.devs).rtag("scratch").getaddr(ctx.devs)
-  args_addr = ka.getaddr(ctx.devs)
+  args_addr = args.getaddr(ctx.devs)

-  user_regs:list = []
-  if data.enable_private_segment_sgpr: user_regs = [scratch_addr | (1 << 63), 0xffffffff, 0x20c14000]
-  if data.enable_dispatch_ptr: user_regs += [args_addr + data.kernargs_segment_size]
-  user_regs += [args_addr]
+  user_regs = []
+  if data.enable_private_segment_sgpr:
+    scratch_hilo = data64_le(scratch_addr)
+    user_regs = [scratch_hilo[0], scratch_hilo[1] | 1 << 31, 0xffffffff, 0x20c14000]
+  if data.enable_dispatch_ptr: user_regs += [*data64_le(args_addr + data.kernargs_segment_size)]
+  user_regs += [*data64_le(args_addr)]

  dispatch_init = ctx.gc.regCOMPUTE_DISPATCH_INITIATOR.encode(
    **({'cs_w32_en': int(data.wave32)} if ctx.target[0] != 9 else {}), force_start_at_000=1, compute_shader_en=1)
  ins = [acquire_mem(ctx, gli=0, gl2=0),
-    wreg(ctx, ctx.gc.regCOMPUTE_PGM_LO, prog_addr >> 8),
+    wreg(ctx, ctx.gc.regCOMPUTE_PGM_LO, *data64_le(prog_addr >> 8)),
    wreg(ctx, ctx.gc.regCOMPUTE_PGM_RSRC1, data.rsrc1, data.rsrc2),
    wreg(ctx, ctx.gc.regCOMPUTE_PGM_RSRC3, data.rsrc3),
    wreg(ctx, ctx.gc.regCOMPUTE_TMPRING_SIZE, ctx.tmpring_size(data.private_segment_size))]
-  ins += [wreg(ctx, ctx.gc.regCOMPUTE_DISPATCH_SCRATCH_BASE_LO, (scratch_addr + data.private_segment_size // ctx.xccs * xcc_id) >> 8)
+  ins += [wreg(ctx, ctx.gc.regCOMPUTE_DISPATCH_SCRATCH_BASE_LO, *data64_le((scratch_addr + data.private_segment_size // ctx.xccs * xcc_id) >> 8))
          for xcc_id in range(ctx.xccs)]
  ins += [wreg(ctx, ctx.gc.regCOMPUTE_RESTART_X, 0, 0, 0),
    wreg(ctx, ctx.gc.regCOMPUTE_USER_DATA_0, *user_regs),
    wreg(ctx, ctx.gc.regCOMPUTE_RESOURCE_LIMITS, ctx.gc.regCOMPUTE_RESOURCE_LIMITS.encode(waves_per_sh=getenv("WAVES_PER_SH"))),
-    wreg(ctx, ctx.gc.regCOMPUTE_START_X, 0, 0, 0, *info.local_size, 0, 0),
+    wreg(ctx, ctx.gc.regCOMPUTE_START_X, 0, 0, 0, *(info.local_size or (1, 1, 1)), 0, 0),
    pkt3(ctx, PM4Ops.DISPATCH_DIRECT, *info.global_size, dispatch_init),
    pkt3(ctx, PM4Ops.EVENT_WRITE, ctx.pm4.EVENT_TYPE(ctx.soc.CS_PARTIAL_FLUSH) | ctx.pm4.EVENT_INDEX(EVENT_INDEX_PARTIAL_FLUSH))]
-  return UOp(Ops.LINEAR, src=tuple(ins))
+  return UOp(Ops.LINEAR, dtypes.void, tuple(ins))

-def pm4_ib(ctx, submit:UOp, lin:UOp) -> UOp|None:
-  # the ring only carries a packet pointing at the ib: the host fence at the start of the batch guarantees the ib is free to reuse
-  if lin.tag is not None or any(w.op in {Ops.CALL, Ops.INS, Ops.LINEAR, Ops.NOOP} for w in lin.src): return None # wait for the flat word linear
-  assert (size_dw:=sum(w.dtype.itemsize for w in lin.src) // 4) < (1 << 20), f"indirect buffer of {size_dw} dwords doesn't fit one packet"
-  pkt = (UOp.const(ctx.pm4.PACKET3(ctx.pm4.PACKET3_INDIRECT_BUFFER, 2), dtypes.uint32), lin.rtag("indirect").getaddr(ctx.devs),
-         UOp.const(size_dw | ctx.pm4.INDIRECT_BUFFER_VALID, dtypes.uint32))
-  return submit.replace(src=(UOp(Ops.LINEAR, src=pkt, arg=lin.arg).rtag(("cmdbuf", ctx.queue)),))
-
-pm_pm4_encode = PatternMatcher([
+pm_pm4_opsel = PatternMatcher([
  (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, name="prg"),), name="call", allow_any_len=True), pm4_program),
-  (UPat(Ops.CUSTOM_FUNCTION, arg="submit_cmdbuf", src=(UPat(Ops.LINEAR, name="lin"),), name="submit"), pm4_ib),

-  (UPat(Ops.INS, arg=("barrier", dtypes.void)), pm4_barrier),
-  (UPat(Ops.INS, arg=("wait", dtypes.void), src=(UPat(name="dst"), UPat(name="val"))), pm4_wait),
-  (UPat(Ops.INS, arg=("timestamp", dtypes.void), src=(UPat(name="dst"),)), pm4_timestamp),
-  (UPat(Ops.INS, arg=("store", dtypes.void), src=(UPat((Ops.BUFFER, Ops.PARAM), name="dst"), UPat(name="val"))), pm4_store),
+  (UPat(Ops.INS, arg="wait", src=(UPat(name="dst"), UPat(name="val"))), pm4_wait),
+  (UPat(Ops.INS, arg="barrier"), pm4_barrier),
+  (UPat(Ops.INS, arg="timestamp", src=(UPat(name="dst"),)), pm4_timestamp),
+  (UPat(Ops.INS, arg="store", src=(UPat((Ops.BUFFER, Ops.PARAM), name="dst"), UPat(name="val"))), pm4_store),
 ])

+def pm4_submit(ctx, lin):
+  # ensure compute queues are allocated
+  for d in (devs:=ctx.devs): q = Device[d].compute_queue
+  ring, wptr, doorbell, put_ptr = (UOp.placeholder((b.size,), b.dtype, 0, device=devs).rtag(f"COMPUTE:0_{name}")
+    for name, b in (("ring", q.ring), ("write_ptr", q.write_ptr), ("doorbell", q.doorbell), ("put_value", q.put_value)))
+
+  # the host fence at the start of the batch guarantees the ib is free to reuse
+  size_dw = sum(len(ins.src) for ins in lin.src)
+  assert size_dw < (1 << 20), f"indirect buffer of {size_dw} dwords doesn't fit one packet"
+
+  ib = UOp.placeholder((size_dw,), dtypes.uint32, next(UOp.unique_num), device=devs, volatile=True).rtag("cmdbuf")
+  cmdbuf = make_cmdbuf(lin, devs, buf=ib)
+
+  # the ring itself only carries a packet pointing at the ib, wrapping the ring
+  put = put_ptr.index(zero:=UOp.const(0, dtypes.int))
+  pkt = (ctx.pm4.PACKET3(ctx.pm4.PACKET3_INDIRECT_BUFFER, 2), *data64_le(cmdbuf.getaddr(devs)), size_dw | ctx.pm4.INDIRECT_BUFFER_VALID)
+  write_pkt = UOp.barrier(*[ring.index(((put + off) % q.ring.size).cast(dtypes.int)).store(UOp.const(x, dtypes.uint32)) for off,x in enumerate(pkt)])
+
+  # advance the put/write pointers past the packet
+  bump_put_ptr = put_ptr.index(zero).store(put + len(pkt))
+  bump_wptr = wptr.index(zero).store(put + len(pkt))
+  flush = UOp.barrier(write_pkt, bump_put_ptr, bump_wptr)
+  return doorbell.after(flush).index(zero).store(put + len(pkt))
+
+pm_pm4_submit = PatternMatcher([(UPat(Ops.LINEAR, name="lin"), pm4_submit)])
+
 # *****************
 # SDMA

+class SDMAOps(FastEnum): COPY = auto(); POLL_REGMEM = auto(); FENCE = auto(); TRAP = auto(); TIMESTAMP = auto()  # noqa: E702
+
 def sdma_copy(ctx, call):
  sz = call.src[2].max_numel() * call.src[2].dtype.itemsize
-  hdr = ctx.sdma.SDMA_OP_COPY | ctx.sdma.SDMA_PKT_COPY_LINEAR_HEADER_SUB_OP(ctx.sdma.SDMA_SUBOP_COPY_LINEAR)
-  return UOp(Ops.LINEAR, src=tuple(x for off in range(0, sz, ctx.max_copy_size) for x in (
-    *(UOp.const(v, dtypes.uint32) for v in (hdr, ctx.sdma.SDMA_PKT_COPY_LINEAR_COUNT_COUNT(min(sz-off, ctx.max_copy_size)-1), 0)),
-    *(a + UOp.const(off, dtypes.uint64) if off else a for a in (call.src[2].getaddr(ctx.devs), call.src[1].getaddr(ctx.devs))))))
+  src_addr, dst_addr = call.src[2].getaddr(ctx.devs), call.src[1].getaddr(ctx.devs)
+  return call.ins(SDMAOps.COPY, src=tuple(UOp.const(x, dtypes.uint32) for off in range(0, sz, ctx.max_copy_size) for x in (
+    ctx.sdma.SDMA_OP_COPY | ctx.sdma.SDMA_PKT_COPY_LINEAR_HEADER_SUB_OP(ctx.sdma.SDMA_SUBOP_COPY_LINEAR),
+    ctx.sdma.SDMA_PKT_COPY_LINEAR_COUNT_COUNT(min(sz-off, ctx.max_copy_size)-1), 0,
+    *data64_le(src_addr+UOp.const(off, dtypes.uint64)), *data64_le(dst_addr+UOp.const(off, dtypes.uint64)))))

-def sdma_wait(ctx, dst, val):
+def sdma_wait(ctx, ins, dst, val):
  op = ctx.sdma.SDMA_OP_POLL_REGMEM | ctx.sdma.SDMA_PKT_POLL_REGMEM_HEADER_FUNC(WAIT_REG_MEM_FUNCTION_GEQ) \
     | ctx.sdma.SDMA_PKT_POLL_REGMEM_HEADER_MEM_POLL(1)
-  return UOp(Ops.LINEAR, src=(
-    UOp.const(op, dtypes.uint32), dst.getaddr(ctx.devs), val.cast(dtypes.uint32), UOp.const(0xffffffff, dtypes.uint32),
-    UOp.const(ctx.sdma.SDMA_PKT_POLL_REGMEM_DW5_INTERVAL(0x04) | ctx.sdma.SDMA_PKT_POLL_REGMEM_DW5_RETRY_COUNT(0xfff), dtypes.uint32)))
+  return ins.ins(SDMAOps.POLL_REGMEM, src=tuple(UOp.const(x, dtypes.uint32) for x in (
+    op, *data64_le(dst.getaddr(ctx.devs)), val, 0xffffffff,
+    ctx.sdma.SDMA_PKT_POLL_REGMEM_DW5_INTERVAL(0x04) | ctx.sdma.SDMA_PKT_POLL_REGMEM_DW5_RETRY_COUNT(0xfff))))

-def sdma_store(ctx, dst, val): # a fence packet then a trap
+def sdma_store(ctx, ins, dst, val):
  op = ctx.sdma.SDMA_OP_FENCE | (ctx.sdma.SDMA_PKT_FENCE_HEADER_MTYPE(3) if ctx.target[0] != 9 else 0)
-  return UOp(Ops.LINEAR, src=(UOp.const(op, dtypes.uint32), dst.getaddr(ctx.devs), val.cast(dtypes.uint32),
-                              UOp.const(ctx.sdma.SDMA_OP_TRAP, dtypes.uint32), UOp.const(0, dtypes.uint32)))
+  return UOp(Ops.LINEAR, src=(
+    ins.ins(SDMAOps.FENCE, src=tuple(UOp.const(x, dtypes.uint32) for x in (op, *data64_le(dst.getaddr(ctx.devs)), val))),
+    ins.ins(SDMAOps.TRAP, src=tuple(UOp.const(x, dtypes.uint32) for x in (ctx.sdma.SDMA_OP_TRAP, 0)))))

-def sdma_timestamp(ctx, dst):
+def sdma_timestamp(ctx, ins, dst):
  op = ctx.sdma.SDMA_OP_TIMESTAMP | ctx.sdma.SDMA_PKT_TIMESTAMP_GET_HEADER_SUB_OP(ctx.sdma.SDMA_SUBOP_TIMESTAMP_GET_GLOBAL)
-  return UOp(Ops.LINEAR, src=(UOp.const(op, dtypes.uint32), dst.getaddr(ctx.devs)))
+  return ins.ins(SDMAOps.TIMESTAMP, src=tuple(UOp.const(x, dtypes.uint32) for x in (op, *data64_le(dst.getaddr(ctx.devs)))))

-pm_sdma_encode = PatternMatcher([
+pm_sdma_opsel = PatternMatcher([
  (UPat(Ops.CALL, src=(UPat(Ops.COPY),), name="call", allow_any_len=True), sdma_copy),

-  (UPat(Ops.INS, arg=("barrier", dtypes.void)), lambda: UOp(Ops.LINEAR)),
-  (UPat(Ops.INS, arg=("wait", dtypes.void), src=(UPat(name="dst"), UPat(name="val"))), sdma_wait),
-  (UPat(Ops.INS, arg=("timestamp", dtypes.void), src=(UPat(name="dst"),)), sdma_timestamp),
-  (UPat(Ops.INS, arg=("store", dtypes.void), src=(UPat((Ops.BUFFER, Ops.PARAM), name="dst"), UPat(name="val"))), sdma_store),
+  (UPat(Ops.INS, arg="barrier"), lambda: UOp(Ops.NOOP, dtypes.void, ())),
+  (UPat(Ops.INS, arg="wait", src=(UPat(name="dst"), UPat(name="val")), name="ins"), sdma_wait),
+  (UPat(Ops.INS, arg="timestamp", src=(UPat(name="dst"),), name="ins"), sdma_timestamp),
+  (UPat(Ops.INS, arg="store", src=(UPat((Ops.BUFFER, Ops.PARAM), name="dst"), UPat(name="val")), name="ins"), sdma_store),
 ])

-# *****************
-# queue submit
+def sdma_submit(cmdbuf, devs):
+  # the cmdbuf to submit + the patch writes that fill it
+  size_dw, zero = cmdbuf.nbytes() // dtypes.uint32.itemsize, UOp.const(0, dtypes.int)

-def _queue_bufs(ctx, q:AMDQueueDesc) -> tuple[UOp, UOp, UOp, UOp]:
-  ring = UOp.placeholder((q.ring.size,), q.ring.dtype, 0, device=ctx.devs, volatile=True).rtag(f"{ctx.queue}_ring")
-  return (ring, *(make_buf(ctx.devs, tag=f"{ctx.queue}_{n}") for n in ("write_ptr", "doorbell", "put_value")))
+  # the sdma queue's ring and its host-side ring/write/put pointers
+  for d in devs: q = Device[d].sdma_queue(0)
+  ring, wptr, doorbell, put_ptr = (UOp.placeholder((b.size,), b.dtype, 0, device=devs).rtag(f"COPY:0_{name}")
+    for name, b in (("ring", q.ring), ("write_ptr", q.write_ptr), ("doorbell", q.doorbell), ("put_value", q.put_value)))

-def pm4_submit(ctx, cmdbuf:UOp) -> UOp:
-  for d in ctx.devs: q = Device[d].compute_queue
-  ring, wptr, doorbell, put = _queue_bufs(ctx, q)
-  p, size_dw = put.after(cmdbuf).index(0).load(), hcq_size_var(cmdbuf) // 4
-  i = UOp.range(size_dw, 10, dtype=dtypes.int, src=(cmdbuf, ring))
-  copy = ring.index(((p + i.cast(p.dtype)) % q.ring.size).cast(dtypes.int)).store(cmdbuf.bitcast(dtypes.uint32).index(i).load()).end(i)
-  next_put = p + size_dw.cast(p.dtype)
-  flush = UOp.barrier(copy, put.index(0).store(next_put), wptr.index(0).store(next_put))
-  return doorbell.after(flush).index(0).store(next_put)
+  # sdma needs the cmdbuf contiguous: if it won't fit before the ring end, restart at 0 and zero the tail
+  put_b = put_ptr.index(zero)
+  tail_off_dw = ((put_b % (q.ring.size * 4)) // 4).cast(dtypes.int)
+  fits = (size_dw <= q.ring.size - tail_off_dw).cast(dtypes.int)
+  start_dw = fits * tail_off_dw
+  zero_amt_dw = (1 - fits) * (q.ring.size - tail_off_dw)

-def sdma_submit(ctx, cmdbuf:UOp) -> UOp:
-  # sdma needs the cmdbuf contiguous in the ring: if it won't fit before the ring end, restart at 0 and zero the tail
-  for d in ctx.devs: q = unwrap(Device[d].sdma_queue(int(ctx.queue.split(":")[1])))
-  (ring, wptr, doorbell, put), rs = _queue_bufs(ctx, q), q.ring.size
-  size_dw = hcq_size_var(cmdbuf) // 4
-  put_b = put.after(cmdbuf).index(0).load()
-  tail = ((put_b % (rs * 4)) // 4).cast(dtypes.int)
-  fits = (size_dw <= rs - tail).cast(dtypes.int)
-  start_dw, zero_amt = fits * tail, (1 - fits) * (rs - tail)
-  zi = UOp.range(zero_amt, 10, dtype=dtypes.int, src=(ring,))
-  zero_tail = ring.index(tail + zi).store(UOp.const(0, dtypes.uint32)).end(zi)
-  i = UOp.range(size_dw, 11, dtype=dtypes.int, src=(cmdbuf, ring))
-  copy = ring.index(start_dw + i).store(cmdbuf.bitcast(dtypes.uint32).index(i).load()).end(i)
-  next_put = put_b + ((zero_amt + size_dw) * 4).cast(put_b.dtype)
-  flush = UOp.barrier(zero_tail, copy, put.index(0).store(next_put), wptr.index(0).store(next_put))
-  return doorbell.after(flush).index(0).store(next_put)
+  # zero the wrapped tail, then copy the cmdbuf into the ring
+  zi = UOp.range(zero_amt_dw, 0, dtype=dtypes.int, src=(cmdbuf,))
+  zero_tail = ring.index(tail_off_dw + zi).store(UOp.const(0, dtypes.uint32)).end(zi)
+  i = UOp.range(UOp.const(size_dw, dtypes.int), 0, dtype=dtypes.int, src=(cmdbuf,))
+  copy_to_ring = ring.index(start_dw + i).store(cmdbuf.index(i).load()).end(i)

-pm_pm4_submit = PatternMatcher([(UPat(Ops.CUSTOM_FUNCTION, arg="submit_cmdbuf", src=(UPat(name="cmdbuf"),)), pm4_submit)])
-pm_sdma_submit = PatternMatcher([(UPat(Ops.CUSTOM_FUNCTION, arg="submit_cmdbuf", src=(UPat(name="cmdbuf"),)), sdma_submit)])
+  # advance the put/write pointers past the zeroed tail and the cmdbuf
+  next_put_b = put_b + ((zero_amt_dw + size_dw) * 4).cast(put_b.dtype)
+  bump_put_ptr = put_ptr.index(zero).store(next_put_b)
+  bump_wptr = wptr.index(zero).store(next_put_b)
+
+  # ring the doorbell once the writes have landed
+  flush = UOp.barrier(zero_tail, copy_to_ring, bump_put_ptr, bump_wptr)
+  return doorbell.after(flush).index(zero).store(next_put_b)
+
+pm_sdma_submit = PatternMatcher([(UPat(Ops.LINEAR, name="lin"),
+  lambda ctx, lin: sdma_submit(make_cmdbuf(lin, ctx.devs), ctx.devs))])
+
+@dataclass(frozen=True)
+class AMDEncodeCtx:  # encode-time constants for one queue: devs (every cmdbuf address resolves into these) + gfx version + packet/ip modules
+  devs: tuple[str, ...]; target: tuple[int, ...]; pm4: Any; sdma: Any; soc: Any  # noqa: E702
+  gc: AMDIP; nbio: AMDIP; xccs: int; max_copy_size: int; tmpring_size: Callable  # noqa: E702
+
+def encode_queue(q:UOp) -> UOp|None:
+  d = Device[(devs:=to_tuple(q.arg[0]))[0]]
+  ctx = AMDEncodeCtx(devs, d.target, d.pm4, d.sdma, d.soc, d.gc, d.nbio, d.xccs, d.max_copy_size, d.tmpring_size)
+  opsel, submit = (pm_pm4_opsel, pm_pm4_submit) if q.arg[1].startswith("COMPUTE") else (pm_sdma_opsel, pm_sdma_submit)
+  return submit.rewrite(graph_rewrite(q, opsel + pm_flatten_linear, walk=True, ctx=ctx, name=f"{q.arg[1]} opsel"), ctx)

@dataclass(frozen=True)
 class AMDProgramData:
@@ -240,10 +262,10 @@ class AMDProgramData:
  private_segment_size:int; kernargs_segment_size:int; kernargs_alloc_size:int
  enable_dispatch_ptr:int; enable_private_segment_sgpr:int

-_amd_program_cache:dict[tuple[bytes, tuple[str, ...]], tuple[AMDProgramData, UOp]] = {}
-def amd_build_program(dev, prg:UOp) -> tuple[AMDProgramData, UOp]:
-  # key on the full device tuple: the same lib can be built for different device sets, each needs its own program buffer
-  if (cached:=_amd_program_cache.get(key:=(lib:=prg.src[3].arg, to_tuple(prg.device)))) is None:
+_amd_program_cache:dict[tuple[bytes,str], tuple[AMDProgramData,bytes]] = {}
+def amd_build_program(prg:UOp) -> UOp:
+  dev = Device[to_tuple(prg.device)[0]] # TODO: rm this
+  if (cached:=_amd_program_cache.get(key:=(lib:=prg.src[3].arg, dev.device))) is None:
    image, sections, relocs = elf_loader(lib)
    rodata = next(sh.header.sh_addr for sh in sections if sh.name == ".rodata")
    for off, sym, typ, addent in relocs:
@@ -260,23 +282,20 @@ def amd_build_program(dev, prg:UOp) -> tuple[AMDProgramData, UOp]:
      wave32=bool(desc.kernel_code_properties & 0x400), private_segment_size=desc.private_segment_fixed_size, kernargs_segment_size=desc.kernarg_size,
      kernargs_alloc_size=desc.kernarg_size + (ctypes.sizeof(hsa.hsa_kernel_dispatch_packet_t) if edp else 0), enable_dispatch_ptr=edp,
      enable_private_segment_sgpr=desc.kernel_code_properties & hsa.AMD_KERNEL_CODE_PROPERTIES_ENABLE_SGPR_PRIVATE_SEGMENT_BUFFER)
-    image = bytes(image).ljust(round_up(len(image), 4), b"\x00") # the program is uploaded as whole dwords
    buf = UOp.placeholder((len(image),), dtypes.uint8, next(UOp.unique_num), device=prg.device).rtag("program")
-    cached = _amd_program_cache[key] = (data, buf.after(buf.store(UOp(Ops.BINARY, src=(), arg=image).bitcast(buf.dtype))))
+    cached = _amd_program_cache[key] = prg.replace(src=(buf.after(make_binary_patch(buf, bytes(image))),), arg=(data, prg.arg))
  return cached

 class AMDAllocator(HCQAllocator['AMDDevice']):
  def __init__(self, dev:AMDDevice):
-    super().__init__(dev, supports_copy_from_disk=dev.has_copy_queue, supports_transfer=dev.has_copy_queue and not dev.is_usb)
+    super().__init__(dev, supports_copy_from_disk=dev.has_sdma_queue, supports_transfer=dev.has_sdma_queue and not dev.is_usb())

-  def _alloc(self, size:int, options:BufferSpec) -> HCQBuffer:
-    return self.dev.iface.alloc(size, host=options.host, uncached=options.uncached, cpu_access=options.cpu_access or not self.dev.has_copy_queue)
+  def _alloc(self, size:int, options:BufferSpec) -> HCQ2Buffer:
+    return self.dev.iface.alloc(size, host=options.host, uncached=options.uncached, cpu_access=options.cpu_access or not self.dev.has_sdma_queue)

  def _do_free(self, opaque, options:BufferSpec): self.dev.iface.free(opaque)

-  def _do_map(self, buf:HCQBuffer): return self.dev.iface.map(buf._base if buf._base is not None else buf)
-
-  def _do_unmap(self, buf:HCQBuffer): self.dev.iface.unmap(buf)
+  def _do_map(self, buf:HCQ2Buffer): return self.dev.iface.map(buf._base if buf._base is not None else buf)

@dataclass
 class AMDQueueDesc:
@@ -369,24 +388,15 @@ class KFDIface:
    return hcqbuf

  def free(self, mem):
-    self._unmap(mem)
-    if mem.va_addr: FileIOInterface.munmap(mem.va_addr, mem.size)
-    kfd.AMDKFD_IOC_FREE_MEMORY_OF_GPU(self.kfd, handle=mem.meta.handle)
-
-  def unmap(self, mem):
-    self._unmap(mem)
-    if getattr(mem, '_owns_kfd_handle', False): kfd.AMDKFD_IOC_FREE_MEMORY_OF_GPU(self.kfd, handle=mem.meta.handle)
-
-  def _unmap(self, mem):
    gpus = (ctypes.c_int32 * 1)(self.gpu_id)
    stm = kfd.AMDKFD_IOC_UNMAP_MEMORY_FROM_GPU(self.kfd, handle=mem.meta.handle, device_ids_array_ptr=ctypes.addressof(gpus), n_devices=1)
    assert stm.n_success == 1
+    if mem.owner == self.dev:
+      if mem.va_addr: FileIOInterface.munmap(mem.va_addr, mem.size)
+      kfd.AMDKFD_IOC_FREE_MEMORY_OF_GPU(self.kfd, handle=mem.meta.handle)

  def map(self, mem):
-    if mem.owner is not None and mem.owner._is_cpu():
-      mapped = self.alloc(mem.size, host=True, cpu_addr=mem.va_addr)
-      mapped._owns_kfd_handle = True
-      return mapped
+    if mem.owner is not None and mem.owner._is_cpu(): return self.alloc(mem.size, host=True, cpu_addr=mem.va_addr)

    c_gpus = (ctypes.c_int32 * 1)(self.gpu_id)
    stm = kfd.AMDKFD_IOC_MAP_MEMORY_TO_GPU(self.kfd, handle=mem.meta.handle, device_ids_array_ptr=ctypes.addressof(c_gpus), n_devices=1)
@@ -458,7 +468,6 @@ class PCIIface(PCIIfaceBase):

  def require_profile_mode(self): return True
  def is_wgp_active(self, xcc, se, sa, wgp) -> bool: return True # TODO: account for WGP disablement on some asics.
-  def unmap(self, mem): self.free(mem)

  def _compute_props(self):
    self.ip_versions = self.dev_impl.ip_ver
@@ -503,7 +512,8 @@ class PCIIface(PCIIfaceBase):
      cq = d.compute_queue
      for b in (cq.put_value, cq.read_ptr, cq.write_ptr): b._buf.view.view(fmt='Q')[0] = 0
      d.iface.dev_impl.gfx.setup_ring(*cq.params)
-      d.signal('timeline')._buf.cpu_view().view(fmt='Q')[0] = d.signal('value', 1, device="CPU")._buf.cpu_view().view(fmt='Q')[0] - 1
+      d.signal('timeline')._buf.cpu_view().mv.cast('Q')[0] = \
+        d.signal('value', 1).as_memoryview(force_zero_copy=True, no_sync=True).cast('Q')[0] - 1

  def sleep(self, timeout):
    if hasattr(self.pci_dev, 'irq_poller') and self.pci_dev.irq_poller is not None and (events_cnt:=len(self.pci_dev.irq_poller.poll(timeout))):
@@ -517,50 +527,31 @@ class PCIIface(PCIIfaceBase):

  def device_fini(self): self.dev_impl.fini()

-class USBIface(PCIIface):
-  def __init__(self, dev, dev_id): # pylint: disable=super-init-not-called
-    if dev_id >= len(visible:=hcq_filter_visible_devices(USB3.list_devices(0xADD1, 0x0001) + USB3.list_devices(0x3801, 0x0001), "AMD")):
-      raise RuntimeError(f"AMD:{dev_id} does not exist ({pluralize('device', len(visible))} available)")
-    self.dev, self.pci_dev, self.vram_bar, self.count = dev, USBPCIDevice("AM", *visible[dev_id]), 0, len(visible)
-    self.dev_impl = AMDev(self.pci_dev)
-    self._compute_props()
-    self.sram = self._dma_region(ctrl_addr=0xf000, sys_addr=0x200000, size=0x80000)
-    self.cq_buf = self._dma_region(ctrl_addr=0xb800, sys_addr=0x822000, size=0x1000) # +12 is the dword that releases an armed read
-    self.usb_handle = unwrap(ctypes.cast(self.pci_dev.usb.usb.handle, ctypes.c_void_p).value)
-
-  def _dma_region(self, ctrl_addr, sys_addr, size):
-    region = self.dev_impl.mm.map_range(vaddr:=self.dev_impl.mm.alloc_vaddr(size=size), size, [(sys_addr, size)], aspace=AddrSpace.SYS, uncached=True)
-    return HCQBuffer(vaddr, size, meta=PCIAllocationMeta(region, has_cpu_mapping=False), view=self.pci_dev.dma_view(ctrl_addr, size), owner=self.dev)
-
-  def alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=False, force_devmem=False, **kwargs) -> HCQBuffer:
-    # everything, even host-style signals, lives in vram: gpu writes into the bridge's own memory collide with an armed 0xF2 read stream
-    return super().alloc(size, host=False, uncached=uncached, cpu_access=cpu_access or host, contiguous=contiguous, force_devmem=True, **kwargs)
-
-  def sleep(self, timeout): pass
-
-  # we don't own the sram region, so the buffer never frees it
-  @functools.cached_property
-  def usb_sram(self) -> Buffer:
-    return Buffer(self.dev.device, (b:=self.sram).size, dtypes.uint8, options=BufferSpec(external_ptr=b.va_addr, nolru=True)).allocate(opaque=b)
-
 def _mock(iface, name=None): return type(name or f"MOCK{iface.__name__}", (iface,), {})

 class AMDDevice(HCQ2Compiled):
+  pm_lower = PatternMatcher([
+    # prep program
+    (UPat(Ops.PROGRAM, src=(UPat(), UPat(), UPat(), UPat(Ops.BINARY)), name="prg"), amd_build_program),
+
+    # encoding of cmdbuf
+    (UPat(Ops.CUSTOM_FUNCTION, arg="submit_cmdbuf", src=(UPat(Ops.LINEAR, name="q"),)), encode_queue),
+  ])
+
  timestamp_divider = 100.0  # AMD GPU clock: ticks/us
  max_scratch_psize = 0
-  pm_encode = {"COMPUTE": pm_pm4_encode, "COPY": pm_sdma_encode}
-  pm_lower = {"COMPUTE": pm_pm4_submit, "COPY": pm_sdma_submit}

-  ifaces = [KFDIface, PCIIface, USBIface, _mock(KFDIface, "MOCKIface"), _mock(KFDIface), _mock(PCIIface), _mock(USBIface)]
+  ifaces = [KFDIface, PCIIface, _mock(KFDIface, "MOCKIface"), _mock(KFDIface), _mock(PCIIface)]

  def device_props(self): return self.iface.props

  def is_am(self) -> bool: return isinstance(self.iface, (PCIIface,))
+  def is_usb(self) -> bool: return False

  def __init__(self, device:str=""):
-    self.iface = self._select_iface(device)
-    self.is_usb = isinstance(self.iface, USBIface)
-    if self.is_usb: self.rt_nbytes = 4 << 20
+    self.device_id = int(device.split(":")[1]) if ":" in device else 0
+
+    self.iface = self._select_iface()

    self.target:tuple[int, ...] = ((trgt:=self.iface.props['gfx_target_version']) // 10000, (trgt // 100) % 100, trgt % 100)
    self.arch = "gfx%d%x%x" % self.target
@@ -585,12 +576,12 @@ class AMDDevice(HCQ2Compiled):

    self.is_aql = getenv("AMD_AQL", int(self.xccs > 1))
    if self.is_aql:
-      self.pm4_ibs = self.iface.alloc(0x2000 if self.is_usb else (16 << 20), uncached=True, cpu_access=True)
+      self.pm4_ibs = self.iface.alloc(0x2000 if self.is_usb() else (16 << 20), uncached=True, cpu_access=True)
      self.pm4_ib_alloc = BumpAllocator(self.pm4_ibs.size, wrap=True)

    self.max_copy_size = 0x40000000 if self.iface.ip_versions[am.SDMA0_HWIP][0] >= 5 else 0x400000
    self.sdma_queues:dict = {}
-    self.has_copy_queue = not getenv("AMD_DISABLE_SDMA")
+    self.has_sdma_queue = True # self.sdma_queue(0) is not None, TODO: think of this

    super().__init__(device, AMDAllocator(self), [HIPRenderer, AMDLLVMRenderer, HIPCCRenderer], None, can_recover=self.is_am(), arch=self.arch)

@@ -598,10 +589,6 @@ class AMDDevice(HCQ2Compiled):
    self.max_private_segment_size = 0
    self.pm_bufferize = PatternMatcher([(UPat(Ops.PARAM, tag="scratch", name="b"), lambda ctx, b: ctx[0].scratch_buffer(b.max_numel()))]) + self.pm_bufferize

-    if self.is_usb:
-      self.pm_bufferize = pm_usb_bufferize + self.pm_bufferize
-      raise NotImplementedError("usb amd is not migrated to sealed submits yet") # a usb pm_lower can override the whole submit graph
-
    self.pmc_enabled:bool = PROFILE > 0 and PMC > 0
    if self.pmc_enabled:
      self.iface.require_profile_mode()
@@ -662,7 +649,7 @@ class AMDDevice(HCQ2Compiled):
    wg_data_size = round_up((vgpr_size_per_cu + sgrp_size_per_cu + lds_size_per_cu + hwreg_size_per_cu) * self.cu_cnt, mmap.PAGESIZE)
    ctl_stack_size = round_up((12 if self.target[0] != 9 else 8) * self.wave_cnt + 8 + 40, mmap.PAGESIZE)
    return self.create_queue(kfd.KFD_IOC_QUEUE_TYPE_COMPUTE_AQL if self.is_aql else kfd.KFD_IOC_QUEUE_TYPE_COMPUTE,
-      0x2000 if self.is_usb else (16 << 20), eop_buffer_size=0x1000,
+      0x2000 if self.is_usb() else (16 << 20), eop_buffer_size=0x1000,
      ctx_save_restore_size=0 if self.is_am() else wg_data_size + ctl_stack_size, ctl_stack_size=ctl_stack_size,
      debug_memory_size=round_up(self.wave_cnt * 32, 64))

@@ -670,7 +657,7 @@ class AMDDevice(HCQ2Compiled):
    if getenv("AMD_DISABLE_SDMA"): return None
    if idx in self.sdma_queues: return self.sdma_queues[idx]
    with contextlib.suppress(OSError):
-      self.sdma_queues[idx] = self.create_queue(kfd.KFD_IOC_QUEUE_TYPE_SDMA, 0x2000 if self.is_usb else (16 << 20), idx=idx)
+      self.sdma_queues[idx] = self.create_queue(kfd.KFD_IOC_QUEUE_TYPE_SDMA, 0x200 if self.is_usb() else (16 << 20), idx=idx)
    return self.sdma_queues.get(idx, None)

  def tmpring_size(self, private_segment_size):
@@ -1,6 +1,8 @@
 from __future__ import annotations
 import functools, pathlib
-from tinygrad import Tensor
+from dataclasses import replace
+from tinygrad import Tensor, dtypes
+from tinygrad.uop.ops import shape_to_shape_arg
 from tinygrad.runtime.support.compiler_amd import HIPCCCompiler

 FP8_MAX = 448.0
@@ -10,7 +12,7 @@ NUM_WG, THREADS_PER_WG = 1024, 256
@functools.cache
 def _local_abs_max_fxn(x_p, device):
  x = Tensor(x_p, device=device)
-  inner = Tensor(x.uop.src[0]) if x.uop.axis is not None else x  # the per-shard view of the flat param
+  inner = Tensor(x.uop.replace(src=(shape_to_shape_arg(x.uop.shard_shape),), arg=replace(x.uop.arg, axis=None))) if x.uop.axis is not None else x
  return (inner.abs().max(),)

 def local_abs_max(x:Tensor) -> Tensor:
@@ -50,7 +50,7 @@ def _custom_quantize_fp8_with_amax(fp8_out:UOp, amax_out:UOp, x:UOp, amax_state:
  else: raise NotImplementedError(f"no atomic max for device {device}")
  amax_idx = amax_out.reshape((1,)).index(UOp.const(0))
  max_val = lds[0].load()
-  atomic = UOp(Ops.CUSTOM, src=(amax_idx, max_val.bitcast(dtypes.int32), max_val, amax_idx.load()), arg=(atomic_arg, dtypes.void))
+  atomic = UOp(Ops.CUSTOM, dtypes.void, (amax_idx, max_val.bitcast(dtypes.int32), max_val, amax_idx.load()), arg=atomic_arg)
  return atomic.end(tid, wg).sink(arg=KernelInfo(f"quantize_fp8_with_amax_{n_elems}", opts_to_apply=()))

@functools.cache
@@ -12,7 +12,7 @@ def _custom_quantize_mxfp4(row_fp4:UOp, row_scale:UOp, col_fp4:UOp, col_scale:UO
  mem = M*N*2 + M*N + M*N//16  # read bf16, write row+col fp4 + e8m0
  outputs = (row_fp4, row_scale, col_fp4, col_scale)
  sink = UOp.sink(*(o.base for o in outputs), x.base,
-                  *(UOp(Ops.CUSTOM, src=(o.base.index(0),), arg=("", dtypes.void)) for o in outputs),
+                  *(UOp(Ops.CUSTOM, dtypes.void, (o.base.index(0),), arg="") for o in outputs),
                  UOp.special(256, "lidx0"), UOp.special(M//128, "gidx0"), UOp.special(N//64, "gidx1"),
                  arg=KernelInfo(name, estimates=Estimates(ops=12*M*N, mem=mem)))
  src = (pathlib.Path(__file__).parent/"quantize_mxfp4.cpp").read_text()
@@ -3,7 +3,7 @@ import os
 # TODO: there is a timing bug without this
 os.environ["AMD_AQL"] = "1"

-from tinygrad import Tensor, Device, GlobalCounters, Context, dtypes
+from tinygrad import Tensor, Device, GlobalCounters, Context
 from tinygrad.helpers import getenv, DEV
 from tinygrad.uop.ops import UOp, Ops, KernelInfo
 from tinygrad.renderer import Estimates
@@ -37,7 +37,7 @@ def launchBenchmark(instruction, vgprIndices, dense=True, accum=False, **kwargs)
    gidx = UOp.special(NUM_WORKGROUPS, "gidx0")
    FLOPs = FLOPS_PER_MATMUL * NUM_WAVES * NUM_WORKGROUPS * INTERNAL_LOOP * INSTRUCTIONS_PER_LOOP
    sink = UOp.sink(A.base, threads, gidx, arg=KernelInfo(inst.op.name.lower(), estimates=Estimates(ops=FLOPs, mem=0)))
-    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=(x, dtypes.void)) for x in insts]))))
+    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
  dummy = Tensor.zeros(1).contiguous().realize()
  out = Tensor.custom_kernel(dummy, fxn=fxn)[0]
  linear = out.schedule_linear()
@@ -5,9 +5,9 @@ from tinygrad.helpers import getenv, DEBUG

 # https://github.com/facebookresearch/llama/blob/1076b9c51c77ad06e9d7ba8a4c6df775741732bd/llama/model.py#L47
 def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0) -> Tensor:
-  freqs = 1.0 / (theta ** (Tensor.arange(0, dim, 2, dtype=dtypes.float32)[:(dim // 2)] / dim))
-  freqs = Tensor.arange(end, dtype=dtypes.float32).unsqueeze(dim=1) * freqs.unsqueeze(dim=0)
-  return Tensor.stack(freqs.cos(), freqs.sin(), dim=-1).cast(dtypes.default_float).reshape(1, end, 1, dim//2, 2)
+  freqs = 1.0 / (theta ** (Tensor.arange(0, dim, 2)[:(dim // 2)] / dim))
+  freqs = Tensor.arange(end).unsqueeze(dim=1) * freqs.unsqueeze(dim=0)
+  return Tensor.stack(freqs.cos(), freqs.sin(), dim=-1).reshape(1, end, 1, dim//2, 2)

 # matches meta, non hugging face weights
 # (a+i*b) * (c+i*d) = (ac-bd) + i*(ad+bc)
@@ -1,12 +1,10 @@
 from tinygrad import Tensor
-import os
 from tinygrad.tensor import _to_np_dtype
 from tinygrad.nn.onnx import OnnxRunner, OnnxValue
 import numpy as np
 import onnxruntime as ort
 ort_options = ort.SessionOptions()
 ort_options.log_severity_level = 3
-ort_options.intra_op_num_threads = os.cpu_count() or 1

 def get_example_inputs(graph_inputs:dict[str, OnnxValue], config={}):
  """
@@ -89,8 +89,7 @@ class TestBeamSearch(unittest.TestCase):
    s.apply_opt(Opt(OptOps.TC, 0, (-1, 0, 1)))
    up = prod([x for x, t in zip(s.full_shape, s.axis_types) if t in (AxisType.UPCAST, AxisType.UNROLL)])
    actions = get_kernel_actions(s, include_0=False, max_up=int(up))
-    upcasted = [s for s in actions.values() if any(o.op is OptOps.SPLIT and o.arg[1] in (AxisType.UPCAST, AxisType.UNROLL)
-                                                  for o in s.applied_opts)]
+    upcasted = [s for s in actions.values() if any(opt.op in (OptOps.UPCAST, OptOps.UNROLL) for opt in s.applied_opts)]
    assert len(upcasted) > 0, f"expected upcast/unroll actions after TC with max_up={up}, but got none"

  def test_max_up(self):
@@ -99,8 +98,8 @@ class TestBeamSearch(unittest.TestCase):
    s = Scheduler(ast, Device[Device.DEFAULT].renderer)
    for max_up in (2, 4):
      actions = get_kernel_actions(s, include_0=False, max_up=max_up)
-      up_opts = [o for s in actions.values() for o in s.applied_opts if o.op is OptOps.SPLIT and o.arg[1] in (AxisType.UPCAST, AxisType.UNROLL)]
-      assert len([opt for opt in up_opts if opt.arg[0] > max_up]) == 0 and len([op for op in up_opts if op.arg[0] <= max_up]) > 0
+      for up_opts in [s.applied_opts for s in actions.values() if any(opt.op in (OptOps.UPCAST, OptOps.UNROLL) for opt in s.applied_opts)]:
+        assert len([opt for opt in up_opts if opt.arg > max_up]) == 0 and len([op for op in up_opts if op.arg <= max_up]) > 0

 if __name__ == '__main__':
  unittest.main()
@@ -1,197 +0,0 @@
-# Runbook: Llama 3 8B Training on DigitalOcean MI350X
-
-## Machine Specs
- 8x MI350X GPUs (gfx950, device ID 75b0), 288GB VRAM each
- 2TB RAM, 192 CPUs, 2TB disk
- ROCm 7.14 at `/opt/rocm` (NOT `/opt/rocm-7.1.1` like the submission scripts assume)
- Python 3.12
-
-## Phase 1: System Setup
-
-### 1.1 Install packages
-```bash
-apt-get update
-apt-get install -y python3-pip python3-venv git tmux rclone clang
-```
-
-### 1.2 Install Python deps
-```bash
-python3 -m pip install --break-system-packages --ignore-installed typing-extensions numpy tqdm wandb tiktoken sentencepiece
-```
-Note: `--ignore-installed typing-extensions` is needed because the base image ships typing-extensions 4.10.0 without a RECORD file, so pip cannot uninstall it.
-
-### 1.3 Install ROCm dev headers
-The base image has ROCm runtime but NOT the HIP dev headers. Need:
-```bash
-apt-get install -y amdrocm-core-dev
-```
-This installs `hip/hip_runtime.h` at `/opt/rocm/core-7.14/include/hip/hip_runtime.h`.
-The symlink `/opt/rocm/include` → `/opt/rocm/core-7.14/include` makes it available at `/opt/rocm/include/hip/hip_runtime.h`.
-
-### 1.4 Configure ROCm comgr
-ROCm 7.14 ships comgr 3.3 at `/opt/rocm/lib/libamd_comgr.so`. tinygrad's DLL loader needs explicit env vars to find it (it searches for `libcomgr.so*` by default, not `libamd_comgr.so*`). Set these in the run command:
-```bash
-export COMGR_PATH=/opt/rocm/lib/libamd_comgr.so
-export COMGR_3_PATH=/opt/rocm/lib/libamd_comgr.so
-```
-Also add ROCm libs to ldconfig so comgr's shared library dependencies resolve:
-```bash
-cat > /etc/ld.so.conf.d/rocm.conf << 'EOF'
-/opt/rocm/lib
-/opt/rocm/lib/llvm/lib
-/opt/rocm/lib/rocm_sysdeps/lib
-EOF
-ldconfig
-```
-
-### 1.5 Install geohot tmux config
-```bash
-curl -sL https://raw.githubusercontent.com/geohot/configuration/master/.tmux.conf -o ~/.tmux.conf
-```
-
-### 1.6 Verify GPU PCI access
-The AM userspace driver accesses the GPUs directly over PCI. Do not load `amdgpu`. `/dev/kfd` is not required.
-```bash
-rmmod amdgpu
-lspci -nnk -d 1002:
-```
-The MI350X devices should not show a `Kernel driver in use: amdgpu`.
-
-## Phase 2: Clone tinygrad
-```bash
-cd /root
-git clone https://github.com/tinygrad/tinygrad.git
-cd tinygrad
-python3 -m pip install --break-system-packages -e .
-```
-
-## Phase 3: Download C4 Dataset
-
-The C4 data is on the MLCommons Cloudflare R2 bucket in Megatron-LM indexed format.
-
-```bash
-rclone config create mlc-training s3 provider=Cloudflare \
-  access_key_id=76ea42eadb867e854061a1806220ee1e \
-  secret_access_key=a53625c4d45e3ca8ac0df8a353ea3a41ffc3292aa25259addd8b7dc5a6ce2936 \
-  endpoint=c2686074cb2caf5cbaf6d134bdba8b47.r2.cloudflarestorage.com
-
-mkdir -p /raid/datasets/c4-8b
-(rclone copy mlc-training:mlcommons-training-wg-public/llama3_1/datasets/c4/llama3_1_8b/ /raid/datasets/c4-8b/ -P && \
-  PYTHONPATH=. python3 examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/buid_dataset_cache.py) \
-  > /root/dataset_cache.log 2>&1 &
-```
-Leave this running and proceed to the beam step while the dataset downloads and its cache builds.
-
-### 3.1 Smoke test (beam search, 2 layers, fake data)
-Always run beam first to validate the pipeline:
-```bash
-tmux new-session -d -s beam 'cd /root/tinygrad && COMGR_PATH=/opt/rocm/lib/libamd_comgr.so COMGR_3_PATH=/opt/rocm/lib/libamd_comgr.so CC=/opt/rocm/core-7.14/lib/llvm/bin/clang DEV=PCI+AMD:HIP ROCM_PATH=/opt/rocm bash examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_beam.sh 2>&1 | tee /root/beam.log'
-```
-
-The beam test runs 10 training steps with 2 layers. Expected results:
- ~0.29s per step after warmup
- ~700K GFLOPS, ~7% MFU (low because only 2 layers)
- ~380 GB VRAM used
- Loss stable at ~12.55 with random init
-
-Files downloaded (~85GB total, ~6 minutes):
- `c4-train.en_6_text_document.bin` (79 GB)
- `c4-train.en_6_text_document.idx` (870 MB)
- `c4-validation-91205-samples.en_text_document.bin` (159 MB)
- `c4-validation-91205-samples.en_text_document.idx` (1.8 MB)
- `LICENSE.txt`, `NOTICE.txt`
-
-**Wait for rclone to fully complete before starting training.** Starting training while the dataset is still downloading will read a truncated .bin file, causing `ValueError: all input arrays must have the same shape` in the dataloader. The stale `.index_cache` and `.blend_cache` files must also be deleted if this happens:
-```bash
-rm -f /raid/datasets/c4-8b/*.index_cache /raid/datasets/c4-8b/*.blend_cache
-```
-
-## Phase 4: wandb Login
-```bash
-wandb login
-```
-Enter API key from https://wandb.ai/authorize
-
-Alternatively, pass the key directly:
-```bash
-wandb login <API_KEY>
-```
-
-## Phase 5: Run Training
-
-Run training in tmux so it survives SSH disconnects:
-```bash
-tmux new-session -d -s train 'cd /root/tinygrad && COMGR_PATH=/opt/rocm/lib/libamd_comgr.so COMGR_3_PATH=/opt/rocm/lib/libamd_comgr.so CC=/opt/rocm/core-7.14/lib/llvm/bin/clang DEV=PCI+AMD:HIP ROCM_PATH=/opt/rocm WANDB=1 bash examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_run.sh 2>&1 | tee /root/train.log'
-```
-Attach with `tmux attach -t train`.
-
-### 5.1 Full training run
-```bash
-tmux new-session -d -s train 'cd /root/tinygrad && COMGR_PATH=/opt/rocm/lib/libamd_comgr.so COMGR_3_PATH=/opt/rocm/lib/libamd_comgr.so CC=/opt/rocm/core-7.14/lib/llvm/bin/clang DEV=PCI+AMD:HIP ROCM_PATH=/opt/rocm WANDB=1 bash examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_run.sh 2>&1 | tee /root/train.log'
-```
-
-## Environment Variable Reference
-
-| Variable | Value | Why |
-|---|---|---|
-| `COMGR_PATH` | `/opt/rocm/lib/libamd_comgr.so` | tinygrad's DLL loader needs explicit path to find comgr 3.3 |
-| `COMGR_3_PATH` | `/opt/rocm/lib/libamd_comgr.so` | comgr 3.x uses a separate `comgr_3` module with its own path var |
-| `CC` | `/opt/rocm/core-7.14/lib/llvm/bin/clang` | System clang doesn't know gfx950; must use ROCm's bundled clang |
-| `DEV` | `PCI+AMD:HIP` | Force HIPRenderer (comgr-based) over HIPCCRenderer (hipcc subprocess) |
-| `ROCM_PATH` | `/opt/rocm` | Script defaults to `/opt/rocm-7.1.1` which doesn't exist |
-| `WANDB` | `1` | Enable wandb logging (off by default) |
-
-## Architecture
-
-| Component | Source file |
-|---|---|
-| Model | `examples/mlperf/models/flat_llama.py` — FlatTransformer, FP8 MXFP4 weights, fused QKV, flash attention |
-| Trainer | `examples/mlperf/model_train.py` → `train_llama3()` |
-| Optimizer | `examples/mlperf/optim.py` — GradAccClipAdamW, master weights, FP8 re-quant |
-| LR schedule | `examples/mlperf/lr_schedulers.py` — CosineAnnealingLRWithWarmup |
-| Dataloader | `examples/mlperf/dataloader.py` — Megatron-LM indexed bin format |
-| ASM GEMM | `extra/gemm/cdna_asm_gemm.py` — gfx950 MFMA assembly, MXFP4 |
-| Flash attention | `extra/thunder/amd/fa.py` |
-| Fused kernels | `extra/llama_kernels/` — rmsnorm, silu, quantize, fused_ce |
-| GPU driver | `tinygrad/runtime/ops_amd.py` — HCQ, using the AM userspace PCI interface |
-| Renderer | `tinygrad/renderer/cstyle.py` — HIPRenderer for gfx950 |
-| comgr compiler | `tinygrad/runtime/support/compiler_amd.py` — HIPCompiler using comgr 3.3 |
-
-## Troubleshooting
-
-### `'hip/hip_runtime.h' file not found`
-Install `amdrocm-core-dev`:
-```bash
-apt-get install -y amdrocm-core-dev
-```
-
-### `'gfx950' is not a recognized processor` + LLVM crash
-System clang doesn't know gfx950. Set `CC=/opt/rocm/core-7.14/lib/llvm/bin/clang`.
-
-### `comgr not available: try setting COMGR_PATH?`
-Add ROCm libs to ldconfig and set `COMGR_PATH` and `COMGR_3_PATH`:
-```bash
-# /etc/ld.so.conf.d/rocm.conf should contain /opt/rocm/lib paths
-ldconfig
-```
-
-### `comgr not available: try setting COMGR_3_PATH?`
-comgr 3.x uses a separate module. Set `COMGR_3_PATH=/opt/rocm/lib/libamd_comgr.so` too.
-
-### `No such file or directory: 'clang'`
-Install clang: `apt-get install -y clang` (for CPU compilation).
-For gfx950 HIP compilation, comgr (not clang) is used — ensure the ROCm 7.14 comgr 3.3 is properly loaded via `COMGR_PATH` and `COMGR_3_PATH`.
-
-## Appendix: KVM Virtualization Observations
-
-### Virtualization detection
-```
-$ systemd-detect-virt
-kvm
-$ lspci -nn | grep AMD
-83:00.0 ... Device [1002:75b0]
-```
-CPU flags include `hypervisor`. `dmesg` shows `Hypervisor detected: KVM`.
-
-### No fan control
-No `fan*` or `pwm*` hwmon entries exist. Only `temp*`, `power*`, `freq*` are exposed. GPU temps read 56-63°C, power ~265W per GPU.
@@ -1,4 +1,4 @@
-import os, subprocess, sys, shlex, pickle
+import os, subprocess, sys, shlex
 from pathlib import Path
 from tinygrad.helpers import temp, getenv

@@ -23,8 +23,5 @@ if __name__ == "__main__":
      # AM_RESET=1 gets a clear trace, does not work on mi300 machines
      subprocess.run([sys.executable, *shlex.split(test)], cwd=EXAMPLES_DIR.parent.parent.parent,
                     env={**os.environ, "DEV":"AMD", "AM_RESET":"1" if not arch.startswith("gfx9") else "0", "VIZ":"-2", "PYTHONPATH":"."})
-      with open(PROFILE_PATH, "rb") as f: events = pickle.load(f)
-      with open(PROFILE_PATH, "wb") as f:
-        pickle.dump([e for e in events if type(e).__name__ in {"ProfilePMCEvent", "ProfileSQTTEvent", "ProfileProgramEvent"}], f)
      PROFILE_PATH.rename(dest:=EXAMPLES_DIR/arch/f"profile_{name}_run_{i}.pkl")
      print(f"saved SQTT trace to {dest}")
--- a/Show More
+++ b/Show More