Module gpu 

Facade for GPU backends: a single call entry point for qtensor/pipeline/ linear_core. Job types and the threshold are canonical HERE; behind the facade dispatch goes to a platform backend:

• gpu_metal (Apple Silicon, unified memory + no-copy buffers);
• gpu_wgpu (C1: Vulkan/DX12/Metal — NVIDIA/Radeon/Intel/Apple, weights resident in VRAM), available under --features gpu.

Runtime selection via CMF_GPU: 1 — native Metal (macOS) or wgpu (other OSes); wgpu — force wgpu (including for the local Metal-via-wgpu parity test). Any backend refusal — false and the honest CPU path, no partial results.

Structs

BatchJob

A single independent batch matvec (GDN projections of one input).

DitBlockArgs

One whole modulated DiT block for dit_block: geometry, norm weights, AdaLN scale/gate vectors (gates pre-tanh’d), a per-token f32 RoPE cos/sin table, and the directory indices of the seven q4t projections. x is in-out [n, hidden].

GraphLayer

Per-layer weights for the whole-token wgpu graph.

GraphW

One weight in the whole-token graph: tensor idx + a codec tag (0=q8_row, 1=q1, 2=q4_tiled, 3=q1t, 4=f32) + per-row scales (q8_row only) + the raw f32 data (kind 4 only — small unquantized projections like GDN in_proj_a/b).

MoeJob

A single MoE-FFN job (an expert with its own weight), executed in one submission: (rows, cols, idx, row_scale) for gate/up/down + prescaled inputs + the down θ-field + the blending weight.

SpecTail

Speculative-verify tail for the batched graph: fold final-norm + lm_head over every batch position and read all k logit rows back; the batch also snapshots the GDN state per position for gdn_spec_restore.

VaeResnetArgs

One VAE resnet block for vae_resnet: norm/conv weights and the channel/shape geometry. shortcut is the 1×1 projection (w, b, k) when in/out channels differ.

Enums

BatchLayout

Which kernel a batched matvec needs. q8 carries row scales in a side buffer; the rest embed them in the payload and differ in stride.

GraphAttn

A layer’s token-mixing op: standard attention or a GDN (linear-attention) block. The surrounding norms + SwiGLU FFN are common to both.

GraphFfn

The FFN of one graph layer: a dense SwiGLU trio, or a routed MoE — router + top-k selection + all selected experts run ON DEVICE (the routing decision depends on the resident hidden state, so a CPU round-trip per layer would forfeit the one-submit design).

GraphPhase

Which half of the run is asking.

OpClass

GPU-eligible op classes, each with an independent probe.

ProbeArm

Probe verdict for one call.

Constants

GPU_MIN_ROWS

Default row threshold: the GPU takes only larger matrices (lm_head class). Below it, the dispatch/readback cost does not pay off on unified memory.

Functions

attn_dropin

Whole attention sub-block on the wgpu token graph (drop-in for qwen_attention): normed hidden in, O-projection out, resident device K/V mirror. false = refusal / not the wgpu backend → CPU path.

backend_available

GPU enabled and initialized on the selected backend? Whether THIS build can bring a GPU up on THIS device: a compiled-in backend plus a live adapter. The mobile FFI exposes it so an app can tell “GPU off” from “GPU impossible” (a CPU-only .so ships no backend at all). Cached after the first call.

bake_precision_strict

Strict-f32 for the bake’s GEMMs (phase A mask training): the mask selects neurons by a gradient signal, and f16 operand rounding on that signal closes the wrong ones. No-op off the wgpu backend.

bake_release

Hand the card back after a bake: drop its resident weights, planes and pools so the ordinary engine (the runtime gate, a serve that follows) starts from a clean budget. No-op off the wgpu backend.

cache_dir_pub

Same directory, for the backends.

chunk_attend

Fused DiT SwiGLU FFN on the device: g=X·W1ᵀ, u=X·W3ᵀ, silu(g)·u, Causal chunk attention on the device: b queries against s0 + b cached keys. wgpu only — Metal’s chunk graph keeps attention inside the resident block and never calls out.

conv1d_gemm

A 1D convolution as a GEMM whose column matrix is expanded on the device instead of being built, transposed and uploaded by the host. yt comes back [out_n x oc]. false = refused, caller runs host.

