Struct WgpuBlockManager 

pub struct WgpuBlockManager {

    pub device: Device,
    pub queue: Queue,
    pub blocks: Vec<WgpuBlock>,
    pub hidden_buf: Buffer,
    pub hidden_buf2: Buffer,
    pub attn_out_buf: Buffer,
    pub ffn_gate_buf: Buffer,
    pub ffn_up_buf: Buffer,
    pub ffn_silu_buf: Buffer,
    pub norm_buf: Buffer,
    pub q_buf: Buffer,
    pub k_buf: Buffer,
    pub v_buf: Buffer,
    pub embed_weight: Buffer,
    pub lm_head_weight: Buffer,
    pub logits_buf: Buffer,
    pub grad_hidden_buf: Buffer,
    pub grad_logits_buf: Buffer,
    pub hidden_size: u32,
    pub intermediate_size: u32,
    pub num_heads: u32,
    pub num_kv_heads: u32,
    pub head_dim: u32,
    pub max_seq_len: u32,
    pub vocab_size: u32,
    pub num_layers: u32,
}

Manages all transformer blocks on GPU + shared buffers.

Fields

device: Device

queue: Queue

blocks: Vec<WgpuBlock>

hidden_buf: Buffer

hidden_buf2: Buffer

attn_out_buf: Buffer

ffn_gate_buf: Buffer

ffn_up_buf: Buffer

ffn_silu_buf: Buffer

norm_buf: Buffer

q_buf: Buffer

k_buf: Buffer

v_buf: Buffer

embed_weight: Buffer

lm_head_weight: Buffer

logits_buf: Buffer

grad_hidden_buf: Buffer

grad_logits_buf: Buffer

hidden_size: u32

intermediate_size: u32

num_heads: u32

num_kv_heads: u32

head_dim: u32

max_seq_len: u32

vocab_size: u32

num_layers: u32

Implementations

impl WgpuBlockManager

pub fn new( device: Device, queue: Queue, hidden_size: u32, intermediate_size: u32, num_heads: u32, num_kv_heads: u32, head_dim: u32, num_layers: u32, vocab_size: u32, max_seq_len: u32, _lora_rank: Option<u32>, _lora_alpha: Option<f32>, ) -> Self

Create a new block manager and upload all transformer weights to GPU.

weights_per_layer is a closure that returns the F32 weights for each layer. This avoids holding all 28 layers in CPU memory simultaneously.

pub fn upload_layer( &mut self, layer_idx: usize, input_norm: &[f32], post_attn_norm: &[f32], w_q: &[f32], w_k: &[f32], w_v: &[f32], w_o: &[f32], w_gate: &[f32], w_up: &[f32], w_down: &[f32], lora_rank: Option<u32>, lora_scale: Option<f32>, )

Upload a single transformer layer’s weights to GPU.

pub fn upload_embeddings(&mut self, embed: &[f32], lm_head: &[f32])

Upload embedding + lm_head weights.

pub fn gpu_memory_bytes(&self) -> u64

Total GPU memory used (approximate, in bytes).

pub fn layer_count(&self) -> usize

Number of uploaded layers.