modeld_v2: Support eGpu

.github/workflows/build-single-tinygrad-model.yaml

@@ -34,6 +34,14 @@ on:
         required: false
         default: true
         type: boolean
+      target_hardware:
+        description: 'Hardware target to compile for'
+        required: false
+        type: choice
+        default: 'qcom'
+        options:
+          - qcom
+          - usbgpu
   workflow_dispatch:
     inputs:
       upstream_branch:
@@ -81,9 +89,17 @@ on:
         description: 'Minimum selector version'
         required: false
         type: string
+      target_hardware:
+        description: 'Hardware target to compile for'
+        required: false
+        type: choice
+        default: 'qcom'
+        options:
+          - qcom
+          - usbgpu
 env:
   RECOMPILED_DIR: recompiled${{ inputs.recompiled_dir }}
-  JSON_FILE: docs/docs/driving_models_v${{ inputs.json_version }}.json
+  JSON_FILE: docs/docs/driving_models_${{ inputs.target_hardware == 'usbgpu' && 'usbgpu_v' || 'v' }}${{ inputs.json_version }}.json
 
 jobs:
   build_model:
@@ -93,6 +109,7 @@ jobs:
       custom_name: ${{ inputs.custom_name || inputs.upstream_branch }}
       is_20hz: ${{ inputs.is_20hz }}
       artifact_suffix: ${{ inputs.artifact_suffix }}
+      target_hardware: ${{ inputs.target_hardware }}
     secrets: inherit
 
   publish_model:

.github/workflows/sunnypilot-build-model.yaml

@@ -191,7 +191,7 @@ jobs:
           if [ "${{ inputs.target_hardware }}" == "usbgpu" ]; then
             echo "USBGPU build"
             export USBGPU=1
-            TG_FLAGS="DEV=AMD USBGPU=1 IMAGE=1 FLOAT16=1 NOLOCALS=1 JIT_BATCH_SIZE=0 OPENPILOT_HACKS=1"
+            TG_FLAGS="DEBUG=2 DEV=USB+AMD:LLVM WARP_DEV=QCOM FLOAT16=1 JIT_BATCH_SIZE=0 GMMU=0"
             OUTPUT_PKL="${{ env.MODELS_DIR }}/big_driving_tinygrad.pkl"
           else
             echo "QCOM build"

openpilot/common/params_keys.h

@@ -197,7 +197,9 @@ inline static std::unordered_map<std::string, ParamKeyAttributes> keys = {
     {"ModelManager_DownloadIndex", {CLEAR_ON_MANAGER_START | CLEAR_ON_ONROAD_TRANSITION, INT}},
     {"ModelManager_Favs", {PERSISTENT | BACKUP, STRING}},
     {"ModelManager_LastSyncTime", {CLEAR_ON_MANAGER_START | CLEAR_ON_OFFROAD_TRANSITION, INT, "0"}},
+    {"ModelManager_LastSyncTime_USBGPU", {CLEAR_ON_MANAGER_START | CLEAR_ON_OFFROAD_TRANSITION, INT, "0"}},
     {"ModelManager_ModelsCache", {PERSISTENT | BACKUP, JSON}},
+    {"ModelManager_ModelsCache_USBGPU", {PERSISTENT | BACKUP, JSON}},
 
     // Neural Network Lateral Control
     {"NeuralNetworkLateralControl", {PERSISTENT | BACKUP, BOOL, "0"}},

openpilot/selfdrive/modeld/compile_modeld.py

@@ -3,7 +3,6 @@ import argparse
 import atexit
 import math
 import os
-import pickle
 import tempfile
 import time
 import shutil
@@ -13,7 +12,6 @@ from collections import namedtuple
 import numpy as np
 
 from openpilot.selfdrive.modeld.helpers import dump_oob, load_oob
-from openpilot.selfdrive.modeld.usbgpu_link import wait_usbgpu_link
 
 def _patch_tinygrad_fetch_fw():
   import hashlib
@@ -31,22 +29,6 @@ def _patch_tinygrad_fetch_fw():
   helpers.fetch_fw = fetch_fw
 _patch_tinygrad_fetch_fw()
 
-def _patch_tinygrad_buffer_reduce():
-  from tinygrad.device import Buffer
-  def __reduce_ex__(self, protocol):
-    buf = None
-    if self._base is not None:
-      return self.__class__, (self.device, self.size, self.dtype, None, None, None, 0, self.base, self.offset, self.is_allocated())
-    if self.device == "NPY":
-      return self.__class__, (self.device, self.size, self.dtype, self._buf, self.options, None, self.uop_refcount)
-    if self.is_allocated():
-      buf = bytearray(self.nbytes)
-      self.copyout(memoryview(buf))
-      if protocol >= 5:
-        buf = pickle.PickleBuffer(buf)
-    return self.__class__, (self.device, self.size, self.dtype, None, self.options, buf, self.uop_refcount)
-  Buffer.__reduce_ex__ = __reduce_ex__
-_patch_tinygrad_buffer_reduce()
 
 from tinygrad.tensor import Tensor
 from tinygrad.helpers import Context
@@ -312,9 +294,6 @@ if __name__ == "__main__":
   p.add_argument('--frame-skip', type=int, required=True)
   args = p.parse_args()
 
-  if 'USB+AMD' in os.environ.get('DEV', ''):
-    wait_usbgpu_link()
-
   model_path = read_file_chunked_to_disk(args.onnx)
   model_w, model_h = args.model_size
 

openpilot/selfdrive/modeld/helpers.py

@@ -6,10 +6,11 @@ import struct
 import tempfile
 from pathlib import Path
 
+from openpilot.common.file_chunker import get_manifest_path
+from openpilot.common.hardware.usb import CHESTNUT_VENDOR_ID, CHESTNUT_PRODUCT_ID, USB_DEVICES_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, usbgpu: bool):
@@ -38,22 +39,21 @@ def dump_oob(obj, f):
 def load_oob(f):
   opcodes = f.read(struct.unpack('<q', f.read(8))[0])
   def buffers():
