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.

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

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

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.

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

discrete

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

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

dit_block

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

enabled

GPU enabled and initialized on the selected backend?

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.

graph_kv_reset

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

graph_race_begin_generation

Called at every generation start (fresh KV). Applies a pending verdict and picks this generation’s arm while racing.

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.

matvec_batch

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

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

probe_arm

Which arm should this GPU-eligible call take? Consult AFTER the eligibility gates (enabled_here / min_rows) so only real candidates alternate.

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.

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

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

Fused DiT SwiGLU FFN on the device: g=X·W1ᵀ, u=X·W3ᵀ, silu(g)·u, y=·W2ᵀ — one command buffer, only X and Y cross the CPU boundary.

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

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_layer

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

vae_conv2d

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

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.

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