feat(dqn-v2): Plan 4 Task 2c.3c.2 — backward infra additions (launch_dw_only_no_bias + saxpy_inplace)

Plan 4 Task 2c.3c.2. Additive only — no production callers (2c.3c.4
wires them).

Two infrastructure additions for the GRN trunk backward chain:

1. launch_dw_only_no_bias on CublasBackwardSet: variant of launch_dw_only
   that skips the bias-grad kernel call. Linear_residual in h_s1 GRN
   block has no bias, so calling launch_dw_only with db=0u64 would
   segfault the bias-grad kernel.

2. saxpy_inplace on CublasGemmSet: y += alpha * x for element-wise
   gradient accumulation. h_s2 GRN's identity residual needs
   d_h_s1 += d_pre_ln_h_s2 after Linear_a_h_s2's backward overwrites
   d_h_s1 with d_x = d_linear_a @ W_a. Implementation reuses the
   existing dqn_saxpy_f32_kernel (already used by the experience
   collector's IQR/ensemble-variance Q-bonus paths) — no new kernel,
   no cuBLAS legacy-handle stream-binding work, kernel handle loaded
   once at CublasGemmSet::new from DQN_UTILITY_CUBIN.

Both methods sit dead-code until 2c.3c.4's wire-up commit.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

crates/ml/src/cuda_pipeline/batched_backward.rs

@@ -1889,6 +1889,41 @@ impl CublasBackwardSet {
         self.launch_bias_grad(stream, dy, db, out_dim, batch)
     }
 
+    /// Compute ONLY weight gradient (no bias gradient, no upstream dX).
+    ///
+    /// Variant of `launch_dw_only` for FC layers that have NO bias term — e.g.
+    /// the GRN block's `Linear_residual` projection (Plan 4 Task 2c.3c.2).
+    /// `launch_dw_only` would segfault here because the 2-phase bias-grad
+    /// kernel dereferences a NULL `db` pointer; this method skips that launch
+    /// entirely and only runs the cuBLAS dW GEMM.
+    ///
+    /// Same convention as `launch_dw_only`:
+    ///   `dW[out, in] += dY^T[out, B] @ X[B, in]`
+    /// using the row-major → col-major dual:
+    ///   `sgemm(N, T, in_dim, out_dim, batch, 1.0, X, in_dim, dY, out_dim, 1.0, dW, in_dim)`.
+    #[allow(clippy::too_many_arguments)]
+    pub(crate) fn launch_dw_only_no_bias(
+        &self,
+        stream: &Arc<CudaStream>,
+        dy: u64,
+        x: u64,
+        dw: u64,
+        out_dim: usize,
+        in_dim: usize,
+        batch: usize,
+    ) -> Result<(), MLError> {
+        self.sgemm_f32(
+            stream,
+            cublas_sys::cublasOperation_t::CUBLAS_OP_N,
+            cublas_sys::cublasOperation_t::CUBLAS_OP_T,
+            in_dim as i32, out_dim as i32, batch as i32,
+            1.0, x, in_dim as i32,
+            dy, out_dim as i32,
+            1.0, dw, in_dim as i32,
+            "dW_only_no_bias",
+        )
+    }
+
     /// Compute only upstream gradient dX = dY @ W^T, with configurable beta.
     ///
     /// `beta=0.0` overwrites dX; `beta=1.0` accumulates into dX.  Used to merge

crates/ml/src/cuda_pipeline/batched_forward.rs

@@ -261,6 +261,13 @@ pub struct CublasGemmSet {
     /// h_s2 GLU output [B, SH2] — target instance scratch.
     _grn_h_s2_glu_buf_tg: CudaSlice<f32>,
     grn_h_s2_glu_ptr_tg: u64,
+
+    // ── Element-wise SAXPY kernel (Plan 4 Task 2c.3c.2) ──────────────────
+    /// `dqn_saxpy_f32_kernel(y, x, alpha, n)` — computes `y[i] += alpha * x[i]`.
+    /// Used by GRN trunk backward to accumulate the h_s2 identity-residual
+    /// gradient (`d_h_s1 += d_pre_ln_h_s2`) after Linear_a's backward
+    /// overwrites `d_h_s1`.
+    saxpy_kernel: CudaFunction,
 }
 
 /// Backward-compatible type alias — existing callers (`GpuDqnTrainer`,
@@ -446,6 +453,22 @@ impl CublasGemmSet {
         let h_s2_linear_b_buf_tg = alloc_grn(batch_size * 2 * shared_h2,  "h_s2_linear_b_tg")?;
         let h_s2_glu_buf_tg      = alloc_grn(batch_size * shared_h2,      "h_s2_glu_tg")?;
 
+        // ── Load SAXPY kernel (Plan 4 Task 2c.3c.2) ──────────────────
+        // `dqn_saxpy_f32_kernel(y, x, alpha, n)` is shared with the
+        // experience collector's IQR/ensemble-variance Q-bonus paths.
+        // Loaded here so the GRN trunk backward (h_s2 identity residual
+        // accumulation) can call `saxpy_inplace` without plumbing another
+        // kernel handle through the trainer's wire-up sequence.
+        let saxpy_kernel = {
+            let saxpy_module = stream_for_grn
+                .context()
+                .load_cubin(super::gpu_dqn_trainer::DQN_UTILITY_CUBIN.to_vec())
+                .map_err(|e| MLError::ModelError(format!("saxpy cubin load: {e}")))?;
+            saxpy_module
+                .load_function("dqn_saxpy_f32_kernel")
+                .map_err(|e| MLError::ModelError(format!("dqn_saxpy_f32_kernel load: {e}")))?
+        };
+
         let h_s1_linear_a_ptr     = h_s1_linear_a_buf.raw_ptr();
         let h_s1_linear_b_ptr     = h_s1_linear_b_buf.raw_ptr();
         let h_s1_residual_ptr     = h_s1_residual_buf.raw_ptr();
@@ -531,6 +554,8 @@ impl CublasGemmSet {
             grn_h_s2_linear_b_ptr_tg: h_s2_linear_b_ptr_tg,
             _grn_h_s2_glu_buf_tg: h_s2_glu_buf_tg,
             grn_h_s2_glu_ptr_tg: h_s2_glu_ptr_tg,
+
+            saxpy_kernel,
         })
     }
 