-    prev = None
     while (h := f.read(8)):
-      if prev is not None:
+      pb = pickle.PickleBuffer(bytearray(struct.unpack('<q', h)[0]))
-        prev.release()
+      f.readinto(pb)
-      buf = bytearray(struct.unpack('<q', h)[0])
+      yield pb
-      f.readinto(buf)
-      prev = pickle.PickleBuffer(buf)
-      yield prev
   return pickle.load(io.BytesIO(opcodes), buffers=buffers())
 
 def usbgpu_present() -> bool:
-  for d in Path("/sys/bus/usb/devices").glob("*"):
+  for d in USB_DEVICES_PATH.glob("*"):
     try:
-      if int((d / "idVendor").read_text(), 16) == USBGPU_VID and \
+      usb_id = (int((d / "idVendor").read_text(), 16), int((d / "idProduct").read_text(), 16))
-          int((d / "idProduct").read_text(), 16) == USBGPU_PID:
+      if usb_id == (CHESTNUT_VENDOR_ID, CHESTNUT_PRODUCT_ID):
         return True
     except Exception:
       pass
   return False
+
+def usbgpu_compiled() -> bool:
+  return Path(get_manifest_path(modeld_pkl_path(usbgpu=True))).is_file()

openpilot/selfdrive/modeld/modeld.py

@@ -2,6 +2,7 @@
 import os
 os.environ['GMMU'] = '0' # for usbgpu fast loading, noop for qcom
 from tinygrad.tensor import Tensor
+import threading
 import time
 import numpy as np
 import openpilot.cereal.messaging as messaging
@@ -18,17 +19,13 @@ from openpilot.common.transformations.camera import DEVICE_CAMERAS
 from openpilot.system.camerad.cameras.nv12_info import get_nv12_info
 from openpilot.common.transformations.model import get_warp_matrix
 from openpilot.selfdrive.controls.lib.desire_helper import DesireHelper
-from openpilot.selfdrive.controls.lib.drive_helpers import get_accel_from_plan, smooth_value, get_curvature_from_plan
+from openpilot.selfdrive.controls.lib.drive_helpers import get_accel_from_plan, should_stop, smooth_value, get_curvature_from_plan
 from openpilot.selfdrive.modeld.parse_model_outputs import Parser
 from openpilot.selfdrive.modeld.compile_modeld import make_input_queues, WARP_INPUTS, POLICY_INPUTS
 from openpilot.selfdrive.modeld.fill_model_msg import fill_model_msg, fill_driving_model_data, fill_pose_msg, PublishState
-from openpilot.common.file_chunker import open_file_chunked, get_manifest_path
+from openpilot.common.file_chunker import open_file_chunked
 from openpilot.selfdrive.modeld.constants import ModelConstants, Plan
-from openpilot.selfdrive.modeld.helpers import usbgpu_present, modeld_pkl_path, get_tg_input_devices, load_oob
+from openpilot.selfdrive.modeld.helpers import usbgpu_present, usbgpu_compiled, modeld_pkl_path, get_tg_input_devices, load_oob
-from openpilot.selfdrive.modeld.usbgpu_link import wait_usbgpu_link
-
-from openpilot.sunnypilot.livedelay.helpers import get_lat_delay
-from openpilot.sunnypilot.modeld_v2.modeld_base import ModelStateBase
 
 PROCESS_NAME = "openpilot.selfdrive.modeld.modeld"
 SEND_RAW_PRED = os.getenv('SEND_RAW_PRED')
@@ -36,13 +33,14 @@ SEND_RAW_PRED = os.getenv('SEND_RAW_PRED')
 LAT_SMOOTH_SECONDS = 0.0
 LONG_SMOOTH_SECONDS = 0.3
 MIN_LAT_CONTROL_SPEED = 0.3
+BIG_MODEL_TIMEOUT = 60
 
 
 def get_action_from_model(model_output: dict[str, np.ndarray], prev_action: log.ModelDataV2.Action,
                           lat_action_t: float, long_action_t: float, v_ego: float) -> log.ModelDataV2.Action:
   if 'action' not in model_output:
     plan = model_output['plan'][0]
