Struct CudaGradWorkspace 

pub struct CudaGradWorkspace { /* private fields */ }

Shared gradient workspace for weight gradients (one per model, NOT per layer).

Contract (C-GRADWS-001)

Backward processes layers sequentially — only one layer’s weight gradients are computed at a time. Sharing this workspace across layers saves (L-1) * per_layer_grad_weight_elements * 4 bytes of VRAM.

For Qwen3-4B: saves 35 * 372 MB = 13.0 GB.

• Precondition: Allocated once before training loop starts
• Postcondition: After backward() for layer i, contains layer i’s weight gradients
• Invariant: Buffer sizes match model config; never reallocated during training

Implementations

impl CudaGradWorkspace

pub fn new(ctx: &Arc<CudaContext>, config: &TransformerConfig) -> Result<Self>

Allocate shared gradient workspace for the given model config.

Called once per training run. GEMM weight gradients are fully overwritten by each backward pass. Norm gradients use atomicAdd accumulation and MUST be zeroed before each rms_norm_backward call (see zero_norm_grads).

pub fn zero_norm_grads(&mut self, zero_buf: &[f32]) -> Result<()>

Zero norm gradient buffers before rms_norm_backward calls.

The BatchedRmsNormBackwardKernel accumulates grad_gamma via atomicAdd, so these buffers MUST be zeroed before each backward pass. Without this, grad_gamma accumulates across steps → exploding norm gradients.