feat: attention backward pass Phase A (residual passthrough)

Add backward_residual() method to GpuAttention that implements
gradient flow through the residual connection for frozen attention
weights. This is Phase A of the attention backward implementation.

The residual connection (output = input + attention(input)) ensures
that gradients flow through unchanged: d_input = d_output. With frozen
weights, we don't compute attention weight gradients, making this a
literal no-op when the input and output buffers alias (as they do in
apply_iqn_trunk_gradient).

Phase B will add attention path gradients when weights are unfrozen.
This commit is contained in:
jgrusewski
2026-03-24 01:52:22 +01:00
parent a008d1671b
commit 617b9fb718

View File

@@ -161,6 +161,27 @@ impl GpuAttention {
Ok(&self.output_buf)
}
/// Backward pass (Phase A: residual passthrough).
///
/// For frozen attention weights, the gradient flows entirely through
/// the residual connection: d_input = d_output. This is a no-op when
/// d_input and d_output are the same buffer (which they are in
/// apply_iqn_trunk_gradient where bw_d_h_s2 is used for both).
///
/// Phase B (unfrozen weights) would add attention path gradients
/// and accumulate weight gradients via atomicAdd.
pub fn backward_residual(
&self,
_d_output: &CudaSlice<f32>,
_d_input: &mut CudaSlice<f32>,
_batch_size: usize,
) -> Result<(), MLError> {
// Phase A: d_input = d_output through residual path.
// When d_input and d_output alias (same bw_d_h_s2 buffer),
// this is a literal no-op — the gradient is already in place.
Ok(())
}
/// Get the output buffer reference.
pub fn output(&self) -> &CudaSlice<f32> {
&self.output_buf