initial usbgpu support (#37906)

* zero ll patched big model * probe in a subprocess so usbgpu lock gets released * compiles * runs * num_jobs gets overwritten, use side effect * poll tg devices * make sure build crashes on missing gpu * fine not to rely on Device.default * seperate tg env for each model runner * comment * Revert "seperate tg env for each model runner" This reverts commit f6470cc4258eaeb3e8e37907ef370871c9af5aa4. * env is shared, gate on flag * no fallback warp dev must be set * build for current device only, unless pc/release * comment * list * listen for plug in * add icon to status bar, read params on every frame (?) * log available devices * try copy out when loading? * Revert "log available devices" This reverts commit e8c52a5d59456d4820ecb13b99a6c46ea1386a20. * Revert "try copy out when loading?" This reverts commit 518f403aa03faeda1950fe3dbce0d9e4c1584455. * don't trigger device probe/caching on modeld prepare * re-export with ll and road edges * dont cache devices in manager process * get USBGPU from params * no usbgpu env * missed one * sconscript don't poll * unconditional env * always explicitely set devices on input tensors * set DEV so amd uses right compiler and iface?? * fix flag * bump tg * rm xdg_cache_home * tg don't bump all the way * missing gmmu=0 at compile time * dm set dev * tg backend * update gitignore * missing import * unused imports * rely on Device.DEFAULT at compile time (already the case bc onnxrunner) * comments * dm warp needs DEV set too * build both smol and big * misc typos * set dev at compile time * don't need * DEV=CPU when getting metadata, ensure we don't grab gpu lock * this would also grab lock * put bool * warp compile always prepare only * missed one * poll ui * missing here * don't force usbgpu at build time * tmp patch fetch_fw * catch all, follow hardwared patterns * simpler * compile make input queues * revert this * group this more readable * rm empty line * make dummy frame using numpy * revert compile make input queues * no compiler at runtime * cleanup * fine to rebuild all on change to device node for now * fix usbgpu_present * fix sconscript * no size in header stream decompress * DEBUG=2 * minimal viable feedback * egpu gray * oops * gotta do this actually * modeld build only depends on modeld devices * don't ship onnx to release? or chunk * don't need * can only set compiler on dev= * none device works, will use default * make linter happy * chunk agnostic onnx input to compile_modeld * chunk big onnx * +x chunker * fix #! * and don't ship chunked onnx to release * firmware now in correct location * better err on missing onnx/chunk * SConscript also need to accept chunked onnx * metadata also need to load maybe chunked * dedupe cmd * this needs to be on cpu * devices are set in the tgflags, we already depend on them * rebuilding on changed order is fine * read file chunked can already load either chunked or not * chunk all big onnx * less confusing * unused import * python device to load onnx bytes * default device for runners, python for metadata * why not * chunked to shm
2026-08-05 00:05:57 +08:00 · 2026-05-19 22:41:57 -07:00
parent 82338fd8c5
commit 52e182611d
15 changed files with 184 additions and 61 deletions
@@ -1,3 +1,5 @@
+#!/usr/bin/env python3
+import sys
 import math
 import os
 from pathlib import Path
@@ -10,10 +12,13 @@ def get_chunk_name(name, idx, num_chunks):
 def get_manifest_path(name):
  return f"{name}.chunkmanifest"

-def get_chunk_paths(path, file_size):
-  num_chunks = math.ceil(file_size / CHUNK_SIZE)
+def _chunk_paths(path, num_chunks):
  return [get_manifest_path(path)] + [get_chunk_name(path, i, num_chunks) for i in range(num_chunks)]

+def get_chunk_targets(path, file_size):
+  num_chunks = math.ceil(file_size / CHUNK_SIZE)
+  return _chunk_paths(path, num_chunks)
+
 def chunk_file(path, targets):
  manifest_path, *chunk_paths = targets
  with open(path, 'rb') as f:
@@ -26,6 +31,13 @@ def chunk_file(path, targets):
  Path(manifest_path).write_text(str(len(chunk_paths)))
  os.remove(path)

+def get_existing_chunks(path):
+  if os.path.isfile(path):
+    return [path]
+  if os.path.isfile(manifest := get_manifest_path(path)):
+    num_chunks = int(Path(manifest).read_text().strip())
+    return _chunk_paths(path, num_chunks)
+  raise FileNotFoundError(path)

 def read_file_chunked(path):
  manifest_path = get_manifest_path(path)
@@ -35,3 +47,9 @@ def read_file_chunked(path):
  if os.path.isfile(path):
    return Path(path).read_bytes()
  raise FileNotFoundError(path)
+
+
+if __name__ == "__main__":
+  path = sys.argv[1]
+  chunk_paths = get_chunk_targets(path, os.path.getsize(path))
+  chunk_file(path, chunk_paths)
@@ -127,5 +127,7 @@ inline static std::unordered_map<std::string, ParamKeyAttributes> keys = {
    {"UpdaterLastFetchTime", {PERSISTENT, TIME}},
    {"UptimeOffroad", {PERSISTENT, FLOAT, "0.0"}},
    {"UptimeOnroad", {PERSISTENT, FLOAT, "0.0"}},
+    {"UsbGpuPresent", {CLEAR_ON_MANAGER_START | CLEAR_ON_OFFROAD_TRANSITION, BOOL}},
+    {"UsbGpuCompiled", {CLEAR_ON_MANAGER_START | CLEAR_ON_OFFROAD_TRANSITION, BOOL}},
    {"Version", {PERSISTENT, STRING}},
 };
@@ -72,7 +72,7 @@ find . -name '*.pyc' -delete
 find . -name 'moc_*' -delete
 find . -name '__pycache__' -delete
 rm -rf .sconsign.dblite Jenkinsfile release/
-rm -f selfdrive/modeld/models/*.onnx
+rm -f selfdrive/modeld/models/*.onnx*

 # Mark as prebuilt release
 touch prebuilt
@@ -45,6 +45,8 @@ cd $TARGET_DIR
 rm -rf .git/modules/
 rm -f panda/board/obj/panda.bin.signed

+find selfdrive/modeld/models -name '*.onnx' -size +95M -exec ./common/file_chunker.py {} \;
+
 # include source commit hash and build date in commit
 GIT_HASH=$(git --git-dir=$SOURCE_DIR/.git rev-parse HEAD)
 GIT_COMMIT_DATE=$(git --git-dir=$SOURCE_DIR/.git show --no-patch --format='%ct %ci' HEAD)
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ec3dcf64cbc34251d8423cb8b3b31d743e93d14002dec43c389a857cb7e8eb17
+size 10875
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7409c53d7c72681c24982fd83b56ce70f80797c9c0f936d9296a5c18557ac472
+size 7279
@@ -1,14 +1,14 @@
 import glob
 import json
 import os
+import sys, subprocess
 from SCons.Script import Value
-from openpilot.common.file_chunker import chunk_file, get_chunk_paths
+from openpilot.common.file_chunker import chunk_file, get_chunk_targets, get_existing_chunks
 from openpilot.common.transformations.camera import _ar_ox_fisheye, _os_fisheye
 from openpilot.common.transformations.model import MEDMODEL_INPUT_SIZE, DM_INPUT_SIZE
 from openpilot.selfdrive.modeld.constants import ModelConstants
-from tinygrad import Device
 from openpilot.system.hardware import HARDWARE, PC
-from openpilot.selfdrive.modeld.helpers import TG_INPUT_DEVICES_PATH
+from openpilot.selfdrive.modeld.helpers import TG_INPUT_DEVICES_PATH, usbgpu_present, modeld_pkl_path

 Import('env', 'arch', 'release')

@@ -36,7 +36,13 @@ def estimate_pickle_max_size(onnx_size):
  return 1.2 * onnx_size  + 10 * 1024 * 1024  # 20% + 10MB is plenty

 # get fastest TG config
-available = set(Device.get_available_devices())
+# probe in subprocess so usbgpu locks gets released on process exit
+def probe_devices():
+  return set(subprocess.run(
+    [sys.executable, '-c', 'from tinygrad import Device\nprint("\\n".join(Device.get_available_devices()))'],
+    capture_output=True, text=True, check=True).stdout.strip().splitlines())
+
+available = probe_devices()
 if 'CUDA' in available:
  tg_backend = 'CUDA'
  tg_flags = f'DEV={tg_backend}'
@@ -44,14 +50,25 @@ elif 'QCOM' in available:
  tg_backend = 'QCOM'
  tg_flags = f'DEV={tg_backend} IMAGE=1 FLOAT16=1 NOLOCALS=1 JIT_BATCH_SIZE=0 OPENPILOT_HACKS=1'
 else:
-  tg_backend = 'CPU' if arch == 'Darwin' else 'CPU:LLVM'
-  tg_flags = f'DEV={tg_backend}'
+  tg_backend = 'CPU'
+  tg_flags = f'DEV=CPU' if arch == 'Darwin' else 'DEV=CPU:LLVM'

 tg_devices = { # which device to put jit inputs to at runtime
-  'selfdrive.modeld.modeld': tg_backend,
-  'selfdrive.modeld.dmonitoringmodeld': tg_backend,
+  'selfdrive.modeld.modeld': {
+    'default': {'WARP_DEV': tg_backend, 'QUEUE_DEV': tg_backend}
+  },
+  'selfdrive.modeld.dmonitoringmodeld': {
+    'default': {'DEV': tg_backend}
+  },
 }

+USBGPU = usbgpu_present() # or release # TODO always build big model on release
+if USBGPU:
+  tg_devices['selfdrive.modeld.modeld']['usbgpu'] = {'WARP_DEV': tg_backend, 'QUEUE_DEV': 'AMD'}
+  usbgpu_tg_flags = f'DEBUG=2 DEV=USB+AMD:LLVM WARP_DEV={tg_backend} FLOAT16=1 JIT_BATCH_SIZE=0 GMMU=0'
+  # the USB+AMD GPU takes an exclusive flock; serialize all targets that touch it
+  usbgpu_lock = File("models/.usb_gpu.lock").abspath
+
 def write_tg_devices(target, source, env):
  with open(str(target[0]), "w") as f:
    json.dump(tg_devices, f)
@@ -59,7 +76,7 @@ def write_tg_devices(target, source, env):

 tg_devices_node = lenv.Command(
  str(TG_INPUT_DEVICES_PATH),
-  [Value(json.dumps(tg_devices, sort_keys=True))],
+  [Value(tg_devices)],
  write_tg_devices,
 )

@@ -68,27 +85,33 @@ mac_brew_string = f'HOME={os.path.expanduser("~")}' if arch == 'Darwin' else ''

 modeld_dir = Dir("#selfdrive/modeld").abspath
 compile_modeld_script = [File(f"{modeld_dir}/compile_modeld.py")]
-driving_onnx_deps = [File(f"models/{m}.onnx").abspath for m in ['driving_vision', 'driving_policy']]
-
 model_w, model_h = MEDMODEL_INPUT_SIZE
 frame_skip = ModelConstants.MODEL_RUN_FREQ // ModelConstants.MODEL_CONTEXT_FREQ
-pkl_path = File("models/driving_tinygrad.pkl").abspath
-camera_res_args = ' '.join(f'{cw}x{ch}' for cw, ch in CAMERA_CONFIGS)
-cmd = (f'{tg_flags} {mac_brew_string} python3 {modeld_dir}/compile_modeld.py '
-       f'--model-size {model_w}x{model_h} '
-       f'--camera-resolutions {camera_res_args} '
-       f'--vision-onnx {File("models/driving_vision.onnx").abspath} '
-       f'--policy-onnx {File("models/driving_policy.onnx").abspath} '
-       f'--output {pkl_path} --frame-skip {frame_skip}')
-onnx_sizes_sum = sum(os.path.getsize(f) for f in driving_onnx_deps)
-chunk_targets = get_chunk_paths(pkl_path, estimate_pickle_max_size(onnx_sizes_sum)*2) # TODO make weight dedupe work on QCOM
-def do_chunk(target, source, env, pkl=pkl_path, chunks=chunk_targets):
-  chunk_file(pkl, chunks)
-lenv.Command(
-  chunk_targets,
-  tinygrad_files + compile_modeld_script + driving_onnx_deps + [Value(camera_res_args), chunker_file, tg_devices_node],
-  [cmd, do_chunk],
-)
+
+for usbgpu in [False, True] if USBGPU else [False]:
+  target_pkl_path = File(modeld_pkl_path(usbgpu)).abspath
+  file_prefix, cmd_flags = ('big_', usbgpu_tg_flags) if usbgpu else ('', tg_flags)
+  driving_onnx_deps = [p for m in [f'{file_prefix}driving_vision', f'{file_prefix}driving_policy']
+                       for p in get_existing_chunks(File(f"models/{m}.onnx").abspath)]
+  camera_res_args = ' '.join(f'{cw}x{ch}' for cw, ch in CAMERA_CONFIGS)
+  cmd = (f'{cmd_flags} {mac_brew_string} python3 {modeld_dir}/compile_modeld.py '
+        f'--model-size {model_w}x{model_h} '
+        f'--camera-resolutions {camera_res_args} '
+        f'--vision-onnx {File(f"models/{file_prefix}driving_vision.onnx").abspath} '
+        f'--policy-onnx {File(f"models/{file_prefix}driving_policy.onnx").abspath} '
+        f'--output {target_pkl_path} --frame-skip {frame_skip}')
+  onnx_sizes_sum = sum(os.path.getsize(f) for f in driving_onnx_deps)
+  size_multiplier = 1 if usbgpu else 2 # TODO make weight dedupe work on QCOM
+  chunk_targets = get_chunk_targets(target_pkl_path, estimate_pickle_max_size(onnx_sizes_sum)*size_multiplier)
+  def do_chunk(target, source, env, pkl=target_pkl_path, chunks=chunk_targets):
+    chunk_file(pkl, chunks)
+  node = lenv.Command(
+    chunk_targets,
+    tinygrad_files + compile_modeld_script + driving_onnx_deps + [Value(camera_res_args), chunker_file],
+    [cmd, do_chunk],
+  )
+  if usbgpu:
+    lenv.SideEffect(usbgpu_lock, node)

 # get model metadata
 fn = File(f"models/dmonitoring_model").abspath
@@ -110,7 +133,7 @@ def tg_compile(flags, model_name):
  fn = File(f"models/{model_name}").abspath
  pkl = fn + "_tinygrad.pkl"
  onnx_path = fn + ".onnx"
-  chunk_targets = get_chunk_paths(pkl, estimate_pickle_max_size(os.path.getsize(onnx_path)))
+  chunk_targets = get_chunk_targets(pkl, estimate_pickle_max_size(os.path.getsize(onnx_path)))
  def do_chunk(target, source, env):
    chunk_file(pkl, chunk_targets)
  return lenv.Command(
@@ -1,5 +1,6 @@
 #!/usr/bin/env python3
 import argparse
+import atexit
 import os
 import pickle
 import time
@@ -7,24 +8,46 @@ from functools import partial
 from collections import namedtuple, defaultdict

 import numpy as np
+
+def _patch_tinygrad_fetch_fw():
+  import hashlib
+  import pathlib
+  import zstandard
+  from tinygrad import helpers
+  _orig = helpers.fetch_fw
+  def fetch_fw(path, name, sha256):
+    p = pathlib.Path(f"/lib/firmware/{path}/{name}.zst")
+    if p.is_file():
+      blob = zstandard.ZstdDecompressor().stream_reader(p.read_bytes()).read()
+      if hashlib.sha256(blob).hexdigest() == sha256:
+        return blob
+    return _orig(path, name, sha256)
+  helpers.fetch_fw = fetch_fw
+_patch_tinygrad_fetch_fw()
+
 from tinygrad.tensor import Tensor
 from tinygrad.helpers import Context
 from tinygrad.device import Device
 from tinygrad.engine.jit import TinyJit

+from openpilot.common.file_chunker import read_file_chunked
+from openpilot.system.hardware.hw import Paths
+

 NV12Frame = namedtuple("NV12Frame", ['width', 'height', 'stride', 'y_height', 'uv_height', 'size'])

 UV_SCALE_MATRIX = np.array([[0.5, 0, 0], [0, 0.5, 0], [0, 0, 1]], dtype=np.float32)
 UV_SCALE_MATRIX_INV = np.linalg.inv(UV_SCALE_MATRIX)

+WARP_DEV = os.getenv('WARP_DEV')
+

 def warp_perspective_tinygrad(src_flat, M_inv, dst_shape, src_shape, stride_pad, border_fill_val=None):
  w_dst, h_dst = dst_shape
  h_src, w_src = src_shape

-  x = Tensor.arange(w_dst).reshape(1, w_dst).expand(h_dst, w_dst).reshape(-1)
-  y = Tensor.arange(h_dst).reshape(h_dst, 1).expand(h_dst, w_dst).reshape(-1)
+  x = Tensor.arange(w_dst, device=WARP_DEV).reshape(1, w_dst).expand(h_dst, w_dst).reshape(-1)
+  y = Tensor.arange(h_dst, device=WARP_DEV).reshape(h_dst, 1).expand(h_dst, w_dst).reshape(-1)

  # inline 3x3 matmul as elementwise to avoid reduce op (enables fusion with gather)
  src_x = M_inv[0, 0] * x + M_inv[0, 1] * y + M_inv[0, 2]
@@ -68,7 +91,7 @@ def make_frame_prepare(nv12: NV12Frame, model_w, model_h):

  def frame_prepare_tinygrad(input_frame, M_inv):
    # UV_SCALE @ M_inv @ UV_SCALE_INV simplifies to elementwise scaling
-    M_inv_uv = M_inv * Tensor([[1.0, 1.0, 0.5], [1.0, 1.0, 0.5], [2.0, 2.0, 1.0]])
+    M_inv_uv = M_inv * Tensor([[1.0, 1.0, 0.5], [1.0, 1.0, 0.5], [2.0, 2.0, 1.0]], device=WARP_DEV)
    # deinterleave NV12 UV plane (UVUV... -> separate U, V)
    uv = input_frame[uv_offset:uv_offset + uv_height * stride].reshape(uv_height, stride)
    with Context(SPLIT_REDUCEOP=0):
@@ -132,14 +155,16 @@ def make_run_policy(vision_runner, policy_runner, nv12: NV12Frame, model_w, mode
  sample_desire_fn = partial(sample_desire, frame_skip=frame_skip)

  def run_policy(img_q, big_img_q, feat_q, desire_q, desire, traffic_convention, tfm, big_tfm, frame, big_frame):
-    tfm = tfm.to(Device.DEFAULT)
-    big_tfm = big_tfm.to(Device.DEFAULT)
+    tfm = tfm.to(WARP_DEV)
+    big_tfm = big_tfm.to(WARP_DEV)
    desire = desire.to(Device.DEFAULT)
    traffic_convention = traffic_convention.to(Device.DEFAULT)
    Tensor.realize(tfm, big_tfm, desire, traffic_convention)

-    img = shift_and_sample(img_q, frame_prepare(frame, tfm).unsqueeze(0), sample_skip_fn)
-    big_img = shift_and_sample(big_img_q, frame_prepare(big_frame, big_tfm).unsqueeze(0), sample_skip_fn)
+    warped_frame = frame_prepare(frame, tfm).unsqueeze(0).to(Device.DEFAULT)
+    warped_big_frame = frame_prepare(big_frame, big_tfm).unsqueeze(0).to(Device.DEFAULT)
+    img = shift_and_sample(img_q, warped_frame, sample_skip_fn)
+    big_img = shift_and_sample(big_img_q, warped_big_frame, sample_skip_fn)

    if prepare_only:
      return img, big_img
@@ -180,8 +205,8 @@ def compile_modeld(nv12: NV12Frame, model_w, model_h, prepare_only, frame_skip,
    n_runs = 1 if testing else 3

    for i in range(n_runs):
-      frame = Tensor.randint(nv12.size, low=0, high=256, dtype='uint8').realize()
-      big_frame = Tensor.randint(nv12.size, low=0, high=256, dtype='uint8').realize()
+      frame = Tensor.randint(nv12.size, low=0, high=256, dtype='uint8', device=WARP_DEV).realize()
+      big_frame = Tensor.randint(nv12.size, low=0, high=256, dtype='uint8', device=WARP_DEV).realize()
      for v in npy.values():
        v[:] = np.random.randn(*v.shape).astype(v.dtype)
      Device.default.synchronize()
@@ -219,6 +244,14 @@ def _parse_size(s):
  return int(w), int(h)


+def read_file_chunked_to_shm(path):
+  shm_path = os.path.join(Paths.shm_path(), os.path.basename(path))
+  atexit.register(lambda: os.path.exists(shm_path) and os.remove(shm_path))
+  with open(shm_path, 'wb') as f:
+    f.write(read_file_chunked(path))
+  return shm_path
+
+
 if __name__ == "__main__":
  from tinygrad.nn.onnx import OnnxRunner
  from openpilot.system.camerad.cameras.nv12_info import get_nv12_info
@@ -235,10 +268,11 @@ if __name__ == "__main__":
  out = defaultdict(dict)
  # init runners once so weights are shared
  from get_model_metadata import make_metadata_dict
-  vision_runner = OnnxRunner(args.vision_onnx)
-  policy_runner = OnnxRunner(args.policy_onnx)
-  out['metadata']['vision'] = make_metadata_dict(args.vision_onnx)
-  out['metadata']['policy'] = make_metadata_dict(args.policy_onnx)
+  vision_path, policy_path = read_file_chunked_to_shm(args.vision_onnx), read_file_chunked_to_shm(args.policy_onnx)
+  vision_runner = OnnxRunner(vision_path)
+  policy_runner = OnnxRunner(policy_path)
+  out['metadata']['vision'] = make_metadata_dict(vision_path)
+  out['metadata']['policy'] = make_metadata_dict(policy_path)

  for cam_w, cam_h in args.camera_resolutions:
    nv12 = NV12Frame(cam_w, cam_h, *get_nv12_info(cam_w, cam_h))
@@ -27,7 +27,7 @@ class ModelState:
  output: np.ndarray

  def __init__(self, cam_w: int, cam_h: int):
-    self.DEV = get_tg_input_devices(PROCESS_NAME)
+    self.DEV = get_tg_input_devices(PROCESS_NAME, usbgpu=False)['DEV']
    with open(METADATA_PATH, 'rb') as f:
      model_metadata = pickle.load(f)
      self.input_shapes = model_metadata['input_shapes']
@@ -3,8 +3,24 @@ from pathlib import Path

 MODELS_DIR = Path(__file__).resolve().parent / 'models'
 TG_INPUT_DEVICES_PATH = MODELS_DIR / 'tg_input_devices.json'
+USBGPU_VID = 0xADD1
+USBGPU_PID = 0x0001


-def get_tg_input_devices(process_name: str) -> dict[str, str]:
+def get_tg_input_devices(process_name: str, usbgpu: bool):
  with open(TG_INPUT_DEVICES_PATH) as f:
-    return json.load(f)[process_name]
+    return json.load(f)[process_name]['default' if not usbgpu else 'usbgpu']
+
+def modeld_pkl_path(usbgpu: bool):
+  prefix = 'big_' if usbgpu else ''
+  return MODELS_DIR / f'{prefix}driving_tinygrad.pkl'
+
+def usbgpu_present() -> bool:
+  for d in Path("/sys/bus/usb/devices").glob("*"):
+    try:
+      if int((d / "idVendor").read_text(), 16) == USBGPU_VID and \
+          int((d / "idProduct").read_text(), 16) == USBGPU_PID:
+        return True
+    except Exception:
+      pass
+  return False
@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 import os
-from openpilot.selfdrive.modeld.helpers import MODELS_DIR, get_tg_input_devices
+os.environ['GMMU'] = '0' # for usbgpu fast loading, noop for qcom
 from tinygrad.tensor import Tensor
 import time
 import pickle
@@ -22,8 +22,9 @@ from openpilot.selfdrive.controls.lib.drive_helpers import get_accel_from_plan,
 from openpilot.selfdrive.modeld.parse_model_outputs import Parser
 from openpilot.selfdrive.modeld.compile_modeld import make_input_queues
 from openpilot.selfdrive.modeld.fill_model_msg import fill_model_msg, fill_pose_msg, PublishState
-from openpilot.common.file_chunker import read_file_chunked
+from openpilot.common.file_chunker import read_file_chunked, get_manifest_path
 from openpilot.selfdrive.modeld.constants import ModelConstants, Plan
+from openpilot.selfdrive.modeld.helpers import usbgpu_present, modeld_pkl_path, get_tg_input_devices


 PROCESS_NAME = "selfdrive.modeld.modeld"
@@ -71,9 +72,10 @@ class FrameMeta:
 class ModelState:
  prev_desire: np.ndarray  # for tracking the rising edge of the pulse

-  def __init__(self, cam_w: int, cam_h: int):
-    self.DEV = get_tg_input_devices(PROCESS_NAME)
-    jits = pickle.loads(read_file_chunked(MODELS_DIR / 'driving_tinygrad.pkl'))
+  def __init__(self, cam_w: int, cam_h: int, usbgpu: bool):
+    input_devices = get_tg_input_devices(PROCESS_NAME, usbgpu)
+    self.WARP_DEV, self.QUEUE_DEV = input_devices['WARP_DEV'], input_devices['QUEUE_DEV']
+    jits = pickle.loads(read_file_chunked(modeld_pkl_path(usbgpu)))
    vision_metadata = jits['metadata']['vision']
    self.vision_input_shapes =  vision_metadata['input_shapes']
    self.vision_input_names = list(self.vision_input_shapes.keys())
@@ -86,7 +88,7 @@ class ModelState:
    self.prev_desire = np.zeros(ModelConstants.DESIRE_LEN, dtype=np.float32)

    self.frame_skip = ModelConstants.MODEL_RUN_FREQ // ModelConstants.MODEL_CONTEXT_FREQ
-    self.input_queues, self.npy = make_input_queues(self.vision_input_shapes, self.policy_input_shapes, self.frame_skip, device=self.DEV)
+    self.input_queues, self.npy = make_input_queues(self.vision_input_shapes, self.policy_input_shapes, self.frame_skip, device=self.QUEUE_DEV)
    self.full_frames : dict[str, Tensor] = {}
    self._blob_cache : dict[int, Tensor] = {}
    self.parser = Parser()
@@ -95,8 +97,8 @@ class ModelState:
    self.warp_enqueue = jits[(cam_w,cam_h)]['warp_enqueue']
    self.warp_enqueue(
      **self.input_queues,
-      frame=Tensor(np.zeros(self.frame_buf_params['img'][3], dtype=np.uint8), device=self.DEV).contiguous().realize(),
-      big_frame=Tensor(np.zeros(self.frame_buf_params['big_img'][3], dtype=np.uint8), device=self.DEV).contiguous().realize())
+      frame=Tensor(np.zeros(self.frame_buf_params['img'][3], dtype=np.uint8), device=self.WARP_DEV).contiguous().realize(),
+      big_frame=Tensor(np.zeros(self.frame_buf_params['big_img'][3], dtype=np.uint8), device=self.WARP_DEV).contiguous().realize())

  def slice_outputs(self, model_outputs: np.ndarray, output_slices: dict[str, slice]) -> dict[str, np.ndarray]:
    parsed_model_outputs = {k: model_outputs[np.newaxis, v] for k,v in output_slices.items()}
@@ -110,7 +112,7 @@ class ModelState:
      # There is a ringbuffer of imgs, just cache tensors pointing to all of them
      cache_key = (key, ptr)
      if cache_key not in self._blob_cache:
-        self._blob_cache[cache_key] = Tensor.from_blob(ptr, (yuv_size,), dtype='uint8', device=self.DEV)
+        self._blob_cache[cache_key] = Tensor.from_blob(ptr, (yuv_size,), dtype='uint8', device=self.WARP_DEV)
      self.full_frames[key] = self._blob_cache[cache_key]

    # Model decides when action is completed, so desire input is just a pulse triggered on rising edge
@@ -143,7 +145,13 @@ class ModelState:
 def main(demo=False):
  cloudlog.warning("modeld init")

-  USBGPU = False
+  _present = usbgpu_present()
+  _compiled = os.path.isfile(get_manifest_path(modeld_pkl_path(usbgpu=True)))
+  USBGPU = _present and _compiled
+  params = Params()
+  params.put_bool("UsbGpuPresent", _present)
+  params.put_bool("UsbGpuCompiled", _compiled)
+
  if not USBGPU:
    # USB GPU currently saturates a core so can't do this yet,
    # also need to move the aux USB interrupts for good timings
@@ -174,7 +182,7 @@ def main(demo=False):

  st = time.monotonic()
  cloudlog.warning("loading model")
-  model = ModelState(vipc_client_main.width, vipc_client_main.height)
+  model = ModelState(vipc_client_main.width, vipc_client_main.height, USBGPU)
  cloudlog.warning(f"models loaded in {time.monotonic() - st:.1f}s, modeld starting")

  # messaging
@@ -1 +0,0 @@
-driving_policy.onnx
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:565e53c38dcd64c50dd3fe4d5ee1530213aeefd66c3f6b67ea6a72a32612a6bf
+size 14061419
@@ -1 +0,0 @@
-driving_vision.onnx
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1f0cab5033fe9e3bc5e174a2e790fa277f7d9fc44c65822d734064d2f899a9a0
+size 296203378
@@ -139,6 +139,8 @@ class MiciHomeLayout(Widget):
    self._version_text = self._get_version_text()

    self._experimental_icon = IconWidget("icons_mici/experimental_mode.png", (48, 48))
+    self._egpu_icon = IconWidget("icons_mici/egpu.png", (50, 37))
+    self._egpu_icon_gray = IconWidget("icons_mici/egpu_gray.png", (50, 37))
    self._mic_icon = IconWidget("icons_mici/microphone.png", (32, 46))
    self._body_icon = IconWidget("icons_mici/body.png", (54, 37))

@@ -148,6 +150,8 @@ class MiciHomeLayout(Widget):
      IconWidget("icons_mici/settings.png", (48, 48), opacity=0.9),
      NetworkIcon(),
      self._experimental_icon,
+      self._egpu_icon,
+      self._egpu_icon_gray,
      self._body_icon,
      self._mic_icon,
    ], spacing=18)
@@ -244,6 +248,8 @@ class MiciHomeLayout(Widget):

    # ***** Center-aligned bottom section icons *****
    self._experimental_icon.set_visible(ui_state.experimental_mode)
+    self._egpu_icon.set_visible(ui_state.usbgpu and ui_state.usbgpu_compiled)
+    self._egpu_icon_gray.set_visible(ui_state.usbgpu and not ui_state.usbgpu_compiled)
    self._mic_icon.set_visible(ui_state.recording_audio)
    self._body_icon.set_visible(ui_state.is_body)

@@ -75,6 +75,8 @@ class UIState:
    self.is_release = self.params.get_bool("IsReleaseBranch")
    self.always_on_dm: bool = self.params.get_bool("AlwaysOnDM")
    self.experimental_mode: bool = self.params.get_bool("ExperimentalMode")
+    self.usbgpu: bool = self.params.get_bool("UsbGpuPresent")
+    self.usbgpu_compiled: bool = self.params.get_bool("UsbGpuCompiled")
    self.started: bool = False
    self.ignition: bool = False
    self.recording_audio: bool = False
@@ -201,6 +203,8 @@ class UIState:
    self.is_metric = self.params.get_bool("IsMetric")
    self.always_on_dm = self.params.get_bool("AlwaysOnDM")
    self.experimental_mode = self.params.get_bool("ExperimentalMode")
+    self.usbgpu = self.params.get_bool("UsbGpuPresent")
+    self.usbgpu_compiled = self.params.get_bool("UsbGpuCompiled")


 class Device: