From 1f36d2fd244a672054fea6c41cabf0e789c19281 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sun, 19 Apr 2026 13:50:32 +0200 Subject: [PATCH] perf: q_denoise_backward cuBLAS pipeline + attn/IQL 2-phase bias grad q_denoise_backward (63.5% GPU time, 169ms/call): decomposed into cuBLAS forward replay + backward GEMMs. 8 cuBLAS GEMMs + elementwise kernels replace 1800-thread serial loop. Target: <1ms/call. attn_bias_grad_reduce + iql_bias_grad_reduce: converted from serial batch loops to 2-phase shared-memory reduction (same pattern as IQN). Co-Authored-By: Claude Opus 4.6 (1M context) --- .../attention_backward_kernel.cu | 46 +- .../src/cuda_pipeline/experience_kernels.cu | 150 ++++ crates/ml/src/cuda_pipeline/gpu_attention.rs | 51 +- .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 730 +++++++++++++++++- .../ml/src/cuda_pipeline/gpu_iql_trainer.rs | 76 +- .../ml/src/cuda_pipeline/iql_value_kernel.cu | 44 +- 6 files changed, 1044 insertions(+), 53 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/attention_backward_kernel.cu b/crates/ml/src/cuda_pipeline/attention_backward_kernel.cu index 3351e607d..1dc199298 100644 --- a/crates/ml/src/cuda_pipeline/attention_backward_kernel.cu +++ b/crates/ml/src/cuda_pipeline/attention_backward_kernel.cu @@ -310,23 +310,45 @@ void attn_bias_add( } /* ===================================================================== - * attn_bias_grad_reduce — Bias gradient reduction for attention projections + * attn_bias_grad_reduce_p1 — Phase 1: block-level shared-memory reduce * - * Reduces d_pre [B, out_dim] (row-major) over the batch dimension - * and scales by inv_batch to produce d_bias [out_dim]. - * - * Grid: (ceil(out_dim/blockDim.x)), Block: (256). + * Reduces d_pre [out_dim, B] (col-major, cuBLAS output) over B. + * Grid: (num_blocks, out_dim), Block: 256, shared: 256*sizeof(float). * ===================================================================== */ extern "C" __global__ -void attn_bias_grad_reduce( +void attn_bias_grad_reduce_p1( const float* __restrict__ d_pre, + float* __restrict__ partials, + int out_dim, int B +) { + extern __shared__ float sdata[]; + int j = blockIdx.y; + int tid = threadIdx.x; + int gid = blockIdx.x * blockDim.x + tid; + float val = (gid < B) ? d_pre[j + (long)gid * out_dim] : 0.0f; + sdata[tid] = val; + __syncthreads(); + for (int s = blockDim.x / 2; s > 0; s >>= 1) { + if (tid < s) sdata[tid] += sdata[tid + s]; + __syncthreads(); + } + if (tid == 0) partials[blockIdx.x * out_dim + j] = sdata[0]; +} + +/* ===================================================================== + * attn_bias_grad_reduce_p2 — Phase 2: warp-shuffle final reduce + scale + * + * Grid: ceil(out_dim/256), Block: 256. + * ===================================================================== */ +extern "C" __global__ +void attn_bias_grad_reduce_p2( + const float* __restrict__ partials, float* __restrict__ d_bias, - int out_dim, int B, float inv_batch + int out_dim, int num_blocks, float inv_batch ) { int j = blockIdx.x * blockDim.x + threadIdx.x; - if (j < out_dim) { - float sum = 0.0f; - for (int b = 0; b < B; b++) { sum += d_pre[j + b * out_dim]; } - d_bias[j] = sum * inv_batch; - } + if (j >= out_dim) return; + float sum = 0.0f; + for (int i = 0; i < num_blocks; i++) sum += partials[i * out_dim + j]; + d_bias[j] = sum * inv_batch; } diff --git a/crates/ml/src/cuda_pipeline/experience_kernels.cu b/crates/ml/src/cuda_pipeline/experience_kernels.cu index 8d623752f..01f1e5faf 100644 --- a/crates/ml/src/cuda_pipeline/experience_kernels.cu +++ b/crates/ml/src/cuda_pipeline/experience_kernels.cu @@ -3998,6 +3998,156 @@ extern "C" __global__ void liquid_tau_rk4_step( } } +/* ================================================================== */ +/* Denoise cuBLAS helper kernels */ +/* ================================================================== */ + +/** + * denoise_build_input — Build col-major input [H=24, B] from row-major Q [B, D=12] + * and row-major var_q [B, D=12]. First D rows = Q values, next D rows = sqrt(var_q) * schedule. + * + * Col-major [H, B] layout: element (h, b) at offset h + b * H. + * Grid: ceil(B/256), Block: 256. One thread per sample. + */ +extern "C" __global__ void denoise_build_input( + const float* __restrict__ q_values, + const float* __restrict__ var_q, + float* __restrict__ input_buf, + int B, int D, float schedule +) { + int b = blockIdx.x * blockDim.x + threadIdx.x; + if (b >= B) return; + int H = 2 * D; + for (int i = 0; i < D; i++) { + input_buf[i + b * H] = q_values[b * D + i]; + input_buf[(D + i) + b * H] = sqrtf(fmaxf(var_q[b * D + i], 0.0f)) * schedule; + } +} + +/** + * denoise_bias_silu — Add bias + SiLU activation. + * Operates on col-major [H, B] buffer from cuBLAS GEMM output. + * Saves pre-activation (with bias) in pre_act for backward. + * Writes SiLU output to h_out. + * + * Col-major [H, B]: element at linear index idx has row h = idx % H. + * Grid: ceil(H*B/256), Block: 256. + */ +extern "C" __global__ void denoise_bias_silu( + float* __restrict__ pre_act, + float* __restrict__ h_out, + const float* __restrict__ bias, + int H, int N +) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= N) return; + int h = idx % H; + float v = pre_act[idx] + bias[h]; + pre_act[idx] = v; + float sig = 1.0f / (1.0f + expf(-v)); + h_out[idx] = v * sig; +} + +/** + * denoise_silu_bwd — SiLU backward: d_pre = d_h * SiLU'(pre_act). + * SiLU'(x) = sigmoid(x) * (1 + x * (1 - sigmoid(x))). + * Grid: ceil(N/256), Block: 256. + */ +extern "C" __global__ void denoise_silu_bwd( + const float* __restrict__ d_h, + const float* __restrict__ pre_act, + float* __restrict__ d_pre, + int N +) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= N) return; + float x = pre_act[idx]; + float sig = 1.0f / (1.0f + expf(-x)); + d_pre[idx] = d_h[idx] * sig * (1.0f + x * (1.0f - sig)); +} + +/** + * denoise_loss_grad — MSE gradient with 0.1x residual scaling. + * Reads row-major [B, D], writes col-major [D, B]. + * d_res[d, b] = 0.1 * 2.0 * (Q_refined[b*D+d] - Q_target[b*D+d]) * inv_B. + * + * Col-major [D, B] output: element (d, b) at offset d + b * D. + * Grid: ceil(D*B/256), Block: 256. + */ +extern "C" __global__ void denoise_loss_grad( + const float* __restrict__ Q_refined, + const float* __restrict__ Q_target, + float* __restrict__ d_res, + int B, int D, float inv_B +) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + int total = D * B; + if (idx >= total) return; + int d = idx % D; + int b = idx / D; + d_res[idx] = 0.2f * (Q_refined[b * D + d] - Q_target[b * D + d]) * inv_B; +} + +/** + * denoise_skip_connection — q_out = q_in + scale * (mlp_out_colmaj + bias). + * q_in/q_out: row-major [B, D]. mlp_out: col-major [D, B]. + * Grid: ceil(B*D/256), Block: 256. + */ +extern "C" __global__ void denoise_skip_connection( + const float* __restrict__ q_in, + const float* __restrict__ mlp_out, + const float* __restrict__ bias, + float* __restrict__ q_out, + int B, int D, float scale +) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= B * D) return; + int b = idx / D; + int d = idx % D; + q_out[idx] = q_in[idx] + scale * (mlp_out[d + b * D] + bias[d]); +} + +/** + * denoise_bias_grad_p1 — Phase 1 block-level reduction for bias gradient. + * Reduces col-major [out_dim, B] over B dimension. + * Grid: (ceil(B/256), out_dim), Block: 256. + */ +extern "C" __global__ void denoise_bias_grad_p1( + const float* __restrict__ d_pre, + float* __restrict__ partials, + int out_dim, int B +) { + int j = blockIdx.y; + __shared__ float sdata[256]; + float sum = 0.0f; + for (int b = blockIdx.x * blockDim.x + threadIdx.x; b < B; b += gridDim.x * blockDim.x) { + sum += d_pre[j + b * out_dim]; + } + sdata[threadIdx.x] = sum; + __syncthreads(); + for (int s = 128; s > 0; s >>= 1) { + if (threadIdx.x < (unsigned)s) sdata[threadIdx.x] += sdata[threadIdx.x + s]; + __syncthreads(); + } + if (threadIdx.x == 0) partials[blockIdx.x * out_dim + j] = sdata[0]; +} + +/** + * denoise_bias_grad_p2 — Phase 2 final reduction across block partials. + * Grid: ceil(out_dim/256), Block: 256. + */ +extern "C" __global__ void denoise_bias_grad_p2( + const float* __restrict__ partials, + float* __restrict__ d_bias, + int out_dim, int num_blocks +) { + int j = blockIdx.x * blockDim.x + threadIdx.x; + if (j >= out_dim) return; + float sum = 0.0f; + for (int i = 0; i < num_blocks; i++) sum += partials[i * out_dim + j]; + d_bias[j] = sum; +} + /* ================================================================== */ /* Kernel: q_denoise_backward */ /* ================================================================== */ diff --git a/crates/ml/src/cuda_pipeline/gpu_attention.rs b/crates/ml/src/cuda_pipeline/gpu_attention.rs index 815feb885..4595a99da 100644 --- a/crates/ml/src/cuda_pipeline/gpu_attention.rs +++ b/crates/ml/src/cuda_pipeline/gpu_attention.rs @@ -130,7 +130,10 @@ pub struct GpuAttention { sdp_bwd_kernel: CudaFunction, layer_norm_fwd_kernel: CudaFunction, layer_norm_bwd_kernel: CudaFunction, - bias_grad_reduce_kernel: CudaFunction, + bias_grad_reduce_p1_kernel: CudaFunction, + bias_grad_reduce_p2_kernel: CudaFunction, + bias_grad_partials_buf: CudaSlice, + bias_grad_num_blocks: u32, grad_norm_phase1_kernel: CudaFunction, grad_norm_phase2_kernel: CudaFunction, adam_kernel: CudaFunction, @@ -229,8 +232,14 @@ impl GpuAttention { .map_err(|e| MLError::ModelError(format!("attn_sdp_bwd kernel load: {e}")))?; let layer_norm_bwd_kernel = bwd_module.load_function("attn_layer_norm_bwd") .map_err(|e| MLError::ModelError(format!("attn_layer_norm_bwd kernel load: {e}")))?; - let bias_grad_reduce_kernel = bwd_module.load_function("attn_bias_grad_reduce") - .map_err(|e| MLError::ModelError(format!("attn_bias_grad_reduce load: {e}")))?; + let bias_grad_reduce_p1_kernel = bwd_module.load_function("attn_bias_grad_reduce_p1") + .map_err(|e| MLError::ModelError(format!("attn_bias_grad_reduce_p1 load: {e}")))?; + let bias_grad_reduce_p2_kernel = bwd_module.load_function("attn_bias_grad_reduce_p2") + .map_err(|e| MLError::ModelError(format!("attn_bias_grad_reduce_p2 load: {e}")))?; + // Partials buffer for 2-phase bias grad reduce + let bias_grad_num_blocks = ((b + 255) / 256) as u32; + let bias_grad_partials_buf = stream.alloc_zeros::(bias_grad_num_blocks as usize * d) + .map_err(|e| MLError::ModelError(format!("attn bias_grad_partials alloc: {e}")))?; let grad_norm_phase1_kernel = bwd_module.load_function("attn_grad_norm_phase1") .map_err(|e| MLError::ModelError(format!("attn_grad_norm_phase1 load: {e}")))?; let grad_norm_phase2_kernel = bwd_module.load_function("attn_grad_norm_phase2") @@ -409,7 +418,10 @@ impl GpuAttention { sdp_bwd_kernel, layer_norm_fwd_kernel, layer_norm_bwd_kernel, - bias_grad_reduce_kernel, + bias_grad_reduce_p1_kernel, + bias_grad_reduce_p2_kernel, + bias_grad_partials_buf, + bias_grad_num_blocks, grad_norm_phase1_kernel, grad_norm_phase2_kernel, adam_kernel, @@ -958,25 +970,44 @@ impl GpuAttention { self.launch_bias_add_ex(out_ptr, bias_ptr, out_dim, batch, &self.stream) } - /// Launch bias gradient reduction kernel with an explicit stream. + /// Launch 2-phase bias gradient reduction with an explicit stream. fn launch_bias_grad_reduce_ex(&self, d_pre_ptr: u64, d_bias_ptr: u64, out_dim: usize, batch: usize, inv_batch: f32, stream: &CudaStream) -> Result<(), MLError> { let out_dim_i32 = out_dim as i32; let batch_i32 = batch as i32; - let blocks = ((out_dim + 255) / 256) as u32; + let num_blocks = self.bias_grad_num_blocks; + let partials_ptr = self.bias_grad_partials_buf.raw_ptr(); + // Phase 1: block-level reduce unsafe { stream - .launch_builder(&self.bias_grad_reduce_kernel) + .launch_builder(&self.bias_grad_reduce_p1_kernel) .arg(&d_pre_ptr) - .arg(&d_bias_ptr) + .arg(&partials_ptr) .arg(&out_dim_i32) .arg(&batch_i32) + .launch(LaunchConfig { + grid_dim: (num_blocks, out_dim as u32, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 256 * 4, + }) + .map_err(|e| MLError::ModelError(format!("attn_bias_grad_reduce_p1: {e}")))?; + } + // Phase 2: final reduce + scale + let num_blocks_i32 = num_blocks as i32; + let p2_blocks = ((out_dim + 255) / 256) as u32; + unsafe { + stream + .launch_builder(&self.bias_grad_reduce_p2_kernel) + .arg(&partials_ptr) + .arg(&d_bias_ptr) + .arg(&out_dim_i32) + .arg(&num_blocks_i32) .arg(&inv_batch) .launch(LaunchConfig { - grid_dim: (blocks, 1, 1), + grid_dim: (p2_blocks, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0, }) - .map_err(|e| MLError::ModelError(format!("attention bias_grad_reduce: {e}")))?; + .map_err(|e| MLError::ModelError(format!("attn_bias_grad_reduce_p2: {e}")))?; } Ok(()) } diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index cb816dcc4..24bb79e0b 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -1397,6 +1397,57 @@ pub struct GpuDqnTrainer { denoise_t_pinned: *mut i32, denoise_t_dev_ptr: u64, + // ── Denoise cuBLAS backward pipeline ────────────────────────────── + /// Col-major scratch [H=24, B] — denoiser MLP input for step 0. + denoise_input0_buf: CudaSlice, + /// Col-major scratch [H=24, B] — denoiser MLP input for step 1. + denoise_input1_buf: CudaSlice, + /// Col-major scratch [H=24, B] — pre-activation (with bias) for step 0. + denoise_pre_h0_buf: CudaSlice, + /// Col-major scratch [H=24, B] — pre-activation (with bias) for step 1. + denoise_pre_h1_buf: CudaSlice, + /// Col-major scratch [H=24, B] — SiLU output for step 0. + denoise_h0_buf: CudaSlice, + /// Col-major scratch [H=24, B] — SiLU output for step 1. + denoise_h1_buf: CudaSlice, + /// Row-major scratch [B, D=12] — Q values after step 0. + denoise_q_step0_buf: CudaSlice, + /// Col-major scratch [D=12, B] — loss gradient (upstream for backward). + denoise_d_res_buf: CudaSlice, + /// Col-major scratch [H=24, B] — d_h / d_pre scratch for backward. + denoise_d_h_buf: CudaSlice, + /// Scratch for bias gradient reduction phase 1 [num_blocks * max(H,D)]. + denoise_bias_grad_partials: CudaSlice, + /// GEMM desc: forward W @ input, C[H=24, B] = W[H,H] @ X[H,B]. TRANSA=T, TRANSB=N. + denoise_gemm_fwd_hh: Mamba2GemmDesc, + /// GEMM desc: forward W2 @ h, C[D=12, B] = W2[D,H] @ h[H,B]. TRANSA=T, TRANSB=N. + denoise_gemm_fwd_dh: Mamba2GemmDesc, + /// GEMM desc: backward dW1 = input @ d_pre^T, output row-major [H,H]. + /// Actually computes dW1^T[H,H]_colmaj = input[H,B] @ d_pre[H,B]^T. + /// TRANSA=N, TRANSB=T, m=H, n=H, k=B. + denoise_gemm_dw1: Mamba2GemmDesc, + /// GEMM desc: backward dW2 = h @ d_res^T, output row-major [D,H]. + /// Actually computes dW2^T[H,D]_colmaj = h[H,B] @ d_res[D,B]^T. + /// TRANSA=N, TRANSB=T, m=H, n=D, k=B. + denoise_gemm_dw2: Mamba2GemmDesc, + /// GEMM desc: backward d_h = W2^T @ d_res, C[H=24, B] = W2^T[H,D] @ d_res[D,B]. + /// W2 stored row-major [D,H] = col-major [H,D]. TRANSA=N (no transpose), TRANSB=N. + denoise_gemm_bwd_dh: Mamba2GemmDesc, + /// Kernel handle: denoise_build_input. + denoise_build_input_kernel: CudaFunction, + /// Kernel handle: denoise_bias_silu. + denoise_bias_silu_kernel: CudaFunction, + /// Kernel handle: denoise_silu_bwd. + denoise_silu_bwd_kernel: CudaFunction, + /// Kernel handle: denoise_loss_grad. + denoise_loss_grad_kernel: CudaFunction, + /// Kernel handle: denoise_skip_connection. + denoise_skip_conn_kernel: CudaFunction, + /// Kernel handle: denoise_bias_grad_p1. + denoise_bias_grad_p1_kernel: CudaFunction, + /// Kernel handle: denoise_bias_grad_p2. + denoise_bias_grad_p2_kernel: CudaFunction, + // ── xLSTM temporal context (mLSTM cell) ──────────────────────────── /// [528] 6 weight matrices × [11, 8]: W_k, W_v, W_q, W_i, W_f, W_o qlstm_weights: CudaSlice, @@ -5796,6 +5847,102 @@ impl GpuDqnTrainer { }; info!("GpuDqnTrainer: q_denoise_backward kernel loaded (1800 params, denoiser training)"); + // ── Denoise cuBLAS backward pipeline ────────────────────────────── + const DENOISE_D: usize = 12; + const DENOISE_H: usize = 24; + let denoise_input0_buf = alloc_f32(&stream, DENOISE_H * b, "denoise_input0")?; + let denoise_input1_buf = alloc_f32(&stream, DENOISE_H * b, "denoise_input1")?; + let denoise_pre_h0_buf = alloc_f32(&stream, DENOISE_H * b, "denoise_pre_h0")?; + let denoise_pre_h1_buf = alloc_f32(&stream, DENOISE_H * b, "denoise_pre_h1")?; + let denoise_h0_buf = alloc_f32(&stream, DENOISE_H * b, "denoise_h0")?; + let denoise_h1_buf = alloc_f32(&stream, DENOISE_H * b, "denoise_h1")?; + let denoise_q_step0_buf = alloc_f32(&stream, DENOISE_D * b, "denoise_q_step0")?; + let denoise_d_res_buf = alloc_f32(&stream, DENOISE_D * b, "denoise_d_res")?; + let denoise_d_h_buf = alloc_f32(&stream, DENOISE_H * b, "denoise_d_h")?; + // Bias grad partials: ceil(B/256) blocks × max(H=24, D=12) = ceil(B/256)*24 + let denoise_bias_num_blocks = (b + 255) / 256; + let denoise_bias_grad_partials = alloc_f32(&stream, denoise_bias_num_blocks * DENOISE_H, "denoise_bias_grad_partials")?; + + // Load denoise cuBLAS helper kernels from experience_kernels cubin + let cpbi_module_denoise_cublas = stream.context().load_cubin(EXPECTED_Q_CUBIN.to_vec()) + .map_err(|e| MLError::ModelError(format!("cpbi cubin (denoise_cublas): {e}")))?; + let denoise_build_input_kernel = cpbi_module_denoise_cublas.load_function("denoise_build_input") + .map_err(|e| MLError::ModelError(format!("denoise_build_input load: {e}")))?; + let denoise_bias_silu_kernel = cpbi_module_denoise_cublas.load_function("denoise_bias_silu") + .map_err(|e| MLError::ModelError(format!("denoise_bias_silu load: {e}")))?; + let denoise_silu_bwd_kernel = cpbi_module_denoise_cublas.load_function("denoise_silu_bwd") + .map_err(|e| MLError::ModelError(format!("denoise_silu_bwd load: {e}")))?; + let denoise_loss_grad_kernel = cpbi_module_denoise_cublas.load_function("denoise_loss_grad") + .map_err(|e| MLError::ModelError(format!("denoise_loss_grad load: {e}")))?; + let denoise_skip_conn_kernel = cpbi_module_denoise_cublas.load_function("denoise_skip_connection") + .map_err(|e| MLError::ModelError(format!("denoise_skip_connection load: {e}")))?; + let denoise_bias_grad_p1_kernel = cpbi_module_denoise_cublas.load_function("denoise_bias_grad_p1") + .map_err(|e| MLError::ModelError(format!("denoise_bias_grad_p1 load: {e}")))?; + let denoise_bias_grad_p2_kernel = cpbi_module_denoise_cublas.load_function("denoise_bias_grad_p2") + .map_err(|e| MLError::ModelError(format!("denoise_bias_grad_p2 load: {e}")))?; + + // cuBLAS GEMM descriptors for denoise MLP backward + let denoise_lt_handle = shared_cublas.lt_handle.0; + let denoise_lt_ws_size = shared_cublas.lt_workspace_size; + // W1 stored row-major [H=24, H=24]. cuBLAS sees it as col-major [H, H] = W1^T. + // Forward: pre_h[H,B] = W1 @ input[H,B]. With W1 row-major = W1^T col-major: + // cuBLAS: C[H,B] = W1_cm^T[H,H] @ X_cm[H,B] → TRANSA=T, TRANSB=N + let denoise_gemm_fwd_hh = create_mamba2_gemm_desc( + denoise_lt_handle, 1, 0, // TRANSA=T, TRANSB=N + DENOISE_H, b, DENOISE_H, // m=24, n=B, k=24 + DENOISE_H, DENOISE_H, DENOISE_H, // lda=24, ldb=24, ldc=24 + denoise_lt_ws_size, "denoise_fwd_hh", + )?; + + // W2 stored row-major [D=12, H=24]. cuBLAS sees it as col-major [H, D] = W2^T. + // Forward: out[D,B] = W2 @ h[H,B]. With W2 row-major = W2^T col-major: + // cuBLAS: C[D,B] = W2_cm^T[D,H] @ h_cm[H,B] → TRANSA=T, TRANSB=N + let denoise_gemm_fwd_dh = create_mamba2_gemm_desc( + denoise_lt_handle, 1, 0, // TRANSA=T, TRANSB=N + DENOISE_D, b, DENOISE_H, // m=12, n=B, k=24 + DENOISE_H, DENOISE_H, DENOISE_D, // lda=24 (stored [H,D]->ld=H), ldb=24, ldc=12 + denoise_lt_ws_size, "denoise_fwd_dh", + )?; + + // Backward dW1: we want row-major dW1[H, H]. + // dW1 = d_pre[H,B] @ input[H,B]^T in math (both col-major). + // For correct row-major output: compute dW1^T col-major = input @ d_pre^T. + // cuBLAS: C[H,H] = input[H,B] @ d_pre[H,B]^T -> TRANSA=N, TRANSB=T, m=H, n=H, k=B + let denoise_gemm_dw1 = create_mamba2_gemm_desc( + denoise_lt_handle, 0, 1, // TRANSA=N, TRANSB=T + DENOISE_H, DENOISE_H, b, // m=24, n=24, k=B + DENOISE_H, DENOISE_H, DENOISE_H, // lda=24, ldb=24, ldc=24 + denoise_lt_ws_size, "denoise_bwd_dw1", + )?; + + // Backward dW2: we want row-major dW2[D, H]. + // dW2 = d_res[D,B] @ h[H,B]^T in math. + // For correct row-major output: compute dW2^T col-major = h @ d_res^T. + // cuBLAS: C[H,D] = h[H,B] @ d_res[D,B]^T -> TRANSA=N, TRANSB=T, m=H, n=D, k=B + let denoise_gemm_dw2 = create_mamba2_gemm_desc( + denoise_lt_handle, 0, 1, // TRANSA=N, TRANSB=T + DENOISE_H, DENOISE_D, b, // m=24, n=12, k=B + DENOISE_H, DENOISE_D, DENOISE_H, // lda=24, ldb=12, ldc=24 + denoise_lt_ws_size, "denoise_bwd_dw2", + )?; + + // Backward d_h = W2^T @ d_res. W2 stored row-major [D,H] = col-major [H,D]. + // We need C[H,B] = W2^T[H,D] @ d_res[D,B]. + // W2 col-major has shape [H,D] with ld=H. No transpose needed. + // cuBLAS: C[H,B] = A[H,D] @ B[D,B] -> TRANSA=N, TRANSB=N, m=H, n=B, k=D + let denoise_gemm_bwd_dh = create_mamba2_gemm_desc( + denoise_lt_handle, 0, 0, // TRANSA=N, TRANSB=N + DENOISE_H, b, DENOISE_D, // m=24, n=B, k=12 + DENOISE_H, DENOISE_D, DENOISE_H, // lda=24 (W2 col-major ld), ldb=12, ldc=24 + denoise_lt_ws_size, "denoise_bwd_dh", + )?; + + info!( + DENOISE_D, DENOISE_H, b, + input_bytes = DENOISE_H * b * 4 * 2, + "GpuDqnTrainer: Denoise cuBLAS backward initialized (10 scratch bufs, 5 GEMM descs, 7 kernels)" + ); + // ── xLSTM temporal context (mLSTM cell, 528 params) ────────────────────── // 6 weight matrices × [11 input_dim, 8 head_dim] = 6 × 88 = 528 floats. // Xavier uniform: scale = sqrt(2 / (fan_in + fan_out)) = sqrt(2 / 19) ≈ 0.3244 @@ -6431,6 +6578,28 @@ impl GpuDqnTrainer { denoise_clip_buf, denoise_t_pinned, denoise_t_dev_ptr, + denoise_input0_buf, + denoise_input1_buf, + denoise_pre_h0_buf, + denoise_pre_h1_buf, + denoise_h0_buf, + denoise_h1_buf, + denoise_q_step0_buf, + denoise_d_res_buf, + denoise_d_h_buf, + denoise_bias_grad_partials, + denoise_gemm_fwd_hh, + denoise_gemm_fwd_dh, + denoise_gemm_dw1, + denoise_gemm_dw2, + denoise_gemm_bwd_dh, + denoise_build_input_kernel, + denoise_bias_silu_kernel, + denoise_silu_bwd_kernel, + denoise_loss_grad_kernel, + denoise_skip_conn_kernel, + denoise_bias_grad_p1_kernel, + denoise_bias_grad_p2_kernel, qlstm_weights, qlstm_c, qlstm_n, @@ -8245,6 +8414,565 @@ impl GpuDqnTrainer { Ok(()) } + /// cuBLAS-accelerated backward through the 2-step diffusion denoiser MLP. + /// + /// Replaces the old `q_denoise_backward` kernel (1800 threads × B=8192 inner loops) + /// with batched cuBLAS GEMMs + lightweight elementwise kernels. + /// + /// Pipeline: + /// Forward replay (save activations): + /// 1. Build input0 [H=24, B] from Q_input + var_q + /// 2. cuBLAS: pre_h0 = W1_0 @ input0 + /// 3. Bias + SiLU → h0, save pre_h0 + /// 4. cuBLAS: out0 = W2_0 @ h0 + /// 5. Skip connection → q_step0 = Q_input + 0.1*(out0 + b2_0) + /// 6. Build input1 [H, B] from q_step0 + var_q + /// 7. cuBLAS: pre_h1 = W1_1 @ input1 + /// 8. Bias + SiLU → h1, save pre_h1 + /// 9. cuBLAS: out1 = W2_1 @ h1 (out1 col-major scratch in d_res_buf) + /// Backward: + /// 10. denoise_loss_grad → d_res [D, B] + /// 11-16. Step 1 backward: dW2_1, db2_1, d_h1, d_pre1, dW1_1, db1_1 + /// 17-22. Step 0 backward: dW2_0, db2_0, d_h0, d_pre0, dW1_0, db1_0 + pub(crate) fn launch_q_denoise_backward_cublas(&mut self, batch_size: usize) -> Result<(), MLError> { + const D: usize = 12; + const H: usize = 24; + const STEP_PARAMS: usize = 900; // W1[576] + b1[24] + W2[288] + b2[12] + let b = batch_size; + + let lt_handle = self.shared_cublas.lt_handle.0; + let lt_ws_ptr = self.shared_cublas.lt_workspace_ptr; + let lt_ws_size = self.shared_cublas.lt_workspace_size; + let cu_stream = self.stream.cu_stream() as cublaslt_sys::cudaStream_t; + let alpha: f32 = 1.0; + let beta: f32 = 0.0; + + let params_ptr = self.denoise_params.raw_ptr(); + let d_params_ptr = self.denoise_grad.raw_ptr(); + let q_input_ptr = self.denoise_q_input_buf.raw_ptr(); + let var_ptr = self.q_var_buf_trainer.raw_ptr(); + let q_refined_ptr = self.q_coord_buf.raw_ptr(); + let q_target_ptr = self.denoise_target_q_buf.raw_ptr(); + + // Weight pointers for step 0 + let w1_0_ptr = params_ptr; + let b1_0_ptr = params_ptr + (576 * 4) as u64; + let w2_0_ptr = params_ptr + (600 * 4) as u64; + let b2_0_ptr = params_ptr + (888 * 4) as u64; + // Weight pointers for step 1 + let w1_1_ptr = params_ptr + (STEP_PARAMS * 4) as u64; + let b1_1_ptr = params_ptr + ((STEP_PARAMS + 576) * 4) as u64; + let w2_1_ptr = params_ptr + ((STEP_PARAMS + 600) * 4) as u64; + let _b2_1_ptr = params_ptr + ((STEP_PARAMS + 888) * 4) as u64; + + // Gradient pointers for step 0 + let dw1_0_ptr = d_params_ptr; + let db1_0_ptr = d_params_ptr + (576 * 4) as u64; + let dw2_0_ptr = d_params_ptr + (600 * 4) as u64; + let db2_0_ptr = d_params_ptr + (888 * 4) as u64; + // Gradient pointers for step 1 + let dw1_1_ptr = d_params_ptr + (STEP_PARAMS * 4) as u64; + let db1_1_ptr = d_params_ptr + ((STEP_PARAMS + 576) * 4) as u64; + let dw2_1_ptr = d_params_ptr + ((STEP_PARAMS + 600) * 4) as u64; + let db2_1_ptr = d_params_ptr + ((STEP_PARAMS + 888) * 4) as u64; + + let blocks_b = ((b as u32 + 255) / 256).max(1); + let bias_num_blocks = ((b + 255) / 256) as i32; + + // ═══════════════════════════════════════════════════════ + // FORWARD REPLAY — save activations for backward + // ═══════════════════════════════════════════════════════ + + // Step 0: schedule = 1.0 - 1/2 = 0.5 + let schedule_0: f32 = 0.5; + let d_val = D as i32; + let input0_ptr = self.denoise_input0_buf.raw_ptr(); + let pre_h0_ptr = self.denoise_pre_h0_buf.raw_ptr(); + let h0_ptr = self.denoise_h0_buf.raw_ptr(); + + // 1. Build input0 [H=24, B] from Q_input + var_q + unsafe { + self.stream.launch_builder(&self.denoise_build_input_kernel) + .arg(&q_input_ptr) + .arg(&var_ptr) + .arg(&input0_ptr) + .arg(&(b as i32)) + .arg(&d_val) + .arg(&schedule_0) + .launch(LaunchConfig { + grid_dim: (blocks_b, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("denoise_build_input step0: {e}")))?; + } + + // 2. cuBLAS: pre_h0 = W1_0 @ input0. W1_0 row-major [H,H] → col-major [H,H]^T. + // C[H,B] = W1^T_cm @ input_cm → TRANSA=T + unsafe { + cublaslt_sys::cublasLtMatmul( + lt_handle, + self.denoise_gemm_fwd_hh.matmul_desc, + &alpha as *const f32 as *const std::ffi::c_void, + w1_0_ptr as *const std::ffi::c_void, + self.denoise_gemm_fwd_hh.a_layout, + input0_ptr as *const std::ffi::c_void, + self.denoise_gemm_fwd_hh.b_layout, + &beta as *const f32 as *const std::ffi::c_void, + pre_h0_ptr as *mut std::ffi::c_void, + self.denoise_gemm_fwd_hh.c_layout, + pre_h0_ptr as *mut std::ffi::c_void, + self.denoise_gemm_fwd_hh.d_layout, + &self.denoise_gemm_fwd_hh.algo as *const _, + lt_ws_ptr as *mut std::ffi::c_void, + lt_ws_size, + cu_stream, + ); + } + + // 3. Bias + SiLU: pre_h0 += b1_0, h0 = SiLU(pre_h0). Saves pre_h0 with bias. + let hb_total = (H * b) as i32; + let h_dim = H as i32; + unsafe { + self.stream.launch_builder(&self.denoise_bias_silu_kernel) + .arg(&pre_h0_ptr) + .arg(&h0_ptr) + .arg(&b1_0_ptr) + .arg(&h_dim) + .arg(&hb_total) + .launch(LaunchConfig { + grid_dim: (((H * b) as u32 + 255) / 256, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("denoise_bias_silu step0: {e}")))?; + } + + // 4. cuBLAS: out0 = W2_0 @ h0. W2_0 row-major [D,H] → TRANSA=T. + // Use denoise_d_res_buf as scratch for out0 [D, B] col-major. + let out0_ptr = self.denoise_d_res_buf.raw_ptr(); // reuse d_res_buf for fwd scratch + unsafe { + cublaslt_sys::cublasLtMatmul( + lt_handle, + self.denoise_gemm_fwd_dh.matmul_desc, + &alpha as *const f32 as *const std::ffi::c_void, + w2_0_ptr as *const std::ffi::c_void, + self.denoise_gemm_fwd_dh.a_layout, + h0_ptr as *const std::ffi::c_void, + self.denoise_gemm_fwd_dh.b_layout, + &beta as *const f32 as *const std::ffi::c_void, + out0_ptr as *mut std::ffi::c_void, + self.denoise_gemm_fwd_dh.c_layout, + out0_ptr as *mut std::ffi::c_void, + self.denoise_gemm_fwd_dh.d_layout, + &self.denoise_gemm_fwd_dh.algo as *const _, + lt_ws_ptr as *mut std::ffi::c_void, + lt_ws_size, + cu_stream, + ); + } + + // 5. Skip connection: q_step0 = Q_input + 0.1 * (out0 + b2_0) + let q_step0_ptr = self.denoise_q_step0_buf.raw_ptr(); + let scale_01: f32 = 0.1; + unsafe { + self.stream.launch_builder(&self.denoise_skip_conn_kernel) + .arg(&q_input_ptr) + .arg(&out0_ptr) + .arg(&b2_0_ptr) + .arg(&q_step0_ptr) + .arg(&(b as i32)) + .arg(&d_val) + .arg(&scale_01) + .launch(LaunchConfig { + grid_dim: (((D * b) as u32 + 255) / 256, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("denoise_skip_connection step0: {e}")))?; + } + + // Step 1: schedule = 1.0 - 2/2 = 0.0 + let schedule_1: f32 = 0.0; + let input1_ptr = self.denoise_input1_buf.raw_ptr(); + let pre_h1_ptr = self.denoise_pre_h1_buf.raw_ptr(); + let h1_ptr = self.denoise_h1_buf.raw_ptr(); + + // 6. Build input1 from q_step0 + var_q + unsafe { + self.stream.launch_builder(&self.denoise_build_input_kernel) + .arg(&q_step0_ptr) + .arg(&var_ptr) + .arg(&input1_ptr) + .arg(&(b as i32)) + .arg(&d_val) + .arg(&schedule_1) + .launch(LaunchConfig { + grid_dim: (blocks_b, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("denoise_build_input step1: {e}")))?; + } + + // 7. cuBLAS: pre_h1 = W1_1 @ input1 + unsafe { + cublaslt_sys::cublasLtMatmul( + lt_handle, + self.denoise_gemm_fwd_hh.matmul_desc, + &alpha as *const f32 as *const std::ffi::c_void, + w1_1_ptr as *const std::ffi::c_void, + self.denoise_gemm_fwd_hh.a_layout, + input1_ptr as *const std::ffi::c_void, + self.denoise_gemm_fwd_hh.b_layout, + &beta as *const f32 as *const std::ffi::c_void, + pre_h1_ptr as *mut std::ffi::c_void, + self.denoise_gemm_fwd_hh.c_layout, + pre_h1_ptr as *mut std::ffi::c_void, + self.denoise_gemm_fwd_hh.d_layout, + &self.denoise_gemm_fwd_hh.algo as *const _, + lt_ws_ptr as *mut std::ffi::c_void, + lt_ws_size, + cu_stream, + ); + } + + // 8. Bias + SiLU for step 1 + unsafe { + self.stream.launch_builder(&self.denoise_bias_silu_kernel) + .arg(&pre_h1_ptr) + .arg(&h1_ptr) + .arg(&b1_1_ptr) + .arg(&h_dim) + .arg(&hb_total) + .launch(LaunchConfig { + grid_dim: (((H * b) as u32 + 255) / 256, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("denoise_bias_silu step1: {e}")))?; + } + + // 9. cuBLAS: out1 = W2_1 @ h1 (we don't need out1 explicitly — Q_refined is already + // computed in the forward pass. But we do need h1 saved for backward.) + // Skip materializing out1 — it's redundant. Q_refined = q_step0 + 0.1*(out1+b2_1) + // is already in q_coord_buf from the forward pass. + + // ═══════════════════════════════════════════════════════ + // BACKWARD + // ═══════════════════════════════════════════════════════ + + // 10. Loss gradient: d_res = 0.1 * 2 * (Q_refined - Q_target) / B + let inv_b: f32 = 1.0 / b as f32; + let d_res_ptr = self.denoise_d_res_buf.raw_ptr(); + unsafe { + self.stream.launch_builder(&self.denoise_loss_grad_kernel) + .arg(&q_refined_ptr) + .arg(&q_target_ptr) + .arg(&d_res_ptr) + .arg(&(b as i32)) + .arg(&d_val) + .arg(&inv_b) + .launch(LaunchConfig { + grid_dim: (((D * b) as u32 + 255) / 256, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("denoise_loss_grad: {e}")))?; + } + + let partials_ptr = self.denoise_bias_grad_partials.raw_ptr(); + let d_h_ptr = self.denoise_d_h_buf.raw_ptr(); + + // ── Step 1 backward ────────────────────────────────────────── + + // 11. dW2_1 = h1 @ d_res^T → output at dw2_1_ptr (row-major [D,H]) + // cuBLAS: C[H,D]_colmaj = h1[H,B] @ d_res[D,B]^T. Col-major [H,D] = row-major [D,H]. + unsafe { + cublaslt_sys::cublasLtMatmul( + lt_handle, + self.denoise_gemm_dw2.matmul_desc, + &alpha as *const f32 as *const std::ffi::c_void, + h1_ptr as *const std::ffi::c_void, + self.denoise_gemm_dw2.a_layout, + d_res_ptr as *const std::ffi::c_void, + self.denoise_gemm_dw2.b_layout, + &beta as *const f32 as *const std::ffi::c_void, + dw2_1_ptr as *mut std::ffi::c_void, + self.denoise_gemm_dw2.c_layout, + dw2_1_ptr as *mut std::ffi::c_void, + self.denoise_gemm_dw2.d_layout, + &self.denoise_gemm_dw2.algo as *const _, + lt_ws_ptr as *mut std::ffi::c_void, + lt_ws_size, + cu_stream, + ); + } + + // 12. db2_1 = sum_b(d_res[d, b]) — bias grad reduce for D=12 + let d_dim = D as i32; + unsafe { + self.stream.launch_builder(&self.denoise_bias_grad_p1_kernel) + .arg(&d_res_ptr) + .arg(&partials_ptr) + .arg(&d_dim) + .arg(&(b as i32)) + .launch(LaunchConfig { + grid_dim: (bias_num_blocks as u32, D as u32, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("denoise_bias_grad_p1 db2_1: {e}")))?; + } + unsafe { + self.stream.launch_builder(&self.denoise_bias_grad_p2_kernel) + .arg(&partials_ptr) + .arg(&db2_1_ptr) + .arg(&d_dim) + .arg(&bias_num_blocks) + .launch(LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("denoise_bias_grad_p2 db2_1: {e}")))?; + } + + // 13. d_h1 = W2_1^T @ d_res. W2_1 row-major [D,H] = col-major [H,D]. + // cuBLAS: C[H,B] = W2_cm[H,D] @ d_res[D,B] → TRANSA=N, TRANSB=N + unsafe { + cublaslt_sys::cublasLtMatmul( + lt_handle, + self.denoise_gemm_bwd_dh.matmul_desc, + &alpha as *const f32 as *const std::ffi::c_void, + w2_1_ptr as *const std::ffi::c_void, + self.denoise_gemm_bwd_dh.a_layout, + d_res_ptr as *const std::ffi::c_void, + self.denoise_gemm_bwd_dh.b_layout, + &beta as *const f32 as *const std::ffi::c_void, + d_h_ptr as *mut std::ffi::c_void, + self.denoise_gemm_bwd_dh.c_layout, + d_h_ptr as *mut std::ffi::c_void, + self.denoise_gemm_bwd_dh.d_layout, + &self.denoise_gemm_bwd_dh.algo as *const _, + lt_ws_ptr as *mut std::ffi::c_void, + lt_ws_size, + cu_stream, + ); + } + + // 14. SiLU backward: d_pre1 = d_h1 * SiLU'(pre_h1). Write into d_h_buf (reuse). + // Actually we need d_pre in a separate buffer since d_h is consumed. + // We can write d_pre over pre_h1 since pre_h1 won't be needed after this. + unsafe { + self.stream.launch_builder(&self.denoise_silu_bwd_kernel) + .arg(&d_h_ptr) + .arg(&pre_h1_ptr) + .arg(&d_h_ptr) // output d_pre1 into d_h_buf (d_h no longer needed) + .arg(&hb_total) + .launch(LaunchConfig { + grid_dim: (((H * b) as u32 + 255) / 256, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("denoise_silu_bwd step1: {e}")))?; + } + // d_pre1 is now in d_h_buf + + // 15. dW1_1 = input1 @ d_pre1^T → output at dw1_1_ptr (row-major [H,H]) + // cuBLAS: C[H,H]_colmaj = input1[H,B] @ d_pre1[H,B]^T + unsafe { + cublaslt_sys::cublasLtMatmul( + lt_handle, + self.denoise_gemm_dw1.matmul_desc, + &alpha as *const f32 as *const std::ffi::c_void, + input1_ptr as *const std::ffi::c_void, + self.denoise_gemm_dw1.a_layout, + d_h_ptr as *const std::ffi::c_void, // d_pre1 + self.denoise_gemm_dw1.b_layout, + &beta as *const f32 as *const std::ffi::c_void, + dw1_1_ptr as *mut std::ffi::c_void, + self.denoise_gemm_dw1.c_layout, + dw1_1_ptr as *mut std::ffi::c_void, + self.denoise_gemm_dw1.d_layout, + &self.denoise_gemm_dw1.algo as *const _, + lt_ws_ptr as *mut std::ffi::c_void, + lt_ws_size, + cu_stream, + ); + } + + // 16. db1_1 = sum_b(d_pre1) — bias grad reduce for H=24 + unsafe { + self.stream.launch_builder(&self.denoise_bias_grad_p1_kernel) + .arg(&d_h_ptr) // d_pre1 + .arg(&partials_ptr) + .arg(&h_dim) + .arg(&(b as i32)) + .launch(LaunchConfig { + grid_dim: (bias_num_blocks as u32, H as u32, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("denoise_bias_grad_p1 db1_1: {e}")))?; + } + unsafe { + self.stream.launch_builder(&self.denoise_bias_grad_p2_kernel) + .arg(&partials_ptr) + .arg(&db1_1_ptr) + .arg(&h_dim) + .arg(&bias_num_blocks) + .launch(LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("denoise_bias_grad_p2 db1_1: {e}")))?; + } + + // ── Step 0 backward ────────────────────────────────────────── + // d_res is the same (skip connection passes gradient through unchanged) + + // 17. dW2_0 = h0 @ d_res^T + unsafe { + cublaslt_sys::cublasLtMatmul( + lt_handle, + self.denoise_gemm_dw2.matmul_desc, + &alpha as *const f32 as *const std::ffi::c_void, + h0_ptr as *const std::ffi::c_void, + self.denoise_gemm_dw2.a_layout, + d_res_ptr as *const std::ffi::c_void, + self.denoise_gemm_dw2.b_layout, + &beta as *const f32 as *const std::ffi::c_void, + dw2_0_ptr as *mut std::ffi::c_void, + self.denoise_gemm_dw2.c_layout, + dw2_0_ptr as *mut std::ffi::c_void, + self.denoise_gemm_dw2.d_layout, + &self.denoise_gemm_dw2.algo as *const _, + lt_ws_ptr as *mut std::ffi::c_void, + lt_ws_size, + cu_stream, + ); + } + + // 18. db2_0 = sum_b(d_res) + unsafe { + self.stream.launch_builder(&self.denoise_bias_grad_p1_kernel) + .arg(&d_res_ptr) + .arg(&partials_ptr) + .arg(&d_dim) + .arg(&(b as i32)) + .launch(LaunchConfig { + grid_dim: (bias_num_blocks as u32, D as u32, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("denoise_bias_grad_p1 db2_0: {e}")))?; + } + unsafe { + self.stream.launch_builder(&self.denoise_bias_grad_p2_kernel) + .arg(&partials_ptr) + .arg(&db2_0_ptr) + .arg(&d_dim) + .arg(&bias_num_blocks) + .launch(LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("denoise_bias_grad_p2 db2_0: {e}")))?; + } + + // 19. d_h0 = W2_0^T @ d_res + unsafe { + cublaslt_sys::cublasLtMatmul( + lt_handle, + self.denoise_gemm_bwd_dh.matmul_desc, + &alpha as *const f32 as *const std::ffi::c_void, + w2_0_ptr as *const std::ffi::c_void, + self.denoise_gemm_bwd_dh.a_layout, + d_res_ptr as *const std::ffi::c_void, + self.denoise_gemm_bwd_dh.b_layout, + &beta as *const f32 as *const std::ffi::c_void, + d_h_ptr as *mut std::ffi::c_void, + self.denoise_gemm_bwd_dh.c_layout, + d_h_ptr as *mut std::ffi::c_void, + self.denoise_gemm_bwd_dh.d_layout, + &self.denoise_gemm_bwd_dh.algo as *const _, + lt_ws_ptr as *mut std::ffi::c_void, + lt_ws_size, + cu_stream, + ); + } + + // 20. SiLU backward step 0: d_pre0 = d_h0 * SiLU'(pre_h0) + unsafe { + self.stream.launch_builder(&self.denoise_silu_bwd_kernel) + .arg(&d_h_ptr) + .arg(&pre_h0_ptr) + .arg(&d_h_ptr) + .arg(&hb_total) + .launch(LaunchConfig { + grid_dim: (((H * b) as u32 + 255) / 256, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("denoise_silu_bwd step0: {e}")))?; + } + + // 21. dW1_0 = input0 @ d_pre0^T + unsafe { + cublaslt_sys::cublasLtMatmul( + lt_handle, + self.denoise_gemm_dw1.matmul_desc, + &alpha as *const f32 as *const std::ffi::c_void, + input0_ptr as *const std::ffi::c_void, + self.denoise_gemm_dw1.a_layout, + d_h_ptr as *const std::ffi::c_void, // d_pre0 + self.denoise_gemm_dw1.b_layout, + &beta as *const f32 as *const std::ffi::c_void, + dw1_0_ptr as *mut std::ffi::c_void, + self.denoise_gemm_dw1.c_layout, + dw1_0_ptr as *mut std::ffi::c_void, + self.denoise_gemm_dw1.d_layout, + &self.denoise_gemm_dw1.algo as *const _, + lt_ws_ptr as *mut std::ffi::c_void, + lt_ws_size, + cu_stream, + ); + } + + // 22. db1_0 = sum_b(d_pre0) + unsafe { + self.stream.launch_builder(&self.denoise_bias_grad_p1_kernel) + .arg(&d_h_ptr) // d_pre0 + .arg(&partials_ptr) + .arg(&h_dim) + .arg(&(b as i32)) + .launch(LaunchConfig { + grid_dim: (bias_num_blocks as u32, H as u32, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("denoise_bias_grad_p1 db1_0: {e}")))?; + } + unsafe { + self.stream.launch_builder(&self.denoise_bias_grad_p2_kernel) + .arg(&partials_ptr) + .arg(&db1_0_ptr) + .arg(&h_dim) + .arg(&bias_num_blocks) + .launch(LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("denoise_bias_grad_p2 db1_0: {e}")))?; + } + + Ok(()) + } + /// Adam optimizer step for denoise_params. Same pattern as step_selectivity_adam. pub(crate) fn step_denoise_adam(&mut self) -> Result<(), MLError> { self.denoise_adam_step += 1; @@ -8710,7 +9438,7 @@ impl GpuDqnTrainer { // Denoise: target + backward + adam self.compute_denoise_target_q(batch_size)?; - self.launch_q_denoise_backward(batch_size)?; + self.launch_q_denoise_backward_cublas(batch_size)?; self.step_denoise_adam()?; // Risk SGD (uses scale_f32_kernel) diff --git a/crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs index 8b452a934..34a990f68 100644 --- a/crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs @@ -184,7 +184,10 @@ pub struct GpuIqlTrainer { silu_bwd_kernel: CudaFunction, expectile_loss_kernel: CudaFunction, bias_add_kernel: CudaFunction, - bias_grad_reduce_kernel: CudaFunction, + bias_grad_reduce_p1_kernel: CudaFunction, + bias_grad_reduce_p2_kernel: CudaFunction, + bias_grad_partials_buf: CudaSlice, + bias_grad_num_blocks: u32, loss_reduce_kernel: CudaFunction, grad_norm_phase1_kernel: CudaFunction, grad_norm_phase2_kernel: CudaFunction, @@ -312,7 +315,11 @@ impl GpuIqlTrainer { let silu_bwd_kernel = load("iql_silu_bwd")?; let expectile_loss_kernel = load("iql_expectile_loss")?; let bias_add_kernel = load("iql_bias_add")?; - let bias_grad_reduce_kernel = load("iql_bias_grad_reduce")?; + let bias_grad_reduce_p1_kernel = load("iql_bias_grad_reduce_p1")?; + let bias_grad_reduce_p2_kernel = load("iql_bias_grad_reduce_p2")?; + let bias_grad_num_blocks = ((b + 255) / 256) as u32; + let bias_grad_partials_buf = stream.alloc_zeros::(bias_grad_num_blocks as usize * h) + .map_err(|e| MLError::ModelError(format!("iql bias_grad_partials alloc: {e}")))?; // Existing kernels (unchanged) let loss_reduce_kernel = load("iql_loss_reduce")?; @@ -431,7 +438,10 @@ impl GpuIqlTrainer { silu_bwd_kernel, expectile_loss_kernel, bias_add_kernel, - bias_grad_reduce_kernel, + bias_grad_reduce_p1_kernel, + bias_grad_reduce_p2_kernel, + bias_grad_partials_buf, + bias_grad_num_blocks, loss_reduce_kernel, grad_norm_phase1_kernel, grad_norm_phase2_kernel, @@ -705,17 +715,29 @@ impl GpuIqlTrainer { ); } - // db3: sum(dv) / B -- scalar bias gradient via bias_grad_reduce + // db3: sum(dv) / B -- scalar bias gradient via 2-phase reduce let db3_ptr = self.grad_buf.raw_ptr() + (b3_off * f32_sz) as u64; + let num_blocks = self.bias_grad_num_blocks; + let partials_ptr = self.bias_grad_partials_buf.raw_ptr(); unsafe { - self.stream.launch_builder(&self.bias_grad_reduce_kernel) + self.stream.launch_builder(&self.bias_grad_reduce_p1_kernel) .arg(&self.dv_buf) - .arg(&db3_ptr) - .arg(&one_i32) // out_dim = 1 (scalar bias) + .arg(&partials_ptr) + .arg(&one_i32) .arg(&batch_size_i32) + .launch(LaunchConfig { grid_dim: (num_blocks, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 256 * 4 }) + .map_err(|e| MLError::ModelError(format!("IQL bias_grad_reduce_p1 b3: {e}")))?; + } + let num_blocks_i32 = num_blocks as i32; + unsafe { + self.stream.launch_builder(&self.bias_grad_reduce_p2_kernel) + .arg(&partials_ptr) + .arg(&db3_ptr) + .arg(&one_i32) + .arg(&num_blocks_i32) .arg(&inv_batch) .launch(LaunchConfig { grid_dim: (1, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 }) - .map_err(|e| MLError::ModelError(format!("IQL bias_grad_reduce b3: {e}")))?; + .map_err(|e| MLError::ModelError(format!("IQL bias_grad_reduce_p2 b3: {e}")))?; } // dh2[H,B] = W3[H,1] @ dv[1,B] @@ -774,18 +796,27 @@ impl GpuIqlTrainer { ); } - // db2: sum columns of dh2_pre -> db2[H] + // db2: sum columns of dh2_pre -> db2[H] via 2-phase reduce let db2_ptr = self.grad_buf.raw_ptr() + (b2_off * f32_sz) as u64; let bias_blocks_h = ((h + 255) / 256) as u32; unsafe { - self.stream.launch_builder(&self.bias_grad_reduce_kernel) + self.stream.launch_builder(&self.bias_grad_reduce_p1_kernel) .arg(&self.dh2_pre_buf) - .arg(&db2_ptr) - .arg(&h_i32) // out_dim = H + .arg(&partials_ptr) + .arg(&h_i32) .arg(&batch_size_i32) + .launch(LaunchConfig { grid_dim: (num_blocks, h as u32, 1), block_dim: (256, 1, 1), shared_mem_bytes: 256 * 4 }) + .map_err(|e| MLError::ModelError(format!("IQL bias_grad_reduce_p1 b2: {e}")))?; + } + unsafe { + self.stream.launch_builder(&self.bias_grad_reduce_p2_kernel) + .arg(&partials_ptr) + .arg(&db2_ptr) + .arg(&h_i32) + .arg(&num_blocks_i32) .arg(&inv_batch) .launch(LaunchConfig { grid_dim: (bias_blocks_h, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 }) - .map_err(|e| MLError::ModelError(format!("IQL bias_grad_reduce b2: {e}")))?; + .map_err(|e| MLError::ModelError(format!("IQL bias_grad_reduce_p2 b2: {e}")))?; } // dh1[H,B] = W2[H,H] @ dh2_pre[H,B] @@ -844,17 +875,26 @@ impl GpuIqlTrainer { ); } - // db1: sum columns of dh1_pre -> db1[H] + // db1: sum columns of dh1_pre -> db1[H] via 2-phase reduce let db1_ptr = self.grad_buf.raw_ptr() + (b1_off * f32_sz) as u64; unsafe { - self.stream.launch_builder(&self.bias_grad_reduce_kernel) + self.stream.launch_builder(&self.bias_grad_reduce_p1_kernel) .arg(&self.dh1_pre_buf) - .arg(&db1_ptr) - .arg(&h_i32) // out_dim = H + .arg(&partials_ptr) + .arg(&h_i32) .arg(&batch_size_i32) + .launch(LaunchConfig { grid_dim: (num_blocks, h as u32, 1), block_dim: (256, 1, 1), shared_mem_bytes: 256 * 4 }) + .map_err(|e| MLError::ModelError(format!("IQL bias_grad_reduce_p1 b1: {e}")))?; + } + unsafe { + self.stream.launch_builder(&self.bias_grad_reduce_p2_kernel) + .arg(&partials_ptr) + .arg(&db1_ptr) + .arg(&h_i32) + .arg(&num_blocks_i32) .arg(&inv_batch) .launch(LaunchConfig { grid_dim: (bias_blocks_h, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 }) - .map_err(|e| MLError::ModelError(format!("IQL bias_grad_reduce b1: {e}")))?; + .map_err(|e| MLError::ModelError(format!("IQL bias_grad_reduce_p2 b1: {e}")))?; } // ── Loss reduce + grad norm + Adam (unchanged) ── diff --git a/crates/ml/src/cuda_pipeline/iql_value_kernel.cu b/crates/ml/src/cuda_pipeline/iql_value_kernel.cu index de2159171..a0ce74e5b 100644 --- a/crates/ml/src/cuda_pipeline/iql_value_kernel.cu +++ b/crates/ml/src/cuda_pipeline/iql_value_kernel.cu @@ -893,24 +893,44 @@ void iql_bias_add( } } -/* iql_bias_grad_reduce — Deterministic bias gradient reduction - * d_bias[j] = inv_batch * sum_{b=0}^{B-1} d_pre[j + b*out_dim] - * Launch: grid=(ceil(out_dim/256),1,1), block=(256,1,1) +/* iql_bias_grad_reduce_p1 — Phase 1: block-level shared-memory reduce + * Grid: (num_blocks, out_dim), Block: 256, shared: 256*sizeof(float) */ extern "C" __global__ -void iql_bias_grad_reduce( +void iql_bias_grad_reduce_p1( const float* __restrict__ d_pre, + float* __restrict__ partials, + int out_dim, + int B +) { + extern __shared__ float sdata[]; + int j = blockIdx.y; + int tid = threadIdx.x; + int gid = blockIdx.x * blockDim.x + tid; + float val = (gid < B) ? d_pre[j + (long)gid * out_dim] : 0.0f; + sdata[tid] = val; + __syncthreads(); + for (int s = blockDim.x / 2; s > 0; s >>= 1) { + if (tid < s) sdata[tid] += sdata[tid + s]; + __syncthreads(); + } + if (tid == 0) partials[blockIdx.x * out_dim + j] = sdata[0]; +} + +/* iql_bias_grad_reduce_p2 — Phase 2: warp-shuffle final reduce + scale + * Grid: ceil(out_dim/256), Block: 256 + */ +extern "C" __global__ +void iql_bias_grad_reduce_p2( + const float* __restrict__ partials, float* __restrict__ d_bias, int out_dim, - int B, + int num_blocks, float inv_batch ) { int j = blockIdx.x * blockDim.x + threadIdx.x; - if (j < out_dim) { - float sum = 0.0f; - for (int b = 0; b < B; b++) { - sum += d_pre[j + b * out_dim]; - } - d_bias[j] = sum * inv_batch; - } + if (j >= out_dim) return; + float sum = 0.0f; + for (int i = 0; i < num_blocks; i++) sum += partials[i * out_dim + j]; + d_bias[j] = sum * inv_batch; }