-    desired_accel, should_stop = get_accel_from_plan(plan[:,Plan.VELOCITY][:,0],
+    desired_accel = get_accel_from_plan(plan[:,Plan.VELOCITY][:,0],
                                         plan[:,Plan.ACCELERATION][:,0],
                                         ModelConstants.T_IDXS,
                                         action_t=long_action_t)
@@ -54,7 +52,7 @@ def get_action_from_model(model_output: dict[str, np.ndarray], prev_action: log.
   else:
     desired_accel = model_output['action'][0,1]
     desired_curvature = model_output['action'][0,0] / (max(1.0, v_ego))**2
-    should_stop = (v_ego < 0.3 and desired_accel < 0.1)
+  stop = should_stop(v_ego, desired_accel)
   desired_accel = smooth_value(desired_accel, prev_action.desiredAcceleration, LONG_SMOOTH_SECONDS)
   if v_ego > MIN_LAT_CONTROL_SPEED:
     desired_curvature = smooth_value(desired_curvature, prev_action.desiredCurvature, LAT_SMOOTH_SECONDS)
@@ -63,7 +61,7 @@ def get_action_from_model(model_output: dict[str, np.ndarray], prev_action: log.
 
   return log.ModelDataV2.Action(desiredCurvature=float(desired_curvature),
                                 desiredAcceleration=float(desired_accel),
-                                shouldStop=bool(should_stop))
+                                shouldStop=bool(stop))
 
 
 class FrameMeta:
@@ -76,12 +74,10 @@ class FrameMeta:
       self.frame_id, self.timestamp_sof, self.timestamp_eof = vipc.frame_id, vipc.timestamp_sof, vipc.timestamp_eof
 
 
-class ModelState(ModelStateBase):
+class ModelState:
   prev_desire: np.ndarray  # for tracking the rising edge of the pulse
 
   def __init__(self, cam_w: int, cam_h: int, usbgpu: bool):
-    ModelStateBase.__init__(self)
-    self.LAT_SMOOTH_SECONDS = LAT_SMOOTH_SECONDS
     input_devices = get_tg_input_devices(PROCESS_NAME, usbgpu)
     self.WARP_DEV, self.QUEUE_DEV = input_devices['WARP_DEV'], input_devices['QUEUE_DEV']
     jits = load_oob(open_file_chunked(modeld_pkl_path(usbgpu)))
@@ -138,16 +134,24 @@ class ModelState(ModelStateBase):
       outputs_dict['raw_pred'] = model_output.copy()
     return outputs_dict
 
+  def warmup(self) -> None:
+    dummy_frames = {k: np.zeros(self.frame_buf_params[k][3], dtype=np.uint8) for k in self.vision_input_names}
+    eye = np.eye(3, dtype=np.float32)
+    dims = {'desire_pulse': ModelConstants.DESIRE_LEN, 'traffic_convention': 2, 'action_t': 2}
+    self.run(dummy_frames, dict.fromkeys(self.vision_input_names, eye), {k: np.zeros(v, dtype=np.float32) for k, v in dims.items()})
+    self.input_queues, self.npy = make_input_queues(self.input_shapes, self.frame_skip, device=self.QUEUE_DEV)
+    self.prev_desire[:] = 0
+    self.full_frames.clear()
+    self._blob_cache.clear()
+
 
 def main(demo=False):
   cloudlog.warning("modeld init")
 
-  _present = usbgpu_present()
+  USBGPU = usbgpu_present() and usbgpu_compiled()
-  _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("UsbGpuLoading", USBGPU)
-  params.put_bool("UsbGpuCompiled", _compiled)
+  params.remove("UsbGpuActive")
 
   config_realtime_process(7, 54)
 
@@ -174,15 +178,33 @@ def main(demo=False):
   if use_extra_client:
     cloudlog.warning(f"connected extra cam with buffer size: {vipc_client_extra.buffer_len} ({vipc_client_extra.width} x {vipc_client_extra.height})")
 
-  if USBGPU:
-    wait_usbgpu_link()
   st = time.monotonic()
   cloudlog.warning("loading model")
-  model = ModelState(vipc_client_main.width, vipc_client_main.height, USBGPU)
+  model = None
+  if USBGPU:
+    big_model = None
+    def load_big():
+      nonlocal big_model
+      try:
+        m = ModelState(vipc_client_main.width, vipc_client_main.height, True)
+        m.warmup()
+        big_model = m
+      except Exception:
+        cloudlog.exception("big model load failed")
+    loader = threading.Thread(target=load_big, daemon=True)
+    loader.start()
+    loader.join(BIG_MODEL_TIMEOUT)
+    model = big_model
+    params.put_bool("UsbGpuActive", model is not None)
+
+  small_model = ModelState(vipc_client_main.width, vipc_client_main.height, False) if model is None or USBGPU else None
+  if model is None:
+    model = small_model
+  params.put_bool("UsbGpuLoading", False)
   cloudlog.warning(f"models loaded in {time.monotonic() - st:.1f}s, modeld starting")
 
   # messaging
-  pm = PubMaster(["modelV2", "drivingModelData", "cameraOdometry", "modelDataV2SP"])
+  pm = PubMaster(["modelV2", "drivingModelData", "cameraOdometry"])
   sm = SubMaster(["deviceState", "carState", "roadCameraState", "liveCalibration", "driverMonitoringState", "carControl", "liveDelay"])
 
   publish_state = PublishState()
@@ -252,9 +274,7 @@ def main(demo=False):
     is_rhd = sm["driverMonitoringState"].isRHD
     frame_id = sm["roadCameraState"].frameId
     v_ego = max(sm["carState"].vEgo, 0.)
-    if sm.frame % 60 == 0:
+    lat_delay = sm["liveDelay"].lateralDelay + LAT_SMOOTH_SECONDS
-      model.lat_delay = get_lat_delay(params, sm["liveDelay"].lateralDelay)
-    lat_delay = model.lat_delay + LAT_SMOOTH_SECONDS
     if sm.updated["liveCalibration"] and sm.seen['roadCameraState'] and sm.seen['deviceState']:
       device_from_calib_euler = np.array(sm["liveCalibration"].rpyCalib, dtype=np.float32)
       dc = DEVICE_CAMERAS[(str(sm['deviceState'].deviceType), str(sm['roadCameraState'].sensor))]
@@ -293,7 +313,17 @@ def main(demo=False):
     }
 
     mt1 = time.perf_counter()
+    try:
       model_output = model.run(bufs, transforms, inputs)
+    except Exception:
+      if not params.get_bool("UsbGpuActive"):
+        raise
+      # fallback to small model
+      cloudlog.exception("big model failed, fall back to small")
+      params.put_bool("UsbGpuActive", False)
+      model = small_model
+      run_count = 0
+      model_output = None
     mt2 = time.perf_counter()
     model_execution_time = mt2 - mt1
 
@@ -301,7 +331,6 @@ def main(demo=False):
       modelv2_send = messaging.new_message('modelV2')
       drivingdata_send = messaging.new_message('drivingModelData')
       posenet_send = messaging.new_message('cameraOdometry')
-      mdv2sp_send = messaging.new_message('modelDataV2SP')
 
       action = get_action_from_model(model_output, prev_action, lat_action_t, long_action_t, v_ego)
       prev_action = action
@@ -316,14 +345,12 @@ def main(demo=False):
       DH.update(sm['carState'], sm['carControl'].latActive, lane_change_prob)
       modelv2_send.modelV2.meta.laneChangeState = DH.lane_change_state
       modelv2_send.modelV2.meta.laneChangeDirection = DH.lane_change_direction
-      mdv2sp_send.modelDataV2SP.laneTurnDirection = DH.lane_turn_direction
 
       fill_driving_model_data(drivingdata_send, modelv2_send)
       fill_pose_msg(posenet_send, model_output, meta_main.frame_id, vipc_dropped_frames, meta_main.timestamp_eof, live_calib_seen)
       pm.send('modelV2', modelv2_send)
       pm.send('drivingModelData', drivingdata_send)
       pm.send('cameraOdometry', posenet_send)
-      pm.send('modelDataV2SP', mdv2sp_send)
     last_vipc_frame_id = meta_main.frame_id
 
 

openpilot/sunnypilot/modeld_v2/SConscript

@@ -1,9 +1,12 @@
 import os
 import glob
+import sys
+import subprocess
 
 from openpilot.common.transformations.camera import _ar_ox_fisheye, _os_fisheye
 from openpilot.common.transformations.model import MEDMODEL_INPUT_SIZE
 from openpilot.common.hardware import HARDWARE, PC
+from openpilot.selfdrive.modeld.helpers import usbgpu_present
 
 Import('env', 'arch', 'release')
 lenv = env.Clone()
@@ -22,14 +25,21 @@ def get_camera_configs():
 
 CAMERA_CONFIGS = get_camera_configs()
 
-tg_flags = {
+def probe_devices():
-  'larch64': 'DEV=QCOM FLOAT16=1 NOLOCALS=1 JIT_BATCH_SIZE=0',
+  return set(subprocess.run(
-  'Darwin': f'DEV=CPU HOME={os.path.expanduser("~")}',
+    [sys.executable, '-c', 'from tinygrad import Device\nprint("\\n".join(Device.get_available_devices()))'],
-}.get(arch, 'DEV=CPU:LLVM')
+    capture_output=True, text=True, check=True).stdout.strip().splitlines())
 
-image_flag = {
+available = probe_devices()
-  'larch64': 'IMAGE=2',
+if 'CUDA' in available:
-}.get(arch, 'IMAGE=0')
+  tg_backend = 'CUDA'
+  tg_flags = f'DEV={tg_backend}'
+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'
+  tg_flags = f'DEV=CPU HOME={os.path.expanduser("~")}' if arch == 'Darwin' else 'DEV=CPU:LLVM'
 
 model_w, model_h = MEDMODEL_INPUT_SIZE
 from openpilot.selfdrive.modeld.constants import ModelConstants
@@ -41,9 +51,20 @@ compile_modeld_script = File("compile_modeld.py").abspath
 upstream_compile_script = File(Dir("#openpilot/selfdrive/modeld").File("compile_modeld.py").abspath)
 script_deps = [File("compile_modeld.py"), upstream_compile_script]
 
+USBGPU = usbgpu_present()
+if USBGPU:
+  usbgpu_tg_flags = f'DEBUG=2 DEV=USB+AMD:LLVM WARP_DEV={tg_backend} FLOAT16=1 JIT_BATCH_SIZE=0 GMMU=0'
+  usbgpu_lock = File("models/.usb_gpu.lock").abspath
+
 def compile_combined(model_type, onnx_args, output_name):
-  output_pkl = File(f"models/{output_name}").abspath
+  for usbgpu in ([False, True] if USBGPU else [False]):
-  cmd = (f'{pythonpath_string} {tg_flags} {image_flag} python3 {compile_modeld_script} '
+    prefix = 'big_' if usbgpu else ('big_' if os.getenv('BIG_INTO_SMALL') else '')
+    final_output_name = prefix + output_name
+    output_pkl = File(f"models/{final_output_name}").abspath
+
+    active_tg_flags = usbgpu_tg_flags if usbgpu else tg_flags
+
+    cmd = (f'{pythonpath_string} {active_tg_flags} python3 {compile_modeld_script} '
            f'--model-type {model_type} '
            f'--model-size {model_w}x{model_h} '
            f'--camera-resolutions {camera_res_args} '
@@ -51,7 +72,9 @@ def compile_combined(model_type, onnx_args, output_name):
            f'--frame-skip {frame_skip} '
            f'--output {output_pkl}')
     onnx_files = [f for f in onnx_args.split() if f.endswith('.onnx')]
-  return lenv.Command(output_pkl, tinygrad_files + script_deps + [File(f) for f in onnx_files if os.path.isfile(f)], cmd)
+    node = lenv.Command(output_pkl, tinygrad_files + script_deps + [File(f) for f in onnx_files if os.path.isfile(f)], cmd)
+    if usbgpu:
+      lenv.SideEffect(usbgpu_lock, node)
 
 # Vision + Policy (stock default model)
 vision_onnx = File("models/driving_vision.onnx").abspath

openpilot/sunnypilot/modeld_v2/compile_modeld.py

@@ -8,10 +8,10 @@ See the LICENSE.md file in the root directory for more details.
 
 import argparse
 import os
-import pickle
+import tempfile
-import time
 from collections import defaultdict
 from functools import partial
+from openpilot.selfdrive.modeld.helpers import dump_oob, load_oob
 import numpy as np
 os.environ['GMMU'] = '0'
 
@@ -76,7 +76,7 @@ def get_policy_npy_shapes(input_shapes: dict, is_supercombo: bool = False) -> tu
 
 
 def generate_queues_and_npy(input_shapes: dict, frame_skip: int, device: str = Device.DEFAULT,
-                            is_supercombo: bool = False, use_packed: bool = True) -> tuple[dict, dict]:
+                            is_supercombo: bool = False) -> tuple[dict, dict]:
   road_key, _ = _detect_vision_keys(input_shapes)
   if not road_key:
     raise ValueError("Vision road key missing from input shapes.")
@@ -92,7 +92,6 @@ def generate_queues_and_npy(input_shapes: dict, frame_skip: int, device: str = D
   desire_shape = input_shapes[desire_key]
   features_buffer = input_shapes.get('features_buffer')
 
-  if use_packed:  # remove packed detection block after all models are recompiled
   npy_arrays = {
     'tfm': np.zeros((3, 3), dtype=np.float32),
     'big_tfm': np.zeros((3, 3), dtype=np.float32)
@@ -119,42 +118,18 @@ def generate_queues_and_npy(input_shapes: dict, frame_skip: int, device: str = D
                        dtype=np.float32), device=device).contiguous().realize()
 
   queues.update({key: Tensor(value, device='NPY').realize() for key, value in npy_arrays.items() if key in ('tfm', 'big_tfm')})
-  else:
-    # TODO-SP: Remove legacy queuing fallback else block after all models are recompiled
-    npy_arrays = {
-      'desire': np.zeros(desire_shape[2], dtype=np.float32),
-      'tfm': np.zeros((3, 3), dtype=np.float32),
-      'big_tfm': np.zeros((3, 3), dtype=np.float32)
-    }
-
-    for key, shape in input_shapes.items():
-      if key not in npy_arrays and 'img' not in key and key not in ('features_buffer', desire_key):
-        npy_arrays[key] = np.zeros(shape, dtype=np.float32)
-
-    queues = {
-      'img_q': Tensor(np.zeros(img_buf_shape, dtype=np.uint8), device=device).contiguous().realize(),
-      'big_img_q': Tensor(np.zeros(img_buf_shape, dtype=np.uint8), device=device).contiguous().realize(),
-      'desire_q': Tensor(np.zeros((frame_skip * desire_shape[1], desire_shape[0], desire_shape[2]),
-                    dtype=np.float32), device=device).contiguous().realize()
-    }
-
-    if features_buffer:
-      queues['feat_q'] = Tensor(np.zeros((frame_skip * (features_buffer[1] - 1) + 1, features_buffer[0], features_buffer[2]),
-                         dtype=np.float32), device=device).contiguous().realize()
-
-    queues.update({key: Tensor(value, device='NPY').realize() for key, value in npy_arrays.items()})
 
   return queues, npy_arrays
 
 
 def make_split_input_queues(vision_input_shapes: dict, policy_input_shapes: dict,
-                            frame_skip: int, device: str = Device.DEFAULT, use_packed: bool = True) -> tuple[dict, dict]:
+                            frame_skip: int, device: str = Device.DEFAULT) -> tuple[dict, dict]:
-  return generate_queues_and_npy({**vision_input_shapes, **policy_input_shapes}, frame_skip, device, is_supercombo=False, use_packed=use_packed)
+  return generate_queues_and_npy({**vision_input_shapes, **policy_input_shapes}, frame_skip, device, is_supercombo=False)
 
 
 def make_supercombo_input_queues(input_shapes: dict, frame_skip: int,
-                                 device: str = Device.DEFAULT, use_packed: bool = True) -> tuple[dict, dict]:
+                                 device: str = Device.DEFAULT) -> tuple[dict, dict]:
-  return generate_queues_and_npy(input_shapes, frame_skip, device, is_supercombo=True, use_packed=use_packed)
+  return generate_queues_and_npy(input_shapes, frame_skip, device, is_supercombo=True)
 
 
 def create_jit_runner(vision_runner, policy_runners: list, nv12: NV12Frame, model_size: tuple[int, int],
@@ -233,29 +208,34 @@ def compile_and_warmup(nv12: NV12Frame, model_size: tuple[int, int], prepare_onl
     raise ValueError("Could not find vision, model, or policy metadata.")
 
   features_slice = feat_meta['output_slices']['hidden_state']
-  WARP_DEV = 'CPU' if "USBGPU" in os.environ else Device.DEFAULT
+  WARP_DEV = os.getenv('WARP_DEV', Device.DEFAULT)
 
   is_supercombo = vision_runner is None
   run_func = create_jit_runner(vision_runner, policy_runners, nv12, model_size, features_slice, frame_skip, all_shapes, prepare_only)
   run_jit = TinyJit(run_func, prune=True)
-  queues, npy_arrays = generate_queues_and_npy(all_shapes, frame_skip, Device.DEFAULT, is_supercombo=is_supercombo)
+  def run_once(seed):
-
+    queues, npy = generate_queues_and_npy(all_shapes, frame_skip, Device.DEFAULT, is_supercombo=is_supercombo)
-  for i in range(3):
+    rng = np.random.default_rng(seed)
-    rng = np.random.default_rng(42 + i)
     frame = Tensor.randint(nv12.size, low=0, high=256, dtype=dtypes.uint8, device=WARP_DEV).realize()
     big_frame = Tensor.randint(nv12.size, low=0, high=256, dtype=dtypes.uint8, device=WARP_DEV).realize()
-    for arr in npy_arrays.values():
+    for value in npy.values():
-      arr[:] = rng.standard_normal(arr.shape).astype(arr.dtype)
+      value[:] = rng.standard_normal(value.shape).astype(value.dtype)
-
     Device.default.synchronize()
-    start_time = time.perf_counter()
+    outs = run_jit(**queues, frame=frame, big_frame=big_frame)
-    run_jit(**queues, frame=frame, big_frame=big_frame)
-    mid_time = time.perf_counter()
     Device.default.synchronize()
-    print(f"  [{i + 1}/3] enqueue {(mid_time - start_time) * 1e3:6.2f} ms -- total {(time.perf_counter() - start_time) * 1e3:6.2f} ms")
+    return [np.copy(value.numpy()) for value in (outs if isinstance(outs, tuple) else [outs])] if outs is not None else []
 
-  # TODO-SP: switch to dump_oob/load_oob on next full recompile of all models
+  for i in range(3):
-  return pickle.loads(pickle.dumps(run_jit)) if not prepare_only else run_jit
+    run_once(42 + i)
+
+  if not prepare_only:
+    baseline = run_once(42)
+    with tempfile.TemporaryFile(dir=".") as f:
+      dump_oob(run_jit, f)
+      f.seek(0)
+      run_jit = load_oob(f)
+    assert all(np.array_equal(baseline, deserialized) for baseline, deserialized in zip(baseline, run_once(42), strict=True)), "OOB pickling regression"
+  return run_jit
 
 
 def _parse_size(size_str: str) -> tuple[int, int]:
@@ -352,8 +332,7 @@ if __name__ == "__main__":
                         vision_runner, policy_runners, output_data['metadata']))
 
   with open(args.output, "wb") as file:
