Enum CudaBlock 

pub enum CudaBlock {
    Fp32(CudaTransformerBlock),
    Nf4(CudaNf4TransformerBlock),
}

Unified enum for CUDA transformer blocks (fp32 or NF4-quantized).

The classify pipeline stores Vec<CudaBlock> and calls forward() without caring which quantization format the frozen weights use.

Variants

Fp32(CudaTransformerBlock)

Standard fp32 weights (full precision, ~16 GB for Qwen3-4B)

Nf4(CudaNf4TransformerBlock)

NF4 quantized weights (~2 GB for Qwen3-4B, ~8x compression)

Implementations

impl CudaBlock

pub fn layer_idx(&self) -> usize

Get the layer index of this block.

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 (only supported for fp32 blocks).

NF4 blocks are frozen — backward is never called when quantize_nf4 is active because gpu_training is set to None.

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

Initialize optimizer state (only supported for fp32 blocks).

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

Download weights from GPU (only supported for fp32 blocks).

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

Optimizer step (only supported for fp32 blocks).

pub fn download_lora_weights( &self, ) -> Result<(Vec<f32>, Vec<f32>, Vec<f32>, Vec<f32>)>

Download LoRA weights from NF4 blocks.

pub fn upload_lora_weights( &mut self, a_q: &[f32], b_q: &[f32], a_v: &[f32], b_v: &[f32], ) -> Result<()>

Upload LoRA weights to NF4 blocks for checkpoint resume (ENT-276).