From b8bc4635632ab5b5015c254c586789e8d6ae295c Mon Sep 17 00:00:00 2001
From: Alex Cheema <alexcheema123@gmail.com>
Date: Sun, 3 May 2026 18:13:12 +0100
Subject: [PATCH 1/4] feat(profilers): add ANE power and profiling

---
 .../src/lib/components/GpuRichBar.svelte      |   17 +-
 .../src/lib/components/TopologyGraph.svelte   |   37 +-
 dashboard/src/lib/stores/app.svelte.ts        |   15 +
 src/exo/api/types/api.py                      |    3 +
 src/exo/shared/apply.py                       |   19 +
 src/exo/shared/types/profiling.py             |   10 +
 src/exo/shared/types/state.py                 |    3 +
 src/exo/utils/info_gatherer/info_gatherer.py  |    2 +
 src/exo/utils/info_gatherer/macmon.py         |    1 +
 src/exo/utils/power_sampler.py                |    4 +
 src/exo/utils/profilers/ane_profiler.py       | 1057 +++++++++++++++++
 src/exo/utils/profilers/profiler_manager.py   |   38 +
 .../profilers/tests/test_ane_profiler.py      |   28 +
 .../profilers/tests/test_apply_profiles.py    |   23 +-
 src/exo/utils/tests/test_macmon_metrics.py    |   31 +
 src/exo/utils/tests/test_power_sampler.py     |   22 +-
 16 files changed, 1295 insertions(+), 15 deletions(-)
 create mode 100644 src/exo/utils/profilers/ane_profiler.py
 create mode 100644 src/exo/utils/profilers/tests/test_ane_profiler.py
 create mode 100644 src/exo/utils/tests/test_macmon_metrics.py

diff --git a/dashboard/src/lib/components/GpuRichBar.svelte b/dashboard/src/lib/components/GpuRichBar.svelte
index 3831617a97..fc397949ea 100644
--- a/dashboard/src/lib/components/GpuRichBar.svelte
+++ b/dashboard/src/lib/components/GpuRichBar.svelte
@@ -1,5 +1,9 @@
 <script lang="ts">
-  import { nodeGpuProfile, topologyData } from "$lib/stores/app.svelte";
+  import {
+    nodeAneProfile,
+    nodeGpuProfile,
+    topologyData,
+  } from "$lib/stores/app.svelte";
 
   interface Props {
     class?: string;
@@ -8,6 +12,7 @@
   let { class: className = "" }: Props = $props();
 
   const profiles = $derived(nodeGpuProfile());
+  const aneProfiles = $derived(nodeAneProfile());
   const topology = $derived(topologyData());
 
   const totalTflops = $derived(
@@ -19,6 +24,12 @@
       0,
     ),
   );
+  const totalAneTflops = $derived(
+    Object.values(aneProfiles).reduce(
+      (sum, p) => sum + (p?.tflopsFp16 ?? 0),
+      0,
+    ),
+  );
   const totalMemoryBytes = $derived(
     Object.values(topology?.nodes ?? {}).reduce(
       (sum, n) => sum + (n.system_info?.memory ?? 0),
@@ -53,6 +64,10 @@
     <span class="stat-value">{formatBandwidth(totalBandwidthGbps)}</span>
     <span class="stat-label">Memory bandwidth</span>
   </div>
+  <div class="stat-block">
+    <span class="stat-value">{formatTflops(totalAneTflops)}</span>
+    <span class="stat-label">ANE compute</span>
+  </div>
   <div class="stat-block">
     <span class="stat-value">{formatMemory(totalMemoryBytes)}</span>
     <span class="stat-label">Memory</span>
diff --git a/dashboard/src/lib/components/TopologyGraph.svelte b/dashboard/src/lib/components/TopologyGraph.svelte
index b5c55f19fe..37bdf15e18 100644
--- a/dashboard/src/lib/components/TopologyGraph.svelte
+++ b/dashboard/src/lib/components/TopologyGraph.svelte
@@ -8,6 +8,7 @@
     nodeThunderboltBridge,
     nodeRdmaCtl,
     nodeIdentities,
+    nodeAneProfile,
     nodeGpuProfile,
     nodeLinkProfiles,
     nodeNetworkRaw,
@@ -44,6 +45,7 @@
   const tbBridgeData = $derived(nodeThunderboltBridge());
   const rdmaCtlData = $derived(nodeRdmaCtl());
   const identitiesData = $derived(nodeIdentities());
+  const aneProfileData = $derived(nodeAneProfile());
   const gpuProfileData = $derived(nodeGpuProfile());
   const linkProfilesData = $derived(nodeLinkProfiles());
   const nodeNetworkData = $derived(nodeNetworkRaw());
@@ -1242,8 +1244,9 @@
           .text(powerText);
       }
 
-      // GPU profile (TFLOPS + memory bandwidth) — only shown when we have a
-      // measurement; otherwise the slot collapses.
+      // Hardware profiles are only shown when we have measurements; otherwise
+      // their slots collapse.
+      const aneProfile = aneProfileData[nodeInfo.id];
       const gpuProfile = gpuProfileData[nodeInfo.id];
 
       // Labels - adapt based on mode
@@ -1317,6 +1320,28 @@
             .attr("fill", "rgba(255,215,0,0.8)")
             .text(`${gpuProfile.memoryBandwidthGbps.toFixed(0)} GB/s`);
         }
+        if (aneProfile) {
+          const aneY = infoY + fontSize * (gpuProfile ? 1.95 : 1.05);
+          const aneText = nodeG
+            .append("text")
+            .attr("x", nodeInfo.x)
+            .attr("y", aneY)
+            .attr("text-anchor", "middle")
+            .attr("font-size", fontSize * 0.8)
+            .attr("font-family", "SF Mono, Monaco, monospace");
+          aneText
+            .append("tspan")
+            .attr("fill", "rgba(96,165,250,0.95)")
+            .text(`ANE ${aneProfile.tflopsFp16.toFixed(1)} TFLOPS`);
+          aneText
+            .append("tspan")
+            .attr("fill", "rgba(179,179,179,0.6)")
+            .text("  ·  ");
+          aneText
+            .append("tspan")
+            .attr("fill", "rgba(255,215,0,0.75)")
+            .text(`${aneProfile.memoryBandwidthGbps.toFixed(0)} GB/s`);
+        }
       } else if (showCompactLabels) {
         // COMPACT MODE: Just name and basic info (4+ nodes)
         const fontSize = Math.max(7, nodeRadius * 0.11);
@@ -1426,10 +1451,15 @@
 
       // Debug mode: Show TB bridge and RDMA status
       if (debugEnabled) {
+        const profileLineCount =
+          (gpuProfile ? 1 : 0) + (aneProfile && showFullLabels ? 1 : 0);
+        const profileDebugOffset =
+          showFullLabels && profileLineCount > 1 ? 12 : 0;
         let debugLabelY =
           nodeInfo.y +
           iconBaseHeight / 2 +
-          (showFullLabels ? 32 : showCompactLabels ? 26 : 22);
+          (showFullLabels ? 32 : showCompactLabels ? 26 : 22) +
+          profileDebugOffset;
         const debugFontSize = showFullLabels ? 9 : 7;
         const debugLineHeight = showFullLabels ? 11 : 9;
 
@@ -1493,6 +1523,7 @@
     const _hoveredNodeId = hoveredNodeId;
     const _filteredNodes = filteredNodes;
     const _highlightedNodes = highlightedNodes;
+    const _ane = aneProfileData;
     const _gpu = gpuProfileData;
     const _links = linkProfilesData;
     const _network = nodeNetworkData;
diff --git a/dashboard/src/lib/stores/app.svelte.ts b/dashboard/src/lib/stores/app.svelte.ts
index 8ee7b4e97c..5a70e5153e 100644
--- a/dashboard/src/lib/stores/app.svelte.ts
+++ b/dashboard/src/lib/stores/app.svelte.ts
@@ -46,6 +46,7 @@ export interface NodeInfo {
     };
     gpu_usage?: [number, number];
     sys_power?: number;
+    ane_power?: number;
   };
   last_macmon_update: number;
   friendly_name?: string;
@@ -100,6 +101,7 @@ interface RawSystemPerformanceProfile {
   gpuUsage?: number;
   temp?: number;
   sysPower?: number;
+  anePower?: number;
   pcpuUsage?: number;
   ecpuUsage?: number;
 }
@@ -133,6 +135,13 @@ export interface RawNodeGpuProfile {
   measuredAt: string;
 }
 
+export interface RawNodeAneProfile {
+  engine: "ane";
+  tflopsFp16: number;
+  memoryBandwidthGbps: number;
+  measuredAt: string;
+}
+
 export interface RawNodeSocketLinkProfile {
   transport: "socket";
   sinkIp: string;
@@ -309,6 +318,8 @@ interface RawStateResponse {
   >;
   // Per-node GPU compute + memory bandwidth profile.
   nodeGpuProfile?: Record<string, RawNodeGpuProfile>;
+  // Per-node ANE compute + streaming bandwidth profile.
+  nodeAneProfile?: Record<string, RawNodeAneProfile>;
   // Per-edge link probe results, keyed source -> sink -> [profiles].
   nodeLinkProfiles?: RawNodeLinkProfiles;
 }
@@ -511,6 +522,7 @@ function transformTopology(
         gpu_usage:
           system?.gpuUsage !== undefined ? [0, system.gpuUsage] : undefined,
         sys_power: system?.sysPower,
+        ane_power: system?.anePower,
       },
       last_macmon_update: Date.now() / 1000,
       friendly_name: identity?.friendlyName,
@@ -634,6 +646,7 @@ class AppStore {
     >
   >({});
   nodeGpuProfile = $state<Record<string, RawNodeGpuProfile>>({});
+  nodeAneProfile = $state<Record<string, RawNodeAneProfile>>({});
   nodeLinkProfiles = $state<RawNodeLinkProfiles>({});
   nodeNetworkRaw = $state<Record<string, RawNodeNetworkInfo>>({});
 
@@ -1406,6 +1419,7 @@ class AppStore {
       this.nodeThunderboltBridge = data.nodeThunderboltBridge ?? {};
       // Profiler outputs
       this.nodeGpuProfile = data.nodeGpuProfile ?? {};
+      this.nodeAneProfile = data.nodeAneProfile ?? {};
       this.nodeLinkProfiles = data.nodeLinkProfiles ?? {};
       // Raw network info — kept so the connection-type inference can use
       // interfaceType, which the topology-shaped `NodeInfo.network_interfaces`
@@ -3671,6 +3685,7 @@ export const nodeThunderboltBridge = () => appStore.nodeThunderboltBridge;
 
 // Profiler outputs
 export const nodeGpuProfile = () => appStore.nodeGpuProfile;
+export const nodeAneProfile = () => appStore.nodeAneProfile;
 export const nodeLinkProfiles = () => appStore.nodeLinkProfiles;
 export const nodeNetworkRaw = () => appStore.nodeNetworkRaw;
 
diff --git a/src/exo/api/types/api.py b/src/exo/api/types/api.py
index 8cfa10dd1a..da9a8a61fc 100644
--- a/src/exo/api/types/api.py
+++ b/src/exo/api/types/api.py
@@ -186,13 +186,16 @@ class NodePowerStats(BaseModel, frozen=True):
     node_id: NodeId
     samples: int
     avg_sys_power: float
+    avg_ane_power: float
 
 
 class PowerUsage(BaseModel, frozen=True):
     elapsed_seconds: float
     nodes: list[NodePowerStats]
     total_avg_sys_power_watts: float
+    total_avg_ane_power_watts: float
     total_energy_joules: float
+    total_ane_energy_joules: float
 
 
 class BenchChatCompletionResponse(ChatCompletionResponse):
diff --git a/src/exo/shared/apply.py b/src/exo/shared/apply.py
index 8abbc5296a..86e94ac00c 100644
--- a/src/exo/shared/apply.py
+++ b/src/exo/shared/apply.py
@@ -33,6 +33,7 @@
 )
 from exo.shared.types.instance_link import InstanceLink, InstanceLinkId
 from exo.shared.types.profiling import (
+    NodeAneProfile,
     NodeGpuProfile,
     NodeIdentity,
     NodeLinkProfile,
@@ -67,6 +68,7 @@
     StaticNodeInformation,
     ThunderboltBridgeInfo,
 )
+from exo.utils.profilers.ane_profiler import AneProfile
 from exo.utils.profilers.gpu_profiler import GpuProfile
 from exo.utils.profilers.link_profiler import RDMALinkProfile, SocketLinkProfile
 
@@ -315,6 +317,11 @@ def apply_node_timed_out(event: NodeTimedOut, state: State) -> State:
         for key, value in state.node_gpu_profile.items()
         if key != event.node_id
     }
+    node_ane_profile = {
+        key: value
+        for key, value in state.node_ane_profile.items()
+        if key != event.node_id
+    }
     # Drop the leaving node both as source (outer key) and as sink (inner key).
     node_link_profiles: dict[NodeId, Mapping[NodeId, Sequence[NodeLinkProfile]]] = {}
     for source_id, sinks in state.node_link_profiles.items():
@@ -348,6 +355,7 @@ def apply_node_timed_out(event: NodeTimedOut, state: State) -> State:
             "node_thunderbolt": node_thunderbolt,
             "node_thunderbolt_bridge": node_thunderbolt_bridge,
             "node_rdma_ctl": node_rdma_ctl,
+            "node_ane_profile": node_ane_profile,
             "node_gpu_profile": node_gpu_profile,
             "node_link_profiles": node_link_profiles,
             "thunderbolt_bridge_cycles": thunderbolt_bridge_cycles,
@@ -467,6 +475,17 @@ def apply_node_gathered_info(event: NodeGatheredInfo, state: State) -> State:
                     measured_at=measured_at,
                 ),
             }
