Struct ComputeEngine 

pub struct ComputeEngine { /* private fields */ }

Compiled WGSL kernels and the methods that dispatch them.

Every method records and submits its own command buffer and returns without waiting; reading a result (e.g. with UnifiedTensor::to_vec) waits for the GPU. Tensors are dense, row-major f32 unless stated otherwise.

Implementations

impl ComputeEngine

pub fn rope( &self, ctx: &WgpuContext, tensor: &UnifiedTensor, head_dim: usize, pos: usize, )

Rotate-half RoPE (theta 10000) applied in place.

tensor is [batch, seq, ...], [seq, hidden] or a single token [hidden]; every head vector of token s gets position pos + s.

pub fn rope_theta( &self, ctx: &WgpuContext, tensor: &UnifiedTensor, head_dim: usize, pos: usize, theta: f32, )

rope with base frequency theta (the rope_theta of a Hugging Face config, e.g. 100000 for SmolLM2).

pub fn flash_attention( &self, ctx: &WgpuContext, q: &UnifiedTensor, k: &UnifiedTensor, v: &UnifiedTensor, out: &UnifiedTensor, causal: bool, )

Single-head attention: q, k, v, out are [seq_len, head_dim].

Runs the v2 kernel on [1, seq_len, 1, head_dim] views of the inputs.

pub fn flash_attention_v2( &self, ctx: &WgpuContext, q: &UnifiedTensor, k: &UnifiedTensor, v: &UnifiedTensor, out: &UnifiedTensor, causal: bool, dropout_prob: f32, seed: u64, mask: Option<&UnifiedTensor>, )

Fused attention forward pass with grouped-query attention, causal masking, an optional additive mask and attention dropout.

• q, out - [batch, seq_q, num_heads, head_dim]
• k, v - [batch, seq_k, num_kv_heads, head_dim]; num_heads must be a multiple of num_kv_heads
• causal - key j is visible to query i only when j <= i
• dropout_prob, seed - attention dropout; pass 0.0 to disable
• mask - additive mask [batch, heads, seq_q, seq_k], where size-1 dimensions broadcast

head_dim must be at most 128. See flash_attention_v2_with for KV-cache offsets and sliding windows.

pub fn flash_attention_v2_with( &self, ctx: &WgpuContext, q: &UnifiedTensor, k: &UnifiedTensor, v: &UnifiedTensor, out: &UnifiedTensor, opts: &AttentionOptions, mask: Option<&UnifiedTensor>, )

flash_attention_v2 with every option.

pub fn flash_attention_backward( &self, ctx: &WgpuContext, q: &UnifiedTensor, k: &UnifiedTensor, v: &UnifiedTensor, out: &UnifiedTensor, grad_out: &UnifiedTensor, grad_q: &UnifiedTensor, grad_k: &UnifiedTensor, grad_v: &UnifiedTensor, causal: bool, )

Single-head attention backward pass for [seq_len, head_dim] tensors. Writes (does not accumulate into) grad_q, grad_k and grad_v.

pub fn flash_attention_v2_backward( &self, ctx: &WgpuContext, q: &UnifiedTensor, k: &UnifiedTensor, v: &UnifiedTensor, out: &UnifiedTensor, grad_out: &UnifiedTensor, grad_q: &UnifiedTensor, grad_k: &UnifiedTensor, grad_v: &UnifiedTensor, causal: bool, )

Backward pass of flash_attention_v2 without mask or dropout. See flash_attention_v2_backward_with.

pub fn flash_attention_v2_backward_with( &self, ctx: &WgpuContext, q: &UnifiedTensor, k: &UnifiedTensor, v: &UnifiedTensor, out: &UnifiedTensor, grad_out: &UnifiedTensor, grad_q: &UnifiedTensor, grad_k: &UnifiedTensor, grad_v: &UnifiedTensor, opts: &AttentionOptions, mask: Option<&UnifiedTensor>, )

Backward pass of flash_attention_v2.

out is the forward output for the same inputs; opts and mask must match the forward call (dropout decisions are regenerated from opts.seed). Writes grad_q, grad_k and grad_v (same shapes as q, k, v); no gradient is produced for the mask.

impl ComputeEngine

pub fn add( &self, ctx: &WgpuContext, a: &UnifiedTensor, b: &UnifiedTensor, out: &UnifiedTensor, )