-    # TODO-SP: switch to dump_oob from openpilot/selfdrive/helpers on next full recompile of all models
+    dump_oob(output_data, file)
-    pickle.dump(output_data, file)
 
   pkl_size = os.path.getsize(args.output)
   print(f"Saved combined JIT to {args.output} ({pkl_size / 1e6:.2f} MB)")

openpilot/sunnypilot/modeld_v2/modeld.py

@@ -10,11 +10,14 @@ import os
 os.environ['GMMU'] = '0'
 from openpilot.common.hardware import TICI
 os.environ['DEV'] = 'QCOM' if TICI else 'CPU'
-USBGPU = "USBGPU" in os.environ
+
+from openpilot.selfdrive.modeld.helpers import usbgpu_present, load_oob
+from openpilot.selfdrive.modeld.usbgpu_link import wait_usbgpu_link
+
+USBGPU = usbgpu_present()
 if USBGPU:
   os.environ['DEV'] = 'AMD'
   os.environ['AMD_IFACE'] = 'USB'
-import pickle
 import time
 import numpy as np
 import openpilot.cereal.messaging as messaging
@@ -108,11 +111,10 @@ class ModelState(ModelStateBase):
 
   def _init_combined(self, pkl_path, cam_w, cam_h, bundle):
     cloudlog.warning(f"loading combined pkl: {pkl_path}")
