Skip to main content

Module gpu

Module gpu 

Source
Expand description

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

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