Module gpu 

Accelerator backends: a trait the model’s linear layers dispatch through, plus a global registry chosen at start-up from the platform detection. Every backend is optional at run time; the CPU kernels remain the reference path and the fallback.

Structs

ConvGeom

Geometry of a k x k convolution with padding k/2 over channels-last input [h*w, cin], producing [oh*ow, n_out]; the weights are in the unfolded [n_out][ky][kx][cin] layout.

Job

One matmul of a batch sharing the same input x: y[t][n_out] = x[t][n_in] * W^T.

NormSpec

Enums

NormKind

A normalization applied to a convolution’s input as it is loaded (so it runs on the device without an extra pass): GroupNorm(groups) over the channels of all pixels, or Layer over each pixel’s (row’s) channels; then * gamma + beta, then SiLU if silu.

WeightRef

A weight matrix in one of the supported storage formats ([n_out, n_in]).

Traits

Backend

Functions

active

exact

See set_exact.

init

Picks and initialises a backend for the detected platform and the user’s preference.

initialized

Whether a backend choice (possibly “none”) was installed already.

install

Installs the backend for this process (call once, before loading models).

moe_offload

See set_moe_offload.

set_exact

Keeps GPU matmuls bit-identical to the CPU: no reduced-precision matrix units (Metal on M5/A19 converts inputs to half with f32 accumulation, a KL divergence of 1e-7 to 1e-3 from the CPU). Process-wide, takes effect on the next dispatch. Default: off, or OXINFER_METAL_MMA=0.

set_moe_offload

Runs MoE expert matmuls on the GPU when the backend can (prefill). Process-wide; affects models loaded afterwards. Default: on, or OXINFER_METAL_MOE=0 for off.

Type Aliases

DeviceWeight

Device-side state for one weight matrix (buffers / handles), owned by the layer.