out = a + b, element-wise; a, b and out have the same length.

pub fn mul( &self, ctx: &WgpuContext, a: &UnifiedTensor, b: &UnifiedTensor, out: &UnifiedTensor, )

out = a * b, element-wise.

pub fn sub( &self, ctx: &WgpuContext, a: &UnifiedTensor, b: &UnifiedTensor, out: &UnifiedTensor, )

out = a - b, element-wise.

pub fn add_backward( &self, ctx: &WgpuContext, grad_out: &UnifiedTensor, grad_a: &UnifiedTensor, grad_b: &UnifiedTensor, )

Gradients of a + b: writes grad_out to both grad_a and grad_b.

pub fn sub_backward( &self, ctx: &WgpuContext, grad_out: &UnifiedTensor, grad_a: &UnifiedTensor, grad_b: &UnifiedTensor, )

Gradients of a - b: grad_a = grad_out, grad_b = -grad_out.

pub fn mul_backward( &self, ctx: &WgpuContext, grad_out: &UnifiedTensor, a: &UnifiedTensor, b: &UnifiedTensor, grad_a: &UnifiedTensor, grad_b: &UnifiedTensor, )

Gradients of a * b: grad_a = grad_out * b, grad_b = grad_out * a.

pub fn silu( &self, ctx: &WgpuContext, input: &UnifiedTensor, out: &UnifiedTensor, )

out = x * sigmoid(x), element-wise.

pub fn silu_backward( &self, ctx: &WgpuContext, grad_out: &UnifiedTensor, input: &UnifiedTensor, grad_in: &UnifiedTensor, )

Gradient of SiLU: grad_in = grad_out * s * (1 + x * (1 - s)) with s = sigmoid(input).

pub fn relu( &self, ctx: &WgpuContext, input: &UnifiedTensor, out: &UnifiedTensor, )

out = max(x, 0), element-wise.

pub fn relu_backward( &self, ctx: &WgpuContext, grad_out: &UnifiedTensor, input: &UnifiedTensor, grad_in: &UnifiedTensor, )

Gradient of ReLU: grad_out where input > 0, else 0.

pub fn silu_gate_fused( &self, ctx: &WgpuContext, up: &UnifiedTensor, gate: &UnifiedTensor, out: &UnifiedTensor, )

Fused SiLU-Gate: out = SiLU(up) * gate Reduces memory bandwidth by 2x compared to separate ops

pub fn silu_gate_backward( &self, ctx: &WgpuContext, up: &UnifiedTensor, gate: &UnifiedTensor, grad_out: &UnifiedTensor, grad_up: &UnifiedTensor, grad_gate: &UnifiedTensor, )

Fused SiLU-Gate backward: computes gradients for up and gate

impl ComputeEngine

pub fn matmul( &self, ctx: &WgpuContext, a: &UnifiedTensor, b: &UnifiedTensor, out: &UnifiedTensor, )

out[M, N] = a[M, K] @ b[K, N].

pub fn matmul_bt( &self, ctx: &WgpuContext, a: &UnifiedTensor, b: &UnifiedTensor, out: &UnifiedTensor, )

out[M, N] = a[M, K] @ b[N, K]^T (e.g. the input gradient of a linear layer).

pub fn matmul_at( &self, ctx: &WgpuContext, a: &UnifiedTensor, b: &UnifiedTensor, out: &UnifiedTensor, )

out[M, N] = a[K, M]^T @ b[K, N] (e.g. the weight gradient of a linear layer).

impl ComputeEngine

pub fn top_k( &self, ctx: &WgpuContext, input: &UnifiedTensor, out_idx: &UnifiedTensor, k: u32, )

Writes the indices of the k largest values of input, largest first, as u32s into out_idx. Runs on a single GPU thread; meant for sampling from one logits vector.

pub fn conv1d( &self, ctx: &WgpuContext, input: &UnifiedTensor, weight: &UnifiedTensor, bias: &UnifiedTensor, out: &UnifiedTensor, )

Depthwise 1-D convolution with kernel size 3 and padding 1: out[d, s] = bias[d] + sum_k weight[d, k] * input[d, s + k - 1], with input/out [dim, seq_len], weight [dim, 3] and bias [dim].

impl ComputeEngine

pub fn rms_norm( &self, ctx: &WgpuContext, input: &UnifiedTensor, weight: &UnifiedTensor, out: &UnifiedTensor, )

