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) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -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;
|
||||
}
|
||||
|
||||
@@ -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 */
|
||||
/* ================================================================== */
|
||||
|
||||
@@ -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<f32>,
|
||||
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::<f32>(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(())
|
||||
}
|
||||
|
||||
@@ -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<f32>,
|
||||
/// Col-major scratch [H=24, B] — denoiser MLP input for step 1.
|
||||
denoise_input1_buf: CudaSlice<f32>,
|
||||
/// Col-major scratch [H=24, B] — pre-activation (with bias) for step 0.
|
||||
denoise_pre_h0_buf: CudaSlice<f32>,
|
||||
/// Col-major scratch [H=24, B] — pre-activation (with bias) for step 1.
|
||||
denoise_pre_h1_buf: CudaSlice<f32>,
|
||||
/// Col-major scratch [H=24, B] — SiLU output for step 0.
|
||||
denoise_h0_buf: CudaSlice<f32>,
|
||||
/// Col-major scratch [H=24, B] — SiLU output for step 1.
|
||||
denoise_h1_buf: CudaSlice<f32>,
|
||||
/// Row-major scratch [B, D=12] — Q values after step 0.
|
||||
denoise_q_step0_buf: CudaSlice<f32>,
|
||||
/// Col-major scratch [D=12, B] — loss gradient (upstream for backward).
|
||||
denoise_d_res_buf: CudaSlice<f32>,
|
||||
/// Col-major scratch [H=24, B] — d_h / d_pre scratch for backward.
|
||||
denoise_d_h_buf: CudaSlice<f32>,
|
||||
/// Scratch for bias gradient reduction phase 1 [num_blocks * max(H,D)].
|
||||
denoise_bias_grad_partials: CudaSlice<f32>,
|
||||
/// 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<f32>,
|
||||
@@ -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)
|
||||
|
||||
@@ -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<f32>,
|
||||
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::<f32>(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) ──
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user