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

perf: parallel backward branches using fork-join on branch_streams

Browse Source 
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-18 16:35:27 +02:00
parent 9fb330da7e
commit 108a0993c7
1 changed files with 647 additions and 181 deletions
Download Patch File Download Diff File Expand all files Collapse all files

828
crates/ml/src/cuda_pipeline/batched_backward.rs
View File

@@ -84,7 +84,7 @@ use std::sync::Arc;
use cudarc::cublas::sys as cublas_sys;
use cudarc::cublaslt::sys as cublaslt_sys;
use cudarc::cublaslt::result as cublaslt_result;
use cudarc::driver::{CudaFunction, CudaSlice, CudaStream, DevicePtr, LaunchConfig, PushKernelArg};
use cudarc::driver::{CudaEvent, CudaFunction, CudaSlice, CudaStream, DevicePtr, LaunchConfig, PushKernelArg};
use crate::MLError;
use super::gpu_dqn_trainer::{GpuDqnTrainConfig, compute_param_sizes, padded_byte_offset};
@@ -197,6 +197,35 @@ pub struct CublasBackwardSet {
/// Used to compute `padded_byte_offset` for gradient buffer offsets.
param_sizes: [usize; super::gpu_dqn_trainer::NUM_WEIGHT_TENSORS],
// ── Multi-stream branch dispatch ──
/// 4 forked CUDA streams for parallel advantage branch backward execution.
/// Each branch submits dW/db GEMMs and KAN backward to its own stream,
/// then the main stream joins all four before the dX (Step 5) accumulation.
branch_streams: [Arc<CudaStream>; 4],
/// Per-branch cublasLt workspace buffers — eliminates workspace contention
/// when 4 branch streams run GEMMs in parallel. Each 32 MB.
_branch_workspace_bufs: [CudaSlice<u8>; 4],
branch_workspace_ptrs: [u64; 4],
/// Reused CUDA events for fork-join synchronization.
trunk_done_event: CudaEvent,
branch_done_events: [CudaEvent; 4],
// ── Per-branch scratch buffers (safe for parallel execution) ──
/// Per-branch GLU value gradient scratch [B, AH] each — avoids cross-branch contention.
branch_d_glu_value: [CudaSlice<f32>; 4],
/// Per-branch GLU gate_pre gradient scratch [B, AH] each.
branch_d_glu_gate: [CudaSlice<f32>; 4],
// ── Per-branch KAN backward intermediate buffers ──
/// Per-element B-spline basis contributions [B*AH, 4] per branch — f32.
branch_kan_d_coeff_per_elem: [CudaSlice<f32>; 4],
/// Per-element span index [B*AH] per branch — i32.
branch_kan_d_coeff_span: [CudaSlice<i32>; 4],
/// Per-element residual_w gradient contribution [B*AH] per branch — f32.
branch_kan_d_residual_per_elem: [CudaSlice<f32>; 4],
// ── Cached GEMM descriptors (created once at init, reused per-call) ──
/// Map from (transa, transb, m, n, k, lda, ldb, ldc) → pre-created descriptors + algo.
gemm_cache: HashMap<BwdGemmKey, CachedBwdGemmDesc>,
@@ -308,6 +337,86 @@ impl CublasBackwardSet {
let kan_d_residual_per_elem = stream.alloc_zeros::<f32>(kan_n)
.map_err(|e| MLError::ModelError(format!("kan_d_residual_per_elem alloc [{kan_n}]: {e}")))?;
// ── Fork 4 branch streams for parallel backward dispatch ────
let branch_streams = [
stream.fork().map_err(|e| MLError::DeviceError(format!("bw fork branch stream 0: {e}")))?,
stream.fork().map_err(|e| MLError::DeviceError(format!("bw fork branch stream 1: {e}")))?,
stream.fork().map_err(|e| MLError::DeviceError(format!("bw fork branch stream 2: {e}")))?,
stream.fork().map_err(|e| MLError::DeviceError(format!("bw fork branch stream 3: {e}")))?,
];
// ── Per-branch cublasLt workspace buffers (32 MB each) ──────
let branch_ws_size: usize = 32 * 1024 * 1024;
let bw_ws_buf_0 = stream.alloc_zeros::<u8>(branch_ws_size)
.map_err(|e| MLError::ModelError(format!("bw branch workspace 0 alloc: {e}")))?;
let bw_ws_buf_1 = stream.alloc_zeros::<u8>(branch_ws_size)
.map_err(|e| MLError::ModelError(format!("bw branch workspace 1 alloc: {e}")))?;
let bw_ws_buf_2 = stream.alloc_zeros::<u8>(branch_ws_size)
.map_err(|e| MLError::ModelError(format!("bw branch workspace 2 alloc: {e}")))?;
let bw_ws_buf_3 = stream.alloc_zeros::<u8>(branch_ws_size)
.map_err(|e| MLError::ModelError(format!("bw branch workspace 3 alloc: {e}")))?;
let branch_workspace_ptrs = [
bw_ws_buf_0.raw_ptr(),
bw_ws_buf_1.raw_ptr(),
bw_ws_buf_2.raw_ptr(),
bw_ws_buf_3.raw_ptr(),
];
let branch_workspace_bufs = [bw_ws_buf_0, bw_ws_buf_1, bw_ws_buf_2, bw_ws_buf_3];
// ── Pre-allocate CUDA events for fork-join ──────────────────
let ctx = stream.context();
let trunk_done_event = ctx.new_event(None)
.map_err(|e| MLError::DeviceError(format!("bw alloc trunk_done_event: {e}")))?;
let branch_done_events = [
ctx.new_event(None).map_err(|e| MLError::DeviceError(format!("bw alloc branch_done_event 0: {e}")))?,
ctx.new_event(None).map_err(|e| MLError::DeviceError(format!("bw alloc branch_done_event 1: {e}")))?,
ctx.new_event(None).map_err(|e| MLError::DeviceError(format!("bw alloc branch_done_event 2: {e}")))?,
ctx.new_event(None).map_err(|e| MLError::DeviceError(format!("bw alloc branch_done_event 3: {e}")))?,
];
// ── Per-branch GLU scratch buffers ──────────────────────────
let b_ah = config.batch_size * config.adv_h;
let alloc_f32 = |n: usize, label: &str| -> Result<CudaSlice<f32>, MLError> {
stream.alloc_zeros::<f32>(n)
.map_err(|e| MLError::ModelError(format!("{label} alloc [{n}]: {e}")))
};
let branch_d_glu_value = [
alloc_f32(b_ah, "branch_d_glu_value_0")?,
alloc_f32(b_ah, "branch_d_glu_value_1")?,
alloc_f32(b_ah, "branch_d_glu_value_2")?,
alloc_f32(b_ah, "branch_d_glu_value_3")?,
];
let branch_d_glu_gate = [
alloc_f32(b_ah, "branch_d_glu_gate_0")?,
alloc_f32(b_ah, "branch_d_glu_gate_1")?,
alloc_f32(b_ah, "branch_d_glu_gate_2")?,
alloc_f32(b_ah, "branch_d_glu_gate_3")?,
];
// ── Per-branch KAN backward intermediate buffers ────────────
let alloc_i32 = |n: usize, label: &str| -> Result<CudaSlice<i32>, MLError> {
stream.alloc_zeros::<i32>(n)
.map_err(|e| MLError::ModelError(format!("{label} alloc [{n}]: {e}")))
};
let branch_kan_d_coeff_per_elem = [
alloc_f32(kan_n * 4, "branch_kan_d_coeff_pe_0")?,
alloc_f32(kan_n * 4, "branch_kan_d_coeff_pe_1")?,
alloc_f32(kan_n * 4, "branch_kan_d_coeff_pe_2")?,
alloc_f32(kan_n * 4, "branch_kan_d_coeff_pe_3")?,
];
let branch_kan_d_coeff_span = [
alloc_i32(kan_n, "branch_kan_d_coeff_span_0")?,
alloc_i32(kan_n, "branch_kan_d_coeff_span_1")?,
alloc_i32(kan_n, "branch_kan_d_coeff_span_2")?,
alloc_i32(kan_n, "branch_kan_d_coeff_span_3")?,
];
let branch_kan_d_residual_per_elem = [
alloc_f32(kan_n, "branch_kan_d_resid_pe_0")?,
alloc_f32(kan_n, "branch_kan_d_resid_pe_1")?,
alloc_f32(kan_n, "branch_kan_d_resid_pe_2")?,
alloc_f32(kan_n, "branch_kan_d_resid_pe_3")?,
];
Ok(Self {
handle: shared,
relu_mask_kernel,
@@ -330,6 +439,16 @@ impl CublasBackwardSet {
branch_1_size: config.branch_1_size,
branch_2_size: config.branch_2_size,
branch_3_size: config.branch_3_size,
branch_streams,
_branch_workspace_bufs: branch_workspace_bufs,
branch_workspace_ptrs,
trunk_done_event,
branch_done_events,
branch_d_glu_value,
branch_d_glu_gate,
branch_kan_d_coeff_per_elem,
branch_kan_d_coeff_span,
branch_kan_d_residual_per_elem,
param_sizes,
gemm_cache,
})
@@ -355,6 +474,8 @@ impl CublasBackwardSet {
///
/// Uses cached descriptors when the shape is pre-cached at init time.
/// Falls back to inline descriptor creation for uncached shapes.
///
/// Uses the shared (main) workspace — NOT safe for branch streams.
fn sgemm_f32(
&self,
stream: &CudaStream,
@@ -367,18 +488,40 @@ impl CublasBackwardSet {
beta: f32,
c: u64, ldc: i32,
label: &str,
) -> Result<(), MLError> {
self.sgemm_f32_ws(stream, transa, transb, m, n, k, alpha, a, lda, b_ptr, ldb, beta, c, ldc, self.handle.lt_workspace_ptr, label)
}
/// F32 GEMM with explicit workspace pointer.
///
/// Branch streams MUST use per-branch workspace pointers to avoid
/// contention with the shared workspace used by the main stream.
#[allow(clippy::too_many_arguments)]
fn sgemm_f32_ws(
&self,
stream: &CudaStream,
transa: cublas_sys::cublasOperation_t,
transb: cublas_sys::cublasOperation_t,
m: i32, n: i32, k: i32,
alpha: f32,
a: u64, lda: i32,
b_ptr: u64, ldb: i32,
beta: f32,
c: u64, ldc: i32,
ws_ptr: u64,
label: &str,
) -> Result<(), MLError> {
let key: BwdGemmKey = (transa as i32, transb as i32, m, n, k, lda, ldb, ldc);
if let Some(cached) = self.gemm_cache.get(&key) {
// ── Fast path: use pre-created descriptors + algo ──
self.sgemm_f32_cached(stream, cached, alpha, a, b_ptr, beta, c, m, n, k, lda, ldb, ldc, label)
self.sgemm_f32_cached_ws(stream, cached, alpha, a, b_ptr, beta, c, m, n, k, lda, ldb, ldc, ws_ptr, label)
} else {
// ── Slow path: inline descriptor creation ──
self.sgemm_f32_uncached(stream, transa, transb, m, n, k, alpha, a, lda, b_ptr, ldb, beta, c, ldc, label)
self.sgemm_f32_uncached_ws(stream, transa, transb, m, n, k, alpha, a, lda, b_ptr, ldb, beta, c, ldc, ws_ptr, label)
}
}
/// Fast path: execute backward GEMM with pre-cached descriptors.
/// Fast path: execute backward GEMM with pre-cached descriptors (shared workspace).
#[allow(clippy::too_many_arguments)]
fn sgemm_f32_cached(
&self,
@@ -392,6 +535,25 @@ impl CublasBackwardSet {
m: i32, n: i32, k: i32,
lda: i32, ldb: i32, ldc: i32,
label: &str,
) -> Result<(), MLError> {
self.sgemm_f32_cached_ws(stream, cached, alpha, a, b_ptr, beta, c, m, n, k, lda, ldb, ldc, self.handle.lt_workspace_ptr, label)
}
/// Fast path: execute backward GEMM with pre-cached descriptors and explicit workspace.
#[allow(clippy::too_many_arguments)]
fn sgemm_f32_cached_ws(
&self,
stream: &CudaStream,
cached: &CachedBwdGemmDesc,
alpha: f32,
a: u64,
b_ptr: u64,
beta: f32,
c: u64,
m: i32, n: i32, k: i32,
lda: i32, ldb: i32, ldc: i32,
ws_ptr: u64,
label: &str,
) -> Result<(), MLError> {
unsafe {
let cu_stream = stream.cu_stream() as cublaslt_sys::cudaStream_t;
@@ -410,7 +572,7 @@ impl CublasBackwardSet {
c as *mut std::ffi::c_void,
cached.d_layout,
&cached.algo as *const cublaslt_sys::cublasLtMatmulAlgo_t,
self.handle.lt_workspace_ptr as *mut std::ffi::c_void,
ws_ptr as *mut std::ffi::c_void,
self.handle.lt_workspace_size,
cu_stream,
);
@@ -420,7 +582,7 @@ impl CublasBackwardSet {
tracing::error!(
m=m, n=n, k=k, lda=lda, ldb=ldb, ldc=ldc,
ws_provided=self.handle.lt_workspace_size,
ws_ptr=self.handle.lt_workspace_ptr,
ws_ptr=ws_ptr,
?e,
"cublasLtMatmul CACHED FAILED for bw {label}"
);
@@ -430,7 +592,7 @@ impl CublasBackwardSet {
Ok(())
}
/// Slow path: inline descriptor creation for uncached backward GEMM shapes.
/// Slow path: inline descriptor creation for uncached backward GEMM shapes (shared workspace).
#[allow(clippy::too_many_arguments)]
fn sgemm_f32_uncached(
&self,
@@ -444,6 +606,25 @@ impl CublasBackwardSet {
beta: f32,
c: u64, ldc: i32,
label: &str,
) -> Result<(), MLError> {
self.sgemm_f32_uncached_ws(stream, transa, transb, m, n, k, alpha, a, lda, b_ptr, ldb, beta, c, ldc, self.handle.lt_workspace_ptr, label)
}
/// Slow path: inline descriptor creation with explicit workspace pointer.
#[allow(clippy::too_many_arguments)]
fn sgemm_f32_uncached_ws(
&self,
stream: &CudaStream,
transa: cublas_sys::cublasOperation_t,
transb: cublas_sys::cublasOperation_t,
m: i32, n: i32, k: i32,
alpha: f32,
a: u64, lda: i32,
b_ptr: u64, ldb: i32,
beta: f32,
c: u64, ldc: i32,
ws_ptr: u64,
label: &str,
) -> Result<(), MLError> {
let f32_type = cublaslt_sys::cudaDataType_t::CUDA_R_32F;
// FAST_TF32: explicit TF32 tensor core path. Required on cuBLAS 13.0+ (SM90)
@@ -531,7 +712,7 @@ impl CublasBackwardSet {
c as *mut std::ffi::c_void,
d_layout,
&heuristic.algo as *const cublaslt_sys::cublasLtMatmulAlgo_t,
self.handle.lt_workspace_ptr as *mut std::ffi::c_void,
ws_ptr as *mut std::ffi::c_void,
self.handle.lt_workspace_size,
cu_stream,
);
@@ -540,7 +721,7 @@ impl CublasBackwardSet {
m=m, n=n, k=k, lda=lda, ldb=ldb, ldc=ldc,
ws_provided=self.handle.lt_workspace_size,
ws_needed=heuristic.workspaceSize,
ws_ptr=self.handle.lt_workspace_ptr,
ws_ptr=ws_ptr,
?e,
"cublasLtMatmul UNCACHED FAILED for bw {label}"
);
@@ -796,6 +977,388 @@ impl CublasBackwardSet {
Ok(())
}
// ══════════════════════════════════════════════════════════════════════════
// Per-branch backward (Steps 1-4: dW/db only, no dX)
// ══════════════════════════════════════════════════════════════════════════
/// KAN spline gate backward using per-branch intermediate buffers.
///
/// Identical to `launch_kan_gate_backward` but uses branch-specific
/// `d_coeff_per_elem`, `d_coeff_span`, and `d_residual_per_elem` buffers,
/// making it safe for concurrent execution on branch streams.
#[allow(clippy::too_many_arguments)]
fn launch_kan_gate_backward_branch(
&self,
stream: &Arc<CudaStream>,
d: usize, // branch index 0..3
d_output: u64, // [B, AH]
gate_pre: u64, // [B, AH]
value: u64, // [B, AH]
spline_coeff: u64, // [AH, 8]
residual_w: u64, // [AH]
d_gate_pre: u64, // [B, AH] output
d_value: u64, // [B, AH] output
d_spline_coeff: u64, // [AH, 8] output (in grad_buf)
d_residual_w: u64, // [AH] output (in grad_buf)
n: usize,
adv_h: usize,
) -> Result<(), MLError> {
let n_i32 = n as i32;
let adv_h_i32 = adv_h as i32;
let blocks = ((n + 255) / 256) as u32;
// Use per-branch intermediate buffers (no cross-branch contention).
let d_coeff_pe_ptr = raw_f32_ptr(&self.branch_kan_d_coeff_per_elem[d], stream);
let d_span_ptr = {
let (ptr, guard) = self.branch_kan_d_coeff_span[d].device_ptr(stream);
let _no_drop = ManuallyDrop::new(guard);
ptr
};
let d_resid_pe_ptr = raw_f32_ptr(&self.branch_kan_d_residual_per_elem[d], stream);
// ── Stage 1: per-element scatter (no atomicAdd) ──
unsafe {
stream
.launch_builder(&self.kan_gate_backward_kernel)
.arg(&d_output)
.arg(&gate_pre)
.arg(&value)
.arg(&spline_coeff)
.arg(&residual_w)
.arg(&d_gate_pre)
.arg(&d_value)
.arg(&d_coeff_pe_ptr)
.arg(&d_span_ptr)
.arg(&d_resid_pe_ptr)
.arg(&n_i32)
.arg(&adv_h_i32)
.launch(LaunchConfig {
grid_dim: (blocks, 1, 1),
block_dim: (256, 1, 1),
shared_mem_bytes: 0,
})
.map_err(|e| MLError::ModelError(format!("kan_gate_backward_kernel branch {d}: {e}")))?;
}
// ── Stage 2: deterministic reduction ──
let total_params = adv_h * 9;
let reduce_blocks = ((total_params + 255) / 256) as u32;
unsafe {
stream
.launch_builder(&self.kan_grad_reduce_kernel)
.arg(&d_coeff_pe_ptr)
.arg(&d_span_ptr)
.arg(&d_resid_pe_ptr)
.arg(&d_spline_coeff)
.arg(&d_residual_w)
.arg(&n_i32)
.arg(&adv_h_i32)
.launch(LaunchConfig {
grid_dim: (reduce_blocks, 1, 1),
block_dim: (256, 1, 1),
shared_mem_bytes: 0,
})
.map_err(|e| MLError::ModelError(format!("kan_grad_reduce_kernel branch {d}: {e}")))?;
}
Ok(())
}
/// Backward Steps 1-4 for a single branch on a branch stream.
///
/// Computes dW/db for the branch output layer, KAN spline backward, and
/// value/gate path weight gradients. All writes go to non-overlapping
/// regions of `grad_buf`, making this safe for parallel execution.
///
/// Does NOT compute Step 5 (dX to d_h_s2 or d_concat). That runs on the
/// main stream after all 4 branches join.
#[allow(clippy::too_many_arguments)]
fn backward_branch_dw(
&self,
stream: &Arc<CudaStream>,
d: usize,
ws_ptr: u64,
// Per-branch data
d_adv_logits_d: u64, // [B, n_d]
save_h_b_d: u64, // [B, AH] (KAN gate output)
scratch_d_h_b_d: u64, // [B, AH] scratch for d_h_bd
glu_gate_pre_d: u64, // [B, AH] saved gate pre-activation
glu_value_d: u64, // [B, AH] saved value path output
// Weight pointers (read-only)
w_out: u64, // W_bdk_out [n_d, AH]
w_ptrs: &[u64; super::gpu_dqn_trainer::NUM_WEIGHT_TENSORS],
// Grad buf offsets
grad_buf_base: u64,
goff_w_bout_d: u64,
goff_b_bout_d: u64,
goff_w_bfc_d: u64,
goff_b_bfc_d: u64,
goff_w_gate_d: u64,
goff_b_gate_d: u64,
goff_kan_coeff_d: u64,
goff_kan_resid_d: u64,
// FC input for dW (branches 1,2,3 may use concat buffers)
fc_input: u64,
fc_in_dim: usize,
// Branch output dim
n_d: usize,
b: usize,
) -> Result<(), MLError> {
// ── Step 1: Branch output layer (no activation — logit layer) ──
// dW += dY^T @ X, db += sum(dY), dX = dY @ W^T → scratch_d_h_b[d]
self.backward_fc_layer_ws(
stream,
d_adv_logits_d,
save_h_b_d,
w_out,
grad_buf_base + goff_w_bout_d,
grad_buf_base + goff_b_bout_d,
scratch_d_h_b_d,
n_d,
self.adv_h,
b,
ws_ptr,
)?;
// Per-branch GLU scratch pointers
let d_glu_value_ptr = raw_f32_ptr(&self.branch_d_glu_value[d], stream);
let d_glu_gate_ptr = raw_f32_ptr(&self.branch_d_glu_gate[d], stream);
// ── Step 2: KAN spline backward ──
let kan_coeff_idx = 42 + d * 2;
let kan_resid_idx = 42 + d * 2 + 1;
self.launch_kan_gate_backward_branch(
stream,
d,
scratch_d_h_b_d,
glu_gate_pre_d,
glu_value_d,
w_ptrs[kan_coeff_idx],
w_ptrs[kan_resid_idx],
d_glu_gate_ptr,
d_glu_value_ptr,
grad_buf_base + goff_kan_coeff_d,
grad_buf_base + goff_kan_resid_d,
b * self.adv_h,
self.adv_h,
)?;
// ── Step 3: Value path weight gradients ──
self.launch_dw_only_ws(
stream,
d_glu_value_ptr,
fc_input,
grad_buf_base + goff_w_bfc_d,
grad_buf_base + goff_b_bfc_d,
self.adv_h,
fc_in_dim,
b,
ws_ptr,
)?;
// ── Step 4: Gate path weight gradients ──
self.launch_dw_only_ws(
stream,
d_glu_gate_ptr,
fc_input,
grad_buf_base + goff_w_gate_d,
grad_buf_base + goff_b_gate_d,
self.adv_h,
fc_in_dim,
b,
ws_ptr,
)?;
Ok(())
}
/// Backward Step 5 for a single branch on the main stream.
///
/// Computes dX (upstream gradient to VSN input / d_h_s2 / concat buffers).
/// Must run on the main stream AFTER all branches finish Steps 1-4 so that
/// d_h_s2 accumulation with beta is ordered correctly.
#[allow(clippy::too_many_arguments)]
fn backward_branch_dx(
&self,
stream: &Arc<CudaStream>,
d: usize,
w_fc: u64,
w_gate: u64,
scratch_d_h_s2: u64,
mag_concat_ptr: u64,
d_mag_concat_ptr: u64,
ord_concat_ptr: u64,
d_ord_concat_ptr: u64,
urg_concat_ptr: u64,
d_urg_concat_ptr: u64,
b: usize,
) -> Result<(), MLError> {
// Read per-branch GLU scratch (produced by Steps 1-4 on branch stream,
// now readable by main stream after join).
let d_glu_value_ptr = raw_f32_ptr(&self.branch_d_glu_value[d], stream);
let d_glu_gate_ptr = raw_f32_ptr(&self.branch_d_glu_gate[d], stream);
if d == 1 && mag_concat_ptr != 0 {
// Magnitude: dX writes to d_mag_concat [B, SH2+3].
self.launch_dx_only(
stream,
d_glu_value_ptr,
w_fc,
d_mag_concat_ptr,
self.adv_h,
self.shared_h2 + 3,
b,
0.0_f32,
)?;
self.launch_dx_only(
stream,
d_glu_gate_ptr,
w_gate,
d_mag_concat_ptr,
self.adv_h,
self.shared_h2 + 3,
b,
1.0_f32,
)?;
} else if d == 2 && ord_concat_ptr != 0 {
// Order: dX writes to d_ord_concat [B, SH2+3].
self.launch_dx_only(
stream,
d_glu_value_ptr,
w_fc,
d_ord_concat_ptr,
self.adv_h,
self.shared_h2 + 3,
b,
0.0_f32,
)?;
self.launch_dx_only(
stream,
d_glu_gate_ptr,
w_gate,
d_ord_concat_ptr,
self.adv_h,
self.shared_h2 + 3,
b,
1.0_f32,
)?;
} else if d == 3 && urg_concat_ptr != 0 {
// Urgency: dX writes to d_urg_concat [B, SH2+3].
self.launch_dx_only(
stream,
d_glu_value_ptr,
w_fc,
d_urg_concat_ptr,
self.adv_h,
self.shared_h2 + 3,
b,
0.0_f32,
)?;
self.launch_dx_only(
stream,
d_glu_gate_ptr,
w_gate,
d_urg_concat_ptr,
self.adv_h,
self.shared_h2 + 3,
b,
1.0_f32,
)?;
} else {
// Direction (d==0) or any branch with no concat: accumulate into d_h_s2 directly.
let beta_s2 = if d == 0 { 0.0_f32 } else { 1.0_f32 };
self.launch_dx_only(
stream,
d_glu_value_ptr,
w_fc,
scratch_d_h_s2,
self.adv_h,
self.shared_h2,
b,
beta_s2,
)?;
self.launch_dx_only(
stream,
d_glu_gate_ptr,
w_gate,
scratch_d_h_s2,
self.adv_h,
self.shared_h2,
b,
1.0_f32,
)?;
}
Ok(())
}
/// Workspace-aware `backward_fc_layer`: dW + db + dX using explicit workspace.
#[allow(clippy::too_many_arguments)]
fn backward_fc_layer_ws(
&self,
stream: &Arc<CudaStream>,
dy: u64, x: u64, w: u64, dw: u64, db: u64, dx: u64,
out_dim: usize, in_dim: usize, batch: usize,
ws_ptr: u64,
) -> Result<(), MLError> {
// dW += dY^T @ X
self.sgemm_f32_ws(
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,
ws_ptr,
"fc_dW_ws",
)?;
self.launch_bias_grad(stream, dy, db, out_dim, batch)?;
// dX = dY @ W^T
if dx != 0 {
self.sgemm_f32_ws(
stream,
cublas_sys::cublasOperation_t::CUBLAS_OP_N,
cublas_sys::cublasOperation_t::CUBLAS_OP_N,
in_dim as i32, batch as i32, out_dim as i32,
1.0, w, in_dim as i32,
dy, out_dim as i32,
0.0, dx, in_dim as i32,
ws_ptr,
"fc_dX_ws",
)?;
}
Ok(())
}
/// Workspace-aware `launch_dw_only`: dW + db using explicit workspace.
#[allow(clippy::too_many_arguments)]
fn launch_dw_only_ws(
&self,
stream: &Arc<CudaStream>,
dy: u64, x: u64, dw: u64, db: u64,
out_dim: usize, in_dim: usize, batch: usize,
ws_ptr: u64,
) -> Result<(), MLError> {
self.sgemm_f32_ws(
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,
ws_ptr,
"dW_only_ws",
)?;
self.launch_bias_grad(stream, dy, db, out_dim, batch)
}
// ══════════════════════════════════════════════════════════════════════════
// Full backward pass
// ══════════════════════════════════════════════════════════════════════════
@@ -876,9 +1439,10 @@ impl CublasBackwardSet {
// GLU saved forward activations for backward pass
glu_gate_pre: &[u64; 4], // [B, AH] each — saved pre-sigmoid gate activations
glu_value: &[u64; 4], // [B, AH] each — saved value path outputs
// GLU backward scratch buffers (reused across branches)
scratch_d_glu_value: u64, // [B, AH] — d_value output from GLU backward
scratch_d_glu_gate: u64, // [B, AH] — d_gate_pre output from GLU backward
// GLU backward scratch buffers (kept for backward compatibility; per-branch
// versions in branch_d_glu_value/branch_d_glu_gate are used for parallel dispatch)
_scratch_d_glu_value: u64, // [B, AH] — superseded by per-branch scratch
_scratch_d_glu_gate: u64, // [B, AH] — superseded by per-branch scratch
) -> Result<(), MLError> {
let b = self.batch_size;
let na = self.num_atoms;
@@ -975,57 +1539,28 @@ impl CublasBackwardSet {
];
// ══════════════════════════════════════════════════════════════════
// BRANCHES (all 4, with KAN spline backward replacing ReLU mask)
// BRANCHES: parallel dW/db (Steps 1-4), sequential dX (Step 5)
//
// Steps 1-4 write to non-overlapping regions of grad_buf and
// per-branch scratch buffers, so they run safely in parallel on
// branch_streams[0..3]. Step 5 (dX) may accumulate into the
// shared scratch_d_h_s2 with beta ordering, so it runs
// sequentially on the main stream after all branches join.
// ══════════════════════════════════════════════════════════════════
// ── Fork: record trunk completion, dispatch Steps 1-4 on branch streams ──
self.trunk_done_event.record(stream)
.map_err(|e| MLError::DeviceError(format!("bw trunk_done_event record: {e}")))?;
for d in 0..4 {
let bs = &self.branch_streams[d];
bs.wait(&self.trunk_done_event)
.map_err(|e| MLError::DeviceError(format!("bw branch stream {d} wait trunk: {e}")))?;
let n_d = branch_n[d];
let w_out = w_ptrs[w_bout_idx[d]]; // W_bdk_out[n_d, AH]
let w_fc = w_ptrs[w_bfc_idx[d]]; // W_bdk_fc[AH, SH2] (value path weight)
let w_gate = w_ptrs[w_gate_idx[d]]; // W_gate_d[AH, SH2] (gate path weight)
let w_out = w_ptrs[w_bout_idx[d]];
// KAN spline weight pointers (from params_buf, read-only)
let kan_coeff_idx = 42 + d * 2; // kan_coeff_d
let kan_resid_idx = 42 + d * 2 + 1; // kan_resid_d
// ── Step 1: Branch output layer (no activation — logit layer) ──
// dY = d_adv_logits[d] [B, n_d]
// X = save_h_b[d] [B, AH] (KAN gate output)
// W = W_bdk_out [n_d, AH]
// dW += dY^T @ X, db += sum(dY), dX_b = dY @ W^T → scratch_d_h_b[d]
self.backward_fc_layer(
stream,
d_adv_logits[d], // dY [B, n_d]
save_h_b[d], // X [B, AH]
w_out, // W [n_d, AH]
grad_buf_base + goff_w_bout[d], // dW
grad_buf_base + goff_b_bout[d], // db
scratch_d_h_b[d], // dX [B, AH] → d_h_bd (gradient into KAN gate output)
n_d, // out_dim
self.adv_h, // in_dim
b,
)?;
// ── Step 2: KAN spline backward (replaces GLU backward) ──────
// d_h_bd → d_value + d_gate_pre + d_coeff + d_resid_w
self.launch_kan_gate_backward(
stream,
scratch_d_h_b[d], // d_output = d_h_bd [B, AH]
glu_gate_pre[d], // saved gate_pre [B, AH]
glu_value[d], // saved value [B, AH]
w_ptrs[kan_coeff_idx], // spline_coeff [AH, 8] (current weights)
w_ptrs[kan_resid_idx], // residual_w [AH] (current weights)
scratch_d_glu_gate, // d_gate_pre output [B, AH]
scratch_d_glu_value, // d_value output [B, AH]
grad_buf_base + goff_kan_coeff[d], // d_spline_coeff [AH, 8] (accumulated in grad_buf)
grad_buf_base + goff_kan_resid[d], // d_residual_w [AH] (accumulated in grad_buf)
b * self.adv_h,
self.adv_h,
)?;
// ── Step 3: Value path weight gradients ────────────────────
// dW_bdf += d_value^T @ vsn_input, db_bdf += sum(d_value)
// Branches 1,2,3 use wider concat input [B, SH2+3] for dW.
// FC input for dW (branches 1,2,3 may use wider concat buffers)
let (fc_input, fc_in_dim) = if d == 1 && mag_concat_ptr != 0 {
(mag_concat_ptr, self.shared_h2 + 3)
} else if d == 2 && ord_concat_ptr != 0 {
@@ -1035,135 +1570,66 @@ impl CublasBackwardSet {
} else {
(save_h_s2, self.shared_h2)
};
self.launch_dw_only(
stream,
scratch_d_glu_value, // dY = d_value [B, AH] — f32
fc_input, // X [B, fc_in_dim] (vsn_input)
grad_buf_base + goff_w_bfc[d], // dW [AH, fc_in_dim]
grad_buf_base + goff_b_bfc[d], // db [AH]
self.adv_h, // out_dim
fc_in_dim, // in_dim
self.backward_branch_dw(
bs,
d,
self.branch_workspace_ptrs[d],
d_adv_logits[d],
save_h_b[d],
scratch_d_h_b[d],
glu_gate_pre[d],
glu_value[d],
w_out,
w_ptrs,
grad_buf_base,
goff_w_bout[d],
goff_b_bout[d],
goff_w_bfc[d],
goff_b_bfc[d],
goff_w_gate[d],
goff_b_gate[d],
goff_kan_coeff[d],
goff_kan_resid[d],
fc_input,
fc_in_dim,
n_d,
b,
)?;
// ── Step 4: Gate path weight gradients ─────────────────────
// dW_gate += d_gate_pre^T @ vsn_input, db_gate += sum(d_gate_pre)
self.launch_dw_only(
self.branch_done_events[d].record(bs)
.map_err(|e| MLError::DeviceError(format!("bw branch_done_event {d} record: {e}")))?;
}
// ── Join: main stream waits for all 4 branches to finish Steps 1-4 ──
for d in 0..4 {
stream.wait(&self.branch_done_events[d])
.map_err(|e| MLError::DeviceError(format!("bw main stream wait branch {d}: {e}")))?;
}
// ── Step 5: dX accumulation runs sequentially on main stream ──
// d_h_s2 requires ordered beta accumulation (d==0 overwrites, d>0 accumulates).
// Concat branches (mag/ord/urg) write to separate buffers, so ordering
// among them does not matter — but they must complete before the trunk
// backward reads d_h_s2.
for d in 0..4 {
let w_fc = w_ptrs[w_bfc_idx[d]];
let w_gate = w_ptrs[w_gate_idx[d]];
self.backward_branch_dx(
stream,
scratch_d_glu_gate, // dY = d_gate_pre [B, AH] — f32
fc_input, // X [B, fc_in_dim] (same vsn_input)
grad_buf_base + goff_w_gate[d], // dW_gate [AH, fc_in_dim]
grad_buf_base + goff_b_gate[d], // db_gate [AH]
self.adv_h, // out_dim
fc_in_dim, // in_dim
d,
w_fc,
w_gate,
scratch_d_h_s2,
mag_concat_ptr,
d_mag_concat_ptr,
ord_concat_ptr,
d_ord_concat_ptr,
urg_concat_ptr,
d_urg_concat_ptr,
b,
)?;
// ── Step 5: Upstream gradient to VSN input ─────────────────
// d_vsn_input = W_bdf^T @ d_value + W_gate^T @ d_gate_pre
// Sum of both paths (value + gate) using beta accumulation.
if d == 1 && mag_concat_ptr != 0 {
// Magnitude: dX writes to d_mag_concat [B, SH2+3].
// Value path: beta=0 (overwrite)
self.launch_dx_only(
stream,
scratch_d_glu_value, // dY = d_value [B, AH]
w_fc, // W_bdf [AH, SH2+3]
d_mag_concat_ptr, // dX [B, SH2+3]
self.adv_h,
self.shared_h2 + 3,
b,
0.0_f32, // overwrite
)?;
// Gate path: beta=1 (accumulate on top of value path)
self.launch_dx_only(
stream,
scratch_d_glu_gate, // dY = d_gate_pre [B, AH]
w_gate, // W_gate [AH, SH2+3]
d_mag_concat_ptr, // dX [B, SH2+3] (accumulate)
self.adv_h,
self.shared_h2 + 3,
b,
1.0_f32, // accumulate
)?;
// Caller will call accumulate_d_h_s2_from_concat(d_mag_concat, d_h_s2, B, 1.0)
// after backward_full returns.
} else if d == 2 && ord_concat_ptr != 0 {
// Order: dX writes to d_ord_concat [B, SH2+3].
// Value path: beta=0 (overwrite)
self.launch_dx_only(
stream,
scratch_d_glu_value,
w_fc,
d_ord_concat_ptr,
self.adv_h,
self.shared_h2 + 3,
b,
0.0_f32,
)?;
// Gate path: beta=1 (accumulate)
self.launch_dx_only(
stream,
scratch_d_glu_gate,
w_gate,
d_ord_concat_ptr,
self.adv_h,
self.shared_h2 + 3,
b,
1.0_f32,
)?;
// Caller will call accumulate_d_h_s2_from_concat(d_ord_concat, d_h_s2, B, 1.0)
} else if d == 3 && urg_concat_ptr != 0 {
// Urgency: dX writes to d_urg_concat [B, SH2+3].
// Value path: beta=0 (overwrite)
self.launch_dx_only(
stream,
scratch_d_glu_value,
w_fc,
d_urg_concat_ptr,
self.adv_h,
self.shared_h2 + 3,
b,
0.0_f32,
)?;
// Gate path: beta=1 (accumulate)
self.launch_dx_only(
stream,
scratch_d_glu_gate,
w_gate,
d_urg_concat_ptr,
self.adv_h,
self.shared_h2 + 3,
b,
1.0_f32,
)?;
// Caller will call accumulate_d_h_s2_from_concat(d_urg_concat, d_h_s2, B, 1.0)
} else {
// Other branches: accumulate into d_h_s2 directly.
let beta_s2 = if d == 0 { 0.0_f32 } else { 1.0_f32 };
// Value path
self.launch_dx_only(
stream,
scratch_d_glu_value, // dY = d_value [B, AH]
w_fc, // W_bdf [AH, SH2]
scratch_d_h_s2, // dX [B, SH2]
self.adv_h,
self.shared_h2,
b,
beta_s2, // d==0: overwrite, d>0: accumulate
)?;
// Gate path: always beta=1 (accumulate on top of value path)
self.launch_dx_only(
stream,
scratch_d_glu_gate, // dY = d_gate_pre [B, AH]
w_gate, // W_gate [AH, SH2]
scratch_d_h_s2, // dX [B, SH2] (accumulate)
self.adv_h,
self.shared_h2,
b,
1.0_f32, // always accumulate
)?;
}
}
// ══════════════════════════════════════════════════════════════════