diff --git a/crates/ml/src/cuda_pipeline/batched_forward.rs b/crates/ml/src/cuda_pipeline/batched_forward.rs index aff323b69..ed938646d 100644 --- a/crates/ml/src/cuda_pipeline/batched_forward.rs +++ b/crates/ml/src/cuda_pipeline/batched_forward.rs @@ -471,7 +471,8 @@ impl CublasGemmSet { // Create RELU_BIAS epilogue variants for hidden layers. // These fuse GEMM + bias-add + ReLU into one kernel, eliminating separate launches. - let relu_bias_shapes: Vec = vec![ + const VSN_HIDDEN_RB: usize = super::gpu_dqn_trainer::VSN_HIDDEN_DIM; + let mut relu_bias_shapes: Vec = vec![ (shared_h1, batch_size, s1_input_dim, s1_ldb), // h_s1 (shared_h2, batch_size, shared_h1, shared_h1), // h_s2 (value_h, batch_size, shared_h2, shared_h2), // h_v @@ -479,6 +480,16 @@ impl CublasGemmSet { (adv_h, batch_size, shared_h2 + branch_0_size, shared_h2 + branch_0_size), // h_bd magnitude wider input (direction-conditioned) (adv_h, batch_size, shared_h2 + 3, shared_h2 + 3), // h_bd order/urgency wider input (OFI-conditioned) ]; + // VSN feature-selection Linear_1 (×6 groups, group_dim varies). + // Phase H Site 3 — fuses per-group GEMM + bias + ReLU into one kernel. + // ldb is the full state_dim_padded so cuBLAS reads the per-group slice + // via pointer offset on the same column-major view as `vsn_forward`. + for &(gb, ge) in &ml_core::state_layout::FEATURE_GROUP_RANGES { + let group_dim = ge - gb; + relu_bias_shapes.push((VSN_HIDDEN_RB, batch_size, group_dim, state_dim_padded)); + } + relu_bias_shapes.sort(); + relu_bias_shapes.dedup(); let mut gemm_cache_relu_bias = HashMap::new(); for &(n, b, k, ldb) in &relu_bias_shapes { match create_cached_fwd_gemm_desc_relu_bias(lt_raw_handle, n, b, k, ldb, lt_ws_size) { @@ -992,42 +1003,66 @@ impl CublasGemmSet { let label_l1 = ["vsn_l1_g0", "vsn_l1_g1", "vsn_l1_g2", "vsn_l1_g3", "vsn_l1_g4", "vsn_l1_g5"][g]; let label_l2 = ["vsn_l2_g0", "vsn_l2_g1", "vsn_l2_g2", "vsn_l2_g3", "vsn_l2_g4", "vsn_l2_g5"][g]; - // Step 1: Linear_1[g]: state[..., gb..ge] @ w1_g^T → h1_g[B, 16] - // Pointer offset selects the per-group column slice; ldb stays - // at state_dim_padded so cuBLAS reads the original [state_dim_padded, B] + // Step 1+2: Linear_1[g] + bias + ReLU. Phase H Site 3: fused via + // cublasLt EPILOGUE_RELU_BIAS — h1_ptr ends up holding the POST-ReLU + // activation (same contract as the legacy two-kernel path). The + // post-ReLU buffer is what 1B-iv's backward needs as the "saved h1" + // (relu_mask_standalone takes the post-ReLU activation as the gating + // signal). No separate DtoD save copy needed. + // + // Pointer offset selects the per-group column slice; ldb stays at + // state_dim_padded so cuBLAS reads the original [state_dim_padded, B] // col-major view starting at the group's first feature. let state_slice_ptr = state_in_ptr + (gb as u64) * f32_size; - self.sgemm_f32_ldb( + let vsn_l1_ws = self.handle.lt_workspace_ptr; + let vsn_l1_wss = self.handle.lt_workspace_size; + if self.sgemm_f32_fused_relu_bias( stream, w1_ptr, state_slice_ptr, h1_ptr, - VSN_HIDDEN, batch_size, group_dim, - state_dim_padded, + b1_ptr, + VSN_HIDDEN, batch_size, group_dim, state_dim_padded, + vsn_l1_ws, vsn_l1_wss, label_l1, - )?; + ).is_err() { + self.sgemm_f32_ldb( + stream, + w1_ptr, + state_slice_ptr, + h1_ptr, + VSN_HIDDEN, batch_size, group_dim, + state_dim_padded, + label_l1, + )?; + self.launch_add_bias_relu_f32_raw(stream, h1_ptr, b1_ptr, VSN_HIDDEN, batch_size)?; + } - // Step 2: fused bias + ReLU. add_bias_relu_f32_kernel writes back - // into h1_ptr, leaving h1_ptr holding the POST-ReLU activation — - // which is exactly what 1B-iv's backward needs as the "saved h1" - // (relu_mask_standalone takes the post-ReLU activation as the gating - // signal). No separate DtoD save copy needed. - self.launch_add_bias_relu_f32_raw(stream, h1_ptr, b1_ptr, VSN_HIDDEN, batch_size)?; - - // Step 3: Linear_2[g]: h1_g[B, 16] @ w2_g^T → linear2_scratch[B, 1] - // w2_g is [1, VSN_HIDDEN_DIM] flat-stored. cuBLAS sees it as - // [VSN_HIDDEN_DIM, 1] col-major (TRANSA=T). Result is [1, B] col-major - // = [B, 1] row-major = a flat B-length f32 array. - self.sgemm_f32( + // Step 3+4: Linear_2[g] + bias_add (no activation — Linear_2 produces a + // logit). Phase H Site 3: fused GEMM+bias via cublasLt EPILOGUE_BIAS; + // falls back to separate GEMM + add_bias on missing-cache. + let vsn_ws = self.handle.lt_workspace_ptr; + let vsn_wss = self.handle.lt_workspace_size; + if self.sgemm_f32_fused_bias( stream, w2_ptr, h1_ptr, linear2_scratch_ptr, - 1, batch_size, VSN_HIDDEN, + b2_ptr, + 1, batch_size, VSN_HIDDEN, VSN_HIDDEN, + vsn_ws, vsn_wss, label_l2, - )?; - // Step 4: bias-add (no activation — Linear_2 produces a logit). - self.launch_add_bias_f32_raw(stream, linear2_scratch_ptr, b2_ptr, 1, batch_size)?; + ).is_err() { + self.sgemm_f32( + stream, + w2_ptr, + h1_ptr, + linear2_scratch_ptr, + 1, batch_size, VSN_HIDDEN, + label_l2, + )?; + self.launch_add_bias_f32_raw(stream, linear2_scratch_ptr, b2_ptr, 1, batch_size)?; + } // Step 5: scatter linear2_scratch[B] → col g of logits_buf[B, num_groups]. // Reuses `strided_scatter` from experience_kernels.cu: @@ -2176,14 +2211,18 @@ impl CublasGemmSet { false, "on_seq", )?; } else { - // Legacy GEMM+bias+ReLU path + // Legacy GEMM+bias+ReLU path — try fused epilogue first. let (fc_input, fc_k) = if d == 1 && mag_concat_ptr != 0 { (mag_concat_ptr, self.mag_concat_dim) } else { (h_s2_ptr, self.shared_h2) }; - self.sgemm_f32(stream, w_ptrs[w_fc_idx], fc_input, branch_h_ptrs[d], self.adv_h, b, fc_k, "h_bd")?; - self.launch_add_bias_relu_f32_raw(stream, branch_h_ptrs[d], w_ptrs[w_fc_idx + 1], self.adv_h, b)?; + if self.sgemm_f32_fused_relu_bias(stream, w_ptrs[w_fc_idx], fc_input, branch_h_ptrs[d], w_ptrs[w_fc_idx + 1], + self.adv_h, b, fc_k, fc_k, ws, wss, "h_bd").is_err() + { + self.sgemm_f32(stream, w_ptrs[w_fc_idx], fc_input, branch_h_ptrs[d], self.adv_h, b, fc_k, "h_bd")?; + self.launch_add_bias_relu_f32_raw(stream, branch_h_ptrs[d], w_ptrs[w_fc_idx + 1], self.adv_h, b)?; + } } let adv_out_ptr = b_logits_ptr + logit_byte_offset; @@ -2747,13 +2786,21 @@ impl CublasGemmSet { let w_fc_idx = w_fc_base[d]; let value_ptr = self.glu_value_ptrs[d]; - if is_branch_stream { - self.sgemm_f32_branch(stream, w_ptrs[w_fc_idx], vsn_input, value_ptr, ah, b, fc_k, d, _label_prefix)?; - } else { - self.sgemm_f32(stream, w_ptrs[w_fc_idx], vsn_input, value_ptr, ah, b, fc_k, _label_prefix)?; + // Phase H Site 3: fused GEMM + bias (no ReLU — GLU provides the gating). + // Workspace selected per-stream: branch streams use the per-branch + // workspace to avoid contention with concurrent branch dispatches. + let glu_ws = if is_branch_stream { self.branch_workspace_ptrs[d] } else { self.handle.lt_workspace_ptr }; + let glu_wss = self.handle.lt_workspace_size; + if self.sgemm_f32_fused_bias(stream, w_ptrs[w_fc_idx], vsn_input, value_ptr, w_ptrs[w_fc_idx + 1], + ah, b, fc_k, fc_k, glu_ws, glu_wss, _label_prefix).is_err() + { + if is_branch_stream { + self.sgemm_f32_branch(stream, w_ptrs[w_fc_idx], vsn_input, value_ptr, ah, b, fc_k, d, _label_prefix)?; + } else { + self.sgemm_f32(stream, w_ptrs[w_fc_idx], vsn_input, value_ptr, ah, b, fc_k, _label_prefix)?; + } + self.launch_add_bias_f32_raw(stream, value_ptr, w_ptrs[w_fc_idx + 1], ah, b)?; } - // Add bias (no ReLU — GLU provides the gating) - self.launch_add_bias_f32_raw(stream, value_ptr, w_ptrs[w_fc_idx + 1], ah, b)?; // ── 4. Gate GEMM: W_gate @ vsn_input → glu_gate_pre [B, AH] + gate bias ── // Plan 4 Task 2c.3a: gate base index 34 +9 → 43. @@ -2761,12 +2808,16 @@ impl CublasGemmSet { let b_gate_idx = 43 + d * 2 + 1; // b_gate_d let gate_ptr = self.glu_gate_pre_ptrs[d]; - if is_branch_stream { - self.sgemm_f32_branch(stream, w_ptrs[w_gate_idx], vsn_input, gate_ptr, ah, b, fc_k, d, _label_prefix)?; - } else { - self.sgemm_f32(stream, w_ptrs[w_gate_idx], vsn_input, gate_ptr, ah, b, fc_k, _label_prefix)?; + if self.sgemm_f32_fused_bias(stream, w_ptrs[w_gate_idx], vsn_input, gate_ptr, w_ptrs[b_gate_idx], + ah, b, fc_k, fc_k, glu_ws, glu_wss, _label_prefix).is_err() + { + if is_branch_stream { + self.sgemm_f32_branch(stream, w_ptrs[w_gate_idx], vsn_input, gate_ptr, ah, b, fc_k, d, _label_prefix)?; + } else { + self.sgemm_f32(stream, w_ptrs[w_gate_idx], vsn_input, gate_ptr, ah, b, fc_k, _label_prefix)?; + } + self.launch_add_bias_f32_raw(stream, gate_ptr, w_ptrs[b_gate_idx], ah, b)?; } - self.launch_add_bias_f32_raw(stream, gate_ptr, w_ptrs[b_gate_idx], ah, b)?; // ── 5. KAN gate: h_bd = kan_spline(gate_pre, coeffs, resid_w) * value ── // Plan 4 Task 2c.3a: KAN coeff/resid base index 42 +9 → 51. diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index b51eb97b5..5f9c77149 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -2,6 +2,31 @@ **Status:** Populated during Plan 1 Task 6 (A.5 orphan audit). Updated on every commit per Invariant 7. +P5T5 Phase H Site 3 — fuse VSN feature-selection + GLU value/gate +GEMM+bias (2026-04-28): migrates the remaining 4 forward sites: +1. **VSN Linear_1** (×6 feature groups) — adds 6 per-group RELU_BIAS + shapes `(VSN_HIDDEN, B, group_dim, state_dim_padded)` to the + relu_bias cache (`group_dim` ∈ {42, 32, 16, 16, 8, 7} from + `FEATURE_GROUP_RANGES`); migrates `vsn_forward` Step 1+2 from + `sgemm_f32_ldb + launch_add_bias_relu_f32_raw` to + `sgemm_f32_fused_relu_bias`. Post-ReLU h1_ptr remains the same + contract for 1B-iv backward (relu_mask_standalone gating signal). +2. **VSN Linear_2** (×6 feature groups, all share shape `(1, B, + VSN_HIDDEN, VSN_HIDDEN)`) — migrates `vsn_forward` Step 3+4 to + `sgemm_f32_fused_bias`. +3. **GLU value head** (`launch_vsn_glu_branch` Step 3) and +4. **GLU gate** (Step 4) — both migrated to `sgemm_f32_fused_bias`, + workspace selected per-stream (per-branch workspace when running on + a branch stream, default workspace otherwise). +Also fixes a partial-refactor leak in `forward_online_raw` — the +sequential branch fallback (when `distinct_branches=false`) was still +on the unfused `sgemm_f32 + launch_add_bias_relu_f32_raw` pair while +its multi-stream sibling already used `sgemm_f32_fused_relu_bias`. +Now both code paths take the fused-first contract. Touched: +`crates/ml/src/cuda_pipeline/batched_forward.rs`. cargo check clean +at 13 warnings; `cargo test --no-run` clean. Layout fingerprint +unchanged (no params buffer changes). + P5T5 Phase H Site 3 — fuse branch FC adv_logits GEMM+bias into BIAS epilogue (2026-04-28): migrates 6 sites — `decoder_forward_only` (1 site, sequential), `forward_online_raw` (2 sites: multi-stream branch