@@ -595,6 +620,47 @@ impl CublasGemmSet {
         self.branch_workspace_ptrs
     }
 
+    /// Element-wise SAXPY: `y[i] += alpha * x[i]` for `i = 0..n`.
+    ///
+    /// Plan 4 Task 2c.3c.2: provides the gradient-accumulation primitive the
+    /// h_s2 GRN block's identity-residual backward needs. Reuses the existing
+    /// `dqn_saxpy_f32_kernel` (loaded once at construction) instead of
+    /// introducing a new cuBLAS legacy-handle stream-bind path.
+    ///
+    /// Both buffers live on the device; `n` is the f32 element count.
+    /// Launch config: 256 threads/block, `ceil(n/256)` blocks (grid-stride
+    /// not needed — kernel is already bounded by `if (i < n)`).
+    pub fn saxpy_inplace(
+        &self,
+        stream: &Arc<CudaStream>,
+        n: usize,
+        alpha: f32,
+        x_ptr: u64,
+        y_ptr: u64,
+    ) -> Result<(), MLError> {
+        if n == 0 {
+            return Ok(());
+        }
+        let blocks = ((n + 255) / 256) as u32;
+        let cfg = LaunchConfig {
+            grid_dim: (blocks, 1, 1),
+            block_dim: (256, 1, 1),
+            shared_mem_bytes: 0,
+        };
+        let n_i32 = n as i32;
+        unsafe {
+            stream
+                .launch_builder(&self.saxpy_kernel)
+                .arg(&y_ptr)
+                .arg(&x_ptr)
+                .arg(&alpha)
+                .arg(&n_i32)
+                .launch(cfg)
+                .map_err(|e| MLError::ModelError(format!("saxpy_inplace launch: {e}")))?;
+        }
+        Ok(())
+    }
+
     // ══════════════════════════════════════════════════════════════════════════
     // Forward pass (online network — saves activations for backward)
     // ══════════════════════════════════════════════════════════════════════════