cpu_scope

Run f with the GPU gates off on this thread (pure-CPU arm).

cur_layer

The layer set_layer last marked on this thread (−1 outside layers).

current_device

The device this thread is pinned to.

default_device

Which GPU this thread’s engine calls address. Multi-card hosts hold one wgpu context PER card (weights, KV mirrors and scratch live inside a context, so per-device contexts give per-device caches for free); this thread-local says which one is current. Default: the process pin (CMF_GPU_ADAPTER) or 0 — so single-card runs behave exactly as they always have.

device_count

How many GPUs this process can address (wgpu adapter count; 1 on Metal, 0 without a backend).

discrete

Is the active backend a discrete card (PCIe VRAM)?

dit_attention

dit_attention_packed

dit_attention_packed_available

Whether dit_attention_packed has an implementation on the backend that is actually selected.

dit_block

dit_block_seg

The same block over a CONCATENATION of independent sequences: attention per segment, everything position-wise batched. wgpu only — the Metal path takes the single-sequence entry above.

dit_chain_supported

Can the selected backend keep the DiT’s hidden state on the device between blocks? Only the wgpu whole-block path; the Metal entry takes and returns host memory every call.

dit_qkv

One whole modulated DiT block on the device — norms, qkv, RoPE, attention, residuals and the SwiGLU FFN in a single command buffer; only x crosses the CPU boundary (in and out). The DiT’s three projections in one submission (wgpu only; the Metal path fuses the whole block instead). False = the caller keeps its three separate calls.

dit_qkv_attention

DiT full bidirectional attention on the device (all heads: scores GEMM → row softmax → P·V → panel unstack, one command buffer). Head-major inputs; out is [n, nh·hd]. Attention from an interleaved qkv panel, splitting into head-major planes ON the device. wgpu only; false elsewhere so the caller keeps its host repack. qkv projection + attention with the panel never leaving the card. wgpu only; false elsewhere and the caller keeps its host chain.

dit_qkv_attn_out

The whole attention half of a DiT block on the card: qkv GEMM, attention, output projection. Only proj comes home.

dit_split_only

dit_state_fetch

Pull the resident hidden state back to the host. For the caller that chained blocks and then hit one the device declined.

enabled

enabled_here

GPU allowed FOR THE CURRENT LAYER: backend is initialized AND the layer falls within CMF_GPU_LAYERS (GPU/CPU layer-split) AND we are not inside a cpu_scope. Op gates call this.

forward_batch_graph

Batched prefill: k contiguous positions through the whole graph in one submit (projections/FFN as GEMMs, attention/GDN looped over scratch). h is [k·hidden] in/out; positions len k. wgpu only.

forward_token_graph

Whole-token decode graph on wgpu: the entire layer stack in ONE submit, hidden resident, one readback. Updates h in place. false = refusal. loop_norm_at: virtual layer indices after which final_norm is applied (Looped Transformer mid-stack norm). Empty for standard models.

fused_block_trusted

Should a FUSED whole-block path trust the device instead of asking the per-op probe? True on native Metal and on discrete wgpu adapters.

fused_dit_block_available

Is a FUSED whole-block device path on offer? The batched-CFG shape (two sequences in one tall batch) and the fused block (one sequence, one command buffer) are alternatives, and the caller picks.

gdn_spec_restore

After a partial speculative acceptance: restore every GDN layer’s device state to the snapshot after batch position slot. wgpu only.

gemm_nt_f32

The backend’s f32 NT GEMM: y[n×m] = x[n×k] · wᵀ[m×k]. Tensor cores where the card has them. Refuses under CMF_BAKE_GPU=0 or strict f32, and for jobs below n·k·m = 4M, where the round trip costs more than the arithmetic saves.

graph_kv_reset

Drop the wgpu token graph’s device K/V mirror for a pipeline.

graph_mark_unsupported

The builder refused for a STRUCTURAL reason — an unsupported weight or layer kind. Callers must NOT report the transient refusals (an unsealed o1 state during prefill, a softcap): those clear on their own and marking them would disable the graph for good.

graph_race_begin_generation

graph_race_first_token_hopeless