-    # TODO-SP: switch to load_oob from openpilot/selfdrive/helpers on next full recompile of all models
+    jits = load_oob(open_file_chunked(pkl_path))
-    jits = pickle.load(open_file_chunked(pkl_path))
 
     self.DEV = Device.DEFAULT
-    self.WARP_DEV = 'CPU' if USBGPU else self.DEV
+    self.WARP_DEV = ('QCOM' if TICI else 'CPU') if USBGPU else self.DEV
     self.QUEUE_DEV = self.DEV
 
     metadata = jits['metadata']
@@ -120,13 +122,6 @@ class ModelState(ModelStateBase):
     self._run_policy = jits[(cam_w, cam_h)]['run_policy']
     self._warp_enqueue = jits[(cam_w, cam_h)]['warp_enqueue']
 
-    # TODO-SP: Remove legacy use_packed detection block after all models are recompiled
-    captured = getattr(self._run_policy, 'captured', None)
-    if captured is not None:
-      use_packed = 'packed_npy_inputs' in getattr(captured, 'expected_names', [])
-    else:
-      use_packed = True
-
     if 'model' in metadata:
       model_metadata = metadata['model']
       self.vision_output_slices = model_metadata['output_slices']
@@ -137,7 +132,7 @@ class ModelState(ModelStateBase):
       from openpilot.sunnypilot.modeld_v2.compile_modeld import make_supercombo_input_queues
       frame_skip = derive_frame_skip({}, model_metadata['input_shapes'])
       self.input_queues, self.numpy_inputs = make_supercombo_input_queues(model_metadata['input_shapes'],
-                                                                          frame_skip, device=self.QUEUE_DEV, use_packed=use_packed)
+                                                                          frame_skip, device=self.QUEUE_DEV)
     else:
       vision_metadata = metadata['vision']
       policy_keys = [k for k in metadata if k != 'vision']