+        case AneProfile():
+            measured_at = datetime.fromisoformat(event.when)
+            update["node_ane_profile"] = {
+                **state.node_ane_profile,
+                event.node_id: NodeAneProfile(
+                    engine=info.engine,
+                    tflops_fp16=info.tflops_fp16,
+                    memory_bandwidth_gbps=info.memory_bandwidth_gbps,
+                    measured_at=measured_at,
+                ),
+            }
         case SocketLinkProfile():
             measured_at = datetime.fromisoformat(event.when)
             new_entry = NodeSocketLinkProfile(
diff --git a/src/exo/shared/types/profiling.py b/src/exo/shared/types/profiling.py
index f5a465e146..fd78624659 100644
--- a/src/exo/shared/types/profiling.py
+++ b/src/exo/shared/types/profiling.py
@@ -63,6 +63,7 @@ class SystemPerformanceProfile(FrozenModel):
     gpu_usage: float = 0.0
     temp: float = 0.0
     sys_power: float = 0.0
+    ane_power: float = 0.0
     pcpu_usage: float = 0.0
     ecpu_usage: float = 0.0
 
@@ -121,6 +122,15 @@ class NodeGpuProfile(FrozenModel):
     measured_at: datetime
 
 
+class NodeAneProfile(FrozenModel):
+    """Measured ANE compute throughput and streaming bandwidth for a node."""
+
+    engine: Literal["ane"]
+    tflops_fp16: float
+    memory_bandwidth_gbps: float
+    measured_at: datetime
+
+
 class NodeSocketLinkProfile(FrozenModel):
     """Per-direction TCP/IP bandwidth + round-trip latency (with jitter).
 
diff --git a/src/exo/shared/types/state.py b/src/exo/shared/types/state.py
index 04b35715c8..d2eb5498dd 100644
--- a/src/exo/shared/types/state.py
+++ b/src/exo/shared/types/state.py
@@ -11,6 +11,7 @@
 from exo.shared.types.profiling import (
     DiskUsage,
     MemoryUsage,
+    NodeAneProfile,
     NodeGpuProfile,
     NodeIdentity,
     NodeLinkProfile,
@@ -63,9 +64,11 @@ class State(FrozenModel):
 
     # Profiler results.
     # node_gpu_profile: per-node GPU TFLOPS / memory bandwidth measurement.
+    # node_ane_profile: per-node ANE TFLOPS / streaming bandwidth measurement.
     # node_link_profiles: per-edge link probe; outer key = source node, inner
     # key = sink node. May contain both a socket and an RDMA profile to the
     # same peer, so the inner value is a sequence.
+    node_ane_profile: Mapping[NodeId, NodeAneProfile] = {}
     node_gpu_profile: Mapping[NodeId, NodeGpuProfile] = {}
     node_link_profiles: Mapping[NodeId, Mapping[NodeId, Sequence[NodeLinkProfile]]] = {}
 
diff --git a/src/exo/utils/info_gatherer/info_gatherer.py b/src/exo/utils/info_gatherer/info_gatherer.py
index a987c22125..760fe750e0 100644
--- a/src/exo/utils/info_gatherer/info_gatherer.py
+++ b/src/exo/utils/info_gatherer/info_gatherer.py
@@ -27,6 +27,7 @@
     ThunderboltIdentifier,
 )
 from exo.utils.channels import Sender
+from exo.utils.profilers.ane_profiler import AneProfile
 from exo.utils.profilers.gpu_profiler import GpuProfile
 from exo.utils.profilers.link_profiler import RDMALinkProfile, SocketLinkProfile
 from exo.utils.pydantic_ext import TaggedModel
@@ -367,6 +368,7 @@ async def _gather_iface_map() -> dict[str, str] | None:
     | MiscData
     | StaticNodeInformation
     | NodeDiskUsage
+    | AneProfile
     | GpuProfile
     | SocketLinkProfile
     | RDMALinkProfile
diff --git a/src/exo/utils/info_gatherer/macmon.py b/src/exo/utils/info_gatherer/macmon.py
index 4fe714e275..81078b645d 100644
--- a/src/exo/utils/info_gatherer/macmon.py
+++ b/src/exo/utils/info_gatherer/macmon.py
@@ -54,6 +54,7 @@ def from_raw(cls, raw: RawMacmonMetrics) -> Self:
                 gpu_usage=raw.gpu_usage[1],
                 temp=raw.temp.gpu_temp_avg,
                 sys_power=raw.sys_power,
+                ane_power=raw.ane_power,
                 pcpu_usage=raw.pcpu_usage[1],
                 ecpu_usage=raw.ecpu_usage[1],
             ),
diff --git a/src/exo/utils/power_sampler.py b/src/exo/utils/power_sampler.py
index b8e985a117..0f1e5c7a58 100644
--- a/src/exo/utils/power_sampler.py
+++ b/src/exo/utils/power_sampler.py
@@ -52,13 +52,17 @@ def result(self) -> PowerUsage:
                     node_id=node_id,
                     samples=n,
                     avg_sys_power=sum(p.sys_power for p in profiles) / n,
+                    avg_ane_power=sum(p.ane_power for p in profiles) / n,
                 )
             )
 
         total_avg_sys = sum(ns.avg_sys_power for ns in node_stats)
+        total_avg_ane = sum(ns.avg_ane_power for ns in node_stats)
         return PowerUsage(
             elapsed_seconds=elapsed,
             nodes=node_stats,
             total_avg_sys_power_watts=total_avg_sys,
+            total_avg_ane_power_watts=total_avg_ane,
             total_energy_joules=total_avg_sys * elapsed,
+            total_ane_energy_joules=total_avg_ane * elapsed,
         )
diff --git a/src/exo/utils/profilers/ane_profiler.py b/src/exo/utils/profilers/ane_profiler.py
new file mode 100644
index 0000000000..68bd9d8d44
--- /dev/null
+++ b/src/exo/utils/profilers/ane_profiler.py
@@ -0,0 +1,1057 @@
+"""Apple Neural Engine compute and streaming bandwidth probe.
+
+The GPU profiler uses MLX workloads because MLX can directly target Metal.
+There is no public ANE Python runtime, so this probe follows the same private
+AppleNeuralEngine.framework path used by the reference ANE inference work:
+build a tiny MIL program, compile it in memory, bind IOSurface inputs/outputs,
+then time `evaluateWithQoS`.
+
+The probe is optional. On non-Darwin hosts, Intel Macs, or machines where the
+private ANE classes are unavailable, `measure()` returns None.
+"""
+
+import ctypes
+import ctypes.util
+import os
+import sys
+import tempfile
+import time
+import uuid
+from collections.abc import Callable, Sequence
+from pathlib import Path
+from typing import Literal, Self, cast, final
+
+import numpy as np
+import numpy.typing as npt
+from anyio import to_thread
+
+from exo.utils.pydantic_ext import TaggedModel
+
+type Pointer = int
+type Float16Array = npt.NDArray[np.float16]
+type UInt8Array = npt.NDArray[np.uint8]
+
+_BYTES_PER_ELEMENT = 2
+
+_COMPUTE_IN_DIM = 2048
+_COMPUTE_OUT_DIM = 2048
+_COMPUTE_SPATIAL = 128
+_COMPUTE_ITERATIONS_PER_PASS = 12
+
+_STREAM_CHANNELS = 4096
+_STREAM_SPATIAL = 2048
+_STREAM_ITERATIONS_PER_PASS = 16
+
+_WARMUP_SECONDS = 0.5
+_MEASUREMENT_PASSES = 5
+_QOS_USER_INITIATED = 21
+
+_objc: ctypes.CDLL | None = None
+_objc_msg_send_address: Pointer | None = None
+_iosurface: ctypes.CDLL | None = None
+
+_k_iosurface_width: Pointer | None = None
+_k_iosurface_height: Pointer | None = None
+_k_iosurface_bytes_per_element: Pointer | None = None
+_k_iosurface_bytes_per_row: Pointer | None = None
+_k_iosurface_alloc_size: Pointer | None = None
+_k_iosurface_pixel_format: Pointer | None = None
+
+_ane_loaded = False
+_ane_descriptor_class: Pointer | None = None
+_ane_in_memory_model_class: Pointer | None = None
+_ane_request_class: Pointer | None = None
+_ane_io_surface_object_class: Pointer | None = None
+
+
+@final
+class AneProfile(TaggedModel):
+    """Wire format for a measured ANE profile, gathered locally on a node."""
+
+    engine: Literal["ane"]
+    tflops_fp16: float
+    memory_bandwidth_gbps: float
+
+    @classmethod
+    async def measure(cls) -> Self | None:
+        if sys.platform != "darwin" or not ane_available():
+            return None
+        return await to_thread.run_sync(cls._measure_blocking)
+
+    @classmethod
+    def _measure_blocking(cls) -> Self:
+        rng = np.random.default_rng(0)
+
+        compute_weights = rng.standard_normal(
+            (_COMPUTE_OUT_DIM, _COMPUTE_IN_DIM)
+        ).astype(np.float16)
+        compute_input = rng.standard_normal((_COMPUTE_IN_DIM, _COMPUTE_SPATIAL)).astype(
+            np.float16
+        )
+
+        stream_input = rng.standard_normal((_STREAM_CHANNELS, _STREAM_SPATIAL)).astype(
+            np.float16
+        )
+
+        compute_kernel = _compile_static_conv_kernel(
+            weights=compute_weights,
+            input_channels=_COMPUTE_IN_DIM,
+            output_channels=_COMPUTE_OUT_DIM,
+            spatial=_COMPUTE_SPATIAL,
+        )
+        compute_kernel.write_input(0, compute_input)
+
+        stream_kernel = _compile_relu_stream_kernel(
+            channels=_STREAM_CHANNELS,
+            spatial=_STREAM_SPATIAL,
+        )
+        stream_kernel.write_input(0, stream_input)
+
+        return cls(
+            engine="ane",
+            tflops_fp16=_measure_compute_tflops(compute_kernel),
+            memory_bandwidth_gbps=_measure_streaming_bandwidth_gbps(stream_kernel),
+        )
+
+
+@final
+class _AneKernel:
+    def __init__(
+        self,
+        *,
+        model: Pointer,
+        request: Pointer,
+        input_surfaces: Sequence[Pointer],
+        output_surfaces: Sequence[Pointer],
+        tmp_dir: Path,
+    ) -> None:
+        self._model = model
+        self._request = request
+        self._input_surfaces = tuple(input_surfaces)
+        self._output_surfaces = tuple(output_surfaces)
+        self._tmp_dir = tmp_dir
+        self._empty_dict = _ns_empty_dict()
+
+    def __del__(self) -> None:
+        try:
+            err_ptr = ctypes.c_void_p(0)
+            unload_fn = ctypes.cast(
+                _require_objc_msg_send_address(),
+                ctypes.CFUNCTYPE(
+                    ctypes.c_bool,
+                    ctypes.c_void_p,
+                    ctypes.c_void_p,
+                    ctypes.c_uint,
+                    ctypes.POINTER(ctypes.c_void_p),
+                ),
+            )
+            unload_fn(
+                ctypes.c_void_p(self._model),
+                ctypes.c_void_p(_sel("unloadWithQoS:error:")),
+                _QOS_USER_INITIATED,
+                ctypes.byref(err_ptr),
+            )
+        except Exception:
+            pass
+
+    def eval(self) -> None:
+        err_ptr = ctypes.c_void_p(0)
+        eval_fn = ctypes.cast(
+            _require_objc_msg_send_address(),
+            ctypes.CFUNCTYPE(
+                ctypes.c_bool,
+                ctypes.c_void_p,
+                ctypes.c_void_p,
+                ctypes.c_uint,
+                ctypes.c_void_p,
+                ctypes.c_void_p,
+                ctypes.POINTER(ctypes.c_void_p),
+            ),
+        )
+        ok = cast(
+            bool,
+            eval_fn(
+                ctypes.c_void_p(self._model),
+                ctypes.c_void_p(_sel("evaluateWithQoS:options:request:error:")),
+                _QOS_USER_INITIATED,
+                ctypes.c_void_p(self._empty_dict),
+                ctypes.c_void_p(self._request),
+                ctypes.byref(err_ptr),
+            ),
+        )
+        if not ok:
+            err_desc = _describe_error(err_ptr.value)
+            raise RuntimeError(f"ANE eval failed: {err_desc}")
+
+    def write_input(self, index: int, values: Float16Array) -> None:
+        _write_iosurface(self._input_surfaces[index], values)
+
+
+def ane_available() -> bool:
+    try:
+        _load_ane_framework()
+        return True
+    except Exception:
+        return False
+
+
+def _measure_compute_tflops(kernel: _AneKernel) -> float:
+    _warm_up_with(kernel.eval, time.perf_counter() + _WARMUP_SECONDS)
+
+    flops_per_iteration = 2 * _COMPUTE_IN_DIM * _COMPUTE_OUT_DIM * _COMPUTE_SPATIAL
+    best_tflops = 0.0
+    for _ in range(_MEASUREMENT_PASSES):
+        start = time.perf_counter()
+        for _ in range(_COMPUTE_ITERATIONS_PER_PASS):
+            kernel.eval()
+        elapsed = time.perf_counter() - start
+        if elapsed <= 0:
+            continue
+        tflops = flops_per_iteration * _COMPUTE_ITERATIONS_PER_PASS / elapsed / 1e12
+        best_tflops = max(best_tflops, tflops)
+    return best_tflops
+
+
+def _measure_streaming_bandwidth_gbps(kernel: _AneKernel) -> float:
+    _warm_up_with(kernel.eval, time.perf_counter() + _WARMUP_SECONDS)
+
+    bytes_per_iteration = _STREAM_CHANNELS * _STREAM_SPATIAL * _BYTES_PER_ELEMENT * 2
+    best_gbps = 0.0
+    for _ in range(_MEASUREMENT_PASSES):
+        start = time.perf_counter()
+        for _ in range(_STREAM_ITERATIONS_PER_PASS):
+            kernel.eval()
+        elapsed = time.perf_counter() - start
+        if elapsed <= 0:
+            continue
+        gbps = bytes_per_iteration * _STREAM_ITERATIONS_PER_PASS / elapsed / 1e9
+        best_gbps = max(best_gbps, gbps)
+    return best_gbps
+
+
+def _warm_up_with(do_op: Callable[[], None], deadline: float) -> None:
+    while time.perf_counter() < deadline:
+        do_op()
+
+
+def _compile_static_conv_kernel(
+    *,
+    weights: Float16Array,
+    input_channels: int,
+    output_channels: int,
+    spatial: int,
+) -> _AneKernel:
+    mil = _mil_static_conv(
+        input_channels=input_channels,
+        output_channels=output_channels,
+        spatial=spatial,
+    )
+    weight_blob = _build_weight_blob_bytes(weights)
+    input_bytes = input_channels * spatial * _BYTES_PER_ELEMENT
+    output_bytes = output_channels * spatial * _BYTES_PER_ELEMENT
+    return _compile_kernel(
+        mil_text=mil,
+        weight_files={"weight.bin": weight_blob},
+        input_sizes=(input_bytes,),
+        output_sizes=(output_bytes,),
+    )
+
+
+def _compile_relu_stream_kernel(*, channels: int, spatial: int) -> _AneKernel:
+    element_count = channels * spatial
+    data_bytes = element_count * _BYTES_PER_ELEMENT
+    return _compile_kernel(
+        mil_text=_mil_relu_stream(channels=channels, spatial=spatial),
+        weight_files={},
+        input_sizes=(data_bytes,),
+        output_sizes=(data_bytes,),
+    )
+
+
+def _compile_kernel(
+    *,
+    mil_text: str,
+    weight_files: dict[str, bytes],
+    input_sizes: Sequence[int],
+    output_sizes: Sequence[int],
+) -> _AneKernel:
+    (
+        descriptor_class,
+        in_memory_model_class,
+        request_class,
+        io_surface_object_class,
+    ) = _load_ane_framework()
+
+    mil_bytes = mil_text.encode("utf-8")
+    mil_data = _ns_data(mil_bytes)
+
+    weight_keys: list[Pointer] = []
+    weight_values: list[Pointer] = []
+    for filename, raw_bytes in weight_files.items():
+        entry = _ns_dict_from_keys_values(
+            [_ns_str("offset"), _ns_str("data")],
+            [_ns_int(0), _ns_data(raw_bytes)],
+        )
+        weight_keys.append(_ns_str(f"@model_path/weights/{filename}"))
+        weight_values.append(entry)
+    weights_dict = (
+        _ns_dict_from_keys_values(weight_keys, weight_values)
+        if weight_keys
+        else _ns_empty_dict()
+    )
+
+    descriptor_fn = ctypes.cast(
+        _require_objc_msg_send_address(),
+        ctypes.CFUNCTYPE(
+            ctypes.c_void_p,
+            ctypes.c_void_p,
+            ctypes.c_void_p,
+            ctypes.c_void_p,
+            ctypes.c_void_p,
+            ctypes.c_void_p,
+        ),
+    )
+    descriptor = _checked_pointer(
+        cast(
+            int | None,
+            descriptor_fn(
+                ctypes.c_void_p(descriptor_class),
+                ctypes.c_void_p(_sel("modelWithMILText:weights:optionsPlist:")),
+                ctypes.c_void_p(mil_data),
+                ctypes.c_void_p(weights_dict),
+                None,
+            ),
+        ),
+        "ANE: modelWithMILText failed",
+    )
+
+    model = _id_id(
+        in_memory_model_class,
+        _sel("inMemoryModelWithDescriptor:"),
+        descriptor,
+        "ANE: inMemoryModelWithDescriptor failed",
+    )
+
+    tmp_dir = _write_model_files(model, mil_bytes, weight_files)
+
+    _compile_and_load_model(model)
+
+    input_surfaces = tuple(_create_iosurface(size) for size in input_sizes)
+    output_surfaces = tuple(_create_iosurface(size) for size in output_sizes)
+    for surface in (*input_surfaces, *output_surfaces):
+        _zero_iosurface(surface)
+
+    wrapped_inputs = _ns_mutable_array(len(input_surfaces))
+    input_indices = _ns_mutable_array(len(input_surfaces))
+    for index, surface in enumerate(input_surfaces):
+        io_object = _id_id(
+            io_surface_object_class,
+            _sel("objectWithIOSurface:"),
+            surface,
+            "ANE: objectWithIOSurface failed for input",
+        )
+        _ns_array_add(wrapped_inputs, io_object)
+        _ns_array_add(input_indices, _ns_int(index))
+
+    wrapped_outputs = _ns_mutable_array(len(output_surfaces))
+    output_indices = _ns_mutable_array(len(output_surfaces))
+    for index, surface in enumerate(output_surfaces):
+        io_object = _id_id(
+            io_surface_object_class,
+            _sel("objectWithIOSurface:"),
+            surface,
+            "ANE: objectWithIOSurface failed for output",
+        )
+        _ns_array_add(wrapped_outputs, io_object)
+        _ns_array_add(output_indices, _ns_int(index))
+
+    request_fn = ctypes.cast(
+        _require_objc_msg_send_address(),
+        ctypes.CFUNCTYPE(
+            ctypes.c_void_p,
+            ctypes.c_void_p,
+            ctypes.c_void_p,
+            ctypes.c_void_p,
+            ctypes.c_void_p,
+            ctypes.c_void_p,
+            ctypes.c_void_p,
+            ctypes.c_void_p,
+            ctypes.c_void_p,
+            ctypes.c_void_p,
+        ),
+    )
+    request = _checked_pointer(
+        cast(
+            int | None,
+            request_fn(
+                ctypes.c_void_p(request_class),
+                ctypes.c_void_p(
+                    _sel(
+                        "requestWithInputs:inputIndices:outputs:outputIndices:"
+                        "weightsBuffer:perfStats:procedureIndex:"
+                    )
+                ),
+                ctypes.c_void_p(wrapped_inputs),
+                ctypes.c_void_p(input_indices),
+                ctypes.c_void_p(wrapped_outputs),
+                ctypes.c_void_p(output_indices),
+                None,
+                None,
+                ctypes.c_void_p(_ns_int(0)),
+            ),
+        ),
+        "ANE: requestWithInputs failed",
+    )
+
+    _retain(model)
+    _retain(request)
+    return _AneKernel(
+        model=model,
+        request=request,
+        input_surfaces=input_surfaces,
+        output_surfaces=output_surfaces,
+        tmp_dir=tmp_dir,
+    )
+
+
+def _compile_and_load_model(model: Pointer) -> None:
+    for selector, options in (
+        ("compileWithQoS:options:error:", _ns_empty_dict()),
+        ("loadWithQoS:options:error:", _load_options()),
+    ):
+        err_ptr = ctypes.c_void_p(0)
+        compile_fn = ctypes.cast(
+            _require_objc_msg_send_address(),
+            ctypes.CFUNCTYPE(
+                ctypes.c_bool,
+                ctypes.c_void_p,
+                ctypes.c_void_p,
+                ctypes.c_uint,
+                ctypes.c_void_p,
+                ctypes.POINTER(ctypes.c_void_p),
+            ),
+        )
+        ok = cast(
+            bool,
+            compile_fn(
+                ctypes.c_void_p(model),
+                ctypes.c_void_p(_sel(selector)),
+                _QOS_USER_INITIATED,
+                ctypes.c_void_p(options),
+                ctypes.byref(err_ptr),
+            ),
+        )
+        if not ok:
+            err_desc = _describe_error(err_ptr.value)
+            raise RuntimeError(f"ANE {selector} failed: {err_desc}")
+
+
+def _load_options() -> Pointer:
+    if not os.environ.get("EXO_ANE_KEEP_MODEL_WIRED"):
+        return _ns_empty_dict()
+
+    number_with_bool = ctypes.cast(
+        _require_objc_msg_send_address(),
+        ctypes.CFUNCTYPE(
+            ctypes.c_void_p,
+            ctypes.c_void_p,
+            ctypes.c_void_p,
+            ctypes.c_bool,
+        ),
+    )
+    yes = _checked_pointer(
+        cast(
+            int | None,
+            number_with_bool(
+                ctypes.c_void_p(_objc_class("NSNumber")),
+                ctypes.c_void_p(_sel("numberWithBool:")),
+                True,
+            ),
+        ),
+        "NSNumber numberWithBool failed",
+    )
+    return _ns_dict_from_keys_values([_ns_str("kANEFKeepModelMemoryWiredKey")], [yes])
+
+
+def _write_model_files(
+    model: Pointer, mil_bytes: bytes, weight_files: dict[str, bytes]
+) -> Path:
+    hex_id = _id_none(model, _sel("hexStringIdentifier"))
+    model_id = _to_cstr(hex_id) if _ns_string_length(hex_id) > 0 else str(uuid.uuid4())
+    tmp_dir = Path(tempfile.gettempdir()) / model_id
+    weights_dir = tmp_dir / "weights"
+    weights_dir.mkdir(parents=True, exist_ok=True)
+    (tmp_dir / "model.mil").write_bytes(mil_bytes)
+    for filename, raw_bytes in weight_files.items():
+        (weights_dir / filename).write_bytes(raw_bytes)
+    return tmp_dir
+
+
+def _write_iosurface(surface: Pointer, values: Float16Array) -> None:
+    raw = cast(UInt8Array, np.ascontiguousarray(values).view(np.uint8).reshape(-1))
+    nbytes = int(raw.nbytes)
+    if nbytes > _iosurface_alloc_size(surface):
+        raise ValueError("input data is larger than the IOSurface allocation")
+
+    _iosurface_lock(surface, readonly=False)
+    try:
+        address = _iosurface_base_address(surface)
+        ctypes.memmove(address, int(raw.ctypes.data), nbytes)
+    finally:
+        _iosurface_unlock(surface, readonly=False)
+
+
+def _zero_iosurface(surface: Pointer) -> None:
+    _iosurface_lock(surface, readonly=False)
+    try:
+        ctypes.memset(
+            _iosurface_base_address(surface), 0, _iosurface_alloc_size(surface)
+        )
+    finally:
+        _iosurface_unlock(surface, readonly=False)
+
+
+def _create_iosurface(nbytes: int) -> Pointer:
+    if nbytes <= 0:
+        nbytes = 4
+    keys = [
+        _require_iosurface_key(_k_iosurface_width, "kIOSurfaceWidth"),
+        _require_iosurface_key(_k_iosurface_height, "kIOSurfaceHeight"),
+        _require_iosurface_key(
+            _k_iosurface_bytes_per_element, "kIOSurfaceBytesPerElement"
+        ),
+        _require_iosurface_key(_k_iosurface_bytes_per_row, "kIOSurfaceBytesPerRow"),
+        _require_iosurface_key(_k_iosurface_alloc_size, "kIOSurfaceAllocSize"),
+        _require_iosurface_key(_k_iosurface_pixel_format, "kIOSurfacePixelFormat"),
+    ]
+    values = [
+        _ns_ulong(nbytes),
+        _ns_int(1),
+        _ns_int(1),
+        _ns_ulong(nbytes),
+        _ns_ulong(nbytes),
+        _ns_int(0),
+    ]
+    dictionary = _ns_dict_from_keys_values(keys, values)
+    surface = _checked_pointer(
+        cast(
+            int | None,
+            _require_iosurface().IOSurfaceCreate(ctypes.c_void_p(dictionary)),
+        ),
+        "IOSurfaceCreate failed",
+    )
+    return surface
+
+
+def _iosurface_lock(surface: Pointer, *, readonly: bool) -> None:
+    flags = 1 if readonly else 0
+    result = cast(
+        int,
+        _require_iosurface().IOSurfaceLock(
+            ctypes.c_void_p(surface), ctypes.c_uint32(flags), None
+        ),
+    )
+    if result != 0:
+        raise RuntimeError(f"IOSurfaceLock failed: {result}")
+
+
+def _iosurface_unlock(surface: Pointer, *, readonly: bool) -> None:
+    flags = 1 if readonly else 0
+    result = cast(
+        int,
+        _require_iosurface().IOSurfaceUnlock(
+            ctypes.c_void_p(surface), ctypes.c_uint32(flags), None
+        ),
+    )
+    if result != 0:
+        raise RuntimeError(f"IOSurfaceUnlock failed: {result}")
+
+
+def _iosurface_base_address(surface: Pointer) -> Pointer:
+    return _checked_pointer(
+        cast(
+            int | None,
+            _require_iosurface().IOSurfaceGetBaseAddress(ctypes.c_void_p(surface)),
+        ),
+        "IOSurfaceGetBaseAddress failed",
+    )
+
+
+def _iosurface_alloc_size(surface: Pointer) -> int:
+    return cast(
+        int,
+        _require_iosurface().IOSurfaceGetAllocSize(ctypes.c_void_p(surface)),
+    )
+
+
+def _build_weight_blob_bytes(values: Float16Array) -> bytes:
+    flat = np.ascontiguousarray(values.reshape(-1), dtype=np.float16)
+    payload = flat.tobytes()
+    header_bytes = 128
+    buf = bytearray(header_bytes + len(payload))
+    buf[0] = 0x01
+    buf[4] = 0x02
+    _pack_uint32_le(buf, 64, 0xDEADBEEF)
+    buf[68] = 0x01
+    _pack_uint32_le(buf, 72, len(payload))
+    _pack_uint32_le(buf, 80, header_bytes)
+    buf[header_bytes:] = payload
+    return bytes(buf)
+
+
+def _pack_uint32_le(buf: bytearray, offset: int, value: int) -> None:
+    buf[offset : offset + 4] = value.to_bytes(4, byteorder="little", signed=False)
+
+
+def _load_ane_framework() -> tuple[Pointer, Pointer, Pointer, Pointer]:
+    global _ane_loaded
+    global _ane_descriptor_class
+    global _ane_in_memory_model_class
+    global _ane_request_class
+    global _ane_io_surface_object_class
+
+    if _ane_loaded:
+        return (
+            _require_pointer(_ane_descriptor_class, "_ANEInMemoryModelDescriptor"),
+            _require_pointer(_ane_in_memory_model_class, "_ANEInMemoryModel"),
+            _require_pointer(_ane_request_class, "_ANERequest"),
+            _require_pointer(_ane_io_surface_object_class, "_ANEIOSurfaceObject"),
+        )
+
+    _load_objc_runtime()
+    _load_iosurface_runtime()
+    ctypes.cdll.LoadLibrary(
+        "/System/Library/PrivateFrameworks/AppleNeuralEngine.framework/"
+        "AppleNeuralEngine"
+    )
+
+    _ane_descriptor_class = _objc_class("_ANEInMemoryModelDescriptor")
+    _ane_in_memory_model_class = _objc_class("_ANEInMemoryModel")
+    _ane_request_class = _objc_class("_ANERequest")
+    _ane_io_surface_object_class = _objc_class("_ANEIOSurfaceObject")
+    _ane_loaded = True
+    return (
+        _ane_descriptor_class,
+        _ane_in_memory_model_class,
+        _ane_request_class,
+        _ane_io_surface_object_class,
+    )
+
+
+def _load_objc_runtime() -> None:
+    global _objc
+    global _objc_msg_send_address
+
+    if _objc is not None and _objc_msg_send_address is not None:
+        return
+    library_name = ctypes.util.find_library("objc")
+    if library_name is None:
+        raise RuntimeError("objc runtime not found")
+
+    objc: ctypes.CDLL = ctypes.cdll.LoadLibrary(library_name)
+    objc.objc_getClass.restype = ctypes.c_void_p
+    objc.objc_getClass.argtypes = [ctypes.c_char_p]
+    objc.sel_registerName.restype = ctypes.c_void_p
+    objc.sel_registerName.argtypes = [ctypes.c_char_p]
+
+    _objc = objc
+    _objc_msg_send_address = _checked_pointer(
+        ctypes.cast(objc.objc_msgSend, ctypes.c_void_p).value,
+        "objc_msgSend not found",
+    )
+
+
+def _load_iosurface_runtime() -> None:
+    global _iosurface
+    global _k_iosurface_width
+    global _k_iosurface_height
+    global _k_iosurface_bytes_per_element
+    global _k_iosurface_bytes_per_row
+    global _k_iosurface_alloc_size
+    global _k_iosurface_pixel_format
+
+    if _iosurface is not None:
+        return
+
+    iosurface: ctypes.CDLL = ctypes.cdll.LoadLibrary(
+        "/System/Library/Frameworks/IOSurface.framework/IOSurface"
+    )
+    iosurface.IOSurfaceCreate.restype = ctypes.c_void_p
+    iosurface.IOSurfaceCreate.argtypes = [ctypes.c_void_p]
+    iosurface.IOSurfaceLock.restype = ctypes.c_int
+    iosurface.IOSurfaceLock.argtypes = [
+        ctypes.c_void_p,
+        ctypes.c_uint32,
+        ctypes.c_void_p,
+    ]
+    iosurface.IOSurfaceUnlock.restype = ctypes.c_int
+    iosurface.IOSurfaceUnlock.argtypes = [
+        ctypes.c_void_p,
+        ctypes.c_uint32,
+        ctypes.c_void_p,
+    ]
+    iosurface.IOSurfaceGetBaseAddress.restype = ctypes.c_void_p
+    iosurface.IOSurfaceGetBaseAddress.argtypes = [ctypes.c_void_p]
+    iosurface.IOSurfaceGetAllocSize.restype = ctypes.c_size_t
+    iosurface.IOSurfaceGetAllocSize.argtypes = [ctypes.c_void_p]
+
+    _iosurface = iosurface
+    _k_iosurface_width = _checked_pointer(
+        ctypes.c_void_p.in_dll(iosurface, "kIOSurfaceWidth").value,
+        "kIOSurfaceWidth not found",
+    )
+    _k_iosurface_height = _checked_pointer(
+        ctypes.c_void_p.in_dll(iosurface, "kIOSurfaceHeight").value,
+        "kIOSurfaceHeight not found",
+    )
+    _k_iosurface_bytes_per_element = _checked_pointer(
+        ctypes.c_void_p.in_dll(iosurface, "kIOSurfaceBytesPerElement").value,
+        "kIOSurfaceBytesPerElement not found",
+    )
+    _k_iosurface_bytes_per_row = _checked_pointer(
+        ctypes.c_void_p.in_dll(iosurface, "kIOSurfaceBytesPerRow").value,
+        "kIOSurfaceBytesPerRow not found",
+    )
+    _k_iosurface_alloc_size = _checked_pointer(
+        ctypes.c_void_p.in_dll(iosurface, "kIOSurfaceAllocSize").value,
+        "kIOSurfaceAllocSize not found",
+    )
+    _k_iosurface_pixel_format = _checked_pointer(
+        ctypes.c_void_p.in_dll(iosurface, "kIOSurfacePixelFormat").value,
+        "kIOSurfacePixelFormat not found",
+    )
+
+
+def _require_objc() -> ctypes.CDLL:
+    if _objc is None:
+        raise RuntimeError("objc runtime was not loaded")
+    return _objc
+
+
+def _require_iosurface() -> ctypes.CDLL:
+    if _iosurface is None:
+        raise RuntimeError("IOSurface runtime was not loaded")
+    return _iosurface
+
+
+def _require_objc_msg_send_address() -> Pointer:
+    return _require_pointer(_objc_msg_send_address, "objc_msgSend")
+
+
+def _require_iosurface_key(value: Pointer | None, name: str) -> Pointer:
+    return _require_pointer(value, name)
+
+
+def _require_pointer(value: Pointer | None, name: str) -> Pointer:
+    if value is None or value == 0:
+        raise RuntimeError(f"{name} is unavailable")
+    return value
+
+
+def _checked_pointer(value: object, context: str) -> Pointer:
+    pointer = cast(int | None, value)
+    if pointer is None or pointer == 0:
+        raise RuntimeError(context)
+    return pointer
+
+
+def _objc_class(name: str) -> Pointer:
+    return _checked_pointer(
+        cast(int | None, _require_objc().objc_getClass(name.encode("utf-8"))),
+        f"Objective-C class {name} not found",
+    )
+
+
+def _sel(name: str) -> Pointer:
+    return _checked_pointer(
+        cast(int | None, _require_objc().sel_registerName(name.encode("utf-8"))),
+        f"Objective-C selector {name} not found",
+    )
+
+
+def _id_none(obj: Pointer, selector: Pointer) -> Pointer:
+    fn = ctypes.cast(
+        _require_objc_msg_send_address(),
+        ctypes.CFUNCTYPE(ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p),
+    )
+    return _checked_pointer(
+        cast(int | None, fn(ctypes.c_void_p(obj), ctypes.c_void_p(selector))),
+        "Objective-C call failed",
+    )
+
+
+def _id_id(obj: Pointer, selector: Pointer, arg: Pointer, context: str) -> Pointer:
+    fn = ctypes.cast(
+        _require_objc_msg_send_address(),
+        ctypes.CFUNCTYPE(
+            ctypes.c_void_p,
+            ctypes.c_void_p,
+            ctypes.c_void_p,
+            ctypes.c_void_p,
+        ),
+    )
+    return _checked_pointer(
+        cast(
+            int | None,
+            fn(ctypes.c_void_p(obj), ctypes.c_void_p(selector), ctypes.c_void_p(arg)),
+        ),
+        context,
+    )
+
+
+def _ns_str(value: str) -> Pointer:
+    fn = ctypes.cast(
+        _require_objc_msg_send_address(),
+        ctypes.CFUNCTYPE(
+            ctypes.c_void_p,
+            ctypes.c_void_p,
+            ctypes.c_void_p,
+            ctypes.c_char_p,
+        ),
+    )
+    return _checked_pointer(
+        cast(
+            int | None,
+            fn(
+                ctypes.c_void_p(_objc_class("NSString")),
+                ctypes.c_void_p(_sel("stringWithUTF8String:")),
+                value.encode("utf-8"),
+            ),
+        ),
+        "NSString stringWithUTF8String failed",
+    )
+
+
+def _ns_int(value: int) -> Pointer:
+    fn = ctypes.cast(
+        _require_objc_msg_send_address(),
+        ctypes.CFUNCTYPE(
+            ctypes.c_void_p,
+            ctypes.c_void_p,
+            ctypes.c_void_p,
+            ctypes.c_int,
+        ),
+    )
+    return _checked_pointer(
+        cast(
+            int | None,
+            fn(
+                ctypes.c_void_p(_objc_class("NSNumber")),
+                ctypes.c_void_p(_sel("numberWithInt:")),
+                value,
+            ),
+        ),
+        "NSNumber numberWithInt failed",
+    )
+
+
+def _ns_ulong(value: int) -> Pointer:
+    fn = ctypes.cast(
+        _require_objc_msg_send_address(),
+        ctypes.CFUNCTYPE(
+            ctypes.c_void_p,
+            ctypes.c_void_p,
+            ctypes.c_void_p,
+            ctypes.c_ulong,
+        ),
+    )
+    return _checked_pointer(
+        cast(
+            int | None,
+            fn(
+                ctypes.c_void_p(_objc_class("NSNumber")),
+                ctypes.c_void_p(_sel("numberWithUnsignedLong:")),
+                value,
+            ),
+        ),
+        "NSNumber numberWithUnsignedLong failed",
+    )
+
+
+def _ns_data(value: bytes) -> Pointer:
+    fn = ctypes.cast(
+        _require_objc_msg_send_address(),
+        ctypes.CFUNCTYPE(
+            ctypes.c_void_p,
+            ctypes.c_void_p,
+            ctypes.c_void_p,
+            ctypes.c_char_p,
+            ctypes.c_ulong,
+        ),
+    )
+    return _checked_pointer(
+        cast(
+            int | None,
+            fn(
+                ctypes.c_void_p(_objc_class("NSData")),
+                ctypes.c_void_p(_sel("dataWithBytes:length:")),
+                value,
+                len(value),
+            ),
+        ),
+        "NSData dataWithBytes failed",
+    )
+
+
+def _ns_empty_dict() -> Pointer:
+    return _id_none(_objc_class("NSDictionary"), _sel("dictionary"))
+
+
+def _ns_mutable_array(capacity: int) -> Pointer:
+    fn = ctypes.cast(
+        _require_objc_msg_send_address(),
+        ctypes.CFUNCTYPE(
+            ctypes.c_void_p,
+            ctypes.c_void_p,
+            ctypes.c_void_p,
+            ctypes.c_ulong,
+        ),
+    )
+    return _checked_pointer(
+        cast(
+            int | None,
+            fn(
+                ctypes.c_void_p(_objc_class("NSMutableArray")),
+                ctypes.c_void_p(_sel("arrayWithCapacity:")),
+                capacity,
+            ),
+        ),
+        "NSMutableArray arrayWithCapacity failed",
+    )
+
+
+def _ns_array_add(array: Pointer, obj: Pointer) -> None:
+    fn = ctypes.cast(
+        _require_objc_msg_send_address(),
+        ctypes.CFUNCTYPE(
+            ctypes.c_void_p,
+            ctypes.c_void_p,
+            ctypes.c_void_p,
+            ctypes.c_void_p,
+        ),
+    )
+    fn(
+        ctypes.c_void_p(array),
+        ctypes.c_void_p(_sel("addObject:")),
+        ctypes.c_void_p(obj),
+    )
+
+
+def _ns_dict_from_keys_values(
+    keys: Sequence[Pointer], values: Sequence[Pointer]
+) -> Pointer:
+    if len(keys) != len(values):
+        raise ValueError("keys and values must have the same length")
+    if not keys:
+        return _ns_empty_dict()
+
+    count = len(keys)
+    key_array_type = ctypes.c_void_p * count
+    value_array_type = ctypes.c_void_p * count
+    fn = ctypes.cast(
+        _require_objc_msg_send_address(),
+        ctypes.CFUNCTYPE(
+            ctypes.c_void_p,
+            ctypes.c_void_p,
+            ctypes.c_void_p,
+            ctypes.POINTER(ctypes.c_void_p),
+            ctypes.POINTER(ctypes.c_void_p),
+            ctypes.c_ulong,
+        ),
+    )
+    return _checked_pointer(
+        cast(
+            int | None,
+            fn(
+                ctypes.c_void_p(_objc_class("NSDictionary")),
+                ctypes.c_void_p(_sel("dictionaryWithObjects:forKeys:count:")),
+                value_array_type(*values),
+                key_array_type(*keys),
+                count,
+            ),
+        ),
+        "NSDictionary dictionaryWithObjects failed",
+    )
+
+
+def _retain(obj: Pointer) -> None:
+    _id_none(obj, _sel("retain"))
+
+
+def _to_cstr(ns_string: Pointer) -> str:
+    fn = ctypes.cast(
+        _require_objc_msg_send_address(),
+        ctypes.CFUNCTYPE(ctypes.c_char_p, ctypes.c_void_p, ctypes.c_void_p),
+    )
+    raw = cast(
+        bytes | None,
+        fn(ctypes.c_void_p(ns_string), ctypes.c_void_p(_sel("UTF8String"))),
+    )
+    if raw is None:
+        return ""
+    return raw.decode("utf-8", errors="replace")
+
+
+def _ns_string_length(ns_string: Pointer) -> int:
+    fn = ctypes.cast(
+        _require_objc_msg_send_address(),
+        ctypes.CFUNCTYPE(ctypes.c_ulong, ctypes.c_void_p, ctypes.c_void_p),
+    )
+    return cast(int, fn(ctypes.c_void_p(ns_string), ctypes.c_void_p(_sel("length"))))
+
+
+def _describe_error(error_pointer: int | None) -> str:
+    if error_pointer is None or error_pointer == 0:
+        return ""
+
+    description = _id_none(error_pointer, _sel("description"))
+    return _to_cstr(description)
+
+
+_MIL_HEADER = (
+    "program(1.3)\n"
+    '[buildInfo = dict<string, string>({{"coremlc-component-MIL", "3510.2.1"}, '
+    '{"coremlc-version", "3505.4.1"}, {"coremltools-component-milinternal", ""}, '
+    '{"coremltools-version", "9.0"}})]\n'
+    "{\n"
+)
+
+_CONV_PARAMS = (
+    '        string pt = const()[name=string("pt"), val=string("valid")];\n'
+    '        tensor<int32, [2]> st = const()[name=string("st"), '
+    "val=tensor<int32, [2]>([1,1])];\n"
+    '        tensor<int32, [4]> pd = const()[name=string("pd"), '
+    "val=tensor<int32, [4]>([0,0,0,0])];\n"
+    '        tensor<int32, [2]> dl = const()[name=string("dl"), '
+    "val=tensor<int32, [2]>([1,1])];\n"
+    '        int32 gr = const()[name=string("gr"), val=int32(1)];\n'
+)
+
+_CONV_ARGS = "dilations=dl, groups=gr, pad=pd, pad_type=pt, strides=st"
+
+
+def _mil_static_conv(*, input_channels: int, output_channels: int, spatial: int) -> str:
+    return (
+        f"{_MIL_HEADER}"
+        f"    func main<ios18>(tensor<fp16, [1, {input_channels}, 1, {spatial}]> x) {{\n"
+        f"        tensor<fp16, [{output_channels}, {input_channels}, 1, 1]> W = const()"
+        f'[name=string("W"), val=tensor<fp16, [{output_channels}, {input_channels}, 1, 1]>'
+        '(BLOBFILE(path=string("@model_path/weights/weight.bin"), '
+        "offset=uint64(64)))];\n"
+        f"{_CONV_PARAMS}"
+        f"        tensor<fp16, [1, {output_channels}, 1, {spatial}]> y = conv("
+        f'{_CONV_ARGS}, weight=W, x=x)[name=string("cv")];\n'
+        "    } -> (y);\n"
+        "}\n"
+    )
+
+
+def _mil_relu_stream(*, channels: int, spatial: int) -> str:
+    return (
+        f"{_MIL_HEADER}"
+        f"    func main<ios18>(tensor<fp16, [1, {channels}, 1, {spatial}]> x) {{\n"
+        f"        tensor<fp16, [1, {channels}, 1, {spatial}]> y = relu(x=x)"
+        '[name=string("relu")];\n'
+        "    } -> (y);\n"
+        "}\n"
+    )
diff --git a/src/exo/utils/profilers/profiler_manager.py b/src/exo/utils/profilers/profiler_manager.py
index 0ac20e7aaa..c616d23b71 100644
--- a/src/exo/utils/profilers/profiler_manager.py
+++ b/src/exo/utils/profilers/profiler_manager.py
@@ -36,6 +36,7 @@
 )
 from exo.utils.channels import Sender
 from exo.utils.info_gatherer.info_gatherer import GatheredInfo
+from exo.utils.profilers.ane_profiler import AneProfile
 from exo.utils.profilers.gpu_profiler import GpuProfile
 from exo.utils.profilers.link_profiler import (
     PROBE_TIMEOUT_SECONDS,
@@ -46,10 +47,12 @@
 from exo.utils.task_group import TaskGroup
 
 GPU_TTL = timedelta(hours=1)
+ANE_TTL = timedelta(hours=1)
 SOCKET_LINK_TTL = timedelta(minutes=5)
 RDMA_LINK_TTL = timedelta(hours=6)
 RECONCILE_TICK_SECONDS = 15.0
 GPU_PROBE_HARD_TIMEOUT_SECONDS = 60.0
+ANE_PROBE_HARD_TIMEOUT_SECONDS = 120.0
 SOCKET_PROBE_HARD_TIMEOUT_SECONDS = 30.0
 RDMA_PROBE_HARD_TIMEOUT_SECONDS = 90.0
 
@@ -69,11 +72,13 @@ class ProfilerManager:
     state_view: Callable[[], State]
     _tg: TaskGroup = field(init=False, default_factory=TaskGroup)
     _gpu_in_flight: bool = field(init=False, default=False)
+    _ane_in_flight: bool = field(init=False, default=False)
     _link_in_flight: set[LinkKey] = field(init=False, default_factory=set)
 
     async def run(self) -> None:
         async with self._tg as tg:
             tg.start_soon(self._reconcile_gpu, RECONCILE_TICK_SECONDS)
+            tg.start_soon(self._reconcile_ane, RECONCILE_TICK_SECONDS)
             tg.start_soon(self._reconcile_links, RECONCILE_TICK_SECONDS)
 
     def shutdown(self) -> None:
@@ -112,6 +117,39 @@ async def _do_gpu_probe(self) -> None:
         finally:
             self._gpu_in_flight = False
 
+    # ----- ANE --------------------------------------------------------------
+
+    async def _reconcile_ane(self, tick_seconds: float) -> None:
+        while True:
+            try:
+                self._maybe_start_ane_probe()
+            except Exception as e:
+                logger.opt(exception=e).warning("ANE reconcile error")
+            await anyio.sleep(tick_seconds)
+
+    def _maybe_start_ane_probe(self) -> None:
+        state = self.state_view()
+        if self._ane_in_flight:
+            return
+        if state.runners:
+            return
+        existing = state.node_ane_profile.get(self.node_id)
+        if existing is not None and not _is_stale(existing.measured_at, ANE_TTL):
+            return
+        self._ane_in_flight = True
+        self._tg.start_soon(self._do_ane_probe)
+
+    async def _do_ane_probe(self) -> None:
+        try:
+            with fail_after(ANE_PROBE_HARD_TIMEOUT_SECONDS):
+                profile = await AneProfile.measure()
+            if profile is not None:
+                await self.info_sender.send(profile)
+        except Exception as e:
+            logger.opt(exception=e).warning("ANE probe failed")
+        finally:
+            self._ane_in_flight = False
+
     # ----- Links ------------------------------------------------------------
 
     async def _reconcile_links(self, tick_seconds: float) -> None:
diff --git a/src/exo/utils/profilers/tests/test_ane_profiler.py b/src/exo/utils/profilers/tests/test_ane_profiler.py
new file mode 100644
index 0000000000..ff9e6f4200
--- /dev/null
+++ b/src/exo/utils/profilers/tests/test_ane_profiler.py
@@ -0,0 +1,28 @@
+import pytest
+
+from exo.utils.profilers.ane_profiler import AneProfile, ane_available
+
+
+@pytest.mark.skipif(
+    not ane_available(),
+    reason="ANE profile requires AppleNeuralEngine.framework",
+)
+async def test_ane_profile_returns_plausible_numbers():
+    profile = await AneProfile.measure()
+    assert profile is not None
+    assert profile.engine == "ane"
+    assert profile.tflops_fp16 > 0
+    assert profile.tflops_fp16 < 1000
+    assert profile.memory_bandwidth_gbps > 0
+    assert profile.memory_bandwidth_gbps < 100_000
+
+
+def test_ane_profile_serializes_with_class_tag():
+    profile = AneProfile(
+        engine="ane",
+        tflops_fp16=12.3,
+        memory_bandwidth_gbps=80.0,
+    )
+    dumped = profile.model_dump()
+    assert "AneProfile" in dumped
+    assert dumped["AneProfile"]["tflops_fp16"] == 12.3
diff --git a/src/exo/utils/profilers/tests/test_apply_profiles.py b/src/exo/utils/profilers/tests/test_apply_profiles.py
index aa8bd5c5be..87bd69b5ea 100644
--- a/src/exo/utils/profilers/tests/test_apply_profiles.py
+++ b/src/exo/utils/profilers/tests/test_apply_profiles.py
@@ -1,5 +1,4 @@
-"""Verify the apply layer wires GpuProfile / SocketLinkProfile / RDMALinkProfile
-through to the granular state mappings."""
+"""Verify the apply layer wires profiler events to granular state mappings."""
 
 from datetime import datetime, timezone
 
@@ -13,6 +12,7 @@
 )
 from exo.shared.types.profiling import NodeSocketLinkProfile
 from exo.shared.types.state import State
+from exo.utils.profilers.ane_profiler import AneProfile
 from exo.utils.profilers.gpu_profiler import GpuProfile
 from exo.utils.profilers.link_profiler import RDMALinkProfile, SocketLinkProfile
 
@@ -43,6 +43,16 @@ def test_apply_gpu_profile_writes_node_gpu_profile():
     assert entry.engine == "mlx"
 
 
+def test_apply_ane_profile_writes_node_ane_profile():
+    state = State()
+    profile = AneProfile(engine="ane", tflops_fp16=12.0, memory_bandwidth_gbps=80.0)
+    new_state = apply(state, _wrap(0, profile))
+    entry = new_state.node_ane_profile[NODE_A]
+    assert entry.tflops_fp16 == 12.0
+    assert entry.memory_bandwidth_gbps == 80.0
+    assert entry.engine == "ane"
+
+
 def test_apply_socket_link_profile_keys_by_source_and_sink():
     state = State()
     profile = SocketLinkProfile(
@@ -156,10 +166,14 @@ def test_apply_node_timed_out_drops_profiles():
         state,
         _wrap(0, GpuProfile(engine="mlx", tflops_fp16=1, memory_bandwidth_gbps=1)),
     )
+    state = apply(
+        state,
+        _wrap(1, AneProfile(engine="ane", tflops_fp16=1, memory_bandwidth_gbps=1)),
+    )
     state = apply(
         state,
         _wrap(
-            1,
+            2,
             SocketLinkProfile(
                 sink_node_id=NODE_B,
                 sink_ip="10.0.0.5",
@@ -172,12 +186,13 @@ def test_apply_node_timed_out_drops_profiles():
     )
 
     timed_out = IndexedEvent(
-        idx=2,
+        idx=3,
         event=NodeTimedOut(event_id=EventId(), node_id=NODE_A),
     )
     state = apply(state, timed_out)
 
     assert NODE_A not in state.node_gpu_profile
+    assert NODE_A not in state.node_ane_profile
     assert NODE_A not in state.node_link_profiles
 
 
diff --git a/src/exo/utils/tests/test_macmon_metrics.py b/src/exo/utils/tests/test_macmon_metrics.py
new file mode 100644
index 0000000000..8003c6b255
--- /dev/null
+++ b/src/exo/utils/tests/test_macmon_metrics.py
@@ -0,0 +1,31 @@
+from exo.utils.info_gatherer.macmon import MacmonMetrics
+
+
+def test_macmon_metrics_preserves_ane_power() -> None:
+    raw_json = """
+    {
+      "timestamp": "2026-01-01T00:00:00Z",
+      "temp": {"cpu_temp_avg": 50.0, "gpu_temp_avg": 55.0},
+      "memory": {
+        "ram_total": 1000,
+        "ram_usage": 400,
+        "swap_total": 200,
+        "swap_usage": 50
+      },
+      "ecpu_usage": [1200, 10.0],
+      "pcpu_usage": [2400, 20.0],
+      "gpu_usage": [800, 30.0],
+      "all_power": 18.0,
+      "ane_power": 2.5,
+      "cpu_power": 4.0,
+      "gpu_power": 6.0,
+      "gpu_ram_power": 1.0,
+      "ram_power": 2.0,
+      "sys_power": 15.0
+    }
+    """
+
+    metrics = MacmonMetrics.from_raw_json(raw_json)
+
+    assert metrics.system_profile.sys_power == 15.0
+    assert metrics.system_profile.ane_power == 2.5
diff --git a/src/exo/utils/tests/test_power_sampler.py b/src/exo/utils/tests/test_power_sampler.py
index 69f4ccee98..9940b3d2b8 100644
--- a/src/exo/utils/tests/test_power_sampler.py
+++ b/src/exo/utils/tests/test_power_sampler.py
@@ -9,8 +9,8 @@
 from exo.utils.power_sampler import PowerSampler
 
 
-def _make_profile(sys_power: float) -> SystemPerformanceProfile:
-    return SystemPerformanceProfile(sys_power=sys_power)
+def _make_profile(sys_power: float, ane_power: float = 0.0) -> SystemPerformanceProfile:
+    return SystemPerformanceProfile(sys_power=sys_power, ane_power=ane_power)
 
 
 NODE_A = NodeId("node-a")
@@ -20,7 +20,7 @@ def _make_profile(sys_power: float) -> SystemPerformanceProfile:
 @pytest.fixture
 def single_node_sampler() -> PowerSampler:
     state: dict[NodeId, SystemPerformanceProfile] = {
-        NODE_A: _make_profile(10.0),
+        NODE_A: _make_profile(10.0, 2.0),
     }
     return PowerSampler(get_node_system=lambda: state)
 
@@ -28,8 +28,8 @@ def single_node_sampler() -> PowerSampler:
 @pytest.fixture
 def multi_node_state() -> dict[NodeId, SystemPerformanceProfile]:
     return {
-        NODE_A: _make_profile(10.0),
-        NODE_B: _make_profile(20.0),
+        NODE_A: _make_profile(10.0, 2.0),
+        NODE_B: _make_profile(20.0, 4.0),
     }
 
 
@@ -44,6 +44,7 @@ async def test_single_sample(single_node_sampler: PowerSampler) -> None:
     assert len(result.nodes) == 1
     assert result.nodes[0].node_id == NODE_A
     assert result.nodes[0].avg_sys_power == 10.0
+    assert result.nodes[0].avg_ane_power == 2.0
     assert result.nodes[0].samples >= 1
     assert result.elapsed_seconds > 0
 
@@ -61,6 +62,7 @@ async def test_multi_node_averaging(
     result = sampler.result()
     assert len(result.nodes) == 2
     assert result.total_avg_sys_power_watts == 30.0
+    assert result.total_avg_ane_power_watts == 6.0
 
 
 async def test_energy_calculation(single_node_sampler: PowerSampler) -> None:
@@ -72,27 +74,31 @@ async def test_energy_calculation(single_node_sampler: PowerSampler) -> None:
 
     result = single_node_sampler.result()
     expected_energy = result.total_avg_sys_power_watts * result.elapsed_seconds
+    expected_ane_energy = result.total_avg_ane_power_watts * result.elapsed_seconds
     assert result.total_energy_joules == expected_energy
+    assert result.total_ane_energy_joules == expected_ane_energy
 
 
 async def test_changing_power_is_averaged() -> None:
     """When power changes mid-sampling, the result should be the average."""
     state: dict[NodeId, SystemPerformanceProfile] = {
-        NODE_A: _make_profile(10.0),
+        NODE_A: _make_profile(10.0, 1.0),
     }
     sampler = PowerSampler(get_node_system=lambda: state, interval=0.05)
 
     async with anyio.create_task_group() as tg:
         tg.start_soon(sampler.run)
         await anyio.sleep(0.15)
-        state[NODE_A] = _make_profile(20.0)
+        state[NODE_A] = _make_profile(20.0, 3.0)
         await anyio.sleep(0.15)
         tg.cancel_scope.cancel()
 
     result = sampler.result()
     avg = result.nodes[0].avg_sys_power
+    avg_ane = result.nodes[0].avg_ane_power
     # Should be between 10 and 20, not exactly either
     assert 10.0 < avg < 20.0
+    assert 1.0 < avg_ane < 3.0
 
 
 async def test_empty_state() -> None:
@@ -108,7 +114,9 @@ async def test_empty_state() -> None:
     result = sampler.result()
     assert len(result.nodes) == 0
     assert result.total_avg_sys_power_watts == 0.0
+    assert result.total_avg_ane_power_watts == 0.0
     assert result.total_energy_joules == 0.0
+    assert result.total_ane_energy_joules == 0.0
 
 
 async def test_result_stops_sampling() -> None:

From 63d095c1a9f4959fe1f98fca330796c972a75cf5 Mon Sep 17 00:00:00 2001
From: Alex Cheema <alexcheema123@gmail.com>
Date: Sun, 3 May 2026 18:36:19 +0100
Subject: [PATCH 2/4] feat(profilers): profile ANE precision modes

---
 .../src/lib/components/GpuRichBar.svelte      |  25 +-
 .../src/lib/components/TopologyGraph.svelte   |  35 +-
 dashboard/src/lib/stores/app.svelte.ts        |  13 +-
 src/exo/shared/apply.py                       |  15 +-
 src/exo/shared/types/profiling.py             |  15 +-
 src/exo/shared/types/state.py                 |   2 +-
 src/exo/utils/profilers/ane_profiler.py       | 402 ++++++++++++++++--
 .../profilers/tests/test_ane_profiler.py      |  43 +-
 .../profilers/tests/test_apply_profiles.py    |  36 +-
 9 files changed, 505 insertions(+), 81 deletions(-)

diff --git a/dashboard/src/lib/components/GpuRichBar.svelte b/dashboard/src/lib/components/GpuRichBar.svelte
index fc397949ea..394d742446 100644
--- a/dashboard/src/lib/components/GpuRichBar.svelte
+++ b/dashboard/src/lib/components/GpuRichBar.svelte
@@ -3,6 +3,7 @@
     nodeAneProfile,
     nodeGpuProfile,
     topologyData,
+    type RawNodeAneProfile,
   } from "$lib/stores/app.svelte";
 
   interface Props {
@@ -24,11 +25,8 @@
       0,
     ),
   );
-  const totalAneTflops = $derived(
-    Object.values(aneProfiles).reduce(
-      (sum, p) => sum + (p?.tflopsFp16 ?? 0),
-      0,
-    ),
+  const totalAneTops = $derived(
+    Object.values(aneProfiles).reduce((sum, p) => sum + getPeakAneTops(p), 0),
   );
   const totalMemoryBytes = $derived(
     Object.values(topology?.nodes ?? {}).reduce(
@@ -42,6 +40,19 @@
     return `${value.toFixed(1)} TFLOPS`;
   }
 
+  function formatTops(value: number): string {
+    if (value >= 1000) return `${(value / 1000).toFixed(2)} POPS`;
+    return `${value.toFixed(1)} TOPS`;
+  }
+
+  function getPeakAneTops(profile: RawNodeAneProfile | undefined): number {
+    if (!profile) return 0;
+    return Math.max(
+      0,
+      ...profile.precisionProfiles.map((p) => p.computeTops ?? 0),
+    );
+  }
+
   function formatBandwidth(value: number): string {
     if (value >= 1000) return `${(value / 1000).toFixed(2)} TB/s`;
     return `${value.toFixed(0)} GB/s`;
@@ -65,8 +76,8 @@
     <span class="stat-label">Memory bandwidth</span>
   </div>
   <div class="stat-block">
-    <span class="stat-value">{formatTflops(totalAneTflops)}</span>
-    <span class="stat-label">ANE compute</span>
+    <span class="stat-value">{formatTops(totalAneTops)}</span>
+    <span class="stat-label">ANE peak</span>
   </div>
   <div class="stat-block">
     <span class="stat-value">{formatMemory(totalMemoryBytes)}</span>
diff --git a/dashboard/src/lib/components/TopologyGraph.svelte b/dashboard/src/lib/components/TopologyGraph.svelte
index 37bdf15e18..a952ae7cc3 100644
--- a/dashboard/src/lib/components/TopologyGraph.svelte
+++ b/dashboard/src/lib/components/TopologyGraph.svelte
@@ -14,6 +14,8 @@
     nodeNetworkRaw,
     nodeThunderbolt,
     type NodeInfo,
+    type RawNodeAnePrecisionProfile,
+    type RawNodeAneProfile,
   } from "$lib/stores/app.svelte";
   import {
     inferRdmaConnectionType,
@@ -72,6 +74,26 @@
     return `${value.toFixed(2)} ms`;
   }
 
+  function formatAneTops(value: number | null | undefined): string {
+    if (value == null || !isFinite(value)) return "—";
+    if (value >= 1000) return `${(value / 1000).toFixed(2)}P`;
+    return `${value.toFixed(value >= 10 ? 0 : 1)}T`;
+  }
+
+  function formatAneProfileLine(profile: RawNodeAneProfile): string {
+    const preferredOrder: RawNodeAnePrecisionProfile["precisionBits"][] = [
+      16, 8, 4, 32,
+    ];
+    const byBits = new Map(
+      profile.precisionProfiles.map((p) => [p.precisionBits, p]),
+    );
+    return preferredOrder
+      .map((bits) => byBits.get(bits))
+      .filter((p): p is RawNodeAnePrecisionProfile => Boolean(p?.supported))
+      .map((p) => `${p.precisionBits}b ${formatAneTops(p.computeTops)}`)
+      .join(" · ");
+  }
+
   interface PairProfileEntry {
     fromId: string;
     toId: string;
@@ -1321,26 +1343,19 @@
             .text(`${gpuProfile.memoryBandwidthGbps.toFixed(0)} GB/s`);
         }
         if (aneProfile) {
+          const aneProfileLine = formatAneProfileLine(aneProfile);
           const aneY = infoY + fontSize * (gpuProfile ? 1.95 : 1.05);
           const aneText = nodeG
             .append("text")
             .attr("x", nodeInfo.x)
             .attr("y", aneY)
             .attr("text-anchor", "middle")
-            .attr("font-size", fontSize * 0.8)
+            .attr("font-size", fontSize * 0.72)
             .attr("font-family", "SF Mono, Monaco, monospace");
           aneText
             .append("tspan")
             .attr("fill", "rgba(96,165,250,0.95)")
-            .text(`ANE ${aneProfile.tflopsFp16.toFixed(1)} TFLOPS`);
-          aneText
-            .append("tspan")
-            .attr("fill", "rgba(179,179,179,0.6)")
-            .text("  ·  ");
-          aneText
-            .append("tspan")
-            .attr("fill", "rgba(255,215,0,0.75)")
-            .text(`${aneProfile.memoryBandwidthGbps.toFixed(0)} GB/s`);
+            .text(aneProfileLine ? `ANE ${aneProfileLine}` : "ANE profiling");
         }
       } else if (showCompactLabels) {
         // COMPACT MODE: Just name and basic info (4+ nodes)
diff --git a/dashboard/src/lib/stores/app.svelte.ts b/dashboard/src/lib/stores/app.svelte.ts
index 5a70e5153e..a5352ca339 100644
--- a/dashboard/src/lib/stores/app.svelte.ts
+++ b/dashboard/src/lib/stores/app.svelte.ts
@@ -135,10 +135,19 @@ export interface RawNodeGpuProfile {
   measuredAt: string;
 }
 
+export interface RawNodeAnePrecisionProfile {
+  precisionBits: 32 | 16 | 8 | 4;
+  weightBits: 32 | 16 | 8 | 4;
+  activationBits: 32 | 16 | 8 | 4;
+  supported: boolean;
+  computeTops?: number | null;
+  memoryBandwidthGbps?: number | null;
+  error?: string | null;
+}
+
 export interface RawNodeAneProfile {
   engine: "ane";
-  tflopsFp16: number;
-  memoryBandwidthGbps: number;
+  precisionProfiles: RawNodeAnePrecisionProfile[];
   measuredAt: string;
 }
 
diff --git a/src/exo/shared/apply.py b/src/exo/shared/apply.py
index 86e94ac00c..4852cff1d0 100644
--- a/src/exo/shared/apply.py
+++ b/src/exo/shared/apply.py
@@ -33,6 +33,7 @@
 )
 from exo.shared.types.instance_link import InstanceLink, InstanceLinkId
 from exo.shared.types.profiling import (
+    NodeAnePrecisionProfile,
     NodeAneProfile,
     NodeGpuProfile,
     NodeIdentity,
@@ -481,8 +482,18 @@ def apply_node_gathered_info(event: NodeGatheredInfo, state: State) -> State:
                 **state.node_ane_profile,
                 event.node_id: NodeAneProfile(
                     engine=info.engine,
-                    tflops_fp16=info.tflops_fp16,
-                    memory_bandwidth_gbps=info.memory_bandwidth_gbps,
+                    precision_profiles=tuple(
+                        NodeAnePrecisionProfile(
+                            precision_bits=profile.precision_bits,
+                            weight_bits=profile.weight_bits,
+                            activation_bits=profile.activation_bits,
+                            supported=profile.supported,
+                            compute_tops=profile.compute_tops,
+                            memory_bandwidth_gbps=profile.memory_bandwidth_gbps,
+                            error=profile.error,
+                        )
+                        for profile in info.precision_profiles
+                    ),
                     measured_at=measured_at,
                 ),
             }
diff --git a/src/exo/shared/types/profiling.py b/src/exo/shared/types/profiling.py
index fd78624659..ab75d805e2 100644
--- a/src/exo/shared/types/profiling.py
+++ b/src/exo/shared/types/profiling.py
@@ -122,12 +122,23 @@ class NodeGpuProfile(FrozenModel):
     measured_at: datetime
 
 
+class NodeAnePrecisionProfile(FrozenModel):
+    """Measured ANE profile for a single weight precision."""
+
+    precision_bits: Literal[32, 16, 8, 4]
+    weight_bits: Literal[32, 16, 8, 4]
+    activation_bits: Literal[32, 16, 8, 4]
+    supported: bool
+    compute_tops: float | None = None
+    memory_bandwidth_gbps: float | None = None
+    error: str | None = None
+
+
 class NodeAneProfile(FrozenModel):
     """Measured ANE compute throughput and streaming bandwidth for a node."""
 
     engine: Literal["ane"]
-    tflops_fp16: float
-    memory_bandwidth_gbps: float
+    precision_profiles: Sequence[NodeAnePrecisionProfile]
     measured_at: datetime
 
 
diff --git a/src/exo/shared/types/state.py b/src/exo/shared/types/state.py
index d2eb5498dd..c9de35086d 100644
--- a/src/exo/shared/types/state.py
+++ b/src/exo/shared/types/state.py
@@ -64,7 +64,7 @@ class State(FrozenModel):
 
     # Profiler results.
     # node_gpu_profile: per-node GPU TFLOPS / memory bandwidth measurement.
-    # node_ane_profile: per-node ANE TFLOPS / streaming bandwidth measurement.
+    # node_ane_profile: per-node ANE compute / streaming bandwidth by precision.
     # node_link_profiles: per-edge link probe; outer key = source node, inner
     # key = sink node. May contain both a socket and an RDMA profile to the
     # same peer, so the inner value is a sequence.
diff --git a/src/exo/utils/profilers/ane_profiler.py b/src/exo/utils/profilers/ane_profiler.py
index 68bd9d8d44..7f2d81056a 100644
--- a/src/exo/utils/profilers/ane_profiler.py
+++ b/src/exo/utils/profilers/ane_profiler.py
@@ -25,17 +25,23 @@
 import numpy.typing as npt
 from anyio import to_thread
 
-from exo.utils.pydantic_ext import TaggedModel
+from exo.utils.pydantic_ext import FrozenModel, TaggedModel
 
 type Pointer = int
 type Float16Array = npt.NDArray[np.float16]
+type Float32Array = npt.NDArray[np.float32]
+type NumericArray = Float16Array | Float32Array
 type UInt8Array = npt.NDArray[np.uint8]
 
-_BYTES_PER_ELEMENT = 2
+_FP16_BYTES_PER_ELEMENT = 2
+_FP32_BYTES_PER_ELEMENT = 4
+
+type AnePrecisionBits = Literal[32, 16, 8, 4]
+_ANE_PRECISION_BITS: tuple[AnePrecisionBits, ...] = (32, 16, 8, 4)
 
 _COMPUTE_IN_DIM = 2048
 _COMPUTE_OUT_DIM = 2048
-_COMPUTE_SPATIAL = 128
+_COMPUTE_SPATIAL = 768
 _COMPUTE_ITERATIONS_PER_PASS = 12
 
 _STREAM_CHANNELS = 4096
@@ -64,13 +70,25 @@
 _ane_io_surface_object_class: Pointer | None = None
 
 
+@final
+class AnePrecisionProfile(FrozenModel):
+    """Per-precision ANE compute and memory-bound probe result."""
+
+    precision_bits: AnePrecisionBits
+    weight_bits: AnePrecisionBits
+    activation_bits: AnePrecisionBits
+    supported: bool
+    compute_tops: float | None = None
+    memory_bandwidth_gbps: float | None = None
+    error: str | None = None
+
+
 @final
 class AneProfile(TaggedModel):
     """Wire format for a measured ANE profile, gathered locally on a node."""
 
     engine: Literal["ane"]
-    tflops_fp16: float
-    memory_bandwidth_gbps: float
+    precision_profiles: Sequence[AnePrecisionProfile]
 
     @classmethod
     async def measure(cls) -> Self | None:
@@ -81,7 +99,60 @@ async def measure(cls) -> Self | None:
     @classmethod
     def _measure_blocking(cls) -> Self:
         rng = np.random.default_rng(0)
+        profiles = tuple(
+            _measure_precision_profile(bits, rng) for bits in _ANE_PRECISION_BITS
+        )
+        return cls(engine="ane", precision_profiles=profiles)
+
+
+def _measure_precision_profile(
+    precision_bits: AnePrecisionBits, rng: np.random.Generator
+) -> AnePrecisionProfile:
+    try:
+        return _measure_supported_precision_profile(precision_bits, rng)
+    except Exception as exc:
+        return AnePrecisionProfile(
+            precision_bits=precision_bits,
+            weight_bits=precision_bits,
+            activation_bits=8 if precision_bits in (8, 4) else precision_bits,
+            supported=False,
+            error=_short_error(exc),
+        )
+
+
+def _measure_supported_precision_profile(
+    precision_bits: AnePrecisionBits, rng: np.random.Generator
+) -> AnePrecisionProfile:
+    activation_bits: AnePrecisionBits = (
+        8 if precision_bits in (8, 4) else precision_bits
+    )
+
+    if precision_bits == 32:
+        compute_weights = rng.standard_normal(
+            (_COMPUTE_OUT_DIM, _COMPUTE_IN_DIM)
+        ).astype(np.float32)
+        compute_input = rng.standard_normal((_COMPUTE_IN_DIM, _COMPUTE_SPATIAL)).astype(
+            np.float32
+        )
+        compute_kernel = _compile_static_conv_kernel(
+            weights=compute_weights,
+            input_channels=_COMPUTE_IN_DIM,
+            output_channels=_COMPUTE_OUT_DIM,
+            spatial=_COMPUTE_SPATIAL,
+            precision_bits=precision_bits,
+        )
+        compute_kernel.write_input(0, compute_input)
 
+        stream_input = rng.standard_normal((_STREAM_CHANNELS, _STREAM_SPATIAL)).astype(
+            np.float32
+        )
+        stream_kernel = _compile_relu_stream_kernel(
+            channels=_STREAM_CHANNELS,
+            spatial=_STREAM_SPATIAL,
+            precision_bits=precision_bits,
+        )
+        stream_kernel.write_input(0, stream_input)
+    elif precision_bits == 16:
         compute_weights = rng.standard_normal(
             (_COMPUTE_OUT_DIM, _COMPUTE_IN_DIM)
         ).astype(np.float16)
@@ -98,21 +169,66 @@ def _measure_blocking(cls) -> Self:
             input_channels=_COMPUTE_IN_DIM,
             output_channels=_COMPUTE_OUT_DIM,
             spatial=_COMPUTE_SPATIAL,
+            precision_bits=precision_bits,
         )
         compute_kernel.write_input(0, compute_input)
 
         stream_kernel = _compile_relu_stream_kernel(
             channels=_STREAM_CHANNELS,
             spatial=_STREAM_SPATIAL,
+            precision_bits=precision_bits,
         )
         stream_kernel.write_input(0, stream_input)
+    else:
+        compute_weights = rng.standard_normal(
+            (_COMPUTE_OUT_DIM, _COMPUTE_IN_DIM)
+        ).astype(np.float32)
+        compute_input = rng.standard_normal((_COMPUTE_IN_DIM, _COMPUTE_SPATIAL)).astype(
+            np.float16
+        )
 
-        return cls(
-            engine="ane",
-            tflops_fp16=_measure_compute_tflops(compute_kernel),
-            memory_bandwidth_gbps=_measure_streaming_bandwidth_gbps(stream_kernel),
+        stream_input = rng.standard_normal((_STREAM_CHANNELS, _STREAM_SPATIAL)).astype(
+            np.float16
         )
 
+        compute_kernel = _compile_quantized_static_conv_kernel(
+            weights=compute_weights,
+            input_channels=_COMPUTE_IN_DIM,
+            output_channels=_COMPUTE_OUT_DIM,
+            spatial=_COMPUTE_SPATIAL,
+            precision_bits=precision_bits,
+        )
+        compute_kernel.write_input(0, compute_input)
+
+        stream_kernel = _compile_quantized_stream_kernel(
+            channels=_STREAM_CHANNELS,
+            spatial=_STREAM_SPATIAL,
+        )
+        stream_kernel.write_input(0, stream_input)
+
+    return AnePrecisionProfile(
+        precision_bits=precision_bits,
+        weight_bits=precision_bits,
+        activation_bits=activation_bits,
+        supported=True,
+        compute_tops=_measure_compute_tops(compute_kernel),
+        memory_bandwidth_gbps=_measure_streaming_bandwidth_gbps(
+            stream_kernel,
+            bytes_per_element=(
+                _FP32_BYTES_PER_ELEMENT
+                if activation_bits == 32
+                else _FP16_BYTES_PER_ELEMENT
+            ),
+        ),
+    )
+
+
+def _short_error(exc: Exception) -> str:
+    message = str(exc).replace("\n", " ")
+    if len(message) <= 220:
+        return message
+    return f"{message[:217]}..."
+
 
 @final
 class _AneKernel:
@@ -183,7 +299,7 @@ def eval(self) -> None:
             err_desc = _describe_error(err_ptr.value)
             raise RuntimeError(f"ANE eval failed: {err_desc}")
 
-    def write_input(self, index: int, values: Float16Array) -> None:
+    def write_input(self, index: int, values: NumericArray) -> None:
         _write_iosurface(self._input_surfaces[index], values)
 
 
@@ -195,11 +311,11 @@ def ane_available() -> bool:
         return False
 
 
-def _measure_compute_tflops(kernel: _AneKernel) -> float:
+def _measure_compute_tops(kernel: _AneKernel) -> float:
     _warm_up_with(kernel.eval, time.perf_counter() + _WARMUP_SECONDS)
 
-    flops_per_iteration = 2 * _COMPUTE_IN_DIM * _COMPUTE_OUT_DIM * _COMPUTE_SPATIAL
-    best_tflops = 0.0
+    operations_per_iteration = 2 * _COMPUTE_IN_DIM * _COMPUTE_OUT_DIM * _COMPUTE_SPATIAL
+    best_tops = 0.0
     for _ in range(_MEASUREMENT_PASSES):
         start = time.perf_counter()
         for _ in range(_COMPUTE_ITERATIONS_PER_PASS):
@@ -207,15 +323,17 @@ def _measure_compute_tflops(kernel: _AneKernel) -> float:
         elapsed = time.perf_counter() - start
         if elapsed <= 0:
             continue
-        tflops = flops_per_iteration * _COMPUTE_ITERATIONS_PER_PASS / elapsed / 1e12
-        best_tflops = max(best_tflops, tflops)
-    return best_tflops
+        tops = operations_per_iteration * _COMPUTE_ITERATIONS_PER_PASS / elapsed / 1e12
+        best_tops = max(best_tops, tops)
+    return best_tops
 
 
-def _measure_streaming_bandwidth_gbps(kernel: _AneKernel) -> float:
+def _measure_streaming_bandwidth_gbps(
+    kernel: _AneKernel, *, bytes_per_element: int
+) -> float:
     _warm_up_with(kernel.eval, time.perf_counter() + _WARMUP_SECONDS)
 
-    bytes_per_iteration = _STREAM_CHANNELS * _STREAM_SPATIAL * _BYTES_PER_ELEMENT * 2
+    bytes_per_iteration = _STREAM_CHANNELS * _STREAM_SPATIAL * bytes_per_element * 2
     best_gbps = 0.0
     for _ in range(_MEASUREMENT_PASSES):
         start = time.perf_counter()
@@ -236,19 +354,23 @@ def _warm_up_with(do_op: Callable[[], None], deadline: float) -> None:
 
 def _compile_static_conv_kernel(
     *,
-    weights: Float16Array,
+    weights: NumericArray,
     input_channels: int,
     output_channels: int,
     spatial: int,
+    precision_bits: Literal[32, 16],
 ) -> _AneKernel:
+    element_type = _floating_mil_type(precision_bits)
     mil = _mil_static_conv(
         input_channels=input_channels,
         output_channels=output_channels,
         spatial=spatial,
+        element_type=element_type,
     )
-    weight_blob = _build_weight_blob_bytes(weights)
-    input_bytes = input_channels * spatial * _BYTES_PER_ELEMENT
-    output_bytes = output_channels * spatial * _BYTES_PER_ELEMENT
+    weight_blob = _build_floating_weight_blob_bytes(weights, precision_bits)
+    bytes_per_element = _bytes_per_precision(precision_bits)
+    input_bytes = input_channels * spatial * bytes_per_element
+    output_bytes = output_channels * spatial * bytes_per_element
     return _compile_kernel(
         mil_text=mil,
         weight_files={"weight.bin": weight_blob},
@@ -257,17 +379,67 @@ def _compile_static_conv_kernel(
     )
 
 
-def _compile_relu_stream_kernel(*, channels: int, spatial: int) -> _AneKernel:
+def _compile_quantized_static_conv_kernel(
+    *,
+    weights: Float32Array,
+    input_channels: int,
+    output_channels: int,
+    spatial: int,
+    precision_bits: Literal[8, 4],
+) -> _AneKernel:
+    weight_blob = _build_quantized_weight_blob_bytes(weights, precision_bits)
+    mil = _mil_quantized_static_conv(
+        input_channels=input_channels,
+        output_channels=output_channels,
+        spatial=spatial,
+        precision_bits=precision_bits,
+    )
+    input_bytes = input_channels * spatial * _FP16_BYTES_PER_ELEMENT
+    output_bytes = output_channels * spatial * _FP16_BYTES_PER_ELEMENT
+    return _compile_kernel(
+        mil_text=mil,
+        weight_files={"weight.bin": weight_blob},
+        input_sizes=(input_bytes,),
+        output_sizes=(output_bytes,),
+    )
+
+
+def _compile_relu_stream_kernel(
+    *, channels: int, spatial: int, precision_bits: Literal[32, 16]
+) -> _AneKernel:
+    element_count = channels * spatial
+    data_bytes = element_count * _bytes_per_precision(precision_bits)
+    return _compile_kernel(
+        mil_text=_mil_relu_stream(
+            channels=channels,
+            spatial=spatial,
+            element_type=_floating_mil_type(precision_bits),
+        ),
+        weight_files={},
+        input_sizes=(data_bytes,),
+        output_sizes=(data_bytes,),
+    )
+
+
+def _compile_quantized_stream_kernel(*, channels: int, spatial: int) -> _AneKernel:
     element_count = channels * spatial
-    data_bytes = element_count * _BYTES_PER_ELEMENT
+    data_bytes = element_count * _FP16_BYTES_PER_ELEMENT
     return _compile_kernel(
-        mil_text=_mil_relu_stream(channels=channels, spatial=spatial),
+        mil_text=_mil_quantized_stream(channels=channels, spatial=spatial),
         weight_files={},
         input_sizes=(data_bytes,),
         output_sizes=(data_bytes,),
     )
 
 
+def _floating_mil_type(precision_bits: Literal[32, 16]) -> Literal["fp32", "fp16"]:
+    return "fp32" if precision_bits == 32 else "fp16"
+
+
+def _bytes_per_precision(precision_bits: Literal[32, 16]) -> int:
+    return _FP32_BYTES_PER_ELEMENT if precision_bits == 32 else _FP16_BYTES_PER_ELEMENT
+
+
 def _compile_kernel(
     *,
     mil_text: str,
@@ -487,7 +659,7 @@ def _write_model_files(
     return tmp_dir
 
 
-def _write_iosurface(surface: Pointer, values: Float16Array) -> None:
+def _write_iosurface(surface: Pointer, values: NumericArray) -> None:
     raw = cast(UInt8Array, np.ascontiguousarray(values).view(np.uint8).reshape(-1))
     nbytes = int(raw.nbytes)
     if nbytes > _iosurface_alloc_size(surface):
@@ -584,18 +756,84 @@ def _iosurface_alloc_size(surface: Pointer) -> int:
     )
 
 
-def _build_weight_blob_bytes(values: Float16Array) -> bytes:
-    flat = np.ascontiguousarray(values.reshape(-1), dtype=np.float16)
+def _build_floating_weight_blob_bytes(
+    values: NumericArray, precision_bits: Literal[32, 16]
+) -> bytes:
+    dtype = np.float32 if precision_bits == 32 else np.float16
+    type_code = 0x02 if precision_bits == 32 else 0x01
+    flat = np.ascontiguousarray(values.reshape(-1), dtype=dtype)
     payload = flat.tobytes()
-    header_bytes = 128
-    buf = bytearray(header_bytes + len(payload))
-    buf[0] = 0x01
-    buf[4] = 0x02
-    _pack_uint32_le(buf, 64, 0xDEADBEEF)
-    buf[68] = 0x01
-    _pack_uint32_le(buf, 72, len(payload))
-    _pack_uint32_le(buf, 80, header_bytes)
-    buf[header_bytes:] = payload
+    buf = _build_multi_weight_blob_bytes([(type_code, payload)])
+    return bytes(buf)
+
+
+def _build_quantized_weight_blob_bytes(
+    values: Float32Array, precision_bits: Literal[8, 4]
+) -> bytes:
+    weights = np.ascontiguousarray(values, dtype=np.float32)
+    abs_weights = cast(Float32Array, np.abs(weights))
+    channel_max = cast(Float32Array, np.max(abs_weights, axis=1, keepdims=True))
+    max_abs = cast(Float32Array, np.maximum(channel_max, np.float32(1e-6)))
+    if precision_bits == 8:
+        scale = np.ascontiguousarray(
+            (max_abs / np.float32(127.0)).reshape(-1, 1, 1, 1),
+            dtype=np.float16,
+        )
+        quantized = np.clip(
+            np.round(weights / max_abs * np.float32(127.0)), -128, 127
+        ).astype(np.int8)
+        data_payload = np.ascontiguousarray(quantized).tobytes()
+        data_type_code = 0x04
+    else:
+        scale = np.ascontiguousarray(
+            (max_abs / np.float32(7.0)).reshape(-1, 1, 1, 1),
+            dtype=np.float16,
+        )
+        quantized = np.clip(
+            np.round(weights / max_abs * np.float32(7.0)), -8, 7
+        ).astype(np.int8)
+        data_payload = _pack_int4_payload(quantized)
+        data_type_code = 0x08
+
+    return _build_multi_weight_blob_bytes(
+        [
+            (data_type_code, data_payload),
+            (0x01, scale.tobytes()),
+        ]
+    )
+
+
+def _pack_int4_payload(values: npt.NDArray[np.int8]) -> bytes:
+    flat = np.ascontiguousarray(values.reshape(-1), dtype=np.int8)
+    if flat.size % 2 != 0:
+        flat = np.pad(flat, (0, 1)).astype(np.int8)
+    nibbles = (flat & 0x0F).astype(np.uint8)
+    packed = nibbles[0::2] | (nibbles[1::2] << 4)
+    return np.ascontiguousarray(packed, dtype=np.uint8).tobytes()
+
+
+def _build_multi_weight_blob_bytes(blobs: Sequence[tuple[int, bytes]]) -> bytes:
+    total_bytes = 64
+    descriptors_and_payloads: list[tuple[bytes, bytes]] = []
+    for type_code, payload in blobs:
+        descriptor = bytearray(64)
+        _pack_uint32_le(descriptor, 0, 0xDEADBEEF)
+        _pack_uint32_le(descriptor, 4, type_code)
+        _pack_uint32_le(descriptor, 8, len(payload))
+        _pack_uint32_le(descriptor, 16, total_bytes + 64)
+        descriptors_and_payloads.append((bytes(descriptor), payload))
+        total_bytes += 64 + len(payload)
+
+    buf = bytearray(total_bytes)
+    _pack_uint32_le(buf, 0, len(blobs))
+    _pack_uint32_le(buf, 4, 0x02)
+
+    offset = 64
+    for descriptor, payload in descriptors_and_payloads:
+        buf[offset : offset + 64] = descriptor
+        offset += 64
+        buf[offset : offset + len(payload)] = payload
+        offset += len(payload)
     return bytes(buf)
 
 
@@ -1030,28 +1268,104 @@ def _describe_error(error_pointer: int | None) -> str:
 _CONV_ARGS = "dilations=dl, groups=gr, pad=pd, pad_type=pt, strides=st"
 
 
-def _mil_static_conv(*, input_channels: int, output_channels: int, spatial: int) -> str:
+def _mil_static_conv(
+    *,
+    input_channels: int,
+    output_channels: int,
+    spatial: int,
+    element_type: Literal["fp32", "fp16"],
+) -> str:
     return (
         f"{_MIL_HEADER}"
-        f"    func main<ios18>(tensor<fp16, [1, {input_channels}, 1, {spatial}]> x) {{\n"
-        f"        tensor<fp16, [{output_channels}, {input_channels}, 1, 1]> W = const()"
-        f'[name=string("W"), val=tensor<fp16, [{output_channels}, {input_channels}, 1, 1]>'
+        f"    func main<ios18>(tensor<{element_type}, [1, {input_channels}, 1, {spatial}]> x) {{\n"
+        f"        tensor<{element_type}, [{output_channels}, {input_channels}, 1, 1]> W = const()"
+        f'[name=string("W"), val=tensor<{element_type}, [{output_channels}, {input_channels}, 1, 1]>'
         '(BLOBFILE(path=string("@model_path/weights/weight.bin"), '
         "offset=uint64(64)))];\n"
         f"{_CONV_PARAMS}"
-        f"        tensor<fp16, [1, {output_channels}, 1, {spatial}]> y = conv("
+        f"        tensor<{element_type}, [1, {output_channels}, 1, {spatial}]> y = conv("
         f'{_CONV_ARGS}, weight=W, x=x)[name=string("cv")];\n'
         "    } -> (y);\n"
         "}\n"
     )
 
 
-def _mil_relu_stream(*, channels: int, spatial: int) -> str:
+def _mil_quantized_static_conv(
+    *,
+    input_channels: int,
+    output_channels: int,
+    spatial: int,
+    precision_bits: Literal[8, 4],
+) -> str:
+    data_type = "int8" if precision_bits == 8 else "int4"
+    scale_blob_offset = 128 + _quantized_data_payload_size(
+        input_channels=input_channels,
+        output_channels=output_channels,
+        precision_bits=precision_bits,
+    )
+    return (
+        f"{_MIL_HEADER}"
+        f"    func main<ios18>(tensor<fp16, [1, {input_channels}, 1, {spatial}]> x) {{\n"
+        f"        tensor<fp16, [{output_channels}, {input_channels}, 1, 1]> W = "
+        f"constexpr_blockwise_shift_scale(data = tensor<{data_type}, "
+        f"[{output_channels}, {input_channels}, 1, 1]>(BLOBFILE(path = "
+        'string("@model_path/weights/weight.bin"), offset = uint64(64))), '
+        f"scale = tensor<fp16, [{output_channels}, 1, 1, 1]>(BLOBFILE(path = "
+        'string("@model_path/weights/weight.bin"), '
+        f'offset = uint64({scale_blob_offset}))))[name = string("Wq")];\n'
+        f"{_mil_quantize_dequantize(channels=input_channels, spatial=spatial)}"
+        f"{_CONV_PARAMS}"
+        f"        tensor<fp16, [1, {output_channels}, 1, {spatial}]> y = conv("
+        f'{_CONV_ARGS}, weight=W, x=dx)[name=string("cv")];\n'
+        "    } -> (y);\n"
+        "}\n"
+    )
+
+
+def _mil_relu_stream(
+    *, channels: int, spatial: int, element_type: Literal["fp32", "fp16"]
+) -> str:
+    return (
+        f"{_MIL_HEADER}"
+        f"    func main<ios18>(tensor<{element_type}, [1, {channels}, 1, {spatial}]> x) {{\n"
+        f"        tensor<{element_type}, [1, {channels}, 1, {spatial}]> y = relu(x=x)"
+        '[name=string("relu")];\n'
+        "    } -> (y);\n"
+        "}\n"
+    )
+
+
+def _mil_quantized_stream(*, channels: int, spatial: int) -> str:
     return (
         f"{_MIL_HEADER}"
         f"    func main<ios18>(tensor<fp16, [1, {channels}, 1, {spatial}]> x) {{\n"
-        f"        tensor<fp16, [1, {channels}, 1, {spatial}]> y = relu(x=x)"
+        f"{_mil_quantize_dequantize(channels=channels, spatial=spatial)}"
+        f"        tensor<fp16, [1, {channels}, 1, {spatial}]> y = relu(x=dx)"
         '[name=string("relu")];\n'
         "    } -> (y);\n"
         "}\n"
     )
+
+
+def _quantized_data_payload_size(
+    *, input_channels: int, output_channels: int, precision_bits: Literal[8, 4]
+) -> int:
+    element_count = input_channels * output_channels
+    if precision_bits == 8:
+        return element_count
+    return (element_count + 1) // 2
+
+
+def _mil_quantize_dequantize(*, channels: int, spatial: int) -> str:
+    return (
+        '        fp16 qs = const()[name = string("qs"), val = fp16(0x1p-7)];\n'
+        '        int8 qz = const()[name = string("qz"), val = int8(0)];\n'
+        '        string qdt = const()[name = string("qdt"), val = string("int8")];\n'
+        f"        tensor<int8, [1, {channels}, 1, {spatial}]> qx = quantize("
+        "input = x, output_dtype = qdt, scale = qs, zero_point = qz)"
+        '[name = string("qx")];\n'
+        '        fp16 dqs = const()[name = string("dqs"), val = fp16(0x1p-7)];\n'
+        '        int8 dqz = const()[name = string("dqz"), val = int8(0)];\n'
+        f"        tensor<fp16, [1, {channels}, 1, {spatial}]> dx = dequantize("
+        'input = qx, scale = dqs, zero_point = dqz)[name = string("dx")];\n'
+    )
diff --git a/src/exo/utils/profilers/tests/test_ane_profiler.py b/src/exo/utils/profilers/tests/test_ane_profiler.py
index ff9e6f4200..944cec256c 100644
--- a/src/exo/utils/profilers/tests/test_ane_profiler.py
+++ b/src/exo/utils/profilers/tests/test_ane_profiler.py
@@ -1,6 +1,10 @@
 import pytest
 
-from exo.utils.profilers.ane_profiler import AneProfile, ane_available
+from exo.utils.profilers.ane_profiler import (
+    AnePrecisionProfile,
+    AneProfile,
+    ane_available,
+)
 
 
 @pytest.mark.skipif(
@@ -11,18 +15,41 @@ async def test_ane_profile_returns_plausible_numbers():
     profile = await AneProfile.measure()
     assert profile is not None
     assert profile.engine == "ane"
-    assert profile.tflops_fp16 > 0
-    assert profile.tflops_fp16 < 1000
-    assert profile.memory_bandwidth_gbps > 0
-    assert profile.memory_bandwidth_gbps < 100_000
+    precision_by_bits = {p.precision_bits: p for p in profile.precision_profiles}
+    assert set(precision_by_bits) == {32, 16, 8, 4}
+    for bits in (16, 8, 4):
+        precision = precision_by_bits[bits]
+        assert precision.supported
+        assert precision.compute_tops is not None
+        assert precision.compute_tops > 0
+        assert precision.compute_tops < 1000
+        assert precision.memory_bandwidth_gbps is not None
+        assert precision.memory_bandwidth_gbps > 0
+        assert precision.memory_bandwidth_gbps < 100_000
+
+    fp32 = precision_by_bits[32]
+    if fp32.supported:
+        assert fp32.compute_tops is not None
+        assert fp32.compute_tops > 0
+    else:
+        assert fp32.error
 
 
 def test_ane_profile_serializes_with_class_tag():
     profile = AneProfile(
         engine="ane",
-        tflops_fp16=12.3,
-        memory_bandwidth_gbps=80.0,
+        precision_profiles=(
+            AnePrecisionProfile(
+                precision_bits=16,
+                weight_bits=16,
+                activation_bits=16,
+                supported=True,
+                compute_tops=12.3,
+                memory_bandwidth_gbps=80.0,
+            ),
+        ),
     )
     dumped = profile.model_dump()
     assert "AneProfile" in dumped
-    assert dumped["AneProfile"]["tflops_fp16"] == 12.3
+    assert dumped["AneProfile"]["precision_profiles"][0]["precision_bits"] == 16
+    assert dumped["AneProfile"]["precision_profiles"][0]["compute_tops"] == 12.3
diff --git a/src/exo/utils/profilers/tests/test_apply_profiles.py b/src/exo/utils/profilers/tests/test_apply_profiles.py
index 87bd69b5ea..c5dc0e41c0 100644
--- a/src/exo/utils/profilers/tests/test_apply_profiles.py
+++ b/src/exo/utils/profilers/tests/test_apply_profiles.py
@@ -12,7 +12,7 @@
 )
 from exo.shared.types.profiling import NodeSocketLinkProfile
 from exo.shared.types.state import State
-from exo.utils.profilers.ane_profiler import AneProfile
+from exo.utils.profilers.ane_profiler import AnePrecisionProfile, AneProfile
 from exo.utils.profilers.gpu_profiler import GpuProfile
 from exo.utils.profilers.link_profiler import RDMALinkProfile, SocketLinkProfile
 
@@ -33,6 +33,30 @@ def _wrap(idx: int, info: object, when: str = WHEN) -> IndexedEvent:
     )
 
 
+def _ane_profile() -> AneProfile:
+    return AneProfile(
+        engine="ane",
+        precision_profiles=(
+            AnePrecisionProfile(
+                precision_bits=16,
+                weight_bits=16,
+                activation_bits=16,
+                supported=True,
+                compute_tops=12.0,
+                memory_bandwidth_gbps=80.0,
+            ),
+            AnePrecisionProfile(
+                precision_bits=8,
+                weight_bits=8,
+                activation_bits=8,
+                supported=True,
+                compute_tops=28.0,
+                memory_bandwidth_gbps=72.0,
+            ),
+        ),
+    )
+
+
 def test_apply_gpu_profile_writes_node_gpu_profile():
     state = State()
     profile = GpuProfile(engine="mlx", tflops_fp16=42.0, memory_bandwidth_gbps=400.0)
@@ -45,12 +69,14 @@ def test_apply_gpu_profile_writes_node_gpu_profile():
 
 def test_apply_ane_profile_writes_node_ane_profile():
     state = State()
-    profile = AneProfile(engine="ane", tflops_fp16=12.0, memory_bandwidth_gbps=80.0)
+    profile = _ane_profile()
     new_state = apply(state, _wrap(0, profile))
     entry = new_state.node_ane_profile[NODE_A]
-    assert entry.tflops_fp16 == 12.0
-    assert entry.memory_bandwidth_gbps == 80.0
     assert entry.engine == "ane"
+    precision_by_bits = {p.precision_bits: p for p in entry.precision_profiles}
+    assert precision_by_bits[16].compute_tops == 12.0
+    assert precision_by_bits[16].memory_bandwidth_gbps == 80.0
+    assert precision_by_bits[8].compute_tops == 28.0
 
 
 def test_apply_socket_link_profile_keys_by_source_and_sink():
@@ -168,7 +194,7 @@ def test_apply_node_timed_out_drops_profiles():
     )
     state = apply(
         state,
-        _wrap(1, AneProfile(engine="ane", tflops_fp16=1, memory_bandwidth_gbps=1)),
+        _wrap(1, _ane_profile()),
     )
     state = apply(
         state,

From 45b43bb41f51d24701bdeb120fd32da10ce7f9cd Mon Sep 17 00:00:00 2001
From: Alex Cheema <alexcheema123@gmail.com>
Date: Sun, 3 May 2026 19:36:52 +0100
Subject: [PATCH 3/4] fix(profilers): verify native ANE quantized compute

---
 .../src/lib/components/TopologyGraph.svelte   |   8 +-
 dashboard/src/lib/stores/app.svelte.ts        |   5 +
 src/exo/shared/apply.py                       |   5 +
 src/exo/shared/types/profiling.py             |   5 +
 src/exo/utils/profilers/ane_profiler.py       | 208 ++++++++++++++----
 .../profilers/tests/test_ane_profiler.py      |  22 ++
 .../profilers/tests/test_apply_profiles.py    |  12 +
 7 files changed, 227 insertions(+), 38 deletions(-)

diff --git a/dashboard/src/lib/components/TopologyGraph.svelte b/dashboard/src/lib/components/TopologyGraph.svelte
index a952ae7cc3..961040a5f9 100644
--- a/dashboard/src/lib/components/TopologyGraph.svelte
+++ b/dashboard/src/lib/components/TopologyGraph.svelte
@@ -80,6 +80,12 @@
     return `${value.toFixed(value >= 10 ? 0 : 1)}T`;
   }
 
+  function formatAnePrecisionLabel(
+    profile: RawNodeAnePrecisionProfile,
+  ): string {
+    return `W${profile.weightBits}A${profile.activationBits} ${formatAneTops(profile.computeTops)}`;
+  }
+
   function formatAneProfileLine(profile: RawNodeAneProfile): string {
     const preferredOrder: RawNodeAnePrecisionProfile["precisionBits"][] = [
       16, 8, 4, 32,
@@ -90,7 +96,7 @@
     return preferredOrder
       .map((bits) => byBits.get(bits))
       .filter((p): p is RawNodeAnePrecisionProfile => Boolean(p?.supported))
-      .map((p) => `${p.precisionBits}b ${formatAneTops(p.computeTops)}`)
+      .map((p) => formatAnePrecisionLabel(p))
       .join(" · ");
   }
 
diff --git a/dashboard/src/lib/stores/app.svelte.ts b/dashboard/src/lib/stores/app.svelte.ts
index a5352ca339..c88169ab42 100644
--- a/dashboard/src/lib/stores/app.svelte.ts
+++ b/dashboard/src/lib/stores/app.svelte.ts
@@ -141,7 +141,12 @@ export interface RawNodeAnePrecisionProfile {
   activationBits: 32 | 16 | 8 | 4;
   supported: boolean;
   computeTops?: number | null;
+  weightOnlyComputeTops?: number | null;
+  singleInstanceComputeTops?: number | null;
+  computeInstances?: number;
   memoryBandwidthGbps?: number | null;
+  activationQuantizationSpeedup?: number | null;
+  nativeQuantizedCompute?: boolean | null;
   error?: string | null;
 }
 
diff --git a/src/exo/shared/apply.py b/src/exo/shared/apply.py
index 4852cff1d0..0002c96495 100644
--- a/src/exo/shared/apply.py
+++ b/src/exo/shared/apply.py
@@ -489,7 +489,12 @@ def apply_node_gathered_info(event: NodeGatheredInfo, state: State) -> State:
                             activation_bits=profile.activation_bits,
                             supported=profile.supported,
                             compute_tops=profile.compute_tops,
+                            weight_only_compute_tops=profile.weight_only_compute_tops,
+                            single_instance_compute_tops=profile.single_instance_compute_tops,
+                            compute_instances=profile.compute_instances,
                             memory_bandwidth_gbps=profile.memory_bandwidth_gbps,
+                            activation_quantization_speedup=profile.activation_quantization_speedup,
+                            native_quantized_compute=profile.native_quantized_compute,
                             error=profile.error,
                         )
                         for profile in info.precision_profiles
diff --git a/src/exo/shared/types/profiling.py b/src/exo/shared/types/profiling.py
index ab75d805e2..a2da314cc5 100644
--- a/src/exo/shared/types/profiling.py
+++ b/src/exo/shared/types/profiling.py
@@ -130,7 +130,12 @@ class NodeAnePrecisionProfile(FrozenModel):
     activation_bits: Literal[32, 16, 8, 4]
     supported: bool
     compute_tops: float | None = None
+    weight_only_compute_tops: float | None = None
+    single_instance_compute_tops: float | None = None
+    compute_instances: int = 1
     memory_bandwidth_gbps: float | None = None
+    activation_quantization_speedup: float | None = None
+    native_quantized_compute: bool | None = None
     error: str | None = None
 
 
diff --git a/src/exo/utils/profilers/ane_profiler.py b/src/exo/utils/profilers/ane_profiler.py
index 7f2d81056a..4dbc4d4b96 100644
--- a/src/exo/utils/profilers/ane_profiler.py
+++ b/src/exo/utils/profilers/ane_profiler.py
@@ -18,8 +18,9 @@
 import time
 import uuid
 from collections.abc import Callable, Sequence
+from concurrent.futures import ThreadPoolExecutor
 from pathlib import Path
-from typing import Literal, Self, cast, final
+from typing import Literal, NamedTuple, Self, cast, final
 
 import numpy as np
 import numpy.typing as npt
@@ -43,6 +44,8 @@
 _COMPUTE_OUT_DIM = 2048
 _COMPUTE_SPATIAL = 768
 _COMPUTE_ITERATIONS_PER_PASS = 12
+_COMPUTE_PARALLEL_INSTANCES = 2
+_NATIVE_QUANTIZATION_SPEEDUP_THRESHOLD = 1.2
 
 _STREAM_CHANNELS = 4096
 _STREAM_SPATIAL = 2048
@@ -79,7 +82,12 @@ class AnePrecisionProfile(FrozenModel):
     activation_bits: AnePrecisionBits
     supported: bool
     compute_tops: float | None = None
+    weight_only_compute_tops: float | None = None
+    single_instance_compute_tops: float | None = None
+    compute_instances: int = 1
     memory_bandwidth_gbps: float | None = None
+    activation_quantization_speedup: float | None = None
+    native_quantized_compute: bool | None = None
     error: str | None = None
 
 
@@ -105,6 +113,12 @@ def _measure_blocking(cls) -> Self:
         return cls(engine="ane", precision_profiles=profiles)
 
 
+class _ComputeMeasurement(NamedTuple):
+    tops: float
+    single_instance_tops: float
+    instances: int
+
+
 def _measure_precision_profile(
     precision_bits: AnePrecisionBits, rng: np.random.Generator
 ) -> AnePrecisionProfile:
@@ -123,9 +137,13 @@ def _measure_precision_profile(
 def _measure_supported_precision_profile(
     precision_bits: AnePrecisionBits, rng: np.random.Generator
 ) -> AnePrecisionProfile:
+    # Core ML accepts int4 weights here, but the ANE activation Q/DQ path that
+    # triggers native low-precision conv uses int8 activations. int4 activation
+    # Q/DQ does not compile through this MIL/private-ANE path, so 4-bit means W4A8.
     activation_bits: AnePrecisionBits = (
         8 if precision_bits in (8, 4) else precision_bits
     )
+    baseline_measurement: _ComputeMeasurement | None = None
 
     if precision_bits == 32:
         compute_weights = rng.standard_normal(
@@ -134,14 +152,18 @@ def _measure_supported_precision_profile(
         compute_input = rng.standard_normal((_COMPUTE_IN_DIM, _COMPUTE_SPATIAL)).astype(
             np.float32
         )
-        compute_kernel = _compile_static_conv_kernel(
-            weights=compute_weights,
-            input_channels=_COMPUTE_IN_DIM,
-            output_channels=_COMPUTE_OUT_DIM,
-            spatial=_COMPUTE_SPATIAL,
-            precision_bits=precision_bits,
+        compute_kernels = tuple(
+            _compile_static_conv_kernel(
+                weights=compute_weights,
+                input_channels=_COMPUTE_IN_DIM,
+                output_channels=_COMPUTE_OUT_DIM,
+                spatial=_COMPUTE_SPATIAL,
+                precision_bits=precision_bits,
+            )
+            for _ in range(_COMPUTE_PARALLEL_INSTANCES)
         )
-        compute_kernel.write_input(0, compute_input)
+        for kernel in compute_kernels:
+            kernel.write_input(0, compute_input)
 
         stream_input = rng.standard_normal((_STREAM_CHANNELS, _STREAM_SPATIAL)).astype(
             np.float32
@@ -164,14 +186,18 @@ def _measure_supported_precision_profile(
             np.float16
         )
 
-        compute_kernel = _compile_static_conv_kernel(
-            weights=compute_weights,
-            input_channels=_COMPUTE_IN_DIM,
-            output_channels=_COMPUTE_OUT_DIM,
-            spatial=_COMPUTE_SPATIAL,
-            precision_bits=precision_bits,
+        compute_kernels = tuple(
+            _compile_static_conv_kernel(
+                weights=compute_weights,
+                input_channels=_COMPUTE_IN_DIM,
+                output_channels=_COMPUTE_OUT_DIM,
+                spatial=_COMPUTE_SPATIAL,
+                precision_bits=precision_bits,
+            )
+            for _ in range(_COMPUTE_PARALLEL_INSTANCES)
         )
-        compute_kernel.write_input(0, compute_input)
+        for kernel in compute_kernels:
+            kernel.write_input(0, compute_input)
 
         stream_kernel = _compile_relu_stream_kernel(
             channels=_STREAM_CHANNELS,
@@ -191,27 +217,64 @@ def _measure_supported_precision_profile(
             np.float16
         )
 
-        compute_kernel = _compile_quantized_static_conv_kernel(
-            weights=compute_weights,
-            input_channels=_COMPUTE_IN_DIM,
-            output_channels=_COMPUTE_OUT_DIM,
-            spatial=_COMPUTE_SPATIAL,
-            precision_bits=precision_bits,
+        compute_kernels = tuple(
+            _compile_quantized_static_conv_kernel(
+                weights=compute_weights,
+                input_channels=_COMPUTE_IN_DIM,
+                output_channels=_COMPUTE_OUT_DIM,
+                spatial=_COMPUTE_SPATIAL,
+                precision_bits=precision_bits,
+                quantize_activations=True,
+            )
+            for _ in range(_COMPUTE_PARALLEL_INSTANCES)
         )
-        compute_kernel.write_input(0, compute_input)
+        for kernel in compute_kernels:
+            kernel.write_input(0, compute_input)
 
         stream_kernel = _compile_quantized_stream_kernel(
             channels=_STREAM_CHANNELS,
             spatial=_STREAM_SPATIAL,
         )
         stream_kernel.write_input(0, stream_input)
+        baseline_kernels = tuple(
+            _compile_quantized_static_conv_kernel(
+                weights=compute_weights,
+                input_channels=_COMPUTE_IN_DIM,
+                output_channels=_COMPUTE_OUT_DIM,
+                spatial=_COMPUTE_SPATIAL,
+                precision_bits=precision_bits,
+                quantize_activations=False,
+            )
+            for _ in range(_COMPUTE_PARALLEL_INSTANCES)
+        )
+        for kernel in baseline_kernels:
+            kernel.write_input(0, compute_input)
+        baseline_measurement = _measure_compute_tops(baseline_kernels)
+
+    measurement = _measure_compute_tops(compute_kernels)
+    activation_quantization_speedup = None
+    native_quantized_compute = None
+    if baseline_measurement is not None:
+        activation_quantization_speedup = (
+            measurement.tops / baseline_measurement.tops
+            if baseline_measurement.tops > 0
+            else 0.0
+        )
+        native_quantized_compute = (
+            activation_quantization_speedup >= _NATIVE_QUANTIZATION_SPEEDUP_THRESHOLD
+        )
 
     return AnePrecisionProfile(
         precision_bits=precision_bits,
         weight_bits=precision_bits,
         activation_bits=activation_bits,
         supported=True,
-        compute_tops=_measure_compute_tops(compute_kernel),
+        compute_tops=measurement.tops,
+        weight_only_compute_tops=(
+            baseline_measurement.tops if baseline_measurement is not None else None
+        ),
+        single_instance_compute_tops=measurement.single_instance_tops,
+        compute_instances=measurement.instances,
         memory_bandwidth_gbps=_measure_streaming_bandwidth_gbps(
             stream_kernel,
             bytes_per_element=(
@@ -220,6 +283,8 @@ def _measure_supported_precision_profile(
                 else _FP16_BYTES_PER_ELEMENT
             ),
         ),
+        activation_quantization_speedup=activation_quantization_speedup,
+        native_quantized_compute=native_quantized_compute,
     )
 
 
@@ -311,23 +376,83 @@ def ane_available() -> bool:
         return False
 
 
-def _measure_compute_tops(kernel: _AneKernel) -> float:
-    _warm_up_with(kernel.eval, time.perf_counter() + _WARMUP_SECONDS)
-
+def _measure_compute_tops(kernels: Sequence[_AneKernel]) -> _ComputeMeasurement:
     operations_per_iteration = 2 * _COMPUTE_IN_DIM * _COMPUTE_OUT_DIM * _COMPUTE_SPATIAL
+    single_instance_tops = _measure_parallel_compute_tops(
+        kernels[:1], operations_per_iteration=operations_per_iteration
+    )
+    best = _ComputeMeasurement(
+        tops=single_instance_tops,
+        single_instance_tops=single_instance_tops,
+        instances=1,
+    )
+    if len(kernels) >= 2:
+        parallel_tops = _measure_parallel_compute_tops(
+            kernels[:2], operations_per_iteration=operations_per_iteration
+        )
+        if parallel_tops > best.tops:
+            best = _ComputeMeasurement(
+                tops=parallel_tops,
+                single_instance_tops=single_instance_tops,
+                instances=2,
+            )
+    return best
+
+
+def _measure_parallel_compute_tops(
+    kernels: Sequence[_AneKernel], *, operations_per_iteration: int
+) -> float:
+    if not kernels:
+        return 0.0
+    _warm_up_kernels(kernels, time.perf_counter() + _WARMUP_SECONDS)
+
     best_tops = 0.0
-    for _ in range(_MEASUREMENT_PASSES):
-        start = time.perf_counter()
-        for _ in range(_COMPUTE_ITERATIONS_PER_PASS):
-            kernel.eval()
-        elapsed = time.perf_counter() - start
-        if elapsed <= 0:
-            continue
-        tops = operations_per_iteration * _COMPUTE_ITERATIONS_PER_PASS / elapsed / 1e12
-        best_tops = max(best_tops, tops)
+    with ThreadPoolExecutor(max_workers=len(kernels)) as executor:
+        for _ in range(_MEASUREMENT_PASSES):
+            start = time.perf_counter()
+            if len(kernels) == 1:
+                _eval_kernel_iterations(kernels[0], _COMPUTE_ITERATIONS_PER_PASS)
+            else:
+                futures = [
+                    executor.submit(
+                        _eval_kernel_iterations,
+                        kernel,
+                        _COMPUTE_ITERATIONS_PER_PASS,
+                    )
+                    for kernel in kernels
+                ]
+                for future in futures:
+                    future.result()
+            elapsed = time.perf_counter() - start
+            if elapsed <= 0:
+                continue
+            tops = (
+                operations_per_iteration
+                * _COMPUTE_ITERATIONS_PER_PASS
+                * len(kernels)
+                / elapsed
+                / 1e12
+            )
+            best_tops = max(best_tops, tops)
     return best_tops
 
 
+def _warm_up_kernels(kernels: Sequence[_AneKernel], deadline: float) -> None:
+    if len(kernels) == 1:
+        _warm_up_with(kernels[0].eval, deadline)
+        return
+    with ThreadPoolExecutor(max_workers=len(kernels)) as executor:
+        while time.perf_counter() < deadline:
+            futures = [executor.submit(kernel.eval) for kernel in kernels]
+            for future in futures:
+                future.result()
+
+
+def _eval_kernel_iterations(kernel: _AneKernel, iterations: int) -> None:
+    for _ in range(iterations):
+        kernel.eval()
+
+
 def _measure_streaming_bandwidth_gbps(
     kernel: _AneKernel, *, bytes_per_element: int
 ) -> float:
@@ -386,6 +511,7 @@ def _compile_quantized_static_conv_kernel(
     output_channels: int,
     spatial: int,
     precision_bits: Literal[8, 4],
+    quantize_activations: bool,
 ) -> _AneKernel:
     weight_blob = _build_quantized_weight_blob_bytes(weights, precision_bits)
     mil = _mil_quantized_static_conv(
@@ -393,6 +519,7 @@ def _compile_quantized_static_conv_kernel(
         output_channels=output_channels,
         spatial=spatial,
         precision_bits=precision_bits,
+        quantize_activations=quantize_activations,
     )
     input_bytes = input_channels * spatial * _FP16_BYTES_PER_ELEMENT
     output_bytes = output_channels * spatial * _FP16_BYTES_PER_ELEMENT
@@ -1296,6 +1423,7 @@ def _mil_quantized_static_conv(
     output_channels: int,
     spatial: int,
     precision_bits: Literal[8, 4],
+    quantize_activations: bool,
 ) -> str:
     data_type = "int8" if precision_bits == 8 else "int4"
     scale_blob_offset = 128 + _quantized_data_payload_size(
@@ -1303,6 +1431,12 @@ def _mil_quantized_static_conv(
         output_channels=output_channels,
         precision_bits=precision_bits,
     )
+    activation_preamble = (
+        _mil_quantize_dequantize(channels=input_channels, spatial=spatial)
+        if quantize_activations
+        else ""
+    )
+    conv_input = "dx" if quantize_activations else "x"
     return (
         f"{_MIL_HEADER}"
         f"    func main<ios18>(tensor<fp16, [1, {input_channels}, 1, {spatial}]> x) {{\n"
@@ -1313,10 +1447,10 @@ def _mil_quantized_static_conv(
         f"scale = tensor<fp16, [{output_channels}, 1, 1, 1]>(BLOBFILE(path = "
         'string("@model_path/weights/weight.bin"), '
         f'offset = uint64({scale_blob_offset}))))[name = string("Wq")];\n'
-        f"{_mil_quantize_dequantize(channels=input_channels, spatial=spatial)}"
+        f"{activation_preamble}"
         f"{_CONV_PARAMS}"
         f"        tensor<fp16, [1, {output_channels}, 1, {spatial}]> y = conv("
-        f'{_CONV_ARGS}, weight=W, x=dx)[name=string("cv")];\n'
+        f'{_CONV_ARGS}, weight=W, x={conv_input})[name=string("cv")];\n'
         "    } -> (y);\n"
         "}\n"
     )
diff --git a/src/exo/utils/profilers/tests/test_ane_profiler.py b/src/exo/utils/profilers/tests/test_ane_profiler.py
index 944cec256c..90b9953b12 100644
--- a/src/exo/utils/profilers/tests/test_ane_profiler.py
+++ b/src/exo/utils/profilers/tests/test_ane_profiler.py
@@ -23,9 +23,20 @@ async def test_ane_profile_returns_plausible_numbers():
         assert precision.compute_tops is not None
         assert precision.compute_tops > 0
         assert precision.compute_tops < 1000
+        assert precision.single_instance_compute_tops is not None
+        assert precision.single_instance_compute_tops > 0
+        assert precision.compute_instances >= 1
         assert precision.memory_bandwidth_gbps is not None
         assert precision.memory_bandwidth_gbps > 0
         assert precision.memory_bandwidth_gbps < 100_000
+    for bits in (8, 4):
+        precision = precision_by_bits[bits]
+        assert precision.weight_bits == bits
+        assert precision.activation_bits == 8
+        assert precision.weight_only_compute_tops is not None
+        assert precision.weight_only_compute_tops > 0
+        assert precision.activation_quantization_speedup is not None
+        assert precision.native_quantized_compute is not None
 
     fp32 = precision_by_bits[32]
     if fp32.supported:
@@ -45,6 +56,9 @@ def test_ane_profile_serializes_with_class_tag():
                 activation_bits=16,
                 supported=True,
                 compute_tops=12.3,
+                weight_only_compute_tops=11.1,
+                single_instance_compute_tops=6.2,
+                compute_instances=2,
                 memory_bandwidth_gbps=80.0,
             ),
         ),
@@ -53,3 +67,11 @@ def test_ane_profile_serializes_with_class_tag():
     assert "AneProfile" in dumped
     assert dumped["AneProfile"]["precision_profiles"][0]["precision_bits"] == 16
     assert dumped["AneProfile"]["precision_profiles"][0]["compute_tops"] == 12.3
+    assert (
+        dumped["AneProfile"]["precision_profiles"][0]["single_instance_compute_tops"]
+        == 6.2
+    )
+    assert (
+        dumped["AneProfile"]["precision_profiles"][0]["weight_only_compute_tops"]
+        == 11.1
+    )
diff --git a/src/exo/utils/profilers/tests/test_apply_profiles.py b/src/exo/utils/profilers/tests/test_apply_profiles.py
index c5dc0e41c0..9fc1556a77 100644
--- a/src/exo/utils/profilers/tests/test_apply_profiles.py
+++ b/src/exo/utils/profilers/tests/test_apply_profiles.py
@@ -43,6 +43,8 @@ def _ane_profile() -> AneProfile:
                 activation_bits=16,
                 supported=True,
                 compute_tops=12.0,
+                single_instance_compute_tops=6.0,
+                compute_instances=2,
                 memory_bandwidth_gbps=80.0,
             ),
             AnePrecisionProfile(
@@ -51,7 +53,12 @@ def _ane_profile() -> AneProfile:
                 activation_bits=8,
                 supported=True,
                 compute_tops=28.0,
+                weight_only_compute_tops=18.0,
+                single_instance_compute_tops=14.0,
+                compute_instances=2,
                 memory_bandwidth_gbps=72.0,
+                activation_quantization_speedup=1.5,
+                native_quantized_compute=True,
             ),
         ),
     )
@@ -75,8 +82,13 @@ def test_apply_ane_profile_writes_node_ane_profile():
     assert entry.engine == "ane"
     precision_by_bits = {p.precision_bits: p for p in entry.precision_profiles}
     assert precision_by_bits[16].compute_tops == 12.0
+    assert precision_by_bits[16].single_instance_compute_tops == 6.0
+    assert precision_by_bits[16].compute_instances == 2
     assert precision_by_bits[16].memory_bandwidth_gbps == 80.0
     assert precision_by_bits[8].compute_tops == 28.0
+    assert precision_by_bits[8].weight_only_compute_tops == 18.0
+    assert precision_by_bits[8].activation_quantization_speedup == 1.5
+    assert precision_by_bits[8].native_quantized_compute
 
 
 def test_apply_socket_link_profile_keys_by_source_and_sink():

From 1a08f699bb6e850ebb122d159d3d0fa335c2eed8 Mon Sep 17 00:00:00 2001
From: Alex Cheema <alexcheema123@gmail.com>
Date: Sun, 3 May 2026 20:00:11 +0100
Subject: [PATCH 4/4] fix(profilers): sweep ANE compute concurrency

---
 src/exo/utils/profilers/ane_profiler.py | 30 ++++++++++++++++---------
 1 file changed, 19 insertions(+), 11 deletions(-)

diff --git a/src/exo/utils/profilers/ane_profiler.py b/src/exo/utils/profilers/ane_profiler.py
index 4dbc4d4b96..f2ccaf6fc9 100644
--- a/src/exo/utils/profilers/ane_profiler.py
+++ b/src/exo/utils/profilers/ane_profiler.py
@@ -44,7 +44,8 @@
 _COMPUTE_OUT_DIM = 2048
 _COMPUTE_SPATIAL = 768
 _COMPUTE_ITERATIONS_PER_PASS = 12
-_COMPUTE_PARALLEL_INSTANCES = 2
+_COMPUTE_PARALLEL_INSTANCES = 4
+_COMPUTE_INSTANCE_COUNTS = (1, 2, 4)
 _NATIVE_QUANTIZATION_SPEEDUP_THRESHOLD = 1.2
 
 _STREAM_CHANNELS = 4096
@@ -378,23 +379,30 @@ def ane_available() -> bool:
 
 def _measure_compute_tops(kernels: Sequence[_AneKernel]) -> _ComputeMeasurement:
     operations_per_iteration = 2 * _COMPUTE_IN_DIM * _COMPUTE_OUT_DIM * _COMPUTE_SPATIAL
-    single_instance_tops = _measure_parallel_compute_tops(
-        kernels[:1], operations_per_iteration=operations_per_iteration
-    )
+    single_instance_tops = 0.0
     best = _ComputeMeasurement(
-        tops=single_instance_tops,
+        tops=0.0,
         single_instance_tops=single_instance_tops,
         instances=1,
     )
-    if len(kernels) >= 2:
-        parallel_tops = _measure_parallel_compute_tops(
-            kernels[:2], operations_per_iteration=operations_per_iteration
+    for instance_count in _COMPUTE_INSTANCE_COUNTS:
+        if len(kernels) < instance_count:
+            continue
+        tops = _measure_parallel_compute_tops(
+            kernels[:instance_count], operations_per_iteration=operations_per_iteration
         )
-        if parallel_tops > best.tops:
+        if instance_count == 1:
+            single_instance_tops = tops
+            best = _ComputeMeasurement(
+                tops=tops,
+                single_instance_tops=single_instance_tops,
+                instances=1,
+            )
+        elif tops > best.tops:
             best = _ComputeMeasurement(
-                tops=parallel_tops,
+                tops=tops,
                 single_instance_tops=single_instance_tops,
-                instances=2,
+                instances=instance_count,
             )
     return best