First decode token of a racing graph generation: hopeless already? (>4x the normal path’s per-token average AND over a second.) Settles the race immediately; the caller discards the graph result and recomputes this token on the normal path.

graph_race_record

Record one decode-token wall time for the racing arm. The first token of each generation is discarded (KV-mirror upload / cold caches on the graph arm; cold mmap on the normal arm).

graph_race_use_graph

Should this decode token try the graph? trusted (discrete adapter, explicit env, or a GDN hybrid whose state lives on the device) skips the race entirely.

graph_unsupported

graph_unsupported_reset

A different model in the same process starts with a clean slate.

matvec_batch

Independent matvecs of one input in a single submission (GDN projections).

matvec_class

Which probe a large matvec belongs to. The head is an order of magnitude bigger than anything else that reaches this gate, and the two populations do not have the same answer.

min_rows

Effective threshold: CMF_GPU_MIN_ROWS overrides. Defaults differ by device class: on a DISCRETE card VRAM bandwidth pays off even for FFN/QKV-class matrices (4096), on unified memory only lm_head-class is worth the dispatch/readback (65536). Field case behind this: a 35B model on an RTX 4090 saw ~0 offload because every layer matrix sat below the old universal 65536.

moe_block

A layer’s MoE-FFN in one submission (amortizing the dispatch cost).

music3_ffn

Music-3’s FFN chain resident on the device — two GEMMs and the GLU between them with no host round trip. false = refused, host runs.

probe_arm

probe_arm_cold_prefers_gpu

probe_deciding

Is the class still collecting samples? (Call sites use this to route cold-weight calls away from the GPU arm during probing.)

probe_record

Record a timed arm sample; on the PROBE_SAMPLES-th clean sample of BOTH arms the class decides for the rest of the process.

probe_set_device

Backends: name the device once at init. The probe cache is keyed by it, because a verdict is a property of THIS silicon and nothing else. First writer wins: a process runs one backend, and on the rare host where two initialize, the one that came up first is the one in use.

q1_force

q1 ops on the native Metal backend skip the probe entirely: the CPU q1 kernel is load-port-bound, the GPU one wins warm — and probe alternation itself cools the device between samples (measured: block times 5.8 ms warm vs 8.8 ms mixed). Other backends keep probing.

q1_matmat

Batched q1 GEMM (prefill). wgpu only — Metal has its own block path.

q1_matvec

q1 matvec: raw f32 activations, tile-embedded scales. Metal only for now (wgpu q1 WGSL is queued); false = CPU fallback.

q1t_matmat

q1t batched GEMM (prefill) — base + overlay on-device (Metal simdgroup or wgpu register-blocked).

q1t_matvec

Ternary (q1t) BASE matvec on the GPU — fills out with the base dot; the caller adds the sparse overlay on the CPU. Metal only for now (wgpu q1t not yet written → CPU fallback).

q2tp_matmat

The same over a two-bit weight plane. Metal has no q2tp kernel, so there it declines and the host takes it.

q4b_matvec

q4_block matvec on the GPU — wgpu only (Metal drives q4_block through the whole-token graph, not a standalone matvec).

q4t_ffn

q4t_matmat

q4t_matvec

Single-token q4_tiled matvec on the device — the lm_head class (a q4t checkpoint’s head is its biggest host matvec, exactly like the q4tp twin above). wgpu holds q4t_mv pipelines only inside the graph encoder — the standalone arm stays an honest refusal until a discrete-GPU q4t model reaches the bench.

q4t_qkv

Fused QKV projection: one upload of the normed chunk, three GEMMs, one readback of Q|K|V back to back. Metal has no twin yet — its chunk graph keeps the whole layer resident and never surfaces QKV.

q4tp_ffn

q4tp_ffn_packed

y=·W2ᵀ — one command buffer, only X and Y cross the CPU boundary. SwiGLU FFN with a row-packed [gate|up] fc1 (MiniMax-H3’s DiT), run end to end on the device. wgpu only: Metal keeps the host loop until its own packed kernel exists.

q4tp_matmat