@@ -156,7 +151,7 @@ class ModelState(ModelStateBase):
       self._vision_input_names = [k for k in vision_input_shapes if 'img' in k]
       frame_skip = derive_frame_skip(vision_input_shapes, policy_input_shapes)
       self.input_queues, self.numpy_inputs = make_split_input_queues(vision_input_shapes, policy_input_shapes,
-                                                                     frame_skip, device=self.QUEUE_DEV, use_packed=use_packed)
+                                                                     frame_skip, device=self.QUEUE_DEV)
 
     self._desire_key = next(key for key in self.numpy_inputs if key.startswith('desire'))
     self._road_key = next(key for key in self._vision_input_names if 'big' not in key)
@@ -189,6 +184,26 @@ class ModelState(ModelStateBase):
       frame=Tensor(np.zeros(yuv_size, dtype=np.uint8), device=self.WARP_DEV).contiguous().realize(),
       big_frame=Tensor(np.zeros(yuv_size, dtype=np.uint8), device=self.WARP_DEV).contiguous().realize())
 
+    if USBGPU:
+      self.warmup()
+
+  def warmup(self) -> None:
+    dummy_frames = {k: np.zeros(self.frame_buf_params[k][3], dtype=np.uint8) for k in self._vision_input_names}
+    transforms = {k: np.eye(3, dtype=np.float32) for k in [self._road_key, self._wide_key] if k}
+
+    dummy_inputs = {}
+    for k, v in self.numpy_inputs.items():
+      if k not in ['tfm', 'big_tfm', 'prev_feat']:
+        dummy_inputs[k] = np.zeros(v.shape, dtype=v.dtype)
+
+    self.run(dummy_frames, transforms, dummy_inputs, prepare_only=False)
+
+    for v in self.numpy_inputs.values():
+      v[:] = 0
+    self.prev_desire[:] = 0
+    self.full_frames.clear()
+    self._blob_cache.clear()
+
 
   @property
   def mlsim(self) -> bool:
