Struct CudaTransformerBlock 

pub struct CudaTransformerBlock { /* private fields */ }

CUDA-accelerated transformer block

All operations run on GPU with minimal CPU<->GPU transfers.

Implementations

impl CudaTransformerBlock

pub fn new( config: &TransformerConfig, layer_idx: usize, ctx: Arc<CudaContext>, input_norm_weight: &[f32], post_attn_norm_weight: &[f32], w_q: &[f32], w_k: &[f32], w_v: &[f32], w_o: &[f32], w_gate: &[f32], w_up: &[f32], w_down: &[f32], max_seq_len: usize, ) -> Result<Self>

Create a new CUDA transformer block from CPU tensors

Uploads all weights to GPU memory.

pub fn set_qk_norm(&mut self, q_norm: &[f32], k_norm: &[f32]) -> Result<()>

Set QK-norm weights (ENT-270). Called after construction when loading Qwen3 models.

pub fn forward( &mut self, input: &GpuBuffer<f32>, output: &mut GpuBuffer<f32>, seq_len: usize, stream: &CudaStream, ) -> Result<()>

Forward pass - all operations on GPU

Arguments

• input - Input tensor on GPU (seq_len * hidden_size)
• output - Output tensor on GPU (seq_len * hidden_size)
• seq_len - Sequence length
• stream - CUDA stream for async execution

pub fn layer_idx(&self) -> usize

Get layer index

pub fn config(&self) -> &TransformerConfig

Get configuration

pub fn backward( &mut self, input: &GpuBuffer<f32>, grad_output: &GpuBuffer<f32>, grad_input: &mut GpuBuffer<f32>, seq_len: usize, stream: &CudaStream, grad_ws: &mut CudaGradWorkspace, ) -> Result<()>

Backward pass - gradient computation on GPU (ENT-151)

Computes gradients for all parameters given upstream gradient.

Arguments

• input - Original input from forward pass (seq_len * hidden_size)
• grad_output - Gradient from upstream layer (seq_len * hidden_size)
• grad_input - Output: gradient w.r.t. input (seq_len * hidden_size)
• seq_len - Sequence length
• stream - CUDA stream for async execution

Returns

Gradients are accumulated into the scratch buffers:

• scratch.grad_input_norm - Gradient for input RMSNorm weight
• scratch.grad_post_attn_norm - Gradient for post-attention RMSNorm weight
• scratch.grad_gate/up/down - Gradients for FFN weights
• scratch.grad_w_q/w_k/w_v/w_o - Gradients for attention projection weights

pub fn init_optimizer_state(&self) -> Result<GpuBlockOptimizerState>

Initialize GPU-resident AdamW optimizer state for all block weights.

Allocates zero-initialized first and second moment buffers for each of the 9 weight tensors (4 attention projections + 3 FFN projections + 2 RMSNorm).

Contract (C-OPTINIT-001)

• Precondition: CUDA context is valid, sufficient GPU memory available
• Postcondition: All m/v buffers are zero-initialized with dimensions matching the corresponding weight tensors
• Invariant: Total GPU memory for optimizer state = 2 × sum(weight_sizes) × 4 bytes

pub fn optimizer_step( &mut self, state: &mut GpuBlockOptimizerState, step: u32, lr: f32, beta1: f32, beta2: f32, eps: f32, weight_decay: f32, stream: &CudaStream, grad_ws: &CudaGradWorkspace, ) -> Result<()>

Run GPU-resident AdamW optimizer step on all block weights.

Updates weights in-place using gradients computed by backward(). All operations run on GPU — zero CPU↔GPU data transfers.

Contract (C-OPTSTEP-001)

• Precondition: backward() completed for this block (scratch grad buffers valid), state initialized via init_optimizer_state(), step > 0
• Postcondition: All 9 weight tensors updated by AdamW rule, m/v states updated with current gradient statistics
• Invariant: Weight dimensions unchanged; no GPU memory allocated or freed

pub fn download_weights(&self) -> Result<BlockWeights>

Download all weight data from GPU to host vectors.

Used to synchronize GPU-updated weights back to CPU model for checkpointing.

Contract (C-DLWEIGHTS-001)

• Precondition: Block weights are valid GPU allocations
• Postcondition: Returned vectors have exact same length and content as GPU buffers
• Invariant: GPU buffers are not modified