Foundation note for the GPU-now track. Custom Metal kernels are a first-class Core AI authoring feature (WWDC 325) and are GPU-only — the ANE cannot run hand-written MSL. Sources:
coreai-torch/docs/guides/custom-metal-kernels.ipynb,coreai-torch/docs/api/TorchMetalKernel.md,coreai-torch/coreai_torch/_torch_metal_kernel.py,coreai/authoring/metal.py,ondevice/_wwdc325_transcript.txt(videoMdlyLT_y3i0),skills/.../model-authoring/references/gpu_rules.md.
A custom Metal kernel lets you write a raw MSL GPU function, wrap it with a PyTorch reference, and have
torch.export + coreai-torch embed it into the .aimodel as a real Core AI op (coreai.metal4_kernel).
It is not a raw-Metal bypass — the MSL travels inside the single .aimodel artifact and runs in the OS
Core AI runtime (WWDC 325, ~line 476–485).
Use it for op fusion — collapse several graph ops (and their dispatch + intermediate-memory traffic) into one kernel dispatch. WWDC 325 (~line 461): “You can take a group of these ops and fuse them into a single operation. This replaces several steps with a single kernel dispatch within the graph.” Core AI already ships fast prepackaged kernels for heavy ops (e.g. SDPA); write custom ones for what it doesn’t cover or to fuse a hot path.
GPU-only — and that is the structural reason the speed track lives on the GPU. The ANE runs only fixed
hardware ops (Conv/LayerNorm/…); it cannot execute arbitrary MSL. So “write fused-int8 kernels” is, by
construction, a GPU strategy — independent of any beta bug. (See compute-units-and-authoring.md.)
from coreai_torch import TorchMetalKernel, TorchConverter, get_decomp_table
from coreai.authoring import MetalParameter # re-exported as coreai_torch.MetalParameter
TorchMetalKernel(
name: str, # kernel id; becomes part of the generated kernel name
input_names: list[str], # names used in the MSL body; count == torch_defn params
result_names: list[str], # output names used in the MSL body (>=1)
src: str, # MSL body ONLY (signature/buffer bindings/#includes auto-generated)
torch_defn: Callable, # PyTorch reference — what torch.export sees, used for shape inference
metal_params: list[MetalParameter] | None = None, # thread attrs, e.g. thread_position_in_grid
helper_src: str | None = None, # extra MSL pasted before the kernel (helpers/typedefs/constants)
template_dtypes: dict[str,str] | None = None, # input_name -> placeholder string in `src`
)
MetalParameter(name: str, dtype: str, attr: str) # e.g. ("gid","uint2","thread_position_in_grid")
(_torch_metal_kernel.py:44-93; metal.py:36-52)
Call it inside an nn.Module like a function, passing dispatch + output shapes per call:
out = kernel(*args,
threads_per_grid=(N,1,1),
threads_per_thread_group=(T,1,1),
result_shapes=[list(out_shape)]) # one entry per result; enables dynamic shapes
result_shapes is how Core AI “bakes in” output-shape-from-input-shape so the kernel works under dynamic
shapes (WWDC 325, ~line 511–518).
converter = TorchConverter()
converter.register_custom_kernels([kernel]) # FIRST
converter.add_exported_program(ep, input_names=[...], output_names=[...]) # THEN
prog = converter.to_coreai(); prog.optimize()
WWDC 325 (~line 519): “I register my custom kernels with the converter, then add the exported program as
before. The metal source gets embedded directly in the asset, a single artifact. The kernel travels with the
model.” This matches register_custom_kernels → builds a register_torch_lowering handler for
coreai_metal_kernels::<name>.default (converter.py:1003-1033).
def torch_add(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: # reference for shape inference
return x + y
vadd = TorchMetalKernel(
"vector_add", input_names=["x","y"], result_names=["output"],
src="output[id] = x[id] + y[id];",
torch_defn=torch_add,
metal_params=[MetalParameter("id","uint","thread_position_in_grid")],
)
class M(nn.Module):
def forward(self, x, y):
return vadd(x, y, threads_per_grid=(x.shape[0],1,1),
threads_per_thread_group=(1,1,1), result_shapes=[list(x.shape)])
ep = torch.export.export(M().eval(), (x,y)).run_decompositions(get_decomp_table())
conv = TorchConverter(); conv.register_custom_kernels([vadd])
conv.add_exported_program(ep, input_names=["x","y"], output_names=["result"])
conv.to_coreai().optimize()
torch.export traces the torch_defn reference; the op is registered under coreai_metal_kernels:: (_torch_metal_kernel.py:265).to_coreai() lowers it to coreai.metal4_kernel via CustomMetalKernel._construct_kernel_op (metal.py:423-479): forces Metal-buffer backing on inputs/results, templates the MSL with the actual dtypes, bakes the result-shape computation..aimodel; the runtime dispatches it on the GPU at the right graph point._ensure_metal_backed, metal.py:347-367) — the GPU can’t read host memory mid-kernel.metal.py PARAMETER_LIMIT=31.bf16→bfloat, f16→half, f32→float, si8→int8_t, ui8→uint8_t, ui32→uint, si32→int, i1→bool, …, metal.py:113-126).template_dtypes: {"A":"TYPE"} replaces the placeholder string "TYPE" in src with input A’s Metal dtype at compile time → one kernel serves multiple dtypes. Keys must be real inputs; placeholders must be unique (metal.py:152-177).torch_defn validation: every param annotated Tensor|int|float|bool, no *args/**kwargs, param count == len(input_names), return Tensor | list[Tensor] | tuple[Tensor,...] with concrete length matching result_names (_torch_metal_kernel.py:141-211).ondevice/_moe_kernel_probe.py). So a kernel can take an index tensor as an INPUT and read only
the rows it points at — W[m, n, e] with e = uint(IDX[slot]) reads only expert-slab e out of
a [E, N, M] tensor (the gather_qmm MoE kernel, models/macos/moe_metal.py → the deferred
MoE-over-read fix; LFM2.5-8B-A1B int8 39→141 tok/s). The torch_defn must stay fake-traceable:
express the gather as torch.index_select(W, 0, idx) (shape-static), NEVER int(idx[i])
(FakeTensor has no concrete value). Same kernel runs on M4 Max AND the iPhone 17 Pro A19 Pro GPU.x across the routed experts,
but the down projection feeds each expert its OWN gated activation — so the kernel’s A must
be [k, K] (one row per slot, A[c, slot]), with x replicated k-wide for gate/up. Treating A
as a single shared [1, K] row silently corrupts down (relative error ~1.3 = garbage).project_macos_speed_state project memory + ondevice/_macos_speed_RESULTS.md.)F.scaled_dot_product_attention) — already fused; don’t hand-roll it.tensorops-quantized-kernels.md.metal4_kernel op is a fusion barrier, and its edges are materialized in the dtype/layout the kernel asks for. A dtype cast on the boundary (e.g. state.float() in, .half() out) is a real coreai.cast op that blows the tensor up — hand large state/activation tensors across in their native dtype and accumulate fp32 in registers.blockwise_shift_scale only fuses into the matmul when its shift is all-zero — so symmetric int4 dequants cheaper than int8, but asymmetric (zero-point) int4 falls off the fast path (~4.6× slower on device). A fused affine-int4 matvec factors the zero point out (acc += s*(dot − z*Σx)) to keep the fast path AND the accuracy; it beats stock-asymmetric 3.3× Mac / ~5× device but only ties/+11% int8 (nibble-unpack ALU caps it). See ssm-hybrid-decode-int4-prefill-findings.md §3.