diff --git a/crates/ml/src/cuda_pipeline/batched_backward.rs b/crates/ml/src/cuda_pipeline/batched_backward.rs index a3170228b..e7f4c234e 100644 --- a/crates/ml/src/cuda_pipeline/batched_backward.rs +++ b/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, + 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 diff --git a/crates/ml/src/cuda_pipeline/batched_forward.rs b/crates/ml/src/cuda_pipeline/batched_forward.rs index be426efa9..d147ae6a9 100644 --- a/crates/ml/src/cuda_pipeline/batched_forward.rs +++ b/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, 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, + 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) // ══════════════════════════════════════════════════════════════════════════ diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index 38416832c..1186cdc6d 100644 --- a/docs/dqn-wire-up-audit.md +++ b/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 |