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].

GpuPause

GpuTrust

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.

cold_epoch

Backends: note a one-off cost (weight upload, buffer-cache fill) so the probe discards this sample. Every buffer creation anywhere bumps this; the graph’s bind-group cache treats any cold event as total invalidation — a stale bind group is silent corruption, a cleared cache is one re-encoded token.

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

gemm_nt_f32_transient

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. gemm_nt_f32 whose w is known to change every call (an accumulation over fresh activations, not a weight): it skips the resident ledger and its per-call fingerprint of the whole operand.

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.

hostprof_encode_done

CMF_GRAPH_HOSTPROF=1: how a graph token’s wall splits between the host encoding the command stream and the tail the GPU still owes after encode. Fifteen GPU-side suspects measured null while the bench counted 17.7k allocations a token — this is the instrument that says whether the thief was on the host all along.

hostprof_total

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.

mm_kill_arm

Disarm / re-arm the contention kill around a phase whose GEMMs are slow for reasons that are not another process (see MM_ARMED).

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.

pause_gpu

Park the device for every thread until the returned guard drops.

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_note_decline

The device arm was chosen and the device refused the work, so there is no time to record. Callers that fall through to the host MUST say so here, or the class can never decide.

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

q8_matmat_2f

GEMM of a prefill batch: pre — prescaled inputs row-major [b, cols], out — row-major [b, rows]. The two-field int8 GEMM with the column field left for the device. wgpu only — Metal’s int8 kernel takes a pre-scaled activation, so the caller keeps that path when this returns false.

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.

stageprof

Per-stage host-encode accumulator for the Metal token loop (CMF_GRAPH_HOSTPROF=1). Stage 0 = GDN-run encode; everything else falls out by subtraction from hostprof’s encode total.

trust_gpu

Take the probe out of the loop until the guard drops.

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_bytes_by

The per-stage split of weight_bytes_dispatched: [misc, dense-ffn, moe, attn, gdn, head].

weight_bytes_dispatched

Active weight bytes dispatched so far (Metal decode path); 0 where the backend does not count. The honest floor’s numerator.

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.