gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

perf(phase-h): fuse VSN + GLU GEMM+bias into RELU_BIAS / BIAS epilogues

Browse Source 
Migrates the remaining 4 forward sites:
- VSN Linear_1 (x6 groups): adds 6 per-group RELU_BIAS shapes to the
  relu_bias cache; migrates sgemm_f32_ldb + launch_add_bias_relu_f32_raw
  pairs to sgemm_f32_fused_relu_bias.
- VSN Linear_2 (x6 groups, single shape (1, B, VSN_HIDDEN, VSN_HIDDEN)):
  migrates to sgemm_f32_fused_bias.
- GLU value head + gate (launch_vsn_glu_branch Steps 3-4): migrates to
  sgemm_f32_fused_bias with per-stream workspace selection.

Also fixes a partial-refactor leak in forward_online_raw — the
sequential branch fallback (distinct_branches=false) was still on the
unfused sgemm_f32 + launch_add_bias_relu_f32_raw pair while its
multi-stream sibling already used the fused epilogue. Now both code
paths take the fused-first contract.

Layout fingerprint unchanged (no params buffer changes).
This commit is contained in:
jgrusewski
2026-04-28 12:44:32 +02:00
parent 19072ee9f9
commit 1c50a4e8a4
2 changed files with 114 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files

127
crates/ml/src/cuda_pipeline/batched_forward.rs
View File

@@ -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<FwdGemmKey> = vec![
const VSN_HIDDEN_RB: usize = super::gpu_dqn_trainer::VSN_HIDDEN_DIM;
let mut relu_bias_shapes: Vec<FwdGemmKey> = 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.

25
docs/dqn-wire-up-audit.md
View File

@@ -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