From 617b9fb71832e85e4f5e1c63aacdb8ab246281de Mon Sep 17 00:00:00 2001
From: jgrusewski <jgrusewski@localhost>
Date: Tue, 24 Mar 2026 01:52:22 +0100
Subject: [PATCH] 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.
---
 crates/ml/src/cuda_pipeline/gpu_attention.rs | 21 ++++++++++++++++++++
 1 file changed, 21 insertions(+)

diff --git a/crates/ml/src/cuda_pipeline/gpu_attention.rs b/crates/ml/src/cuda_pipeline/gpu_attention.rs
index 62db95108..f60d6a5af 100644
--- a/crates/ml/src/cuda_pipeline/gpu_attention.rs
+++ b/crates/ml/src/cuda_pipeline/gpu_attention.rs
@@ -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