Struct WgslForwardPass 

pub struct WgslForwardPass { /* private fields */ }

PMAT-324: WGSL transformer forward pass shaders. GPU-resident transformer layer state. All buffers persist across tokens — only input/output change per step.

Implementations

impl WgslForwardPass

pub fn rmsnorm_shader() -> &'static str

Get the shader sources for external inspection/testing

pub fn silu_mul_shader() -> &'static str

pub fn residual_shader() -> &'static str

pub fn rope_shader() -> &'static str

pub fn new( device: Device, queue: Queue, hidden_dim: usize, num_heads: usize, num_kv_heads: usize, head_dim: usize, intermediate_dim: usize, ) -> Self

PMAT-325: Create a new WGSL forward pass context.

Compiles all shader pipelines and allocates persistent intermediate buffers. Call once at model init. All GPU resources persist until dropped.

pub fn upload_weight(&mut self, name: &str, data: &[f32])

Upload a weight matrix (call once per layer at init). PMAT-342: Bias weights (name contains “bias”) are stored CPU-side.

pub fn upload_q4k_weight(&mut self, name: &str, data: &[u8])

GH-560: Upload raw Q4K weight bytes for fused dequant+GEMV on GPU.

pub fn init_kv_cache(&mut self, num_layers: usize)

GH-560: Initialize per-layer KV cache buffers on GPU.

pub fn weight_count(&self) -> usize

Number of uploaded weight buffers.

pub fn weight_buffer(&self, name: &str) -> Option<&Buffer>

Access a dequantized weight buffer by name (e.g. “layer.0.down_proj”). Used by backward pass for gradient propagation through frozen base weights.

pub fn device_ref(&self) -> &Device

Reference to the wgpu device.

pub fn queue_ref(&self) -> &Queue

Reference to the wgpu queue.

pub fn hidden_buffer(&self) -> &Buffer

Reference to the hidden state buffer (for writing input).

pub fn q_buffer(&self) -> &Buffer

Reference to Q buffer (for LoRA addmm after Q projection).

pub fn k_buffer(&self) -> &Buffer

Reference to K buffer.

pub fn v_buffer(&self) -> &Buffer

Reference to V buffer.

pub fn gpu_residual_add( &self, a: &Buffer, b: &Buffer, output: &Buffer, len: u32, )

Elementwise add: output = a + b. Dispatches residual add shader.

pub fn gpu_rmsnorm(&self, weight: &Buffer, output: &Buffer, _seq_len: u32)

Apply RMSNorm on GPU: normed = rmsnorm(hidden_buf, weight) → output_buf. Contract: gpu-output-norm-v1 / gpu_resident — hidden state never leaves GPU.

pub fn download_hidden(&self, len: usize) -> Vec<f32>

Download hidden state from GPU.

pub fn total_vram_bytes(&self) -> usize

Total VRAM used by all buffers (bytes).

pub fn forward_model( &self, token_id: u32, position: usize, num_layers: usize, token_embedding: &[f32], output_norm_weight: &[f32], lm_head_weight: &[f32], vocab_size: usize, eps: f32, kv_caches: &mut Vec<(Vec<f32>, Vec<f32>)>, ) -> Result<Vec<f32>, String>

PMAT-336: Full model forward — embedding + all layers + output norm + LM head.

Returns logits [vocab_size] for the given token at the given position. Embedding lookup and final LM head are CPU-side (not yet GPU-accelerated). PMAT-344: Added kv_caches for multi-token context

pub fn forward_layer( &self, hidden: &mut [f32], layer_prefix: &str, _position: usize, kv_cache_k: &mut Vec<f32>, kv_cache_v: &mut Vec<f32>, ) -> Result<(), String>

PMAT-325: Execute one transformer layer — 14 passes, 1 submit, 1 readback.

Input: hidden state [hidden_dim] on CPU. Output: updated hidden state [hidden_dim] on CPU. All intermediate computation stays GPU-resident. PMAT-344: KV cache parameters for multi-token context

pub fn encode_forward_layer_training( &self, encoder: &mut CommandEncoder, seq_len: u32, layer_prefix: &str, saved: &LayerActivations, lora: Option<&QkvLoRA<'_>>, ) -> Result<(), String>

Training forward pass for a single transformer layer.

Unlike forward_layer (M=1 decode), this processes the full sequence at once (M=seq_len) and keeps everything on GPU. No CPU readback.

Saves norm_output (pre-projection activations) for backward pass.

Arguments

• seq_len: number of tokens in the sequence
• layer_prefix: e.g. “model.layers.0”
• saved_norm_attn: OUTPUT — saved pre-attention norm for backward (wgpu::Buffer, [seq×hidden])
• saved_norm_ffn: OUTPUT — saved pre-FFN norm for backward (wgpu::Buffer, [seq×hidden])

Forward one layer into an EXISTING encoder (no submit). Caller batches multiple layers into one encoder, submits once.

pub fn forward_layer_traced( &self, seq_len: u32, layer_prefix: &str, saved: &LayerActivations, lora: Option<&QkvLoRA<'_>>, ) -> Result<(), String>

Run one layer with per-operation GPU timing (submit+poll between each op group). Contract: forward-pass-perf-v1 / bottleneck_identified

pub fn alloc_layer_activations(&self, seq_len: u32) -> LayerActivations

Allocate saved activations for one layer.

pub fn forward_layer_training( &self, seq_len: u32, layer_prefix: &str, ) -> Result<LayerActivations, String>

Forward one layer with its own encoder + submit (original API, kept for compat).

pub fn forward_all_layers_training( &self, seq_len: u32, num_layers: usize, ) -> Result<Vec<LayerActivations>, String>

Forward ALL layers in one encoder submit. 28 layers → 1 GPU sync.

pub fn encode_broadcast_bias( &self, encoder: &mut CommandEncoder, buf: &Buffer, bias: &[f32], seq_len: u32, )

Encode causal multi-head attention on GPU. Q: [seq_len, num_heads * head_dim], K/V: [seq_len, num_kv_heads * head_dim] Output written to q_buf (reused as attn output). PMAT-509: Add broadcast bias to a [seq_len, dim] buffer. bias has shape [dim], applied to each of seq_len rows.