docs/dqn-wire-up-audit.md

@@ -287,6 +287,8 @@ Plan 4 Task 2c.3b (2026-04-25): encoder forward swap to GRN (forward path runtim
 
 Plan 4 Task 2c.3c.1 (2026-04-24): `GrnBlock::backward_raw_phase1` + `backward_raw_phase2` added to `gpu_grn.rs` — Rust-side dispatcher split of the existing monolithic `backward()`. Same composition rationale as 2c.3b's forward phase split: the trunk encoder's cuBLAS `Linear_b` backward GEMM must run BETWEEN GLU backward and ELU backward (it produces the `d_elu_out` gradient that ELU backward consumes). Phase 1 sequences `grn_layernorm_backward_dx` → `_dgamma_dbeta_p1` → `_dgamma_dbeta_p2` → `grn_glu_backward`, writing `d_pre_LN` (= `d_glu_out` = `d_residual` via implicit pointer aliasing — pure pass-through, no kernel) and `d_linear_b_out [B, 2*H]` for the caller's Linear_b backward GEMM, plus accumulating `d_gamma`/`d_beta`. Phase 2 runs only `grn_elu_backward`, taking caller-provided `d_elu_out` (output of caller's Linear_b backward) and writing `d_linear_a_out [B, H]` for the caller's Linear_a backward GEMM. Both methods take raw `u64` device pointers matching the forward phase methods' style (graph-capture friendly — bypasses cudarc's `device_ptr` event-tracking machinery). The existing monolithic `backward()` is left untouched (`#[allow(dead_code)]`) for surface-area minimization; will be deleted with 2c.5 cleanup if no production callers emerge. No CUDA kernel changes. **Additive only — ZERO production callers in this commit**; the three backward panic gates (`backward_full`, `apply_iqn_trunk_gradient`, `apply_ensemble_diversity_backward`) remain in place until 2c.3c.4 wires the new phase methods. cargo check clean at 11 warnings (baseline preserved); cargo build compiles 80 cubins (unchanged — no new kernel). +223 LOC. No new module / kernel / ISV slot / Orphan row.
 
+Plan 4 Task 2c.3c.2 (2026-04-24): backward infrastructure additions — two helper methods that 2c.3c.4's GRN trunk backward wire-up requires. (1) `CublasBackwardSet::launch_dw_only_no_bias` — variant of `launch_dw_only` that runs only the cuBLAS dW GEMM and skips the 2-phase bias-grad reduction. The h_s1 GRN block's `Linear_residual` projection has no bias term (matches `gpu_grn.rs::has_residual_projection=true`'s 7-tensor layout: w_a, b_a, w_b, b_b, w_residual, gamma, beta — note no `b_residual`); calling the existing `launch_dw_only` with `db=0u64` would dereference NULL inside `bias_grad_reduce_f32_p1`. Same row-major→col-major sgemm convention as `launch_dw_only` (label `"dW_only_no_bias"` for cublasLt diagnostics), `pub(crate)` visibility matching its sibling. +35 LOC. (2) `CublasGemmSet::saxpy_inplace(stream, n, alpha, x_ptr, y_ptr)` — element-wise `y[i] += alpha * x[i]` for the h_s2 GRN's identity-residual gradient accumulation: after Linear_a_h_s2's backward writes `d_h_s1 = d_linear_a_h_s2 @ W_a_h_s2` (overwrite, β=0), the residual path needs `d_h_s1 += d_pre_ln_h_s2` to capture the identity-residual gradient. Implementation chose **option 2** (existing `dqn_saxpy_f32_kernel` from `dqn_utility_kernels.cu`) over option 1 (cuBLAS `cublasSaxpy_v2` legacy handle): the legacy cuBLAS handle is per-stream-bound state on `PerStreamCublasHandles`, so adopting it for backward saxpy would require either re-binding (race risk vs forward's binding) or a new shared handle slot — versus the existing kernel which is already loaded by `GpuExperienceCollector` for IQR/ensemble-variance Q-bonus saxpy and is graph-capturable via the standard `launch_builder` path. New `saxpy_kernel: CudaFunction` field on `CublasGemmSet`, loaded in `CublasGemmSet::new` from `DQN_UTILITY_CUBIN` (same cubin already include_bytes!'d by `gpu_dqn_trainer.rs`); +66 LOC including field, init block, and the `pub fn saxpy_inplace` method (256 threads/block, `ceil(n/256)` blocks, no grid-stride needed — kernel is already bounded by `if (i < n)`). **Additive only — ZERO production callers in this commit**; both methods sit dead-code until 2c.3c.4's wire-up commit. cargo check clean at 11 warnings (baseline preserved); cargo build unchanged — no new `.cu` files (saxpy uses the existing `dqn_saxpy_f32_kernel` symbol). No new module / kernel / ISV slot / Orphan row.
+
 Plan 4 Task 2c.3a (2026-04-24): GRN trunk param-tensor reshuffle (build-clean, runtime-broken intentionally). `compute_param_sizes()` flat layout reshuffled from 86 → 95 tensors: the 4 legacy trunk Linear tensors (`w_s1`, `b_s1`, `w_s2`, `b_s2` at indices [0..4)) are replaced with **13 GRN tensors** at indices [0..13) — 7 for h_s1 GRN block (`w_a`, `b_a`, `w_b`, `b_b`, `w_residual`, `gamma`, `beta`; SD→SH1 with `Linear_residual` GEMM since dimensions differ) plus 6 for h_s2 GRN block (`w_a`, `b_a`, `w_b`, `b_b`, `gamma`, `beta`; SH1→SH2 with identity residual since SH1==SH2). Every tensor index ≥ 4 in the OLD layout shifts +9 in the NEW layout. `NUM_WEIGHT_TENSORS` 86→95, `FIRST_ISV_TENSOR` 68→77. `layout_fingerprint_seed()` updated to mirror the new tensor names + positions; new `LAYOUT_FINGERPRINT_CURRENT = 0xcf3a24b0a1f70057` (was `0xa504d3c2f275b8af`). All 93 `padded_byte_offset` call sites + ancillary index references migrated in lockstep per `feedback_no_partial_refactor.md`. `xavier_init_params_buf` extended with init paths for the 13 new tensors: Linear matrices (`w_a`, `w_b`, `w_residual`) get Xavier; LayerNorm γ initialised to 1.0; β + biases stay zero. Trunk-forward / trunk-backward / IQN-trunk / ensemble-diversity-backward callers gated by explicit `panic!("Plan 4 Task 2c.3a: trunk param layout migrated to GRN, …")` at function entry — `BatchedForward::encoder_forward_only`, `BatchedForward::forward_target_raw`, `BatchedForward::forward_online_f32`, `BatchedBackward::backward_full`, `GpuDqnTrainer::apply_iqn_trunk_gradient`, `GpuDqnTrainer::apply_ensemble_diversity_backward`. This makes accidental runtime execution fail loudly rather than producing silent garbage from running legacy `Linear→ReLU→Linear` GEMMs against GRN-shaped param tensors. Spectral-norm descriptor (13 matrices) keeps slots [0]/[1] mapped to `w_a_h_s1`/`w_a_h_s2` (shapes match legacy W_s1/W_s2 exactly — the GRN's first Linear has the same shape as the original Linear), so the existing spectral-norm constraint transfers cleanly to the GRN's first Linear; Linear_b / Linear_residual are NOT yet spectral-normed (Task 2c.3c follow-up). **Smoke intentionally NOT run** — build-clean is the validation; runtime would hit the panic at the first encoder forward (the desired behaviour, no need to verify explicitly). Task 2c.3b swaps in the GRN forward (gpu_grn::GrnBlock::forward) at all panic-gated forward callers in place; Task 2c.3c does the same for backward callers and runs the smoke test. cargo check clean at 11 warnings (baseline preserved); cargo build compiles all 58 cubins. No new module / kernel / ISV slot in this commit — pure structural reshuffle + assertion gates.
 
 | Classification | Count |