RMSNorm over the last dimension with eps = 1e-5: out = x / sqrt(mean(x^2) + eps) * weight.

pub fn rms_norm_eps( &self, ctx: &WgpuContext, input: &UnifiedTensor, weight: &UnifiedTensor, out: &UnifiedTensor, eps: f32, )

RMSNorm over the last dimension with a caller-chosen eps.

pub fn rms_norm_residual( &self, ctx: &WgpuContext, x: &UnifiedTensor, residual: &UnifiedTensor, weight: &UnifiedTensor, out: &UnifiedTensor, )

Fused RMSNorm + Residual: out = RMSNorm(x + residual) * weight

pub fn softmax( &self, ctx: &WgpuContext, input: &UnifiedTensor, out: &UnifiedTensor, )

Softmax over the last dimension of input.

pub fn softmax_backward( &self, ctx: &WgpuContext, grad_out: &UnifiedTensor, final_out: &UnifiedTensor, grad_in: &UnifiedTensor, )

Gradient of softmax over the last dimension, given the forward output final_out: grad_in = s * (grad_out - sum(grad_out * s)).

impl ComputeEngine

pub fn dequantize( &self, ctx: &WgpuContext, input: &UnifiedTensor, out: &UnifiedTensor, )

Unpacks bytes to floats: element i of out is byte i of input (read as packed u32s) converted to f32, without scaling.

pub fn affine_dequantize( &self, ctx: &WgpuContext, weight: &Buffer, scales: &Buffer, biases: &Buffer, out: &UnifiedTensor, )

8-bit affine dequantisation with groups of 64 (MLX layout): out[i] = byte(weight, i) * scales[i / 64] + biases[i / 64].

impl ComputeEngine

pub fn sgd_update( &self, ctx: &WgpuContext, weight: &UnifiedTensor, grad: &UnifiedTensor, lr: f32, )

In-place SGD step: weight -= lr * grad.

pub fn adamw_update( &self, ctx: &WgpuContext, weight: &UnifiedTensor, grad: &UnifiedTensor, m: &UnifiedTensor, v: &UnifiedTensor, params: &AdamParams, )

In-place AdamW step, updating weight and the moment buffers m and v.

pub fn mse_loss( &self, ctx: &WgpuContext, pred: &UnifiedTensor, target: &UnifiedTensor, out: &UnifiedTensor, )

Per-element squared error out = 0.5 * (pred - target)^2 (not reduced).

pub fn mse_loss_backward( &self, ctx: &WgpuContext, pred: &UnifiedTensor, target: &UnifiedTensor, grad_in: &UnifiedTensor, )

Gradient of mse_loss: grad_in = pred - target.

pub fn cross_entropy_loss( &self, ctx: &WgpuContext, logits: &UnifiedTensor, target: u32, loss_out: &UnifiedTensor, )

Cross-entropy loss of one logits vector: loss_out[0] = -log(softmax(logits)[target]).

pub fn cross_entropy_backward( &self, ctx: &WgpuContext, logits: &UnifiedTensor, target: u32, grad_out: &UnifiedTensor, )

Cross-entropy loss backward: gradient = softmax(logits) - one_hot(target)

pub fn gradient_norm_squared( &self, ctx: &WgpuContext, input: &UnifiedTensor, out: &UnifiedTensor, )

Compute squared elements (for gradient norm computation)

pub fn clip_grad_by_norm( &self, ctx: &WgpuContext, grad: &UnifiedTensor, scale: f32, )

Clip gradients by scaling: grad *= min(1.0, max_norm / actual_norm)

pub fn gradient_accumulate( &self, ctx: &WgpuContext, accum_grad: &UnifiedTensor, new_grad: &UnifiedTensor, )

Gradient accumulation: accum_grad += new_grad Essential for micro-batching on memory-constrained devices

pub fn gradient_scale( &self, ctx: &WgpuContext, grad: &UnifiedTensor, scale: f32, )

Gradient scaling: grad *= scale_factor Used to average gradients after accumulation

impl ComputeEngine

pub fn new(ctx: &WgpuContext) -> Self

Compiles every kernel for ctx’s device.

pub fn copy_tensor( &self, ctx: &WgpuContext, src: &UnifiedTensor, dst: &UnifiedTensor, dst_offset: usize, )

Copies all of src into dst, starting at byte offset dst_offset.