Batched q4t GEMM on the device (imagegen DiT prefill shapes). Metal: q4t_mul_mm decodes the mmap-resident tiles inside the GEMM’s K loop. wgpu (Vulkan/DX12 → NVIDIA/AMD/Intel/Adreno/Mali): the register-blocked WGSL twin, weights cached in VRAM.

q4tp_matvec

Single-token q4tp matvec on the device — the lm_head class. Through the DEDICATED matvec kernel: the batched GEMM at b=1 measured 11.73 ms against the host’s 9.51 on the release head, so the route that was supposed to save eleven milliseconds a token lost its own probe instead.

q8_matmat

GEMM of a prefill batch: pre — prescaled inputs row-major [b, cols], out — row-major [b, rows].

q8_matvec_range

q8_row/q8_2f matvec, rows [row0, row0+rows). xs — prescaled by the θ-field.

q8_resident_or_upload

Probing helper: true — tensor idx’s quant weights are ALREADY device-resident (a clean GPU sample is possible now); false — they were not (the upload starts within the VRAM budget, so a later call finds them warm) or the tensor cannot go to the GPU at all. Keeps the probe from billing a full cold dispatch+readback to a sample it will discard anyway. The verdict needs only a couple of warm tensors, so probe-driven uploads are capped — the losing-GPU machine should not pay for uploading the whole layer stack it will never use; if the GPU wins, the rest uploads lazily on demand, in the same first-touch order.

set_cache_dir

Loader: name a directory this process can write to. First call wins.

set_current_device

Pin this thread to a device. Server slots call it once per request; the worker pool propagates it into its threads, so a dispatch begun on card 1 does not finish on card 0.

set_layer

Pipeline: mark the current layer (or −1 outside layers) for layer-split.

upload_bytes

Device weight bytes uploaded so far (wgpu; 0 on other backends). Steady-state windows must show a ZERO delta — growth mid-benchmark means eviction/re-upload and disqualifies the number.

vae_attention_packed

vae_attention_packed_layout

vae_conv2d

VAE conv2d on the device (implicit GEMM — the CPU path pays for a multi-GB im2col matrix at high resolutions).

vae_conv2d_coop

The convolution as a GEMM on the matrix units. false = refused.

vae_qkv_attn_out

The VAE decoder’s attention half on the card. Only proj returns.

vae_resnet

One whole VAE resnet block on the device (norm+silu → conv ×2 → shortcut → add, one command buffer).

vae_upsample_conv

Nearest-2× upsample fused with the following conv — the small pre-upsample image is what crosses the CPU boundary.

vram_budget

Weight budget of the current GPU in bytes; 0 when there is none and u64::MAX on unified memory (where the OS pages shared RAM and the question “does the model fit the card” has no separate answer).

weight_is_resident

Which arm should this GPU-eligible call take? Consult AFTER the eligibility gates (enabled_here / min_rows) so only real candidates alternate. While a class is still probing, a call whose weights are NOT yet on the card should take the GPU arm anyway: the upload is work the next step needs regardless, and the sample it produces is discarded as cold — so handing that call to the CPU arm buys nothing and costs a host GEMM. Measured on a diffusion stack, where every layer is touched once per step and therefore EVERY first-step GPU sample is cold: one projection drew the CPU arm for the whole first step, 9.8 s against the 2.8 s it costs once the weights are warm.

wgpu_active

Default-on condition for the wgpu whole-token graph: the wgpu backend on a DISCRETE adapter. NOT plain enabled() (macOS/Metal must not pay a per-token layer scan for a graph its backend refuses), and NOT integrated adapters: the graph’s ~300 barriered dispatches per token are cheap on desktop immediate-mode GPUs but tiled mobile GPUs (Adreno/Mali) drain the pipeline at every barrier — field report: 0.2 tok/s on-graph vs 15 tok/s on the CPU. On integrated adapters the per-op probe path arbitrates each op class against the CPU instead; CMF_GPU_WGPU_GRAPH=1 still forces the graph anywhere. Is the wgpu backend active at all (any adapter)? Eligibility gate for the whole-token graph — whether it actually RUNS is decided by wgpu_graph_default (trusted on discrete) or the generation race.

wgpu_graph_default

wgpu_graph_on

The one place that decides whether the whole-token graph runs.

with_device

Run f with this thread pinned to dev, restoring the previous pin.