@@ -265,11 +280,6 @@ class ModelState(ModelStateBase):
       buf[0, :-1] = buf[0, 1:]
       buf[0, -1, :] = outputs['desired_curvature'][0, :] if not self.mlsim else 0
 
-    # TODO-SP: This is a hack to prevent GPU corruption by calculating in CPU space, it can be removed on next recompile
-    if 'prev_feat' not in self.numpy_inputs and 'feat_q' in self.input_queues:
-      feat_val = self.input_queues['feat_q'].numpy()
-      self.input_queues['feat_q'].assign(feat_val).realize()
-
     return outputs
 
   def get_action_from_model(self, model_output: dict[str, np.ndarray], prev_action: log.ModelDataV2.Action,
@@ -306,6 +316,9 @@ def main(demo=False):
   setproctitle(PROCESS_NAME)
   config_realtime_process(7, 54)
 
+  if USBGPU:
+    wait_usbgpu_link()
+
   # visionipc clients
   while True:
     available_streams = VisionIpcClient.available_streams("camerad", block=False)
@@ -340,6 +353,9 @@ def main(demo=False):
   publish_state = PublishState()
   params = Params()
 
+  params.put_bool("UsbGpuPresent", USBGPU)
+  params.put_bool("UsbGpuCompiled", USBGPU)
+
   # setup filter to track dropped frames
   frame_dropped_filter = FirstOrderFilter(0., 10., 1. / model.constants.MODEL_FREQ)
   frame_id = 0

openpilot/sunnypilot/models/fetcher.py

@@ -13,6 +13,7 @@ from openpilot.common.params import Params
 from openpilot.common.swaglog import cloudlog
 from openpilot.common.hardware.hw import Paths
 from openpilot.sunnypilot.models.helpers import is_bundle_version_compatible
+from openpilot.selfdrive.modeld.helpers import usbgpu_present
 
 from openpilot.cereal import custom
 
@@ -103,11 +104,11 @@ class ModelParser:
 class ModelCache:
   """Handles caching of model data to avoid frequent remote fetches"""
 
-  def __init__(self, params: Params, cache_timeout: int = int(3600 * 1e9)):
+  def __init__(self, params: Params, cache_timeout: int = int(3600 * 1e9), suffix: str = ""):
     self.params = params
     self.cache_timeout = cache_timeout
-    self._LAST_SYNC_KEY = "ModelManager_LastSyncTime"
+    self._LAST_SYNC_KEY = f"ModelManager_LastSyncTime{suffix}"
-    self._CACHE_KEY = "ModelManager_ModelsCache"
+    self._CACHE_KEY = f"ModelManager_ModelsCache{suffix}"
 
   def _is_expired(self) -> bool:
     """Checks if the cache has expired"""
@@ -139,24 +140,28 @@ class ModelCache:
 
 class ModelFetcher:
   """Handles fetching and caching of model data from remote source"""
-  MODEL_URL = "https://raw.githubusercontent.com/sunnypilot/sunnypilot-models/refs/heads/gh-pages/docs/driving_models_v18.json"
 
   def __init__(self, params: Params):
     self.params = params
-    self.model_cache = ModelCache(params)
     self.model_parser = ModelParser()
+    if usbgpu_present():
+      self.model_cache = ModelCache(params, suffix="_USBGPU")
+      self.model_url = "https://raw.githubusercontent.com/sunnypilot/sunnypilot-models/refs/heads/gh-pages/docs/driving_models_usbgpu_v18.json"
+    else:
+      self.model_cache = ModelCache(params)
+      self.model_url = "https://raw.githubusercontent.com/sunnypilot/sunnypilot-models/refs/heads/gh-pages/docs/driving_models_v18.json"
 
   def _fetch_and_cache_models(self) -> list[custom.ModelManagerSP.ModelBundle] | None:
     """Fetches fresh model data from remote and updates cache.
     Returns None on transport errors. Raises on 404 and other fatal HTTP errors.
     """
     try:
-      response = requests.get(self.MODEL_URL, timeout=10)
+      response = requests.get(self.model_url, timeout=10)
 
       # Explicitly handle 404 differently
       if response.status_code == 404:
-        cloudlog.error(f"Models URL returned 404 Not Found: {self.MODEL_URL}")
+        cloudlog.error(f"Models URL returned 404 Not Found: {self.model_url}")
-        raise HTTPError(f"404 Not Found: {self.MODEL_URL}", response=response)
+        raise HTTPError(f"404 Not Found: {self.model_url}", response=response)
 
       # Raise for any other 4xx/5xx
       response.raise_for_status()

openpilot/sunnypilot/models/tests/test_tinygrad_ref.py

@@ -1,11 +1,12 @@
 import requests
 
+from openpilot.common.params import Params
 from openpilot.sunnypilot.models.tinygrad_ref import get_tinygrad_ref
 from openpilot.sunnypilot.models.fetcher import ModelFetcher
 
-
 def fetch_tinygrad_ref():
-  response = requests.get(ModelFetcher.MODEL_URL, timeout=10)
+  fetcher = ModelFetcher(Params())
+  response = requests.get(fetcher.model_url, timeout=10)
   response.raise_for_status()
   json_data = response.json()
   return json_data.get("tinygrad_ref")