files

docs/MODEL_REBUILD.md

@@ -159,7 +159,21 @@ All four files must be updated together.
 
 ## Manifest
 
-Generate the base manifest after compilation, then namespace the release artifacts as v23:
+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:
 
 ```bash
 python3 scripts/model_rebuild_pipeline.py manifest \
@@ -175,9 +189,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 v23 only;
-the manifest is fetched from `Models/model_names_v23.json`, while v22 remains
-available for devices that have not updated yet.
+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.
 
 After importing newly compiled sources, normalize the release namespace before
 copying files into either resource repository:
@@ -205,4 +219,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 v23 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 v25 artifact, StarPilot switches to that built-in model.

scripts/model_compiler.py

@@ -56,14 +56,13 @@ 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)

scripts/model_release.py

@@ -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 = "v24"
+MANIFEST_VERSION = "v25"
 DEFAULT_BEHAVIOR_VERSION = "v16"
 DEVICE_ROOT = "/data/openpilot"
 REPOSITORY_FILE_LIMIT = 100_000_000

scripts/tests/test_model_release.py

@@ -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_v24.json"
+  manifest = tmp_path / "model_names_v25.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},
-    "v24",
+    "v25",
   )
   payload = json.loads(path.read_text())
   assert len(payload["models"]) == 2

selfdrive/controls/tests/test_conditional_experimental_mode.py

@@ -1168,7 +1168,6 @@ def test_starpilot_planner_updates_cem_with_current_frame_state(monkeypatch):
 
   try:
     monkeypatch.setattr(starpilot_planner_module, "calculate_road_curvature", lambda model, v_ego: (0.01, 1.0))
-    monkeypatch.setattr(starpilot_planner_module, "extract_curve_profile", lambda model: ([], []))
     monkeypatch.setattr(planner.starpilot_acceleration, "update", lambda *args, **kwargs: None)
     monkeypatch.setattr(planner.starpilot_events, "update", lambda *args, **kwargs: None)
     monkeypatch.setattr(planner.starpilot_vcruise, "update", lambda *args, **kwargs: 0.0)

selfdrive/controls/tests/test_curve_speed_controller.py

@@ -1,506 +0,0 @@
-import numpy as np
-import pytest
-
-from types import SimpleNamespace
-
-from openpilot.common.realtime import DT_MDL
-from openpilot.starpilot.common.starpilot_variables import DEFAULT_LATERAL_ACCELERATION
-from openpilot.starpilot.controls.lib.curve_speed_controller import (
-  CSC_APPROACH_DECEL,
-  CSC_COMFORT_MARGIN,
-  CSC_COUNT_CAP,
-  CSC_EGO_HEADROOM,
-  CSC_FARFIELD_GAIN,
-  CSC_LAT_ACCEL_MAX,
-  CSC_MIN_SPEED,
-  MAX_CURVATURE,
-  PRIOR_CURVATURE_BP,
-  PRIOR_LAT_ACCEL_V,
-  CSC_NUDGE,
-  CSC_NUDGE_WEIGHT,
-  CSC_OVERRIDE_WATCH_TIME,
-  CSC_TARGET_UP_RATE,
-  CSC_TRAINING_SETTLE_TIME,
-  CurveSpeedController,
-  weighted_isotonic,
-)
-
-
-class FakeParams:
-  def __init__(self, values=None):
-    self.values = dict(values or {})
-
-  def get(self, *args, **kwargs):
-    key = args[0] if args else None
-    return self.values.get(key)
-
-  def put_nonblocking(self, key, value):
-    self.values[key] = value
-
-
-def make_controller(curve_profile=None, curvature_data=None, weather_id=0, reduce_lat=0.0, road_curvature=0.02, driving_in_curve=False):
-  if curve_profile is None:
-    curve_profile = (np.zeros(33), np.linspace(0.0, 300.0, 33))
-
-  planner = SimpleNamespace(
-    params=FakeParams({"CurvatureData": curvature_data} if curvature_data is not None else None),
-    curve_profile=curve_profile,
-    starpilot_weather=SimpleNamespace(weather_id=weather_id, reduce_lateral_acceleration=reduce_lat),
-    road_curvature=road_curvature,
-    driving_in_curve=driving_in_curve,
-    tracking_lead=False,
-    lateral_acceleration=0.0,
-  )
-  controller = CurveSpeedController(SimpleNamespace(starpilot_planner=planner))
-  return planner, controller
-
-
-def make_sm(*, gas=False, brake=False, long_active=True, blinker=False, accel_pressed=False):
-  return {
-    "carControl": SimpleNamespace(longActive=long_active),
-    "carState": SimpleNamespace(gasPressed=gas, brakePressed=brake, leftBlinker=blinker, rightBlinker=False),
-    "starpilotCarState": SimpleNamespace(accelPressed=accel_pressed),
-    "onroadEvents": [],
-  }
-
-
-def single_apex_profile(curvature, distance):
-  distances = np.linspace(0.0, max(distance * 1.5, 1.0), 33)
-  curvatures = np.zeros(33)
-  index = int(np.argmin(np.abs(distances - distance)))
-  distances[index] = distance
-  curvatures[index] = curvature
-  return curvatures, distances
-
-
-def converge(controller, v_ego, v_cruise, frames=600):
-  for _ in range(frames):
-    controller.update_target(v_ego, v_cruise)
-  return controller.target
-
-
-def envelope_speed(controller, curvature, distance):
-  curve_speed = max(float(np.sqrt(controller.lat_accel_for_curvature(curvature) / curvature)), CSC_MIN_SPEED)
-  return float(np.sqrt(curve_speed**2 + 2.0 * CSC_APPROACH_DECEL * distance))
-
-
-def test_straight_road_target_is_cruise_speed():
-  _, controller = make_controller()
-
-  controller.update_target(30.0, 30.0)
-
-  assert controller.target == pytest.approx(30.0)
-
-
-def test_distant_apex_does_not_constrain_until_braking_is_due():
-  # derived from the shipped decel so retuning it doesn't silently invalidate the case
-  _, probe = make_controller()
-  curve_speed = max(float(np.sqrt(probe.lat_accel_for_curvature(0.02) / 0.02)), CSC_MIN_SPEED)
-  beyond_braking = 1.3 * (30.0**2 - curve_speed**2) / (2 * CSC_APPROACH_DECEL)
-
-  _, controller = make_controller(curve_profile=single_apex_profile(0.02, beyond_braking))
-
-  target = converge(controller, 30.0, 30.0)
-
-  assert target == pytest.approx(30.0)
-
-
-def test_apex_in_braking_range_constrains_to_kinematic_envelope():
-  _, controller = make_controller(curve_profile=single_apex_profile(0.02, 150.0))
-
-  target = converge(controller, 30.0, 30.0)
-
-  assert target == pytest.approx(envelope_speed(controller, 0.02, 150.0), abs=0.1)
-  assert target < 30.0
-
-
-def test_exit_recovery_rises_immediately_without_freeze():
-  planner, controller = make_controller(curve_profile=single_apex_profile(0.03, 20.0))
-  low_target = converge(controller, 15.0, 30.0)
-  assert low_target < 20.0
-
-  planner.curve_profile = (np.zeros(33), np.linspace(0.0, 300.0, 33))
-  controller.update_target(15.0, 30.0)
-  assert controller.target > low_target          # rises on the very next frame, no freeze
-  assert controller.target - low_target == pytest.approx(CSC_TARGET_UP_RATE * DT_MDL)
-
-  # and it clears the car by the headroom within the time the up-rate needs
-  frames = int((15.0 + CSC_EGO_HEADROOM - controller.target) / (CSC_TARGET_UP_RATE * DT_MDL)) + 1
-  for _ in range(frames):
-    controller.update_target(15.0, 30.0)
-  assert controller.target >= 15.0 + CSC_EGO_HEADROOM
-
-  recovered = converge(controller, 15.0, 30.0)
-  assert recovered == pytest.approx(30.0)
-
-
-def test_upward_jitter_in_the_envelope_is_rate_limited():
-  # a sweeper the envelope only grazes: raw_target flicks between a mild cap and the
-  # set speed. The target must not chase the jumps, or the glow strobes.
-  planner, controller = make_controller(curve_profile=single_apex_profile(0.002, 40.0))
-  steady = converge(controller, 30.0, 32.0)
-  assert steady < 32.0
-
-  flat = (np.zeros(33), np.linspace(0.0, 300.0, 33))
-  grazing = planner.curve_profile
-  peak = steady
-  for i in range(40):
-    planner.curve_profile = flat if i % 2 else grazing
-    controller.update_target(30.0, 32.0)
-    assert controller.target - peak <= CSC_TARGET_UP_RATE * DT_MDL + 1e-6
-    peak = controller.target
-
-
-def test_firm_distant_curvature_is_corrected_for_the_model_under_read():
-  # the model reads ~0.81x actual at range, so a firm distant bend binds later than it should
-  distance = 90.0
-  _, plain = make_controller(curve_profile=single_apex_profile(0.0045, distance))
-  _, probe = make_controller()
-  corrected = probe._correct_far_field(*single_apex_profile(0.0045, distance))
-
-  assert corrected.max() == pytest.approx(0.0045 * CSC_FARFIELD_GAIN)
-  assert converge(plain, 30.0, 30.0) < envelope_speed(plain, 0.0045, distance) + 1e-6
-
-
-def test_weak_or_near_readings_are_left_alone():
-  _, probe = make_controller()
-
-  # too weak to carry usable magnitude at range
-  weak = probe._correct_far_field(*single_apex_profile(0.002, 90.0))
-  assert weak.max() == pytest.approx(0.002)
-
-  # firm, but close enough that the model is already accurate
-  near = probe._correct_far_field(*single_apex_profile(0.0045, 10.0))
-  assert near.max() == pytest.approx(0.0045)
-
-
-def test_far_field_correction_brings_the_slowdown_forward():
-  profile = single_apex_profile(0.0045, 120.0)
-  _, controller = make_controller(curve_profile=profile)
-
-  corrected = converge(controller, 30.0, 30.0)
-  raw_curvatures, distances = profile
-  uncorrected = float(np.sqrt(
-    max(np.sqrt(controller.lat_accel_for_curvature(0.0045) / 0.0045), CSC_MIN_SPEED) ** 2
-    + 2.0 * CSC_APPROACH_DECEL * 120.0))
-
-  assert corrected < uncorrected      # binds sooner than the model's own reading would
-
-
-def test_fresh_activation_seeds_at_envelope_not_cruise():
-  _, controller = make_controller(curve_profile=(np.full(33, 0.05), np.linspace(0.0, 60.0, 33)))
-
-  controller.update_target(6.0, 30.0)
-
-  assert controller.target < 15.0
-
-
-def test_target_never_trails_accelerating_car_when_unconstrained():
-  planner, controller = make_controller(curve_profile=single_apex_profile(0.03, 20.0))
-  converge(controller, 15.0, 30.0)
-
-  planner.curve_profile = (np.zeros(33), np.linspace(0.0, 300.0, 33))
-  v_ego = 15.0
-  caught_up = None
-  for frame in range(200):
-    v_ego = min(v_ego + 2.0 * DT_MDL, 30.0)
-    controller.update_target(v_ego, 30.0)
-    # the target climbs faster than the car can, so once it is ahead it stays ahead
-    if controller.target >= v_ego:
-      caught_up = caught_up if caught_up is not None else frame
-    assert caught_up is None or controller.target >= min(30.0, v_ego) - 1e-6
-
-  assert caught_up is not None and caught_up * DT_MDL < 2.0
-  assert controller.target == pytest.approx(30.0)
-
-
-def test_target_does_not_ratchet_down_with_ego_speed():
-  _, controller = make_controller(curve_profile=single_apex_profile(0.02, 150.0))
-  target = converge(controller, 30.0, 30.0)
-  assert target > CSC_MIN_SPEED  # a real curve speed, not floored
-
-  controller.update_target(14.0, 30.0)
-
-  assert controller.target == pytest.approx(target, abs=0.2)
-
-
-def test_sharp_curve_target_floors_at_min_speed():
-  _, controller = make_controller(curve_profile=(np.full(33, 0.1), np.linspace(0.0, 100.0, 33)))
-
-  target = converge(controller, 15.0, 30.0)
-
-  assert target == pytest.approx(CSC_MIN_SPEED, abs=0.05)
-
-
-def test_weather_reduces_curve_speed():
-  _, dry = make_controller(curve_profile=single_apex_profile(0.01, 0.0))
-  _, wet = make_controller(curve_profile=single_apex_profile(0.01, 0.0), weather_id=1, reduce_lat=0.2)
-
-  dry_target = converge(dry, 20.0, 30.0)
-  wet_target = converge(wet, 20.0, 30.0)
-
-  assert wet_target < dry_target
-  assert wet_target == pytest.approx(dry_target * np.sqrt(0.8), abs=0.1)
-
-
-def test_prior_gives_higher_lat_accel_for_sharper_curves():
-  _, controller = make_controller()
-
-  assert controller.learned_lat_accel(0.001) == pytest.approx(1.5, abs=0.05)
-  assert controller.learned_lat_accel(MAX_CURVATURE) > controller.learned_lat_accel(0.001)
-  assert controller.learned_lat_accel(MAX_CURVATURE) == pytest.approx(
-    float(np.interp(MAX_CURVATURE, PRIOR_CURVATURE_BP, PRIOR_LAT_ACCEL_V)), abs=0.05)
-  assert controller.lateral_acceleration == pytest.approx(DEFAULT_LATERAL_ACCELERATION)
-
-
-def test_comfort_margin_matches_the_learned_habit():
-  # margin is fixed at 1.0 -- CSC targets exactly the driver's own learned comfort
-  _, controller = make_controller()
-
-  assert CSC_COMFORT_MARGIN == pytest.approx(1.0)
-  assert controller.lat_accel_for_curvature(0.01) == pytest.approx(controller.learned_lat_accel(0.01))
-
-
-def test_binding_distance_reports_the_constraining_point():
-  _, controller = make_controller(curve_profile=single_apex_profile(0.02, 150.0))
-  converge(controller, 30.0, 30.0)
-
-  assert controller.binding_distance == pytest.approx(150.0, abs=1.0)
-
-
-def test_binding_distance_is_zero_when_unconstrained():
-  _, controller = make_controller()
-  converge(controller, 30.0, 30.0)
-
-  assert controller.binding_distance == 0.0
-
-
-def test_heavily_sampled_bucket_dominates_prior():
-  _, controller = make_controller(curvature_data={"0.05": {"average": 3.0, "count": 100000}})
-
-  assert controller.learned_lat_accel(0.05) == pytest.approx(3.0, abs=0.05)
-  assert controller.learned_lat_accel(0.08) >= controller.learned_lat_accel(0.05)
-
-
-def test_learned_curve_stays_monotonic_despite_low_outlier_bucket():
-  _, controller = make_controller(curvature_data={"0.05": {"average": 0.5, "count": 100000}})
-
-  assert controller.learned_lat_accel(0.05) >= controller.learned_lat_accel(0.03)
-
-
-def test_dense_bucket_is_not_overridden_by_sparse_neighbour():
-  # real device data: a running maximum ratcheted the 80-sample bucket up to the 20-sample neighbour
-  _, dense_low = make_controller(curvature_data={
-    "0.003": {"average": 1.95, "count": 20},
-    "0.005": {"average": 1.38, "count": 80},
-  })
-  _, dense_high = make_controller(curvature_data={
-    "0.003": {"average": 1.95, "count": 80},
-    "0.005": {"average": 1.38, "count": 20},
-  })
-
-  assert dense_low.learned_lat_accel(0.005) < 1.95      # not ratcheted to the sparse neighbour
-  assert dense_low.learned_lat_accel(0.005) >= dense_low.learned_lat_accel(0.003)
-  # whichever side is better sampled should pull the fit: swapping the counts must raise it
-  assert dense_high.learned_lat_accel(0.005) > dense_low.learned_lat_accel(0.005)
-
-
-def test_weighted_isotonic_pools_violators_by_weight():
-  fitted = weighted_isotonic(np.array([1.0, 3.0, 1.2]), np.array([1.0, 1.0, 1000.0]))
-
-  assert np.all(np.diff(fitted) >= -1e-9)
-  assert fitted[-1] == pytest.approx(1.2, abs=0.02)
-
-
-def test_weighted_isotonic_leaves_sorted_input_untouched():
-  values = np.array([1.0, 1.5, 2.0, 2.5])
-  fitted = weighted_isotonic(values, np.ones(4))
-
-  assert fitted == pytest.approx(values)
-
-
-def test_legacy_off_grid_curvature_data_merges_into_buckets():
-  _, controller = make_controller(curvature_data={
-    "0.0203": {"average": 2.5, "count": 10},
-    "0.02": {"average": 2.0, "count": 10},
-  })
-
-  assert controller.curvature_data["0.02"]["count"] == 20
-  assert controller.curvature_data["0.02"]["average"] == pytest.approx(2.25)
-
-
-def test_training_update_step_is_capped_by_ema_count():
-  planner, controller = make_controller(curvature_data={"0.02": {"average": 2.0, "count": 10000}}, driving_in_curve=True)
-  planner.lateral_acceleration = 3.0
-  controller.training_timer = CSC_TRAINING_SETTLE_TIME
-
-  controller.log_data(10.0, make_sm(long_active=False))
-
-  data = controller.curvature_data["0.02"]
-  assert data["count"] == 10001
-  assert data["average"] == pytest.approx((2.0 * CSC_COUNT_CAP + 3.0) / (CSC_COUNT_CAP + 1))
-
-
-def test_no_passive_training_right_after_csc_limited_speed():
-  planner, controller = make_controller(curve_profile=single_apex_profile(0.03, 20.0), driving_in_curve=True)
-  planner.lateral_acceleration = 3.0
-  converge(controller, 15.0, 30.0)
-  assert controller.training_quiet_timer > 0.0
-
-  controller.training_timer = CSC_TRAINING_SETTLE_TIME
-  controller.log_data(10.0, make_sm(long_active=False))
-  assert "0.02" not in controller.curvature_data
-  assert not controller.enable_training
-
-  controller.training_quiet_timer = 0.0
-  controller.training_timer = CSC_TRAINING_SETTLE_TIME
-  controller.log_data(10.0, make_sm(long_active=False))
-  assert controller.curvature_data["0.02"]["count"] == 1
-
-
-def test_training_settles_within_a_couple_of_seconds():
-  # a real drive rarely holds every eligibility condition for a whole model horizon,
-  # so the settle time has to be short enough that ordinary curves still teach it
-  planner, controller = make_controller(driving_in_curve=True)
-  planner.lateral_acceleration = 2.4
-  sm = make_sm(long_active=False)
-
-  for _ in range(int(CSC_TRAINING_SETTLE_TIME / DT_MDL) - 2):
-    controller.log_data(10.0, sm)
-  assert "0.02" not in controller.curvature_data
-
-  for _ in range(3):
-    controller.log_data(10.0, sm)
-  assert controller.curvature_data["0.02"]["count"] >= 1
-
-
-def test_brief_ineligibility_does_not_restart_the_settle_timer():
-  planner, controller = make_controller(driving_in_curve=True)
-  planner.lateral_acceleration = 2.4
-  sm = make_sm(long_active=False)
-  for _ in range(int(CSC_TRAINING_SETTLE_TIME / DT_MDL) + 1):
-    controller.log_data(10.0, sm)
-  trained = controller.curvature_data["0.02"]["count"]
-
-  # a lead flickers into the tracker for two frames, then leaves
-  planner.tracking_lead = True
-  controller.log_data(10.0, sm)
-  controller.log_data(10.0, sm)
-  planner.tracking_lead = False
-
-  controller.log_data(10.0, sm)
-  assert controller.curvature_data["0.02"]["count"] == trained + 1
-
-
-def test_sustained_ineligibility_still_drains_the_settle_timer():
-  planner, controller = make_controller(driving_in_curve=True)
-  planner.lateral_acceleration = 2.4
-  engaged = make_sm(long_active=True)
-  manual = make_sm(long_active=False)
-  for _ in range(int(CSC_TRAINING_SETTLE_TIME / DT_MDL) + 1):
-    controller.log_data(10.0, manual)
-
-  for _ in range(int(2 * CSC_TRAINING_SETTLE_TIME / DT_MDL)):
-    controller.log_data(10.0, engaged)
-  assert controller.training_timer == pytest.approx(0.0)
-
-  controller.log_data(10.0, manual)
-  assert not controller.enable_training
-
-
-def settle_override(controller, sm=None, frames=None):
-  """Run the post-override watch out so the pseudo-sample is committed."""
-  sm = sm if sm is not None else make_sm()
-  for _ in range(frames if frames is not None else int(CSC_OVERRIDE_WATCH_TIME / DT_MDL) + 1):
-    controller.handle_override(20.0, False, sm)
-
-
-def test_gas_override_nudges_bucket_up_once_per_episode():
-  _, controller = make_controller()
-  prior = controller.learned_lat_accel(0.02)
-  controller.target = 10.0
-
-  controller.handle_override(20.0, True, make_sm(gas=True))
-  controller.handle_override(20.0, True, make_sm(gas=True))
-  assert "0.02" not in controller.curvature_data  # still watching what the driver holds
-
-  settle_override(controller)
-  assert controller.curvature_data["0.02"]["count"] == CSC_NUDGE_WEIGHT
-  assert controller.curvature_data["0.02"]["average"] > prior
-
-  controller.handle_override(20.0, False, make_sm())
-  controller.target = 10.0
-  controller.handle_override(20.0, True, make_sm(gas=True))
-  settle_override(controller)
-  assert controller.curvature_data["0.02"]["count"] == 2 * CSC_NUDGE_WEIGHT
-
-
-def test_override_learns_the_cornering_the_driver_actually_held():
-  # the whole point: a fixed step needs several rejections to close a real disagreement,
-  # so record what they demonstrated instead
-  planner, observed = make_controller(driving_in_curve=True)
-  observed.target = 10.0
-  observed.handle_override(20.0, True, make_sm(gas=True))
-  planner.lateral_acceleration = 2.9          # they hold the curve much harder than CSC wanted
-  settle_override(observed, make_sm(gas=True))
-
-  _, stepped = make_controller(driving_in_curve=True)
-  stepped._apply_nudge(CSC_NUDGE)             # what the old fixed-step path would have recorded
-
-  assert observed.curvature_data["0.02"]["average"] == pytest.approx(2.9)
-  assert observed.curvature_data["0.02"]["average"] > stepped.curvature_data["0.02"]["average"]
-  assert observed.learned_lat_accel(0.02) > stepped.learned_lat_accel(0.02)
-
-
-def test_override_on_a_straight_still_registers_the_fixed_step():
-  planner, controller = make_controller()
-  prior = controller.learned_lat_accel(0.02)
-  controller.target = 10.0
-
-  controller.handle_override(20.0, True, make_sm(gas=True))
-  planner.lateral_acceleration = 0.0          # never reached a corner
-  settle_override(controller)
-
-  assert controller.curvature_data["0.02"]["average"] == pytest.approx(prior + CSC_NUDGE)
-
-
-def test_res_button_nudges_bucket_up_even_at_target_speed():
-  _, controller = make_controller()
-  prior = controller.learned_lat_accel(0.02)
-  controller.target = 20.0  # car tracking the target, so the gas-press condition would not fire
-
-  controller.handle_override(20.0, True, make_sm(), accel_button=True)
-  settle_override(controller)
-
-  assert controller.curvature_data["0.02"]["count"] == CSC_NUDGE_WEIGHT
-  assert controller.curvature_data["0.02"]["average"] > prior
-
-
-def test_brake_override_nudges_bucket_down():
-  _, controller = make_controller(driving_in_curve=True)
-  prior = controller.learned_lat_accel(0.02)
-
-  controller.handle_override(20.0, True, make_sm(brake=True))
-
-  assert controller.curvature_data["0.02"]["count"] == CSC_NUDGE_WEIGHT
-  assert controller.curvature_data["0.02"]["average"] < prior
-
-
-def test_calibrated_lateral_acceleration_param_is_written_on_flush():
-  planner, controller = make_controller(curvature_data={"0.02": {"average": 2.8, "count": 5000}})
-
-  assert "CalibratedLateralAcceleration" not in planner.params.values
-  controller.flush_data()
-
-  assert planner.params.values["CalibratedLateralAcceleration"] > DEFAULT_LATERAL_ACCELERATION
-  assert controller.lateral_acceleration == planner.params.values["CalibratedLateralAcceleration"]
-
-
-def test_stale_param_from_a_previous_build_is_republished_without_training():
-  # a stale value must not survive a restart just because this drive never trained
-  planner, controller = make_controller(curvature_data={"0.02": {"average": 2.8, "count": 5000}})
-  planner.params.values["CalibratedLateralAcceleration"] = 3.71
-
-  controller.log_data(0.0, make_sm())  # standstill: ineligible -> flush path
-
-  assert planner.params.values["CalibratedLateralAcceleration"] <= CSC_LAT_ACCEL_MAX

selfdrive/controls/tests/test_starpilot_vcruise.py

@@ -4,8 +4,7 @@ import pytest
 
 from openpilot.common.constants import CV
 from openpilot.common.realtime import DT_MDL
-from openpilot.starpilot.common.starpilot_variables import PLANNER_TIME
-from openpilot.starpilot.controls.lib.curve_speed_controller import CSC_GLOW_HOLD_TIME, CSC_GLOW_ON_DELTA
+from openpilot.starpilot.controls.lib.curve_speed_controller import CSC_MAX_DECEL_RATE, CurveSpeedController, MIN_TRAINING_TIME
 from openpilot.starpilot.controls.lib.starpilot_vcruise import (
   FORCE_STOP_CAP_SLACK_M,
   FORCE_STOP_TURN_VETO_STOP_SEEN_HOLD_TIME,
@@ -54,7 +53,6 @@ def make_vcruise(*, red_light=False, raw_model_stopped=False, forcing_stop=False
     raw_model_stopped=raw_model_stopped,
     road_curvature=road_curvature,
     road_curvature_detected=False,
-    lateral_acceleration=0.0,
   )
   vcruise = StarPilotVCruise(planner)
   vcruise.forcing_stop = forcing_stop
@@ -84,12 +82,12 @@ def make_sm(*, standstill=True, min_steer_speed=0.0, car_fingerprint=""):
   }
 
 
-def update_vcruise(vcruise, sm, toggles, *, now, v_ego=0.0, v_cruise=20.0, controls_enabled=True):
+def update_vcruise(vcruise, sm, toggles, *, now, v_ego=0.0, controls_enabled=True):
   return vcruise.update(
     controls_enabled=controls_enabled,
     now=now,
     time_validated=True,
-    v_cruise=v_cruise,
+    v_cruise=20.0,
     v_ego=v_ego,
     sm=sm,
     starpilot_toggles=toggles,
@@ -147,56 +145,30 @@ def test_santa_fe_force_stop_tune_only_applies_to_that_car():
   assert get_force_stop_low_speed_hold(other) is None
 
 
-def test_curve_speed_controller_blinker_releases_the_cap_but_keeps_the_plan():
+def test_curve_speed_controller_holds_target_through_brief_detector_dropout():
   planner, vcruise = make_vcruise()
   sm = make_sm(standstill=False)
   toggles = make_toggles()
   toggles.curve_speed_controller = True
 
-  calls = []
-
-  def set_curve_target(_v_ego, _v_cruise):
-    calls.append(_v_ego)
+  def set_curve_target(_v_ego):
+    vcruise.csc.target_set = True
     vcruise.csc.target = 14.0
 
   vcruise.csc.update_target = set_curve_target
+  planner.road_curvature_detected = True
   result = update_vcruise(vcruise, sm, toggles, now=10.0, v_ego=20.0)
   assert result == pytest.approx(14.0)
   assert vcruise.csc_controlling_speed
 
-  # the cap lifts so CSC can't fight the lane change, but the envelope keeps planning
-  # so the curve doesn't have to be re-discovered from the set speed afterwards
-  sm["carState"].leftBlinker = True
+  planner.road_curvature_detected = False
   result = update_vcruise(vcruise, sm, toggles, now=10.25, v_ego=20.0)
-  assert result == pytest.approx(20.0)
-  assert not vcruise.csc_controlling_speed
-  assert len(calls) == 2  # still planning, so nothing has to be rediscovered
-
-  # blinker off: the plan is already current, so the cap comes straight back
-  sm["carState"].leftBlinker = False
-  result = update_vcruise(vcruise, sm, toggles, now=10.5, v_ego=20.0)
   assert result == pytest.approx(14.0)
   assert vcruise.csc_controlling_speed
 
-
-def test_curve_speed_controller_reseeds_after_a_real_dropout():
-  planner, vcruise = make_vcruise()
-  sm = make_sm(standstill=False)
-  toggles = make_toggles()
-  toggles.curve_speed_controller = True
-
-  def set_curve_target(_v_ego, _v_cruise):
-    vcruise.csc.target = 14.0
-
-  vcruise.csc.update_target = set_curve_target
-  update_vcruise(vcruise, sm, toggles, now=11.0, v_ego=20.0)
-  assert vcruise.csc_controlling_speed
-
-  # disengaging is a real dropout, not a momentary veto -- that still resets
-  sm["carControl"].longActive = False
-  update_vcruise(vcruise, sm, toggles, now=11.05, v_ego=20.0)
+  result = update_vcruise(vcruise, sm, toggles, now=10.8, v_ego=20.0)
+  assert result == pytest.approx(20.0)
   assert not vcruise.csc_controlling_speed
-  assert vcruise.csc.seed_pending
 
 
 def test_curve_speed_controller_releases_immediately_when_disabled():
@@ -205,19 +177,53 @@ def test_curve_speed_controller_releases_immediately_when_disabled():
   toggles = make_toggles()
   toggles.curve_speed_controller = True
 
-  def set_curve_target(_v_ego, _v_cruise):
+  def set_curve_target(_v_ego):
+    vcruise.csc.target_set = True
     vcruise.csc.target = 14.0
 
   vcruise.csc.update_target = set_curve_target
+  planner.road_curvature_detected = True
   update_vcruise(vcruise, sm, toggles, now=20.0, v_ego=20.0)
   assert vcruise.csc_controlling_speed
 
+  planner.road_curvature_detected = False
   toggles.curve_speed_controller = False
   result = update_vcruise(vcruise, sm, toggles, now=20.1, v_ego=20.0)
   assert result == pytest.approx(20.0)
   assert not vcruise.csc_controlling_speed
 
 
+def test_curve_speed_controller_does_not_compete_with_force_stop():
+  planner, vcruise = make_vcruise(red_light=True, road_curvature=0.001)
+  sm = make_sm(standstill=False)
+  toggles = make_toggles()
+  toggles.curve_speed_controller = True
+  planner.road_curvature_detected = True
+  vcruise.csc.target_set = True
+  vcruise.csc.target = 12.0
+
+  update_vcruise(vcruise, sm, toggles, now=25.0, v_ego=20.0)
+
+  assert not vcruise.csc_controlling_speed
+  assert not vcruise.csc.target_set
+
+
+def test_curve_speed_controller_learns_through_a_signaled_curve():
+  planner, vcruise = make_vcruise(road_curvature=0.02)
+  sm = make_sm(standstill=False)
+  sm["carControl"].longActive = False
+  sm["carState"].leftBlinker = True
+  planner.driving_in_curve = True
+  planner.lateral_acceleration = 2.4
+  vcruise.csc.training_timer = MIN_TRAINING_TIME
+
+  vcruise.csc.log_data(20.0, sm)
+
+  assert vcruise.csc.enable_training
+  assert vcruise.csc.curvature_data["0.02"]["count"] == 1
+
+
+
 def test_curve_speed_controller_can_be_limited_to_driving_without_a_lead():
   planner, vcruise = make_vcruise()
   sm = make_sm(standstill=False)
@@ -225,10 +231,12 @@ def test_curve_speed_controller_can_be_limited_to_driving_without_a_lead():
   toggles.curve_speed_controller = True
   toggles.csc_no_lead = True
 
-  def set_curve_target(_v_ego, _v_cruise):
+  def set_curve_target(_v_ego):
+    vcruise.csc.target_set = True
     vcruise.csc.target = 14.0
 
   vcruise.csc.update_target = set_curve_target
+  planner.road_curvature_detected = True
 
   result = update_vcruise(vcruise, sm, toggles, now=30.0, v_ego=20.0)
   assert result == pytest.approx(14.0)
@@ -246,8 +254,10 @@ def test_curve_speed_controller_stays_enabled_with_a_lead_by_default():
   toggles = make_toggles()
   toggles.curve_speed_controller = True
   planner.starpilot_following.following_lead = True
+  planner.road_curvature_detected = True
 
-  def set_curve_target(_v_ego, _v_cruise):
+  def set_curve_target(_v_ego):
+    vcruise.csc.target_set = True
     vcruise.csc.target = 14.0
 
   vcruise.csc.update_target = set_curve_target
@@ -271,7 +281,7 @@ def test_curve_speed_controller_learns_when_speed_is_manually_controlled(long_ac
   planner.driving_in_curve = True
   planner.road_curvature_detected = True
   planner.lateral_acceleration = 2.4
-  vcruise.csc.training_timer = PLANNER_TIME
+  vcruise.csc.training_timer = MIN_TRAINING_TIME
 
   update_vcruise(vcruise, sm, toggles, now=50.0, v_ego=20.0)
 
@@ -289,7 +299,7 @@ def test_curve_speed_controller_learns_when_longitudinal_override_event_is_activ
   planner.driving_in_curve = True
   planner.road_curvature_detected = True
   planner.lateral_acceleration = 2.4
-  vcruise.csc.training_timer = PLANNER_TIME
+  vcruise.csc.training_timer = MIN_TRAINING_TIME
 
   update_vcruise(vcruise, sm, toggles, now=50.0, v_ego=20.0)
 
@@ -303,7 +313,7 @@ def test_curve_speed_controller_persists_data_after_leaving_curve():
   sm["carControl"].longActive = False
   planner.driving_in_curve = True
   planner.lateral_acceleration = 2.4
-  vcruise.csc.training_timer = PLANNER_TIME
+  vcruise.csc.training_timer = MIN_TRAINING_TIME
 
   vcruise.csc.log_data(20.0, sm)
   assert not any(key == "CurvatureData" for key, _ in planner.params.writes)
@@ -314,276 +324,54 @@ def test_curve_speed_controller_persists_data_after_leaving_curve():
   assert any(key == "CurvatureData" for key, _ in planner.params.writes)
 
 
-def test_csc_res_press_cancels_for_episode_and_rearms():
-  planner, vcruise = make_vcruise()
+def test_curve_speed_controller_publishes_live_values_to_memory_params():
+  planner, vcruise = make_vcruise(road_curvature=0.02)
   sm = make_sm(standstill=False)
-  toggles = make_toggles()
-  toggles.curve_speed_controller = True
-
-  curve_target = {"v": 14.0}
-
-  def set_curve_target(_v_ego, _v_cruise):
-    vcruise.csc.target = curve_target["v"]
-
-  vcruise.csc.update_target = set_curve_target
-  result = update_vcruise(vcruise, sm, toggles, now=60.0, v_ego=20.0)
-  assert result == pytest.approx(14.0)
-  assert vcruise.csc_controlling_speed
-
-  sm["starpilotCarState"].accelPressed = True
-  result = update_vcruise(vcruise, sm, toggles, now=60.05, v_ego=20.0)
-  assert result == pytest.approx(20.0)
-  assert not vcruise.csc_controlling_speed
-  assert vcruise.csc_override
-
-  # latches for the rest of the episode, not just while pressed
-  sm["starpilotCarState"].accelPressed = False
-  result = update_vcruise(vcruise, sm, toggles, now=60.1, v_ego=20.0)
-  assert result == pytest.approx(20.0)
-  assert vcruise.csc_override
-
-  # curve ends -> re-arms
-  curve_target["v"] = 20.0
-  update_vcruise(vcruise, sm, toggles, now=60.15, v_ego=20.0)
-  assert not vcruise.csc_override
-
-  curve_target["v"] = 14.0
-  result = update_vcruise(vcruise, sm, toggles, now=60.2, v_ego=20.0)
-  assert result == pytest.approx(14.0)
-  assert vcruise.csc_controlling_speed
-
-
-def test_csc_res_press_does_not_latch_when_csc_was_not_active():
-  planner, vcruise = make_vcruise()
-  sm = make_sm(standstill=False)
-  toggles = make_toggles()
-  toggles.curve_speed_controller = True
-
-  def set_curve_target(_v_ego, _v_cruise):
-    vcruise.csc.target = 14.0
-
-  vcruise.csc.update_target = set_curve_target
-  # press before CSC ever limited: suspends it while held, but must not latch a cancel
-  sm["starpilotCarState"].accelPressed = True
-  result = update_vcruise(vcruise, sm, toggles, now=70.0, v_ego=20.0)
-  assert result == pytest.approx(20.0)
-  assert not vcruise.csc_override
-
-  sm["starpilotCarState"].accelPressed = False
-  result = update_vcruise(vcruise, sm, toggles, now=70.05, v_ego=20.0)
-  assert result == pytest.approx(14.0)
-  assert vcruise.csc_controlling_speed
-
-
-def test_csc_res_press_defers_to_slc_confirmation():
-  planner, vcruise = make_vcruise()
-  sm = make_sm(standstill=False)
-  toggles = make_toggles()
-  toggles.curve_speed_controller = True
-
-  def set_curve_target(_v_ego, _v_cruise):
-    vcruise.csc.target = 14.0
-
-  vcruise.csc.update_target = set_curve_target
-  update_vcruise(vcruise, sm, toggles, now=80.0, v_ego=20.0)
-  assert vcruise.csc_controlling_speed
-
-  # confirming a speed limit must not also cancel the curve slowdown
-  vcruise.slc.speed_limit_changed_timer = 1.0
-  vcruise.slc.unconfirmed_speed_limit = 25.0
-  sm["starpilotCarState"].accelPressed = True
-  update_vcruise(vcruise, sm, toggles, now=80.05, v_ego=20.0)
-  assert not vcruise.csc_override
-
-  sm["starpilotCarState"].accelPressed = False
-  result = update_vcruise(vcruise, sm, toggles, now=80.1, v_ego=20.0)
-  assert result == pytest.approx(14.0)
-  assert vcruise.csc_controlling_speed
-
-
-def test_curve_speed_controller_glow_ignores_a_trivial_graze():
-  planner, vcruise = make_vcruise()
-  sm = make_sm(standstill=False)
-  toggles = make_toggles()
-  toggles.curve_speed_controller = True
-
-  # a long gentle bend where the envelope only shaves a little: the target hovers either
-  # side of the threshold for the whole curve, so a low bar strobes the glow
-  def set_curve_target(_v_ego, _v_cruise):
-    vcruise.csc.target = 20.0 - (CSC_GLOW_ON_DELTA / 2.0)
-
-  vcruise.csc.update_target = set_curve_target
-  result = update_vcruise(vcruise, sm, toggles, now=160.0, v_ego=20.0)
-
-  assert result < 20.0                      # the cap is still applied
-  assert not vcruise.csc_controlling_speed  # it just isn't worth announcing
-
-
-def test_curve_speed_controller_glow_holds_through_a_brief_release():
-  planner, vcruise = make_vcruise()
-  sm = make_sm(standstill=False)
-  toggles = make_toggles()
-  toggles.curve_speed_controller = True
-
-  curve_target = {"v": 14.0}
-
-  def set_curve_target(_v_ego, _v_cruise):
-    vcruise.csc.target = curve_target["v"]
-
-  vcruise.csc.update_target = set_curve_target
-  update_vcruise(vcruise, sm, toggles, now=130.0, v_ego=20.0)
-  assert vcruise.csc_controlling_speed
-
-  # one curve routinely lets go and re-engages; the glow must ride through it
-  curve_target["v"] = 20.0
-  now = 130.0
-  for _ in range(int((CSC_GLOW_HOLD_TIME - 0.2) / DT_MDL)):
-    now += DT_MDL
-    update_vcruise(vcruise, sm, toggles, now=now, v_ego=20.0)
-    assert vcruise.csc_controlling_speed
-
-  curve_target["v"] = 14.0
-  now += DT_MDL
-  update_vcruise(vcruise, sm, toggles, now=now, v_ego=20.0)
-  assert vcruise.csc_controlling_speed
-
-
-def test_curve_speed_controller_glow_clears_once_the_release_sticks():
-  planner, vcruise = make_vcruise()
-  sm = make_sm(standstill=False)
-  toggles = make_toggles()
-  toggles.curve_speed_controller = True
-
-  curve_target = {"v": 14.0}
-
-  def set_curve_target(_v_ego, _v_cruise):
-    vcruise.csc.target = curve_target["v"]
-
-  vcruise.csc.update_target = set_curve_target
-  update_vcruise(vcruise, sm, toggles, now=140.0, v_ego=20.0)
-  assert vcruise.csc_controlling_speed
-
-  curve_target["v"] = 20.0
-  now = 140.0
-  for _ in range(int(CSC_GLOW_HOLD_TIME / DT_MDL) + 1):
-    now += DT_MDL
-    update_vcruise(vcruise, sm, toggles, now=now, v_ego=20.0)
-  assert not vcruise.csc_controlling_speed
-
-
-def test_curve_speed_controller_keeps_the_cap_when_signalling_mid_curve():
-  planner, vcruise = make_vcruise()
-  sm = make_sm(standstill=False)
-  toggles = make_toggles()
-  toggles.curve_speed_controller = True
-
-  def set_curve_target(_v_ego, _v_cruise):
-    vcruise.csc.target = 14.0
-
-  vcruise.csc.update_target = set_curve_target
-  result = update_vcruise(vcruise, sm, toggles, now=150.0, v_ego=20.0)
-  assert result == pytest.approx(14.0)
-
-  # a lane change taken inside a curve must not hand the speed back
+  sm["carControl"].longActive = False
   planner.driving_in_curve = True
-  sm["carState"].leftBlinker = True
-  result = update_vcruise(vcruise, sm, toggles, now=150.05, v_ego=20.0)
-  assert result == pytest.approx(14.0)
-  assert vcruise.csc_controlling_speed
+  planner.lateral_acceleration = 2.4
+  vcruise.csc.training_timer = MIN_TRAINING_TIME
 
-  # on a straight it still yields, so CSC can't fight the manoeuvre
-  planner.driving_in_curve = False
-  result = update_vcruise(vcruise, sm, toggles, now=150.1, v_ego=20.0)
-  assert result == pytest.approx(20.0)
-  assert not vcruise.csc_controlling_speed
+  vcruise.csc.log_data(20.0, sm)
+
+  assert any(key == "CalibratedLateralAcceleration" for key, _ in planner.params_memory.writes)
+  assert any(key == "CalibrationProgress" for key, _ in planner.params_memory.writes)
+  assert planner.params_memory.values["CalibrationProgress"] > 0.0
 
 
-def test_curve_speed_controller_glow_lights_when_the_car_arrives_at_the_cap_from_below():
-  planner, vcruise = make_vcruise()
-  sm = make_sm(standstill=False)
-  toggles = make_toggles()
-  toggles.curve_speed_controller = True
+def test_curve_speed_controller_ramps_toward_curve_speed_at_bounded_rate():
+  planner = SimpleNamespace(
+    params=FakeParams(),
+    road_curvature=0.004,
+    time_to_curve=2.0,
+    starpilot_weather=SimpleNamespace(weather_id=0, reduce_lateral_acceleration=0.0),
+  )
+  controller = CurveSpeedController(SimpleNamespace(starpilot_planner=planner))
+  controller.lateral_acceleration = 2.0
+  controller.target_set = True
+  controller.target = 30.0
 
-  # accelerating out of a slow zone into a curve: the target is never under v_ego, but it
-  # is still the only thing stopping the car from reaching the set speed
-  def set_curve_target(_v_ego, _v_cruise):
-    vcruise.csc.target = 22.0
+  controller.update_target(30.0)
 
-  vcruise.csc.update_target = set_curve_target
-
-  update_vcruise(vcruise, sm, toggles, now=120.0, v_ego=15.0, v_cruise=32.0)
-  assert not vcruise.csc_controlling_speed          # still climbing, CSC isn't holding it yet
-
-  result = update_vcruise(vcruise, sm, toggles, now=120.05, v_ego=22.0, v_cruise=32.0)
-  assert result == pytest.approx(22.0)
-  assert vcruise.csc_controlling_speed              # arrived at the cap, and it binds
+  assert controller.target == pytest.approx(30.0 - CSC_MAX_DECEL_RATE * DT_MDL)
+  assert controller.target > (controller.lateral_acceleration / planner.road_curvature) ** 0.5
 
 
-def test_curve_speed_controller_glow_stays_off_while_the_target_is_above_v_ego():
-  planner, vcruise = make_vcruise()
-  sm = make_sm(standstill=False)
-  toggles = make_toggles()
-  toggles.curve_speed_controller = True
+def test_curve_speed_controller_does_not_slow_for_curve_speed_above_ego():
+  planner = SimpleNamespace(
+    params=FakeParams(),
+    road_curvature=0.001,
+    time_to_curve=2.0,
+    starpilot_weather=SimpleNamespace(weather_id=0, reduce_lateral_acceleration=0.0),
+  )
+  controller = CurveSpeedController(SimpleNamespace(starpilot_planner=planner))
+  controller.lateral_acceleration = 2.0
+  controller.target_set = True
+  controller.target = 28.0
 
-  # a highway sweeper trims the target well under the set speed but never under v_ego,
-  # so the car keeps accelerating and the driver feels nothing
-  def set_curve_target(_v_ego, _v_cruise):
-    vcruise.csc.target = 26.0
-
-  vcruise.csc.update_target = set_curve_target
-  result = update_vcruise(vcruise, sm, toggles, now=90.0, v_ego=20.0, v_cruise=30.0)
-
-  assert result == pytest.approx(26.0)
-  assert not vcruise.csc_controlling_speed
-
-
-def test_curve_speed_controller_glow_holds_through_the_recovery_ramp():
-  planner, vcruise = make_vcruise()
-  sm = make_sm(standstill=False)
-  toggles = make_toggles()
-  toggles.curve_speed_controller = True
-
-  curve_target = {"v": 14.0}
-
-  def set_curve_target(_v_ego, _v_cruise):
-    vcruise.csc.target = curve_target["v"]
-
-  vcruise.csc.update_target = set_curve_target
-  update_vcruise(vcruise, sm, toggles, now=100.0, v_ego=20.0)
-  assert vcruise.csc_controlling_speed
-
-  # past the apex the target climbs back above v_ego while the car is still cornering
-  curve_target["v"] = 18.0
-  update_vcruise(vcruise, sm, toggles, now=100.05, v_ego=15.0)
-  assert vcruise.csc_controlling_speed
-
-  # fully released, but the glow only clears once the release has stuck
-  curve_target["v"] = 20.0
-  now = 100.1
-  update_vcruise(vcruise, sm, toggles, now=now, v_ego=17.0)
-  assert vcruise.csc_controlling_speed
-
-  for _ in range(int(CSC_GLOW_HOLD_TIME / DT_MDL) + 1):
-    now += DT_MDL
-    update_vcruise(vcruise, sm, toggles, now=now, v_ego=17.0)
-  assert not vcruise.csc_controlling_speed
-
-
-def test_curve_speed_controller_hysteresis_keeps_glow_off_for_marginal_targets():
-  planner, vcruise = make_vcruise()
-  sm = make_sm(standstill=False)
-  toggles = make_toggles()
-  toggles.curve_speed_controller = True
-
-  def set_curve_target(_v_ego, _v_cruise):
-    vcruise.csc.target = 19.7
-
-  vcruise.csc.update_target = set_curve_target
-  result = update_vcruise(vcruise, sm, toggles, now=50.0, v_ego=20.0)
-
-  assert result == pytest.approx(19.7)
-  assert not vcruise.csc_controlling_speed
+  controller.update_target(30.0)
 
+  assert controller.target == pytest.approx(30.0)
 
 
 def test_active_slc_control_target_applies_offset_and_cluster_diff():

selfdrive/modeld/modeld.py

@@ -7,6 +7,10 @@ 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

selfdrive/modeld/models/driving_tinygrad.pkl.sha256

@@ -1 +1 @@
-a77db33c2e2d6a7570dc2a4a70c2b877429ee8bd9ca5dfeda74b5a41231aaff9  driving_tinygrad.pkl
+31902b114b7fb8455af694d83333a86a44112b83be662e064ffbd67e8daafe72  driving_tinygrad.pkl

selfdrive/ui/onroad/starpilot/starpilot_border.py

@@ -58,8 +58,6 @@ def _csc_state():
 
   plan = sm["starpilotPlan"]
   params = ui_state.ui_params
-  # A pending speed limit flashes the speed limit sign, not the border -- it has no reason
-  # to blank this, and doing so hid real curve slowdowns for the whole confirmation window.
   if not params.get_bool("ShowCSCStatus"):
     return None
 

starpilot/assets/model_manager.py

@@ -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 = ("v24",)
+MANIFEST_CANDIDATES = ("v25",)
 MODEL_NAMESPACE_SUFFIX = "3"
 DEFAULT_MODEL_KEY = "rdf43"
 LOCAL_MODEL_PREFIX = "local-"

starpilot/assets/tests/test_model_pipeline.py

@@ -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_v24_is_the_only_manifest_candidate():
-  assert MANIFEST_CANDIDATES == ("v24",)
+def test_v25_is_the_only_manifest_candidate():
+  assert MANIFEST_CANDIDATES == ("v25",)
 
 
-def test_v24_manifest_is_loaded_from_models_checkout():
-  assert ModelManager._manifest_paths("v24") == ("Models/model_names_v24.json",)
+def test_v25_manifest_is_loaded_from_models_checkout():
+  assert ModelManager._manifest_paths("v25") == ("Models/model_names_v25.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("v24") == (
-    "model_names_v24.json",
-    "manifests/model_names_v24.json",
+  assert ModelManager._hf_manifest_paths("v25") == (
+    "model_names_v25.json",
+    "manifests/model_names_v25.json",
   )
 
 

starpilot/common/starpilot_utilities.py

@@ -138,23 +138,6 @@ def calculate_road_curvature(modelData, v_ego):
   return float(predicted_lateral_acc / max(v_ego, 1)**2), max(time_to_curve, 1)
 
 
-PROFILE_MIN_SPEED = 3.0     # m/s — model points planned near standstill have unusable curvature
-PROFILE_MAX_CURVATURE = 0.1
-
-
-def extract_curve_profile(modelData):
-  orientation_rate = np.abs(np.array(modelData.orientationRate.z))
-  velocity = np.array(modelData.velocity.x)
-  distances = np.array(modelData.position.x)
-
-  # k = psi_dot / v per point, against the model's own planned speed so its
-  # slowdowns don't inflate the curvature
-  curvatures = orientation_rate / np.clip(velocity, PROFILE_MIN_SPEED, None)
-  curvatures = np.where(velocity < PROFILE_MIN_SPEED, 0.0, np.minimum(curvatures, PROFILE_MAX_CURVATURE))
-
-  return curvatures, distances
-
-
 def clean_model_name(name):
   return name.replace("(Default)", "").strip()
 

starpilot/controls/lib/curve_speed_controller.py

@@ -2,97 +2,19 @@
 import numpy as np
 
 from openpilot.common.constants import CV
-from openpilot.common.filter_simple import FirstOrderFilter
 from openpilot.common.realtime import DT_MDL
 
-from openpilot.starpilot.common.starpilot_variables import (
-  CITY_SPEED_LIMIT,
-  CRUISING_SPEED,
-  DEFAULT_LATERAL_ACCELERATION,
-  PLANNER_TIME,
-)
+from openpilot.starpilot.common.starpilot_variables import CITY_SPEED_LIMIT, CRUISING_SPEED, DEFAULT_LATERAL_ACCELERATION, PLANNER_TIME
 
 CALIBRATION_PROGRESS_THRESHOLD = 10 / DT_MDL
+MIN_TRAINING_TIME = 5.0
 CSC_MIN_SPEED = CITY_SPEED_LIMIT * CV.MPH_TO_MS
-
-# braking distance is (v^2 - v_curve^2) / (2 * this), so lower starts the slowdown
-# sooner and spreads it further.
-CSC_APPROACH_DECEL = 0.3
-CSC_TARGET_UP_RATE = 3.0
-CSC_TARGET_DOWN_RATE = 2.5
-CSC_TARGET_FILTER_RC = 0.4
-CSC_EGO_HEADROOM = 2.0            # target never trails below v_ego, so CSC can't drag re-acceleration
-CSC_RELEASE_DEBOUNCE = 0.25       # s the envelope must stay clear before that floor applies
-CSC_ACTIVE_ON_DELTA = 0.5
-CSC_ACTIVE_OFF_DELTA = 0.25
-CSC_GLOW_ON_DELTA = 1.0           # ~2.2 mph; separate from CSC_ACTIVE_ON_DELTA (training) so a trivial graze doesn't light the glow
-CSC_GLOW_HOLD_TIME = 3.0          # s the cap must stay released before the glow clears, so it doesn't flicker on/off across one curve
-
-CSC_COUNT_CAP = 600               # EMA floor: samples beyond this stop shrinking the update step
-CSC_PRIOR_COUNT = 100             # bucket count at which learned data and the prior have equal weight
-CSC_LAT_ACCEL_MIN = 1.2
-CSC_LAT_ACCEL_MAX = 3.2
-CSC_NUDGE = 0.15
-CSC_NUDGE_WEIGHT = 20             # counts a single override pseudo-sample is worth
-CSC_OVERRIDE_WATCH_TIME = 6.0     # s to keep watching what the driver holds after they reject a cut
-CSC_TRAINING_QUIET_TIME = 5.0     # blocks passive samples after CSC limited speed, so it can't learn its own cap
-CSC_TRAINING_SETTLE_TIME = 2.0    # driver-owned seconds before a sample counts, so it isn't openpilot's leftover speed
-CSC_COMFORT_MARGIN = 1.0          # 1.0 = matches the driver's own learned cornering, no extra cushion
-
-# The model under-reads curvature at range: measured 0.81x actual beyond ~75 m. That holds
-# only where the reading is already firm -- weak distant readings carry no usable magnitude
-# (0.40x median with a 14:1 spread), so scaling those would amplify noise, not signal.
-CSC_FARFIELD_MIN_CURVATURE = 0.004   # ~R 250 m; at this strength range readings were 85%+ reliable
-CSC_FARFIELD_MIN_DISTANCE = 30.0     # inside this the model is already accurate
-CSC_FARFIELD_GAIN = 1.23             # 1 / 0.81
-
-# Buckets are spaced geometrically, not linearly: comfort is a speed and v = sqrt(a/k), so equal
-# steps in k give wildly uneven speed resolution. Regridding is safe -- _normalize_curvature_data
-# re-buckets stored keys on load.
-MIN_CURVATURE = 0.0005            # R 2000 m — gentler than this never constrains anything
-MAX_CURVATURE = 0.02              # R 50 m — already well below the CSC_MIN_SPEED floor
-CURVATURE_BUCKETS = 24            # keeps every bucket under ~7 mph wide without over-thinning the data
-ROUNDING_PRECISION = 6
-CURVATURE_GRID = MIN_CURVATURE * np.power(MAX_CURVATURE / MIN_CURVATURE,
-                                          np.arange(CURVATURE_BUCKETS) / (CURVATURE_BUCKETS - 1))
-LOG_CURVATURE_GRID = np.log(CURVATURE_GRID)
-
-# Drivers accept more lateral acceleration in sharp slow corners than in highway sweepers.
-PRIOR_CURVATURE_BP = [0.001, 0.003, 0.01, 0.03, 0.1]
-PRIOR_LAT_ACCEL_V = [1.5, 1.8, 2.2, 2.6, 2.9]
-
-
-def weighted_isotonic(values, weights):
-  """Weighted non-decreasing fit (pool adjacent violators).
-
-  Keeps comfort from falling as curves tighten, without letting a sparse bucket
-  overrule a well-sampled neighbour the way a running maximum would.
-  """
-  block_values: list[float] = []
-  block_weights: list[float] = []
-  block_sizes: list[int] = []
-
-  for value, weight in zip(values, weights, strict=True):
-    block_values.append(float(value))
-    block_weights.append(float(weight))
-    block_sizes.append(1)
-
-    while len(block_values) > 1 and block_values[-2] > block_values[-1]:
-      merged_weight = block_weights[-2] + block_weights[-1]
-      merged_value = ((block_values[-2] * block_weights[-2]) + (block_values[-1] * block_weights[-1])) / merged_weight
-      block_values.pop()
-      block_weights.pop()
-      merged_size = block_sizes.pop()
-      block_values[-1] = merged_value
-      block_weights[-1] = merged_weight
-      block_sizes[-1] += merged_size
-
-  fitted = np.empty(len(values))
-  index = 0
-  for value, size in zip(block_values, block_sizes, strict=True):
-    fitted[index:index + size] = value
-    index += size
-  return fitted
+CSC_MAX_DECEL_RATE = 1.5
+MAX_CURVATURE = 0.1
+MIN_CURVATURE = 0.001
+PERCENTILE = 90
+ROUNDING_PRECISION = 5
+STEP = 0.001
 
 
 def is_user_overriding_longitudinal(sm):
@@ -121,45 +43,26 @@ class CurveSpeedController:
     self.starpilot_planner = StarPilotVCruise.starpilot_planner
 
     self.enable_training = False
-    self.nudge_applied = False
-
-    self.override_watch_key = None
-    self.override_watch_peak = 0.0
-    self.override_watch_timer = 0.0
+    self.target_set = False
 
     self.training_timer = 0.0
     self.persistence_timer = 0.0
-    self.training_quiet_timer = 0.0
     self.data_dirty = False
 
-    self.target = 0.0
-    self.binding_distance = 0.0
-    self.release_timer = 0.0
-    self.target_filter = FirstOrderFilter(0.0, CSC_TARGET_FILTER_RC, DT_MDL, initialized=False)
-    self.seed_pending = True
-
-    self._long_active_prev = False
-
     curvature_data = self.starpilot_planner.params.get("CurvatureData")
     self.curvature_data = self._normalize_curvature_data(curvature_data)
 
-    # built through the bucketer so the keys are byte-identical to what training writes
-    self.required_curvatures = [self._bucket_curvature(curvature) for curvature in CURVATURE_GRID]
+    self.required_curvatures = [str(round(road_curvature, ROUNDING_PRECISION)) for road_curvature in np.arange(MIN_CURVATURE, MAX_CURVATURE + STEP, STEP)]
 
-    self.rebuild_lat_accel_curve()
-    # publish on the first flush even if this drive never trains, or the readout
-    # keeps showing whatever a previous build left behind
-    self.data_dirty = True
-    # the Settings screen reads the memory param for a live readout -- seed it now, or it
-    # shows nothing until the first disk flush
-    self._publish_live_values()
+    self.update_lateral_acceleration()
+    self._publish_calibration_progress(persist=True)
 
   @staticmethod
   def _bucket_curvature(road_curvature):
-    clipped_curvature = float(np.clip(abs(road_curvature), MIN_CURVATURE, MAX_CURVATURE))
-    # nearest in log space, so a bucket is a constant speed step rather than a constant radius one
-    bucket_index = int(np.argmin(np.abs(LOG_CURVATURE_GRID - np.log(clipped_curvature))))
-    return str(round(float(CURVATURE_GRID[bucket_index]), ROUNDING_PRECISION))
+    clipped_curvature = float(np.clip(road_curvature, MIN_CURVATURE, MAX_CURVATURE))
+    bucket_index = round((clipped_curvature - MIN_CURVATURE) / STEP)
+    bucketed_curvature = MIN_CURVATURE + (bucket_index * STEP)
+    return str(round(bucketed_curvature, ROUNDING_PRECISION))
 
   @classmethod
   def _normalize_curvature_data(cls, curvature_data):
@@ -197,6 +100,17 @@ class CurveSpeedController:
 
     return normalized
 
+  def _persist_data(self):
+    if not self.data_dirty:
+      return
+
+    progress = self._calibration_progress()
+    self.starpilot_planner.params.put_nonblocking("CalibrationProgress", progress)
+    self.starpilot_planner.params.put_nonblocking("CurvatureData", self.curvature_data)
+    self._put_memory_param("CalibrationProgress", progress)
+    self.data_dirty = False
+    self.persistence_timer = 0.0
+
   def _calibration_progress(self):
     progress = 0.0
     for key in self.required_curvatures:
@@ -204,48 +118,31 @@ class CurveSpeedController:
         progress += min(self.curvature_data[key]["count"] / CALIBRATION_PROGRESS_THRESHOLD, 1.0)
     return (progress / len(self.required_curvatures)) * 100
 
-  def _publish_live_values(self, progress=None):
-    # memory-only and cheap, so this can run every frame training touches the data --
-    # it's what the on-device Settings screen reads for a live readout between disk flushes
-    params_memory = getattr(self.starpilot_planner, "params_memory", None)
-    if params_memory is None:
-      return
-    if progress is None:
-      progress = self._calibration_progress()
-    params_memory.put_nonblocking("CalibratedLateralAcceleration", self.lateral_acceleration)
-    params_memory.put_nonblocking("CalibrationProgress", progress)
-
-  def _persist_data(self):
-    if not self.data_dirty:
-      return
-
+  def _publish_calibration_progress(self, persist=False):
     progress = self._calibration_progress()
-    self.starpilot_planner.params.put_nonblocking("CalibratedLateralAcceleration", self.lateral_acceleration)
-    self.starpilot_planner.params.put_nonblocking("CalibrationProgress", progress)
-    self.starpilot_planner.params.put_nonblocking("CurvatureData", self.curvature_data)
-    self._publish_live_values(progress)
-    self.data_dirty = False
-    self.persistence_timer = 0.0
+    if persist:
+      self.starpilot_planner.params.put_nonblocking("CalibrationProgress", progress)
+    self._put_memory_param("CalibrationProgress", progress)
+
+  def _put_memory_param(self, key, value):
+    params_memory = getattr(self.starpilot_planner, "params_memory", None)
+    if params_memory is not None:
+      params_memory.put_nonblocking(key, value)
 
   def flush_data(self):
     self._persist_data()
 
   def log_data(self, v_ego, sm):
-    self.training_quiet_timer = max(self.training_quiet_timer - DT_MDL, 0.0)
-
     eligible = (
       v_ego > CRUISING_SPEED and
       not self.starpilot_planner.tracking_lead and
-      is_manual_speed_control(sm) and
-      self.training_quiet_timer <= 0.0
+      is_manual_speed_control(sm)
     )
     self.enable_training = False
 
     if not eligible:
       self.flush_data()
-      # decay instead of resetting: a lead flickering in and out of the tracker used to
-      # cost the full re-arm, which left almost nothing to learn from on a real drive
-      self.training_timer = max(self.training_timer - DT_MDL, 0.0)
+      self.training_timer = 0.0
       self.persistence_timer = 0.0
       return
 
@@ -254,9 +151,8 @@ class CurveSpeedController:
       self.persistence_timer += DT_MDL
 
     in_curve = (
-      self.training_timer >= CSC_TRAINING_SETTLE_TIME and
-      self.starpilot_planner.driving_in_curve and
-      not (sm["carState"].leftBlinker or sm["carState"].rightBlinker)
+      self.training_timer >= MIN_TRAINING_TIME and
+      self.starpilot_planner.driving_in_curve
     )
     if in_curve:
       lateral_acceleration = abs(self.starpilot_planner.lateral_acceleration)
@@ -264,11 +160,11 @@ class CurveSpeedController:
 
       if road_curvature in self.curvature_data:
         data = self.curvature_data[road_curvature]
-        # capped so an established bucket still tracks a change in driving style
-        effective_count = min(data["count"], CSC_COUNT_CAP)
+        average = data["average"]
+        count = data["count"]
         self.curvature_data[road_curvature] = {
-          "average": ((data["average"] * effective_count) + lateral_acceleration) / (effective_count + 1),
-          "count": data["count"] + 1
+          "average": ((average * count) + lateral_acceleration) / (count + 1),
+          "count": count + 1
         }
       else:
         self.curvature_data[road_curvature] = {
@@ -277,8 +173,8 @@ class CurveSpeedController:
         }
 
       self.data_dirty = True
-      self.rebuild_lat_accel_curve()
-      self._publish_live_values()
+      self.update_lateral_acceleration()
+      self._publish_calibration_progress()
       self.enable_training = True
 
       if self.persistence_timer >= PLANNER_TIME:
@@ -286,166 +182,30 @@ class CurveSpeedController:
     elif self.data_dirty:
       self.flush_data()
 
-  def handle_override(self, v_ego, was_controlling, sm, accel_button=False):
-    long_active = bool(sm["carControl"].longActive)
-    long_dropped = self._long_active_prev and not long_active
-    self._long_active_prev = long_active
-
-    self._update_override_watch(sm)
-
-    if not was_controlling:
-      self.nudge_applied = False
-      return
-
-    if self.nudge_applied:
-      return
-
-    if accel_button or (sm["carState"].gasPressed and self.target < v_ego - 0.5):
-      # Watch what the driver actually holds instead of stepping by a fixed amount -- CSC is
-      # suspended while overridden, so their cornering now measures their real comfort.
-      self.override_watch_key = self._bucket_curvature(abs(self.starpilot_planner.road_curvature))
-      self.override_watch_peak = abs(self.starpilot_planner.lateral_acceleration)
-      self.override_watch_timer = CSC_OVERRIDE_WATCH_TIME
-      self.nudge_applied = True
-    elif (getattr(sm["carState"], "brakePressed", False) or long_dropped) and self.starpilot_planner.driving_in_curve:
-      self._apply_nudge(-CSC_NUDGE)
-
-  def _update_override_watch(self, sm):
-    if self.override_watch_key is None:
-      return
-
-    lateral_acceleration = abs(self.starpilot_planner.lateral_acceleration)
-    if lateral_acceleration > self.override_watch_peak:
-      # credit the bucket the peak actually happened in, not the one at the button press
-      self.override_watch_peak = lateral_acceleration
-      self.override_watch_key = self._bucket_curvature(abs(self.starpilot_planner.road_curvature))
-
-    self.override_watch_timer -= DT_MDL
-    if self.override_watch_timer > 0.0 and (is_user_overriding_longitudinal(sm) or
-                                            self.starpilot_planner.driving_in_curve):
-      return
-
-    key = self.override_watch_key
-    self.override_watch_key = None
-    # floored at the old fixed step, so a rejection that never reaches a corner still counts
-    # and this path can only ever raise the bucket
-    self._record_pseudo_sample(key, max(self.override_watch_peak,
-                                        self.learned_lat_accel(float(key)) + CSC_NUDGE))
-
-  def _apply_nudge(self, offset):
-    key = self._bucket_curvature(abs(self.starpilot_planner.road_curvature))
-    # relative to the learned value, not the margined one, or repeated overrides walk the bucket down
-    self._record_pseudo_sample(key, self.learned_lat_accel(float(key)) + offset)
-    self.nudge_applied = True
-
-  def _record_pseudo_sample(self, key, sample):
-    sample = float(np.clip(sample, CSC_LAT_ACCEL_MIN, CSC_LAT_ACCEL_MAX))
-
-    data = self.curvature_data.get(key, {"average": sample, "count": 0})
-    effective_count = min(data["count"], CSC_COUNT_CAP)
-    total = effective_count + CSC_NUDGE_WEIGHT
-    self.curvature_data[key] = {
-      "average": ((data["average"] * effective_count) + (sample * CSC_NUDGE_WEIGHT)) / total,
-      "count": data["count"] + CSC_NUDGE_WEIGHT,
-    }
-
-    self.rebuild_lat_accel_curve()
-    self.data_dirty = True
-    self.flush_data()
-
-  def rebuild_lat_accel_curve(self):
-    grid_k = np.array([float(key) for key in self.required_curvatures])
-    prior = np.interp(grid_k, PRIOR_CURVATURE_BP, PRIOR_LAT_ACCEL_V)
-
-    blended = prior.copy()
-    counts = np.zeros(len(grid_k))
-    for i, key in enumerate(self.required_curvatures):
-      data = self.curvature_data.get(key)
-      if data:
-        confidence = data["count"] / (data["count"] + CSC_PRIOR_COUNT)
-        blended[i] = confidence * data["average"] + (1.0 - confidence) * prior[i]
-        counts[i] = data["count"]
-
-    blended = np.clip(blended, CSC_LAT_ACCEL_MIN, CSC_LAT_ACCEL_MAX)
-    blended = weighted_isotonic(blended, counts + CSC_PRIOR_COUNT)
-
-    self._curve_k = grid_k
-    self._curve_a = blended
-
-    if counts.sum() > 0:
-      self.lateral_acceleration = float(np.average(blended, weights=counts))
+  def update_lateral_acceleration(self):
+    if self.curvature_data:
+      all_samples = [data["average"] for data in self.curvature_data.values()]
+      self.lateral_acceleration = float(np.percentile(all_samples, PERCENTILE))
     else:
       self.lateral_acceleration = DEFAULT_LATERAL_ACCELERATION
 
-  def learned_lat_accel(self, curvature):
-    """Comfort level learned for this curvature, before any control margin."""
-    return float(np.interp(abs(curvature), self._curve_k, self._curve_a))
+    self.starpilot_planner.params.put_nonblocking("CalibratedLateralAcceleration", self.lateral_acceleration)
+    self._put_memory_param("CalibratedLateralAcceleration", self.lateral_acceleration)
 
-  def lat_accel_for_curvature(self, curvature):
-    lat_accel = np.interp(np.abs(curvature), self._curve_k, self._curve_a) * CSC_COMFORT_MARGIN
+  def update_target(self, v_ego):
+    lateral_acceleration = self.lateral_acceleration
+    if self.starpilot_planner.starpilot_weather.weather_id != 0:
+      lateral_acceleration -= self.lateral_acceleration * self.starpilot_planner.starpilot_weather.reduce_lateral_acceleration
 
-    weather = self.starpilot_planner.starpilot_weather
-    if weather.weather_id != 0:
-      lat_accel = lat_accel * (1.0 - weather.reduce_lateral_acceleration)
-
-    return lat_accel
-
-  @staticmethod
-  def _correct_far_field(curvatures, distances):
-    """Undo the model's known under-read of distant curvature, where the reading is firm."""
-    firm = (curvatures >= CSC_FARFIELD_MIN_CURVATURE) & (distances >= CSC_FARFIELD_MIN_DISTANCE)
-    return np.minimum(np.where(firm, curvatures * CSC_FARFIELD_GAIN, curvatures), MAX_CURVATURE)
-
-  def reset(self, v_cruise):
-    self.target = float(v_cruise)
-    self.release_timer = 0.0
-    self.target_filter.x = float(v_cruise)
-    self.target_filter.initialized = True
-    self.seed_pending = True
-
-  def update_target(self, v_ego, v_cruise):
-    if not self.target_filter.initialized:
-      self.reset(v_cruise)
-
-    curvatures, distances = self.starpilot_planner.curve_profile
-    if len(curvatures) == 0:
-      raw_target = float(v_cruise)
-      self.binding_distance = 0.0
+    if self.target_set:
+      csc_speed = (lateral_acceleration / abs(self.starpilot_planner.road_curvature))**0.5
+      csc_speed = max(float(csc_speed), CSC_MIN_SPEED)
+      if csc_speed >= v_ego:
+        self.target = v_ego
+      else:
+        time_to_curve = max(float(self.starpilot_planner.time_to_curve), DT_MDL)
+        decel_rate = float(np.clip((v_ego - csc_speed) / time_to_curve, 0.0, CSC_MAX_DECEL_RATE))
+        self.target = float(np.clip(self.target - decel_rate * DT_MDL, csc_speed, v_ego))
     else:
-      curvatures = self._correct_far_field(curvatures, distances)
-      lat_accel = self.lat_accel_for_curvature(curvatures)
-      point_speeds = np.sqrt(lat_accel / np.maximum(curvatures, 1e-4))
-      point_speeds = np.maximum(point_speeds, CSC_MIN_SPEED)
-      allowed_speeds = np.sqrt(point_speeds**2 + 2.0 * CSC_APPROACH_DECEL * np.maximum(distances, 0.0))
-      binding_index = int(np.argmin(allowed_speeds))
-      raw_target = min(float(allowed_speeds[binding_index]), float(v_cruise))
-      self.binding_distance = float(distances[binding_index]) if raw_target < v_cruise else 0.0
-
-    # a fresh activation starts at the envelope, or it spends seconds ramping down
-    # toward a curve it already sees (engaging or launching into a turn)
-    if self.seed_pending:
-      seed = min(float(v_cruise), max(raw_target, v_ego + CSC_EGO_HEADROOM))
-      self.target = seed
-      self.target_filter.x = seed
-      self.seed_pending = False
-
-    if raw_target >= v_ego:
-      self.release_timer += DT_MDL
-    else:
-      self.release_timer = 0.0
-
-    # The headroom aim goes through the rate limiter with everything else; applying it
-    # after the clamp let every upward jitter in raw_target reach the target unsmoothed.
-    filtered = self.target_filter.update(raw_target)
-    self.target = float(np.clip(max(filtered, min(raw_target, v_ego + CSC_EGO_HEADROOM)),
-                                self.target - CSC_TARGET_DOWN_RATE * DT_MDL,
-                                self.target + CSC_TARGET_UP_RATE * DT_MDL))
-
-    # Once the envelope really has released, the target must not sit under the car or it
-    # drags re-acceleration. Debounced, because a single jittery frame doing this yanks a
-    # legitimate cut back up to v_ego and strobes the glow on sweepers.
-    if self.release_timer >= CSC_RELEASE_DEBOUNCE:
-      self.target = max(self.target, min(raw_target, v_ego))
-
-    if self.target < v_cruise - CSC_ACTIVE_ON_DELTA:
-      self.training_quiet_timer = CSC_TRAINING_QUIET_TIME
+      self.target_set = True
+      self.target = v_ego

starpilot/controls/lib/starpilot_vcruise.py

@@ -6,13 +6,7 @@ from openpilot.common.constants import CV
 from openpilot.common.realtime import DT_MDL
 
 from openpilot.starpilot.common.starpilot_variables import CITY_SPEED_LIMIT, CRUISING_SPEED
-from openpilot.starpilot.controls.lib.curve_speed_controller import (
-  CSC_ACTIVE_OFF_DELTA,
-  CSC_GLOW_HOLD_TIME,
-  CSC_GLOW_ON_DELTA,
-  CurveSpeedController,
-  is_manual_speed_control,
-)
+from openpilot.starpilot.controls.lib.curve_speed_controller import CurveSpeedController, is_manual_speed_control
 from openpilot.starpilot.controls.lib.speed_limit_controller import SpeedLimitController
 from openpilot.selfdrive.controls.lib.longitudinal_vehicle_tunes import (
   get_force_stop_distance_bias,
@@ -22,6 +16,7 @@ from openpilot.selfdrive.controls.lib.longitudinal_vehicle_tunes import (
 )
 
 CSC_MIN_SPEED = CITY_SPEED_LIMIT * CV.MPH_TO_MS
+CSC_CURVE_RELEASE_HOLD_TIME = 0.75
 OVERRIDE_FORCE_STOP_TIMER = 10
 STANDSTILL_FORCE_STOP_CLEAR_TIME = 0.75
 # Open-loop — green is undetectable at standstill, so this only needs to cover the
@@ -207,9 +202,8 @@ class StarPilotVCruise:
     self._nav_instruction_state = {}
     self._applied_slc_control_target = 0.0
     self.csc_controlling_speed = False
-    self.csc_glow_release_timer = 0.0
-    self.csc_override = False
     self.csc_target = 0.0
+    self.csc_curve_last_seen_at = None
 
   def _update_nav_instruction_state(self):
     raw = self.starpilot_planner.params_memory.get("NavInstructionState") or {}
@@ -577,62 +571,28 @@ class StarPilotVCruise:
       starpilot_toggles.curve_speed_controller and
       (not getattr(starpilot_toggles, "csc_no_lead", False) or not following_lead)
     )
-    # The blinker veto is for lane changes/turns, not for an already-real curve -- releasing it
-    # there let the car accelerate into the bend, then claw the speed back once the blinker cleared.
-    csc_blinker_on = ((sm["carState"].leftBlinker or sm["carState"].rightBlinker) and
-                      not self.starpilot_planner.driving_in_curve)
-    csc_was_controlling = self.csc_controlling_speed
-    # a pending SLC confirmation owns the accel button
-    slc_confirmation_pending = self.slc.speed_limit_changed_timer > DT_MDL and self.slc.unconfirmed_speed_limit >= 1
-    csc_accel_button = bool(sm["starpilotCarState"].accelPressed) and not slc_confirmation_pending
+    csc_curve_detected = csc_available and self.starpilot_planner.road_curvature_detected
+    if csc_curve_detected:
+      self.csc.update_target(v_ego)
 
-    # Latched outside the availability branch: the press itself suspends CSC this frame, so
-    # latching inside it would never see the press, and the slowdown would return on release.
-    if csc_was_controlling and csc_accel_button:
-      self.csc_override = True
-    if not (long_control_active and starpilot_toggles.curve_speed_controller):
-      self.csc_override = False
-
-    if csc_available and not csc_blinker_on:
-      self.csc.update_target(v_ego, v_cruise)
-
-      if self.csc_override and self.csc.target > v_cruise - CSC_ACTIVE_OFF_DELTA:
-        self.csc_override = False
-
-      if self.csc_override:
-        self.csc_controlling_speed = False
-        self.csc_glow_release_timer = 0.0
-        self.csc_target = v_cruise
-      else:
-        self.csc_target = self.csc.target
-        # A low target alone means nothing until the car has actually reached it (slowed down
-        # to it, or accelerated up into it). Release still waits for the set speed, so the glow
-        # spans the hold and the recovery, not just the braking.
-        if self.csc_target < v_cruise - CSC_GLOW_ON_DELTA and v_ego >= self.csc_target - CSC_ACTIVE_OFF_DELTA:
-          self.csc_controlling_speed = True
-          self.csc_glow_release_timer = 0.0
-        elif self.csc_target > v_cruise - CSC_ACTIVE_OFF_DELTA:
-          # hold through a brief release: one curve routinely lets go and re-engages
-          self.csc_glow_release_timer += DT_MDL
-          if self.csc_glow_release_timer >= CSC_GLOW_HOLD_TIME:
-            self.csc_controlling_speed = False
-        else:
-          self.csc_glow_release_timer = 0.0
-    elif csc_available:
-      # Release the cap so CSC can't fight the lane change, but keep planning -- resetting here
-      # threw the braking plan away and re-planned from the set speed with the curve closer.
-      self.csc.update_target(v_ego, v_cruise)
-      self.csc_controlling_speed = False
-      self.csc_glow_release_timer = 0.0
-      self.csc_target = v_cruise
+      self.csc_controlling_speed = True
+      self.csc_target = self.csc.target
+      self.csc_curve_last_seen_at = now
     else:
-      self.csc.reset(v_cruise)
-      self.csc_controlling_speed = False
-      self.csc_glow_release_timer = 0.0
-      self.csc_target = v_cruise
+      csc_release_hold = bool(
+        csc_available and
+        self.csc_controlling_speed and
+        self.csc_curve_last_seen_at is not None and
+        self._elapsed_seconds(now, self.csc_curve_last_seen_at) < CSC_CURVE_RELEASE_HOLD_TIME
+      )
+      if not csc_release_hold:
+        self.csc.log_data(v_ego, sm)
 
-    self.csc.handle_override(v_ego, csc_was_controlling, sm, accel_button=csc_accel_button)
-    self.csc.log_data(v_ego, sm)
+        self.csc_controlling_speed = False
+        self.csc.target_set = False
+        self.csc_curve_last_seen_at = None
+
+        self.csc_target = v_cruise
 
     # Pfeiferj's Speed Limit Controller
     self.slc.starpilot_toggles = starpilot_toggles

starpilot/controls/starpilot_planner.py

@@ -20,7 +20,7 @@ from openpilot.selfdrive.controls.lib.lead_behavior import (
 from openpilot.selfdrive.controls.lib.longitudinal_mpc_lib.long_mpc import A_CHANGE_COST, DANGER_ZONE_COST, J_EGO_COST, STOP_DISTANCE
 from openpilot.selfdrive.controls.lib.longitudinal_vehicle_tunes import get_lead_follow_jerk_scale
 
-from openpilot.starpilot.common.starpilot_utilities import calculate_lane_width, calculate_road_curvature, extract_curve_profile
+from openpilot.starpilot.common.starpilot_utilities import calculate_lane_width, calculate_road_curvature
 from openpilot.starpilot.common.starpilot_variables import CRUISING_SPEED, MINIMUM_LATERAL_ACCELERATION, PLANNER_TIME, THRESHOLD
 from openpilot.starpilot.controls.lib.conditional_chill_mode import ConditionalChillMode
 from openpilot.starpilot.controls.lib.conditional_experimental_mode import ConditionalExperimentalMode
@@ -214,7 +214,6 @@ class StarPilotPlanner:
     self.model_stopped = self.raw_model_stopped or self.starpilot_vcruise.forcing_stop
 
     self.road_curvature, self.time_to_curve = calculate_road_curvature(sm["modelV2"], v_ego)
-    self.curve_profile = extract_curve_profile(sm["modelV2"])
 
     self.road_curvature_detected = (1 / abs(self.road_curvature))**0.5 < v_ego > CRUISING_SPEED and not (sm["carState"].leftBlinker or sm["carState"].rightBlinker)
 
@@ -329,9 +328,6 @@ class StarPilotPlanner:
     starpilotPlan.cscControllingSpeed = self.starpilot_vcruise.csc_controlling_speed
     starpilotPlan.cscSpeed = float(self.starpilot_vcruise.csc_target)
     starpilotPlan.cscTraining = self.starpilot_vcruise.csc.enable_training
-    starpilotPlan.cscOverridden = self.starpilot_vcruise.csc_override
-    starpilotPlan.cscLearnedLatAccel = float(self.starpilot_vcruise.csc.learned_lat_accel(self.road_curvature))
-    starpilotPlan.cscBindingDistance = float(self.starpilot_vcruise.csc.binding_distance)
 
     starpilotPlan.desiredFollowDistance = int(self.starpilot_following.desired_follow_distance)
     starpilotPlan.disableThrottle = (

starpilot/system/the_galaxy/assets/components/tools/device_settings.css

@@ -384,7 +384,6 @@
   padding: 0.2rem 0.6rem;
 }
 
-/* read-only: no border, since there is nothing here to click or edit */
 .ds-row-readout {
   background-color: transparent;
   border: none;

starpilot/system/the_galaxy/assets/components/tools/device_settings.js

@@ -484,11 +484,11 @@ function formatSliderValue(val, stepStr, precisionInt, key) {
 
 function formatReadoutValue(p) {
   const raw = state.values[p.key]
-  const v = parseFloat(raw)
-  if (raw === undefined || raw === null || Number.isNaN(v)) return "--"
+  const value = parseFloat(raw)
+  if (raw === undefined || raw === null || Number.isNaN(value)) return "--"
 
   const precision = p.precision !== undefined && p.precision !== null ? Number(p.precision) : 2
-  const formatted = Number(v.toFixed(Math.max(0, precision))).toString()
+  const formatted = Number(value.toFixed(Math.max(0, precision))).toString()
   return p.unit ? `${formatted}${p.unit}` : formatted
 }
 

starpilot/system/the_galaxy/the_galaxy.py

@@ -5995,16 +5995,6 @@ def setup(app):
     result["VehicleParked"] = _get_vehicle_parked()
     result["AlphaLongitudinalAvailable"] = _get_alpha_longitudinal_available()
     result["HasRivianAngleHarness"] = _get_has_rivian_angle_harness()
-    # read-only: excluded from allowed_keys (and so from the write paths) but still
-    # worth surfacing as a display-only readout
-    try:
-      result["CalibratedLateralAcceleration"] = _get_current_param_value("CalibratedLateralAcceleration", float, defaults_lookup)
-    except Exception:
-      result["CalibratedLateralAcceleration"] = None
-    try:
-      result["CalibrationProgress"] = _get_current_param_value("CalibrationProgress", float, defaults_lookup)
-    except Exception:
-      result["CalibrationProgress"] = None
 
     for key in ("CalibratedLateralAcceleration", "CalibrationProgress"):
       try:

system/hardware/fan_controller.py

@@ -7,7 +7,9 @@ from openpilot.common.swaglog import cloudlog
 from openpilot.common.pid import PIDController
 from openpilot.system.hardware import HARDWARE
 
-OFFSET = 0
+# raise fan setpoint on tici/tizi to reduce noise
+# after raising LMH threshold in AGNOS 18.1 to prevent CPU throttling
+OFFSET = 0 if HARDWARE.get_device_type() == "mici" else 5
 
 class BaseFanController(ABC):
   @abstractmethod

tinygrad_repo/.github/actions/setup-tinygrad/action.yml

@@ -4,7 +4,7 @@ inputs:
   python-version:
     description: 'Python version to use'
     required: false
-    default: '' # if you don't set a version, the native python version will be used
+    default: '3.14'
   key:
     description: 'Key for the python cache'
     required: false
@@ -41,12 +41,12 @@ inputs:
     description: "Install LLVM?"
     required: false
     default: 'false'
-  tinydreno:
-    description: "Install tinydreno"
+  qemu:
+    description: "Install qemu?"
     required: false
     default: 'false'
-  qemu:
-    description: "Install qemu"
+  ninja:
+    description: "Install ninja?"
     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 .venv
+        uv venv --allow-existing --python ${{ inputs.python-version }} "$VIRTUAL_ENV"
         DEPS="${{ inputs.deps }}"
-        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/
+        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/
     - name: Install dependencies in venv (without extra)
       if: inputs.deps == ''
       shell: bash
       run: |
-        uv venv .venv
-        uv pip install --python .venv -e . ${{ inputs.pydeps }}
+        uv venv --allow-existing --python ${{ inputs.python-version }} "$VIRTUAL_ENV"
+        uv pip install --python "$VIRTUAL_ENV" -e . ${{ inputs.pydeps }}
     - name: Prune uv cache
       if: github.event_name != 'pull_request'
       shell: bash
@@ -125,16 +125,15 @@ runs:
     - name: Configure venv
       shell: bash
       run: |
-        echo "VIRTUAL_ENV=${{ github.workspace }}/.venv" >> "$GITHUB_ENV"
         if [[ "$RUNNER_OS" == "Windows" ]]; then
-          echo "${{ github.workspace }}/.venv/Scripts" >> "$GITHUB_PATH"
+          echo "$VIRTUAL_ENV/Scripts" >> "$GITHUB_PATH"
         else
-          echo "${{ github.workspace }}/.venv/bin" >> "$GITHUB_PATH"
+          echo "$VIRTUAL_ENV/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')
+      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true')
       shell: bash
       run: |
         sudo mkdir -p /var/cache/apt/archives
@@ -162,7 +161,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')
+      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true')
       id: apt-pkgs
       shell: bash
       run: |
@@ -187,25 +186,37 @@ 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') && github.event_name == 'pull_request'
+      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'
       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') && github.event_name != 'pull_request'
+      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'
       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')
+      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'
       shell: bash
       run: |
         sudo apt -qq update || true
@@ -277,12 +288,6 @@ 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'

tinygrad_repo/.github/workflows/autogen.yml

@@ -35,15 +35,15 @@ jobs:
         key: 'autogen'
         amd: 'true'
         llvm: 'true'
-        pydeps: 'pyyaml mako'
+        deps: 'autogen'
     - 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 "comgr.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 "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_3, hsa, hip, amd_gpu, sqtt, rocprof, amdgpu_kd, amdgpu_drm"
+        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.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"
+        python3 -c "from tinygrad.runtime.autogen import mlx5, bnxt"
         python3 -c "from tinygrad.runtime.autogen import ggml_common"
         REGEN=1 python3 -c "from tinygrad.runtime.autogen import libclang"
     - name: Check for differences
@@ -102,42 +102,3 @@ 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

tinygrad_repo/.github/workflows/benchmark.yml

@@ -88,13 +88,13 @@ jobs:
       fail-fast: false
       matrix:
         dev: ['METAL', 'AMD', 'NV']
-    timeout-minutes: 60
+    timeout-minutes: 30
     defaults:
       run:
         shell: bash -e -o pipefail {0}
     env:
       DEV: ${{ matrix.dev }}
-      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
+      HCQ2: '0'
     if: github.repository_owner == 'tinygrad'
     steps:
     - name: Checkout Code
@@ -102,16 +102,11 @@ jobs:
     - name: Setup (AMD)
       if: ${{ matrix.dev == 'AMD' }}
       run: |
-        ./extra/amdpci/setup_python_cap.sh
-        ./extra/hcq/hcq_smi.py amd rmmod
-        ./extra/hcq/hcq_smi.py amd kill_pids
+        ./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: 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
+      run: lsof -tQ /dev/nvidia* | { xargs -r kill -9 || true; }
     - name: setup staging db
       if: github.ref == 'refs/heads/update_benchmark_staging'
       run: |
@@ -121,18 +116,14 @@ 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.6
-      # qwen3.6:35b-a3b doesn't fit on mac
+    - name: Run qwen3.8
+      # qwen3.8:27b doesn't fit on mac
       if: ${{ matrix.dev != 'METAL' }}
-      run: BENCHMARK_LOG=qwen36_35b-a3b JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 -m tinygrad.llm -m qwen3.6:35b-a3b --benchmark --warmup
+      run: BENCHMARK_LOG=qwen38_27b JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 -m tinygrad.llm -m qwen3.8:27b --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
 
@@ -143,13 +134,13 @@ jobs:
       fail-fast: false
       matrix:
         dev: ['METAL', 'AMD', 'NV']
-    timeout-minutes: 60
+    timeout-minutes: 10
     defaults:
       run:
         shell: bash -e -o pipefail {0}
     env:
       DEV: ${{ matrix.dev }}
-      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
+      HCQ2: '0'
     if: github.repository_owner == 'tinygrad'
     steps:
     - name: Checkout Code
@@ -157,12 +148,11 @@ jobs:
     - name: Setup (AMD)
       if: ${{ matrix.dev == 'AMD' }}
       run: |
-        ./extra/amdpci/setup_python_cap.sh
-        ./extra/hcq/hcq_smi.py amd rmmod
-        ./extra/hcq/hcq_smi.py amd kill_pids
+        ./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: sudo lsof -tQ /dev/nvidia* | { xargs -r sudo kill -9 || true; }
+      run: lsof -tQ /dev/nvidia* | { xargs -r kill -9 || true; }
     - name: setup staging db
       if: github.ref == 'refs/heads/update_benchmark_staging'
       run: |
@@ -182,10 +172,6 @@ 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
 
@@ -196,13 +182,13 @@ jobs:
       fail-fast: false
       matrix:
         dev: ['AMD', 'NV']
-    timeout-minutes: 60
+    timeout-minutes: 5
     defaults:
       run:
         shell: bash -e -o pipefail {0}
     env:
       DEV: ${{ matrix.dev }}
-      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
+      HCQ2: '0'
     if: github.repository_owner == 'tinygrad'
     steps:
     - name: Checkout Code
@@ -210,12 +196,11 @@ jobs:
     - name: Setup (AMD)
       if: ${{ matrix.dev == 'AMD' }}
       run: |
-        ./extra/amdpci/setup_python_cap.sh
-        ./extra/hcq/hcq_smi.py amd rmmod
-        ./extra/hcq/hcq_smi.py amd kill_pids
+        ./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: sudo lsof -tQ /dev/nvidia* | { xargs -r sudo kill -9 || true; }
+      run: lsof -tQ /dev/nvidia* | { xargs -r kill -9 || true; }
     - name: Symlink models and datasets
       run: |
         mkdir -p extra/datasets
@@ -227,15 +212,8 @@ 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
 
@@ -246,13 +224,13 @@ jobs:
       fail-fast: false
       matrix:
         dev: ['METAL', 'AMD', 'NV']
-    timeout-minutes: 60
+    timeout-minutes: 15
     defaults:
       run:
         shell: bash -e -o pipefail {0}
     env:
       DEV: ${{ matrix.dev }}
-      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
+      HCQ2: '0'
     if: github.repository_owner == 'tinygrad'
     steps:
     - name: Checkout Code
@@ -260,12 +238,11 @@ jobs:
     - name: Setup (AMD)
       if: ${{ matrix.dev == 'AMD' }}
       run: |
-        ./extra/amdpci/setup_python_cap.sh
-        ./extra/hcq/hcq_smi.py amd rmmod
-        ./extra/hcq/hcq_smi.py amd kill_pids
+        ./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: sudo lsof -tQ /dev/nvidia* | { xargs -r sudo kill -9 || true; }
+      run: lsof -tQ /dev/nvidia* | { xargs -r kill -9 || true; }
     - name: setup staging db
       if: github.ref == 'refs/heads/update_benchmark_staging'
       run: |
@@ -285,6 +262,59 @@ 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' }}"]
@@ -292,7 +322,7 @@ jobs:
       fail-fast: false
       matrix:
         dev: ['METAL', 'AMD', 'NV']
-    timeout-minutes: 60
+    timeout-minutes: 10
     defaults:
       run:
         shell: bash -e -o pipefail {0}
@@ -305,12 +335,11 @@ jobs:
     - name: Setup (AMD)
       if: ${{ matrix.dev == 'AMD' }}
       run: |
-        ./extra/amdpci/setup_python_cap.sh
-        ./extra/hcq/hcq_smi.py amd rmmod
-        ./extra/hcq/hcq_smi.py amd kill_pids
+        ./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: sudo lsof -tQ /dev/nvidia* | { xargs -r sudo kill -9 || true; }
+      run: lsof -tQ /dev/nvidia* | { xargs -r kill -9 || true; }
     - name: setup staging db
       if: github.ref == 'refs/heads/update_benchmark_staging'
       run: |
@@ -374,7 +403,7 @@ jobs:
       run: python test/external/external_benchmark_multitensor_allreduce.py
     - name: HEVC Decode Benchmark
       if: ${{ matrix.dev == 'NV' }}
-      run: VALIDATE=1 MAX_FRAMES=100 ASSERT_FPS=1400 JITBEAM=1 PYTHONPATH=. python3 extra/hevc/decode.py
+      run: IGNORE_BEAM_CACHE=1 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:
@@ -387,7 +416,7 @@ jobs:
   testusbgpu:
     name: UsbGPU Benchmark
     runs-on: [self-hosted, macOS]
-    timeout-minutes: 10
+    timeout-minutes: 3
     defaults:
       run:
         shell: bash -e -o pipefail {0}
@@ -402,32 +431,66 @@ jobs:
         rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
     - name: Kill stale pids
       run: |
-        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
+        ./extra/hcq/hcq_smi.py amd kill_pids --sudoless
+        ./extra/hcq/hcq_smi.py nv kill_pids --sudoless
     - name: UsbGPU boot time
-      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
+      run: 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: sudo -E PYTHONDONTWRITEBYTECODE=1 PYTHONPATH=. GMMU=0 DEV=USB+AMD python3.11 test/test_tiny.py
+      run: GMMU=0 DEV=USB+AMD python3.11 test/test_tiny.py
     - name: UsbGPU copy speeds
-      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
+      run: SIZE=64000000 PYTHONPATH=. GMMU=0 DEV=USB+AMD python3.11 test/external/external_test_usb_asm24.py TestDevCopySpeeds
     - name: UsbGPU (USB4/TB) install script
-      run: PYTHONPATH=. sh extra/setup_tinygpu_osx.sh
+      run: sh extra/setup_tinygpu_osx.sh
     - name: UsbGPU (USB4/TB) boot time
-      run: PYTHONPATH=. DEBUG=3 DEV=PCI+NV:NAK time python3.11 test/test_tiny.py TestTiny.test_plus
+      run: DEBUG=3 DEV=PCI+NV:NAK time python3.11 test/test_tiny.py TestTiny.test_plus
     - name: UsbGPU (USB4/TB) tiny tests
-      run: PYTHONPATH=. DEV=PCI+NV:NAK python3.11 test/test_tiny.py
+      run: DEV=PCI+NV:NAK python3.11 test/test_tiny.py
 
-  testcommalatest:
-    name: comma Benchmark (0.11.2)
-    runs-on: [self-hosted, Linux, comma]
-    timeout-minutes: 12
+  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
     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
@@ -438,45 +501,10 @@ 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: 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: 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: Run process replay tests
       uses: ./.github/actions/process-replay
 
@@ -489,15 +517,6 @@ 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
@@ -510,8 +529,8 @@ jobs:
     - name: benchmark MobileNetV2 on DSP
       run: |
         # generate quantized weights
-        ln -s /data/home/tiny/tinygrad/extra/datasets/imagenet extra/datasets/imagenet
-        ln -s /data/home/tiny/tinygrad/testsig-*.so .
+        ln -s ~/tinygrad/extra/datasets/imagenet extra/datasets/imagenet
+        ln -s ~/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
@@ -521,7 +540,7 @@ jobs:
   testcommausbgpubenchmark:
     name: UsbGPU Benchmark (comma)
     runs-on: [self-hosted, Linux, comma4]
-    timeout-minutes: 20
+    timeout-minutes: 10
     defaults:
       run:
         shell: bash -e -o pipefail {0}
@@ -541,7 +560,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=64e6 PYTHONPATH=. GMMU=0 DEV=USB+AMD python3 test/external/external_test_usb_asm24.py TestDevCopySpeeds
+      run: SIZE=64000000 PYTHONPATH=. GMMU=0 DEV=USB+AMD python3 test/external/external_test_usb_asm24.py TestDevCopySpeeds
 
   driverbenchmarks:
     name: PCI Driver Benchmark (DEV=${{ matrix.dev }})
@@ -550,7 +569,7 @@ jobs:
       fail-fast: false
       matrix:
         dev: ['AMD', 'NV']
-    timeout-minutes: 20
+    timeout-minutes: 5
     defaults:
       run:
         shell: bash -e -o pipefail {0}
@@ -562,9 +581,8 @@ jobs:
       uses: actions/checkout@v6
     - name: Setup
       run: |
-        ./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
+        ./extra/hcq/hcq_smi.py ${{ matrix.dev }} rmmod --expect
+        ./extra/hcq/hcq_smi.py ${{ matrix.dev }} kill_pids --sudoless
         mkdir -p extra/datasets
         ln -s /raid/datasets/imagenet extra/datasets/imagenet
     - name: setup staging db
@@ -599,6 +617,9 @@ 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
@@ -621,12 +642,12 @@ jobs:
   llvmspeed:
     name: LLVM Speed
     runs-on: [self-hosted, Linux, tinyboxrandom]
-    timeout-minutes: 20
+    timeout-minutes: 10
     if: github.repository_owner == 'tinygrad'
     steps:
     - name: Checkout Code
       uses: actions/checkout@v6
     - name: Speed Test
-      run: DEV=CPU:LLVM THREADS=0 python3 test/speed/external_test_speed_v_torch.py
+      run: DEV=CPU:LLVM python3 test/speed/external_test_speed_v_torch.py
     - name: Speed Test (BEAM=2)
-      run: BEAM=2 DEV=CPU:LLVM THREADS=0 python3 test/speed/external_test_speed_v_torch.py
+      run: IGNORE_BEAM_CACHE=1 BEAM=2 DEV=CPU:LLVM python3 test/speed/external_test_speed_v_torch.py

tinygrad_repo/.github/workflows/docs.yml

@@ -8,7 +8,7 @@ permissions:
   contents: write
 jobs:
   deploy:
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-24.04
     steps:
       - uses: actions/checkout@v6
       - name: Configure Git Credentials

tinygrad_repo/.github/workflows/platform.yml

@@ -166,7 +166,7 @@ jobs:
         uses: ./.github/actions/setup-tinygrad
         with:
           key: windows-${{ matrix.dev }}-minimal
-          deps: testing_unit
+          deps: testing_minimal
           pydeps: ${{ matrix.dev == 'WEBGPU' && 'dawn-python' || '' }}
       - name: Set env
         shell: bash
@@ -179,35 +179,3 @@ 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

tinygrad_repo/.github/workflows/python-publish.yml

@@ -10,7 +10,7 @@ on:
 
 jobs:
   deploy:
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-24.04
     steps:
     - uses: actions/checkout@v6
     - name: Set up Python

tinygrad_repo/.github/workflows/szdiff.yml

@@ -10,7 +10,7 @@ concurrency:
 jobs:
   checkbranch:
     name: Check PR Branch status
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-24.04
     outputs:
       branchstat: ${{ steps.brstat.outputs.stat}}
     steps:
@@ -44,7 +44,7 @@ jobs:
     permissions:
       contents: read
       pull-requests: write
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-24.04
     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-latest
+    runs-on: ubuntu-24.04
     needs: checkbranch
     if: needs.checkbranch.outputs.branchstat == 'true'
     steps:

tinygrad_repo/.github/workflows/test.yml

@@ -21,7 +21,7 @@ concurrency:
 jobs:
   docs:
     name: Docs
-    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' }}
+    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: 10
     env:
       CHECK_OOB: 0
@@ -31,8 +31,7 @@ jobs:
     - name: Setup Environment
       uses: ./.github/actions/setup-tinygrad
       with:
-        deps: docs
-        pydeps: "capstone torch"
+        deps: "docs testing_minimal"
     - name: Build wheel and show size
       run: |
         uv build --wheel
@@ -61,7 +60,7 @@ jobs:
 
   torchbackend:
     name: Torch Backend Tests
-    runs-on: *linux
+    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
@@ -73,10 +72,7 @@ jobs:
         deps: testing_unit
         pydeps: "pillow torchvision expecttest"
         llvm: 'true'
-    - name: Install ninja
-      run: |
-        sudo apt update || true
-        sudo apt install -y --no-install-recommends ninja-build
+        ninja: 'true'
     - name: Test ResNet-18
       run: DEBUG=2 python3 extra/torch_backend/example.py
     - name: Test one op in torch tests
@@ -86,9 +82,26 @@ 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: *linux
+    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
@@ -126,7 +139,7 @@ jobs:
 
   linter:
     name: Linters
-    runs-on: *linux
+    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: 10
 
     steps:
@@ -157,7 +170,7 @@ jobs:
 
   nulltest:
     name: Null Tests
-    runs-on: *linux
+    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:
@@ -191,7 +204,7 @@ jobs:
 
   unittest:
     name: Unit Tests
-    runs-on: *linux
+    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:
@@ -220,7 +233,7 @@ jobs:
     - name: Run process replay tests
       uses: ./.github/actions/process-replay
     - name: Repo line count <= 26000 lines
-      run: MAX_LINE_COUNT=26000 python sz.py
+      run: MAX_LINE_COUNT=26500 python sz.py
 
   spec:
     strategy:
@@ -228,7 +241,7 @@ jobs:
       matrix:
         group: [1, 2]
     name: SPEC=2 (${{ matrix.group }})
-    runs-on: *linux
+    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
@@ -244,7 +257,7 @@ jobs:
 
   fuzzing:
     name: Fuzzing
-    runs-on: *linux
+    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: 10
     steps:
     - name: Checkout Code
@@ -260,7 +273,7 @@ jobs:
 
   testopenclimage:
     name: CL IMAGE Tests
-    runs-on: *linux
+    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
@@ -280,7 +293,7 @@ jobs:
 
   testopenpilot:
     name: openpilot Compile Tests
-    runs-on: *linux
+    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
@@ -309,7 +322,7 @@ jobs:
 
   testonnxcpu:
     name: ONNX (CPU) Tests
-    runs-on: *linux
+    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: 20
 
     steps:
@@ -328,7 +341,7 @@ jobs:
 
   testoptim:
     name: Optimization Tests
-    runs-on: *linux
+    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: 20
     steps:
       - name: Checkout Code
@@ -360,7 +373,7 @@ jobs:
 
   testllm:
     name: Test LLM
-    runs-on: *linux
+    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
     env:
       CHECK_OOB: 0
@@ -377,17 +390,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 --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'
+            ::: $'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'
                 # NOTE: qwen is dumb and only knows about female chickens
 
 # ****** Models Tests ******
 
   testmodels:
     name: Models
-    runs-on: *linux
+    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
@@ -407,7 +420,7 @@ jobs:
 
   testdsp:
     name: Linux (DSP)
-    runs-on: *linux
+    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
@@ -435,7 +448,7 @@ jobs:
           - 'WEBGPU'
 
     name: Linux (DEV=${{ matrix.dev }})
-    runs-on: *linux
+    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: 20
     steps:
     - name: Checkout Code
@@ -461,7 +474,7 @@ jobs:
 
   testamdasm:
     name: AMD ASM IDE
-    runs-on: *linux
+    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: 20
     env:
       DEV: MOCKKFD+AMD
@@ -491,7 +504,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: PROFILE=1 SQTT=1 python3 -m pytest -n=auto test/amd/test_sqtt_profiler.py
+        run: VIZ=-2 python3 -m pytest -n=auto test/amd/test_sqtt_profiler.py
       - name: Run AMD emulated tests on NULL backend
         env:
           AMD: 0
@@ -507,7 +520,7 @@ jobs:
 
   hcq2:
     name: hcq2
-    runs-on: *linux
+    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: 5
     steps:
       - name: Checkout Code
@@ -521,16 +534,15 @@ 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 test/unit/test_multitensor.py \
-            TestMultiTensor.test_simple_add TestMultiTensor.test_shard_reduce \
-            TestMultiTensor.test_backward_sum TestMultiTensor.test_matmul_shard_0_0
+        run: HCQ_RUNTIME_DEV=PYTHON HCQ2=1 DEV=MOCKKFD+AMD FORWARD_ONLY=1 PYTHONPATH=. python -m pytest -n=auto test/backend/test_multitensor.py
       - 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: *linux
+    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
     env:
       DEV: MOCKPCI+AMD
@@ -566,7 +578,7 @@ jobs:
         arch: [gfx1100, gfx1201, gfx950]
 
     name: Linux (${{ matrix.backend }} ${{ matrix.arch }})
-    runs-on: *linux
+    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
     env:
       DEV: MOCKKFD+AMD:${{ matrix.backend == 'amdllvm' && 'LLVM' || '' }}:${{ matrix.arch }}
@@ -589,7 +601,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 test/opt/test_tensor_cores.py --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 --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
@@ -604,7 +616,7 @@ jobs:
         backend: [ptx, nv]
 
     name: Linux (${{ matrix.backend }})
-    runs-on: *linux
+    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: 20
     env:
       FORWARD_ONLY: 1
@@ -638,10 +650,17 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        backend: [ir3, nak]
-    name: Compile-only (${{ matrix.backend }})
-    runs-on: *linux
+        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' }}
     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
@@ -650,21 +669,15 @@ jobs:
         with:
           key: compile-${{ matrix.backend }}
           deps: "testing_unit mesa"
-      - name: Set env
+          qemu: ${{ contains(matrix.dev, 'QCOMCL') }}
+      - name: Test IMAGE
         shell: bash
-        run: printf "NULL_ALLOW_COPYOUT=1\n${{ matrix.backend == 'ir3' && 'DEV=NULL:IR3:a630' || matrix.backend == 'nak' && 'DEV=NULL:NAK:sm_120' }}" >> $GITHUB_ENV
+        if: contains(matrix.dev, 'a630')
+        run: DEBUG=7 python3 test/backend/test_ops.py TestOps.test_gemm | grep isam
       - 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)
-        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
+          # 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

tinygrad_repo/.gitignore

@@ -69,3 +69,4 @@ mutants
 dagre/
 graphlib/
 uv.lock
+pi_session_window0.jsonl

tinygrad_repo/.pre-commit-config.yaml

@@ -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/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
+        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
         language: system
         always_run: true
         pass_filenames: false

tinygrad_repo/AGENTS.md

@@ -4,3 +4,4 @@
 - 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.

tinygrad_repo/README.md

@@ -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)
+x = Tensor.eye(3).clone()  # clone to make it a buffer
 y = Tensor([[2.0,0,-2.0]])
 z = y.matmul(x).sum()
 z.backward()

tinygrad_repo/TINYGRAD_COMMIT

@@ -1 +1 @@
-8611fe22a7fcc7d1928bbde19ded66277cb12f3e
+f6fc4e3f2c3db5fae1e19cbfbc3ad9fc579a12ae

tinygrad_repo/conftest.py

@@ -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", 300)), os.kill, args=(os.getpid(), signal.SIGABRT))
+  t = threading.Timer(int(os.getenv("TEST_TIMEOUT", 90)), os.kill, args=(os.getpid(), signal.SIGABRT))
   t.start()
   try: yield
   finally:

tinygrad_repo/docs/abstractions4.py

@@ -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) for x in self.instructions]))))
+      return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=(x, dtypes.void)) for x in self.instructions]))))
 
   CU_COUNT = 32
   LANES = 64

tinygrad_repo/docs/developer/speed.md

@@ -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. 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.
+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.
 
 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.
 

tinygrad_repo/examples/mlperf/dataloader.py

@@ -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, NUM_CPU_THREADS
+from tinygrad.helpers import getenv, prod, Context, round_up, tqdm, OSX, CPU_COUNT
 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(NUM_CPU_THREADS.value):
+    for _ in range(CPU_COUNT):
       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(NUM_CPU_THREADS.value, len(train_files))
+  cycle_length = min(CPU_COUNT, 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(NUM_CPU_THREADS.value):
+    for _ in range(CPU_COUNT):
       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(NUM_CPU_THREADS.value):
+    for _ in range(CPU_COUNT):
       proc = Process(
         target=load_retinanet_data,
         args=(base_dir, val, queue_in, queue_out, imgs, boxes, labels),

tinygrad_repo/examples/mlperf/model_train.py

@@ -1667,15 +1667,14 @@ 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, clip_grads
+  from examples.mlperf.optim import GradAccClipAdamW, GradAccClipAdamWGroup, fclip_grads
 
   BENCHMARK = getenv("BENCHMARK")
 
   config = {}
   BASEDIR            = config["BASEDIR"]                = Path(getenv("BASEDIR", "/raid/datasets/c4-8b/"))
   BS                 = config["BS"]                     = getenv("BS", 16)
-  grad_acc           = config["GRADIENT_ACC_STEPS"]     = getenv("GRADIENT_ACC_STEPS", 1)
-  GBS                = config["GLOBAL_BATCH_SIZE"]      = BS * grad_acc
+  GBS                = config["GLOBAL_BATCH_SIZE"]      = BS
   SEED               = config["SEED"]                   = getenv("SEED", 5760)
   DATA_SEED          = config["DATA_SEED"]              = getenv("DATA_SEED", SEED)
   SEQLEN             = config["SEQLEN"]                 = getenv("SEQLEN", 8192)
@@ -1737,13 +1736,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=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),
+    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),
   )
 
   for p in optim.params:
-    grad_dtype = dtypes.bfloat16 if p.dtype == FP8_DTYPE else p.dtype
-    p.grad = p.zeros_like(dtype=grad_dtype).contiguous()
+    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)
   grads = [p.grad for p in optim.params]
 
   from extra.gemm.cdna_asm_gemm import _mx_block_scale
@@ -1770,31 +1769,32 @@ def train_gptoss():
 
   @TinyJit
   @Context(TRAINING=1)
-  def minibatch(tokens:Tensor):
+  def step(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)
-    loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
+    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:])
 
     for g, new_g in zip(grads, loss.gradient(*optim.params)):
       apply_grad(g, new_g.uop)
 
-    loss_cpu = loss.flatten().float().to("CPU")
-    return loss_cpu.realize(*grads)
+    Tensor.realize(loss, *grads)
 
-  @TinyJit
-  def optim_step():
-    grad_norm = clip_grads(grads, grad_acc, 1.0)
-    optim.fstep(grads, grad_norm)
+    clipped_grads, grad_norm = fclip_grads(grads, 1.0)
+    optim.fstep(clipped_grads, grad_norm)
     scheduler.step()
 
-    for g in grads: g.assign(0)
-
+    loss_cpu = loss.flatten().float().to("CPU")
     lr_cpu = optim.lr.float().to("CPU")
     grad_norm_cpu = grad_norm.float().to("CPU")
-    Tensor.realize(lr_cpu, grad_norm_cpu, *grads, *fp8_inv_scales)
+    Tensor.realize(loss_cpu, lr_cpu, grad_norm_cpu, *grads, *fp8_inv_scales)
 
-    return lr_cpu, grad_norm_cpu
+    return loss_cpu, lr_cpu, grad_norm_cpu
 
   @TinyJit
   @Context(TRAINING=0)
@@ -1843,30 +1843,20 @@ def train_gptoss():
       profile_marker(f"train @ {i}")
       st = 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
+      ist = time.perf_counter()
 
-      gt = time.perf_counter()
-      ret = optim_step()
-      lr, grad_norm = ret[0].item(), ret[1].item()
+      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()
       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
@@ -1876,7 +1866,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, {gbs_time:.3f} s gbs, {optim_time:.3f} s optim, {data_time:.3f} s data, {loss:.4f} loss, " \
+          f"{i:5} {step_time:.3f} s step, {dev_time:.3f} s dev, {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())))
 
@@ -1886,8 +1876,6 @@ 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,

tinygrad_repo/examples/mlperf/models/gpt_oss.py

@@ -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
+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.moe_gemm import grouped_mx_gemm
 from extra.gemm.moe_routing import route, dispatch, combine
 
@@ -61,7 +61,25 @@ 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, w_q:Tensor, w_scale:Tensor) -> Tensor:
+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)
   l_shape = x.shape[:-1]
   if ASM_GEMM:
     from extra.gemm.cdna_asm_gemm import asm_gemm, can_use_asm_gemm, mx_pack
@@ -146,6 +164,7 @@ 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)
 
@@ -174,20 +193,30 @@ 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
-    x_normed, rrms = rmsnorm(x, self.norm_eps)
-    qkv = matmul_mx(x_normed * attention_norm, wqkv, wqkv_scale) + wqkv_bias
+
+    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]
+
     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)
 
-    if sliding:
-      attn = self._sliding_attention(xq, xk, xv, sinks)
-    elif getenv("HK_FLASH_ATTENTION"):
+    fa_saves = []
+    if 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, _, 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:
+      attn = self._sliding_attention(xq, xk, xv, sinks)
     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)
@@ -199,13 +228,19 @@ 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, [x_normed, rrms, attn]
+    return out, [attn] + norm_saves + fa_saves
 
   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):
-    x_normed, rrms = rmsnorm(x, self.norm_eps)
-    inp = x_normed * ffn_norm
+    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
+
     logits = inp.float() @ gate.float().T + gate_bias.float()
     dim, inter = self.dim, self.intermediate_size
 
@@ -220,6 +255,7 @@ 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)
@@ -263,7 +299,11 @@ 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)
 
-    logits = self.norm(h) @ self.output.T
+    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
     return logits
 
 def _get_pads(uop:UOp) -> list[UOp]:
@@ -274,14 +314,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(grad_buf.uop + new_grad))
+    grad_buf.uop = grad_buf.uop.after(grad_buf.uop.store(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(buf_slice + pad.src[0].cast(cur.dtype)))
+      cur = cur.after(buf_slice.store(pad.src[0].cast(cur.dtype)))
     else:
       cur = cur.after(cur.store(cur + pad.cast(cur.dtype)))
   grad_buf.uop = cur

tinygrad_repo/examples/mlperf/optim.py

@@ -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, dtypes.default_float, tuple(Tensor.rand(*shape, device=d).uop for d in x.device)))
+    noise = Tensor(UOp(Ops.MSTACK, src=tuple(Tensor.rand(*shape, device=d).uop for d in x.device)))
   else:
     noise = x.rand_like()
   noise = (noise * 0xFFFF).cast(dtypes.uint32)
@@ -27,6 +27,11 @@ 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)

tinygrad_repo/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/buid_dataset_cache.py

@@ -0,0 +1,13 @@
+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)

tinygrad_repo/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_beam.sh

@@ -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="bfloat16"
+export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="float32"
 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=4096 MAX_STEPS=1200000
+export LR="1e-3" END_LR="1e-4" WARMUP_SAMPLES=2048 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=1
+export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
 
 export FAKEDATA=${FAKEDATA:-1} BENCHMARK=${BENCHMARK:-10}
 if [ -z "$FULL_LAYERS" ]; then

tinygrad_repo/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_run.sh

@@ -1,8 +1,8 @@
 #!/usr/bin/env bash
 
 export PYTHONPATH="."
-export PATH="/opt/rocm-7.1.1/bin:$PATH"
-export ROCM_PATH="/opt/rocm-7.1.1"
+export ROCM_PATH=${ROCM_PATH:-/opt/rocm-7.1.1}
+export PATH="$ROCM_PATH/bin:$PATH"
 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="bfloat16"
+export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="float32"
 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=4096 MAX_STEPS=1200000
+export LR="1e-3" END_LR="1e-4" WARMUP_SAMPLES=2048 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=1
+export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
 
 python3 examples/mlperf/model_train.py

tinygrad_repo/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/run_and_time.sh

@@ -17,7 +17,7 @@ export USE_ATOMICS=1
 export ASM_GEMM=1
 export WQKV=1
 export MASTER_WEIGHTS=1
-export FP8=1
+export MXFP4=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="bfloat16"
+export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="float32"
 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=1
+export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
 
 export LOGMLPERF=1
 

tinygrad_repo/examples/mlperf/training_submission_v6.1/tinycorp/benchmarks/gpt_oss/implementations/tinybox_8xMI350X/dev_beam.sh

@@ -11,6 +11,7 @@ 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}

tinygrad_repo/examples/mlperf/training_submission_v6.1/tinycorp/benchmarks/gpt_oss/implementations/tinybox_8xMI350X/dev_run.sh

@@ -11,6 +11,7 @@ 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}

tinygrad_repo/examples/openpilot/compile3.py

@@ -107,14 +107,21 @@ 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()
-    out = run(**inputs)
-    mt = time.perf_counter()
-    val = out.numpy()
+    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()
     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")
@@ -160,12 +167,6 @@ 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)
@@ -181,6 +182,3 @@ 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)

tinygrad_repo/extra/amdpci/am_smi.py

@@ -84,7 +84,8 @@ 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):
+  def __init__(self, dev_filter=None):
+    self.dev_filter = dev_filter
     self.devs = []
     self.opened_pcidevs = []
     self.opened_pci_resources = {}
@@ -135,6 +136,7 @@ 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)
 
@@ -406,7 +408,7 @@ if __name__ == "__main__":
 
   try:
     if not args.list: os.system('clear')
-    smi_ctx = SMICtx()
+    smi_ctx = SMICtx(args.dev)
     while True:
       smi_ctx.rescan_devs()
       smi_ctx.draw(args.list)

tinygrad_repo/extra/bench_log.py

@@ -35,7 +35,7 @@ class WallTimeEvent:
     return self
   def __exit__(self, *_):
     self.time = time.monotonic() - self.start
-    _events[self.event]["wall"].append(self.time)
+    _events[self.event]["wall"].append((self.time, BENCHMARK_LOG.value))
     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)
+    _events[self.event]["kernel"].append((GlobalCounters.time_sum_s - self.start, BENCHMARK_LOG.value))
     return False
 
 def log_event_instant(event:InstantBenchEvent, value:float):
-  _events[event].append(value)
+  _events[event].append((value, BENCHMARK_LOG.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, run):
-    point = Point(BENCHMARK_LOG.value).tag("id", run_id).tag("index", i)
+  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)
     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 in enumerate(values):
-            point = _create_point(run_id, i, attempt, ref, commit, f"{event.value}_{event_type}", value, run)
+          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)
             points.append(point)
       else:
-        for i, value in enumerate(_events[event]):
-          point = _create_point(run_id, i, attempt, ref, commit, event.value, value, run)
+        for i, (value, log_name) in enumerate(_events[event]):
+          point = _create_point(run_id, i, attempt, ref, commit, event.value, value, log_name, run)
           points.append(point)
 
     write_options = WriteOptions(write_type=WriteType.synchronous, retry_interval=5000, max_retries=5, max_retry_delay=30000, exponential_base=2)

tinygrad_repo/extra/bnxt_driver/bnxtdev.py

@@ -0,0 +1,238 @@
+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

tinygrad_repo/extra/bnxt_driver/connect.py

@@ -0,0 +1,118 @@
+#!/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()

tinygrad_repo/extra/bnxt_driver/loopback.py

@@ -0,0 +1,46 @@
+#!/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)

tinygrad_repo/extra/export_model.py

@@ -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 b.op is not Ops.BIND]
+    arg_uops = [b for b in call.src[1:] if not b.is_bound_var]
     prg = to_program(call.src[0], Device[arg_uops[0].device].renderer)
     info = prg.arg
     functions[info.function_name] = prg.src[2].arg
@@ -241,8 +241,7 @@ 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"
 
-  # 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)
+  with Context(JIT=2): 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}

tinygrad_repo/extra/gemm/amd_asm_matmul.py

@@ -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) for x in insts]))))
+    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=(x, dtypes.void)) for x in insts]))))
   c = Tensor.custom_kernel(a, b, c, fxn=asm_kernel)[2]
   linear = c.schedule_linear()
 

tinygrad_repo/extra/gemm/cdna_asm_gemm.py

@@ -122,9 +122,10 @@ 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"custom_mxfp4_gemm_{M}_{N}_{K}", estimates=Estimates(ops=2*M*N*K)))
+                  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)))
   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) for x in insts))))
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple(UOp(Ops.INS, arg=(x, dtypes.void)) 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
@@ -214,7 +215,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, dtype=dtypes.float32).cast(C.dtype)
+  red = mul.reduce(k, arg=Ops.ADD).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}'))
 

tinygrad_repo/extra/gemm/max_matmul.py

@@ -113,7 +113,6 @@ 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 = {

tinygrad_repo/extra/gemm/moe_gemm.py

@@ -1,10 +1,32 @@
 import functools, pathlib
 from tinygrad import Tensor, dtypes
-from tinygrad.uop.ops import UOp, Ops, KernelInfo
+from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType
+from tinygrad.helpers import getenv
 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
@@ -58,7 +80,8 @@ 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]
-  out = out.sum(0) if is_multi else out.squeeze(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)
   return out.reshape(n_experts, N, K)
 
 def mx_pack_3d(e8:Tensor) -> Tensor:

tinygrad_repo/extra/gemm/moe_routing.py

@@ -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, dtypes.void, (gtab.index(g, row, j), val), arg=atomic_str)
+    atomic = UOp(Ops.CUSTOM, src=(gtab.index(g, row, j), val), arg=(atomic_str, dtypes.void))
     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)

tinygrad_repo/extra/gemm/rdna4_asm_matmul.py

@@ -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) for x in insts]))))
+    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=(x, dtypes.void)) for x in insts]))))
 
   c = Tensor.custom_kernel(a, b, c, fxn=asm_kernel)[2]
   linear = c.schedule_linear()

tinygrad_repo/extra/gemm/tinygrad_nv_matmul.py

@@ -1,6 +1,7 @@
 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
 
@@ -13,17 +14,17 @@ if __name__ == "__main__":
   C = A.matmul(B)
   if getenv("GEMV"):
     opts = [
-      Opt(op=OptOps.UNROLL, axis=0, amt=8),
-      Opt(op=OptOps.GROUP, axis=0, amt=32),
+      Opt(op=OptOps.SPLIT, axis=1, arg=(8, AxisType.UNROLL)),
+      Opt(op=OptOps.SPLIT, axis=1, arg=(32, AxisType.GROUP_REDUCE)),
     ]
   else:
     opts = [
       Opt(op=OptOps.TC, axis=0, amt=0),
-      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),
+      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)),
     ]
   linear = C.schedule_linear()
   call = linear.src[-1]

tinygrad_repo/extra/gemm/triton_nv_matmul.py

@@ -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 s.op is not Ops.BIND]
+  bufs = [s.buffer for s in last_call.src[1:] if not s.is_bound_var]
 
   src = compiled.asm["ptx"]
   # specify the shared memory here so we don't need to do it dynamically

tinygrad_repo/extra/gptoss_kernels/rmsnorm/__init__.py

@@ -0,0 +1,29 @@
+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))

tinygrad_repo/extra/hcq/hcq_smi.py

@@ -16,9 +16,12 @@ 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", "ast"] if args.backend == "nv" else ["amdgpu"]
+  modules = ["nvidia_drm", "nvidia_modeset", "nvidia_uvm", "nvidia"] 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)
@@ -60,17 +63,19 @@ 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', 'lsof', temp(f'{prefix}_{dev}.lock')]).decode('utf-8').strip().split('\n')[1].split()[1]
+        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]
         except subprocess.CalledProcessError: break
 
         print(f"Killing process {pid} (which uses {dev})")
-        subprocess.run(['sudo', 'kill', '-9', pid], check=True)
+        subprocess.run((['sudo'] if use_sudo else []) + ['kill', '-9', pid], check=True)
       except subprocess.CalledProcessError as e:
         print(f"Failed to kill process for device {dev}: {e}", file=sys.stderr)
 
@@ -79,6 +84,7 @@ 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")
@@ -91,17 +97,20 @@ 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", help="NVIDIA GPUs")
+  nv_parser = backend_subparsers.add_parser("nv", aliases=["NV"], help="NVIDIA GPUs")
+  nv_parser.set_defaults(backend="nv")
   nv_commands = nv_parser.add_subparsers(dest="command", required=True)
   add_common_commands(nv_commands)
 
-  amd_parser = backend_subparsers.add_parser("amd", help="AMD GPUs")
+  amd_parser = backend_subparsers.add_parser("amd", aliases=["AMD"], help="AMD GPUs")
+  amd_parser.set_defaults(backend="amd")
   amd_commands = amd_parser.add_subparsers(dest="command", required=True)
   add_common_commands(amd_commands)
 

tinygrad_repo/extra/hcq2/ops_amd2.py

@@ -1,33 +1,32 @@
 from __future__ import annotations
-from typing import cast, Any, Callable
-import os, ctypes, struct, hashlib, functools, importlib, mmap, errno, array, contextlib, sys, weakref, itertools, collections, atexit
+from typing import cast
+import os, ctypes, struct, functools, importlib, mmap, errno, contextlib, sys, itertools, atexit
 assert sys.platform != 'win32'
 from dataclasses import dataclass
-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.runtime.support.hcq2 import HCQ2Compiled, HCQAllocator, make_buf, hcq_size_var
 from tinygrad.uop.ops import sint, UOp
-from tinygrad.device import Compiled, BufferSpec, Buffer, Device
+from tinygrad.device import BufferSpec, Buffer, Device
 from tinygrad.dtype import dtypes
-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.helpers import getenv, round_up, data64_le, DEBUG, PROFILE, lo32, hi32
+from tinygrad.helpers import 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, sqtt, amdgpu_kd, amdgpu_drm
+from tinygrad.runtime.autogen import kfd, hsa, 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
+from tinygrad.runtime.support.usb import USB3, pm_usb_bufferize
 from tinygrad.runtime.support.memory import AddrSpace, BumpAllocator
-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
+from tinygrad.runtime.ops_amd import SQTT, PMC
+from tinygrad.runtime.ops_amd import EVENT_INDEX_PARTIAL_FLUSH, 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_runtime, pm_flatten_linear
+from tinygrad.engine.realize import get_call_arg_uops, get_call_var_uops
 from tinygrad.uop import FastEnum, auto
-from tinygrad.uop.ops import Ops, UPat, PatternMatcher, graph_rewrite
+from tinygrad.uop.ops import Ops, UPat, PatternMatcher
 
 # *****************
 # PM4
@@ -36,9 +35,11 @@ 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.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)))
+  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)))
 
 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')
@@ -52,7 +53,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, *(data64_le(mem) if mem is not None else (reg, reg_done)), value, mask, 4)
+  return pkt3(ctx, PM4Ops.WAIT_REG_MEM, wrm_info_dw, *((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:
@@ -83,16 +84,18 @@ 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
-  return pkt3(ctx, PM4Ops.RELEASE_MEM, event_dw | cache_flags_dw, memsel_dw, *data64_le(address), *data64_le(value), ctxid)
+  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)
 
 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, dtypes.void, (
+  return UOp(Ops.LINEAR, src=(
     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, mem=dst.getaddr(ctx.devs))
+def pm4_wait(ctx, dst, val): return wait_reg_mem(ctx, val.cast(dtypes.uint32), mem=dst.getaddr(ctx.devs))
 
 def pm4_barrier(ctx): return memory_barrier(ctx)
 
@@ -106,155 +109,130 @@ def pm4_timestamp(ctx, dst):
                      ctx.pm4.int_sel__mec_release_mem__none)
 
 def pm4_program(ctx, call, prg):
-  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 = args.getaddr(ctx.devs)
+  data, lib = amd_build_program(ctx.dev, prg)
+  info = prg.arg
 
-  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)]
+  # 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
+  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)
+
+  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]
 
   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, *data64_le(prog_addr >> 8)),
+    wreg(ctx, ctx.gc.regCOMPUTE_PGM_LO, 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, *data64_le((scratch_addr + data.private_segment_size // ctx.xccs * xcc_id) >> 8))
+  ins += [wreg(ctx, ctx.gc.regCOMPUTE_DISPATCH_SCRATCH_BASE_LO, (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 or (1, 1, 1)), 0, 0),
+    wreg(ctx, ctx.gc.regCOMPUTE_START_X, 0, 0, 0, *info.local_size, 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, dtypes.void, tuple(ins))
+  return UOp(Ops.LINEAR, src=tuple(ins))
 
-pm_pm4_opsel = PatternMatcher([
+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([
   (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="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),
+  (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),
 ])
 
-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
-  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)))))
+  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))))))
 
-def sdma_wait(ctx, ins, dst, val):
+def sdma_wait(ctx, 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 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, 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=(
-    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)))))
+    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)))
 
-def sdma_timestamp(ctx, ins, dst):
+def sdma_store(ctx, dst, val): # a fence packet then a trap
+  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)))
+
+def sdma_timestamp(ctx, dst):
   op = ctx.sdma.SDMA_OP_TIMESTAMP | ctx.sdma.SDMA_PKT_TIMESTAMP_GET_HEADER_SUB_OP(ctx.sdma.SDMA_SUBOP_TIMESTAMP_GET_GLOBAL)
-  return ins.ins(SDMAOps.TIMESTAMP, src=tuple(UOp.const(x, dtypes.uint32) for x in (op, *data64_le(dst.getaddr(ctx.devs)))))
+  return UOp(Ops.LINEAR, src=(UOp.const(op, dtypes.uint32), dst.getaddr(ctx.devs)))
 
-pm_sdma_opsel = PatternMatcher([
+pm_sdma_encode = PatternMatcher([
   (UPat(Ops.CALL, src=(UPat(Ops.COPY),), name="call", allow_any_len=True), sdma_copy),
 
-  (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),
+  (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),
 ])
 
-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)
+# *****************
+# queue submit
 
-  # 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 _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")))
 
-  # 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 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)
 
-  # 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)
+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)
 
-  # 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)
+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)])
 
 @dataclass(frozen=True)
 class AMDProgramData:
@@ -262,10 +240,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,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:
+_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:
     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:
@@ -282,20 +260,23 @@ def amd_build_program(prg:UOp) -> 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] = prg.replace(src=(buf.after(make_binary_patch(buf, bytes(image))),), arg=(data, prg.arg))
+    cached = _amd_program_cache[key] = (data, buf.after(buf.store(UOp(Ops.BINARY, src=(), arg=image).bitcast(buf.dtype))))
   return cached
 
 class AMDAllocator(HCQAllocator['AMDDevice']):
   def __init__(self, dev:AMDDevice):
-    super().__init__(dev, supports_copy_from_disk=dev.has_sdma_queue, supports_transfer=dev.has_sdma_queue and not dev.is_usb())
+    super().__init__(dev, supports_copy_from_disk=dev.has_copy_queue, supports_transfer=dev.has_copy_queue and not dev.is_usb)
 
-  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 _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 _do_free(self, opaque, options:BufferSpec): self.dev.iface.free(opaque)
 
-  def _do_map(self, buf:HCQ2Buffer): return self.dev.iface.map(buf._base if buf._base is not None else buf)
+  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)
 
 @dataclass
 class AMDQueueDesc:
@@ -388,15 +369,24 @@ 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(): return self.alloc(mem.size, host=True, cpu_addr=mem.va_addr)
+    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
 
     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)
@@ -468,6 +458,7 @@ 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
@@ -512,8 +503,7 @@ 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().mv.cast('Q')[0] = \
-        d.signal('value', 1).as_memoryview(force_zero_copy=True, no_sync=True).cast('Q')[0] - 1
+      d.signal('timeline')._buf.cpu_view().view(fmt='Q')[0] = d.signal('value', 1, device="CPU")._buf.cpu_view().view(fmt='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))):
@@ -527,31 +517,50 @@ 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, _mock(KFDIface, "MOCKIface"), _mock(KFDIface), _mock(PCIIface)]
+  ifaces = [KFDIface, PCIIface, USBIface, _mock(KFDIface, "MOCKIface"), _mock(KFDIface), _mock(PCIIface), _mock(USBIface)]
 
   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.device_id = int(device.split(":")[1]) if ":" in device else 0
-
-    self.iface = self._select_iface()
+    self.iface = self._select_iface(device)
+    self.is_usb = isinstance(self.iface, USBIface)
+    if self.is_usb: self.rt_nbytes = 4 << 20
 
     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
@@ -576,12 +585,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_sdma_queue = True # self.sdma_queue(0) is not None, TODO: think of this
+    self.has_copy_queue = not getenv("AMD_DISABLE_SDMA")
 
     super().__init__(device, AMDAllocator(self), [HIPRenderer, AMDLLVMRenderer, HIPCCRenderer], None, can_recover=self.is_am(), arch=self.arch)
 
@@ -589,6 +598,10 @@ 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()
@@ -649,7 +662,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))
 
@@ -657,7 +670,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, 0x200 if self.is_usb() else (16 << 20), idx=idx)
+      self.sdma_queues[idx] = self.create_queue(kfd.KFD_IOC_QUEUE_TYPE_SDMA, 0x2000 if self.is_usb else (16 << 20), idx=idx)
     return self.sdma_queues.get(idx, None)
 
   def tmpring_size(self, private_segment_size):

tinygrad_repo/extra/llama_kernels/__init__.py

@@ -1,8 +1,6 @@
 from __future__ import annotations
 import functools, pathlib
-from dataclasses import replace
-from tinygrad import Tensor, dtypes
-from tinygrad.uop.ops import shape_to_shape_arg
+from tinygrad import Tensor
 from tinygrad.runtime.support.compiler_amd import HIPCCCompiler
 
 FP8_MAX = 448.0
@@ -12,7 +10,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.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
+  inner = Tensor(x.uop.src[0]) if x.uop.axis is not None else x  # the per-shard view of the flat param
   return (inner.abs().max(),)
 
 def local_abs_max(x:Tensor) -> Tensor:

tinygrad_repo/extra/llama_kernels/quantize_fp8_delayed/__init__.py

@@ -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, dtypes.void, (amax_idx, max_val.bitcast(dtypes.int32), max_val, amax_idx.load()), arg=atomic_arg)
+  atomic = UOp(Ops.CUSTOM, src=(amax_idx, max_val.bitcast(dtypes.int32), max_val, amax_idx.load()), arg=(atomic_arg, dtypes.void))
   return atomic.end(tid, wg).sink(arg=KernelInfo(f"quantize_fp8_with_amax_{n_elems}", opts_to_apply=()))
 
 @functools.cache

tinygrad_repo/extra/llama_kernels/quantize_mxfp4/__init__.py

@@ -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, dtypes.void, (o.base.index(0),), arg="") for o in outputs),
+                  *(UOp(Ops.CUSTOM, src=(o.base.index(0),), arg=("", dtypes.void)) 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()

tinygrad_repo/extra/mmapeak/mmapeak.py

@@ -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
+from tinygrad import Tensor, Device, GlobalCounters, Context, dtypes
 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) for x in insts]))))
+    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=(x, dtypes.void)) for x in insts]))))
   dummy = Tensor.zeros(1).contiguous().realize()
   out = Tensor.custom_kernel(dummy, fxn=fxn)[0]
   linear = out.schedule_linear()

tinygrad_repo/extra/models/llama.py

@@ -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)[:(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)
+  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)
 
 # matches meta, non hugging face weights
 # (a+i*b) * (c+i*d) = (ac-bd) + i*(ad+bc)

tinygrad_repo/extra/onnx_helpers.py

@@ -1,10 +1,12 @@
 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={}):
   """

tinygrad_repo/extra/optimization/test_beam_search.py

@@ -89,7 +89,8 @@ 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(opt.op in (OptOps.UPCAST, OptOps.UNROLL) for opt in s.applied_opts)]
+    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)]
     assert len(upcasted) > 0, f"expected upcast/unroll actions after TC with max_up={up}, but got none"
 
   def test_max_up(self):
@@ -98,8 +99,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)
-      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
+      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
 
 if __name__ == '__main__':
   unittest.main()

tinygrad_repo/extra/runbook_digitalocean_mi350x.md

@@ -0,0 +1,197 @@
+# 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.

tinygrad_repo/extra/sqtt/examples/generate_examples.py

@@ -1,4 +1,4 @@
-import os, subprocess, sys, shlex
+import os, subprocess, sys, shlex, pickle
 from pathlib import Path
 from tinygrad.helpers import temp, getenv
 
@@ -23,5 +23,8 @@ 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}")