Struct LoadedModel 

pub struct LoadedModel { /* private fields */ }

A CudaModel with all kernels loaded into GPU memory and parameter buffers allocated, ready to run inference (and, with the training feature, training steps).

Implementations

impl LoadedModel

pub fn param_info(&self) -> impl Iterator<Item = (usize, &[Vec<usize>])>

Iterate over every node that has parameter buffers.

Yields (node_idx, param_shapes) where param_shapes[i] is the concrete shape of parameter slot i (e.g. [out_features, in_features] for a weight matrix). Use load_param_f32(node_idx, i, data) to upload values.

pub fn param_info_named( &self, ) -> impl Iterator<Item = (String, usize, usize)> + '_

Iterate over every named parameter slot.

Yields (full_key, node_idx, param_idx) where full_key is the dotted safetensors key (e.g. "model.0.conv.weight"), built by joining the node name from the graph with the slot name from the runtime op. Nodes without a name or without parameters are skipped.

pub fn load_param_f32( &mut self, node_idx: usize, param_idx: usize, data: &[f32], ) -> Result<()>

Copy f32 parameter data into a node’s pre-allocated device buffer.

node_idx — the DAG node index. param_idx — which parameter slot (0 = weight, 1 = bias, …). data — host f32 slice; must match the buffer element count exactly.

pub fn read_param_grad_f32( &self, node_idx: usize, param_idx: usize, ) -> Result<Vec<f32>>

Copy the accumulated parameter gradient (dL/dParam) back to host as f32.

Call after backward and before zero_grad.

pub fn node_parents(&self, idx: usize) -> &[usize]

Return the parent node indices for a DAG node.

pub fn forward( &self, device: &CudaDevice<'_>, batch_size: usize, inputs: &[TensorRef], ) -> Result<TensorRef>

Run a single forward pass through the loaded model.

device — the CUDA device context. batch_size — concrete value for dynamic (None) batch dimensions. inputs — device tensors matched to Input nodes in topological order.

Returns the TensorRef of the last DAG node. Intermediate output buffers are allocated per-call and freed when the returned TensorRef is dropped (caller owns the final buffer; all intermediate ones are freed at the end).

pub fn forward_train( &self, device: &CudaDevice<'_>, batch_size: usize, inputs: &[TensorRef], ) -> Result<(TensorRef, ActivationCache)>

Run a forward pass and retain ALL intermediate activation buffers.

Returns (final_output, activation_cache) where activation_cache[i] is the output tensor of node i. Call drop(cache) after backward to release the device buffers.

pub fn zero_grad(&mut self)

Zero all parameter gradient buffers. Call before each backward pass.

pub fn adamw_step( &mut self, device: &CudaDevice<'_>, kernel: &AdamwKernel, lr: f32, beta1: f32, beta2: f32, eps: f32, weight_decay: f32, ) -> Result<()>

Apply an AdamW update to all parameters using the accumulated gradient buffers.

kernel — pre-compiled adamw_step PTX (compile with AdamwStep::new(1024) from teeny_kernels::nn::optim::adam).

pub fn terminal_node_indices(&self) -> Vec<usize>

Indices of all DAG nodes that have no children (sinks / output nodes).

For single-output models this returns one element. YOLO26 returns two: the boxes node and the scores node.

pub fn node_name(&self, idx: usize) -> Option<&str>

Terminal node indices sorted by output tensor element count (ascending).

For YOLO26 this reliably gives [boxes_idx, scores_idx]: boxes has 4·A elements in the channel dim while scores has nc·A, and nc > 4 for all practical detection models. Return the name for a node index, if one was recorded during compilation.

pub fn terminal_node_indices_sorted_by_size(&self) -> Vec<usize>

Terminal (no-consumer) node indices, sorted by output element count (largest last).

pub fn backward( &mut self, device: &CudaDevice<'_>, batch_size: usize, grad_output: TensorRef, cache: &ActivationCache, ) -> Result<()>

Backward pass seeded from a single output node (common case for single-output models).

grad_output — dL/d(model_output), provided by the loss backward. cache — the activation cache returned by forward_train.

pub fn backward_multi( &mut self, device: &CudaDevice<'_>, batch_size: usize, seed_grads: &[(usize, TensorRef)], cache: &ActivationCache, ) -> Result<()>

Backward pass seeded from multiple output nodes (e.g. YOLO26 boxes + scores).

seed_grads — list of (node_idx, grad_tensor) pairs, one per output node. cache — the activation cache returned by forward_train.

pub fn capture_graph( &self, device: &CudaDevice<'_>, batch_size: usize, input_shapes: &[Vec<usize>], output_node_indices: &[usize], ) -> Result<CudaGraphModel>

Capture a fixed-batch CUDA graph for low-overhead repeated inference.

All device buffers (intermediate activations, scratch pads, TMA padding) are pre-allocated once. The kernel sequence is then recorded via cuStreamBeginCapture_v2 / cuStreamEndCapture and instantiated as a CUgraphExec. Subsequent CudaGraphModel::run calls replay it with a single cuGraphLaunch + cuStreamSynchronize.

Parameters

• batch_size — resolves dynamic (None) shape dimensions.
• input_shapes — concrete shape per Input node, in topological order.
• output_node_indices — which DAG node indices to read back as outputs. Use LoadedModel::terminal_node_indices_sorted_by_size to obtain them.

Constraints

• Inference-only (no backward pass).
• Fixed topology: re-capture if batch_size or input shapes change.
• f32 inputs and outputs assumed for the convenience CudaGraphModel::run API.