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

perf: iqn_bias_grad_reduce — 2-phase warp-parallel (16ms×130 → <0.2ms×130)

Browse Source 
Replace serial 3-thread kernel (each looping 524K iters) with 2-phase
block-parallel reduction: Phase 1 uses shared memory block reduce across
ceil(BQ/256) blocks per neuron; Phase 2 warp-reduces block partials.
Preallocated scratch buffer [ceil(BQ/256), max_out_dim] avoids runtime alloc.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-19 12:33:35 +02:00
parent d751d2ad4f
commit 3f94802d8f
2 changed files with 123 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files

89
crates/ml/src/cuda_pipeline/gpu_iqn_head.rs
View File

@@ -191,7 +191,8 @@ pub struct GpuIqnHead {
hadamard_sigmoid_bwd_kernel: CudaFunction,
quantile_huber_loss_kernel: CudaFunction,
bias_add_kernel: CudaFunction,
bias_grad_reduce_kernel: CudaFunction,
bias_grad_reduce_p1_kernel: CudaFunction,
bias_grad_reduce_p2_kernel: CudaFunction,
h_s2_tile_kernel: CudaFunction,
d_h_s2_reduce_kernel: CudaFunction,
cos_tile_kernel: CudaFunction,
@@ -267,6 +268,10 @@ pub struct GpuIqnHead {
d_embed_pre_buf: CudaSlice<f32>,
/// d_h_s2_tiled from hadamard_sigmoid_bwd [H, B*Q]
d_h_s2_tiled_buf: CudaSlice<f32>,
/// Scratch for 2-phase bias_grad_reduce: [ceil(BQ/256), max_out_dim]
bias_grad_partials_buf: CudaSlice<f32>,
/// Number of phase-1 blocks along the BQ axis (ceil(BQ/256)).
bias_grad_num_blocks: usize,
// ── Per-step buffers ─────────────────────────────────────────────
/// Pre-sampled τ values for online network [B, N]
@@ -437,6 +442,11 @@ impl GpuIqnHead {
let d_embed_pre_buf = alloc_f32(&stream, h * bq, "iqn_d_embed_pre")?;
let d_h_s2_tiled_buf = alloc_f32(&stream, h * bq, "iqn_d_h_s2_tiled")?;
// 2-phase bias_grad_reduce scratch: [ceil(BQ/256), max_out_dim]
// max_out_dim = h (hidden_dim, used for embed layer — the largest bias)
let bias_grad_num_blocks = (bq + 255) / 256;
let bias_grad_partials_buf = alloc_f32(&stream, bias_grad_num_blocks * h, "iqn_bias_grad_partials")?;
// ── Per-step buffers ────────────────────────────────────────────
// Fixed τ midpoints (QR-DQN style): τ_i = (2i + 1) / (2N)
let online_taus;
@@ -554,7 +564,8 @@ impl GpuIqnHead {
hadamard_sigmoid_bwd_kernel: kernels.hadamard_sigmoid_bwd,
quantile_huber_loss_kernel: kernels.quantile_huber_loss,
bias_add_kernel: kernels.bias_add,
bias_grad_reduce_kernel: kernels.bias_grad_reduce,
bias_grad_reduce_p1_kernel: kernels.bias_grad_reduce_p1,
bias_grad_reduce_p2_kernel: kernels.bias_grad_reduce_p2,
h_s2_tile_kernel: kernels.h_s2_tile,
d_h_s2_reduce_kernel: kernels.d_h_s2_reduce,
cos_tile_kernel: kernels.cos_tile,
@@ -587,6 +598,8 @@ impl GpuIqnHead {
d_embed_buf,
d_embed_pre_buf,
d_h_s2_tiled_buf,
bias_grad_partials_buf,
bias_grad_num_blocks,
online_taus,
target_taus,
cos_features,
@@ -1200,20 +1213,41 @@ impl GpuIqnHead {
)?;
grad_off += bs * h;
// db_bk: bias gradient reduce
// db_bk: 2-phase bias gradient reduce
// Phase 1: grid(ceil(BQ/256), bs, 1), block(256) — partial sums per block
// Phase 2: grid(bs, 1, 1), block(256) — final reduce + scale
let db_ptr = self.grad_buf.raw_ptr() + grad_off as u64 * f32_sz;
let bs_i32 = bs as i32;
let bq_i32 = bq as i32;
let p1_blocks = self.bias_grad_num_blocks as u32;
let num_blocks_i32 = self.bias_grad_num_blocks as i32;
let partials_ptr = self.bias_grad_partials_buf.raw_ptr();
unsafe {
effective_stream
.launch_builder(&self.bias_grad_reduce_kernel)
.launch_builder(&self.bias_grad_reduce_p1_kernel)
.arg(&dq_ptr)
.arg(&db_ptr)
.arg(&partials_ptr)
.arg(&bs_i32)
.arg(&bq_i32)
.launch(LaunchConfig {
grid_dim: (p1_blocks, bs as u32, 1),
block_dim: (256, 1, 1),
shared_mem_bytes: 256 * 4,
})
.map_err(|e| MLError::ModelError(format!("IQN bias_grad_reduce_p1 branch_{i}: {e}")))?;
effective_stream
.launch_builder(&self.bias_grad_reduce_p2_kernel)
.arg(&partials_ptr)
.arg(&db_ptr)
.arg(&num_blocks_i32)
.arg(&bs_i32)
.arg(&inv_bq)
.launch(elem_config(bs))
.map_err(|e| MLError::ModelError(format!("IQN bias_grad_reduce branch_{i}: {e}")))?;
.launch(LaunchConfig {
grid_dim: (bs as u32, 1, 1),
block_dim: (256, 1, 1),
shared_mem_bytes: 0,
})
.map_err(|e| MLError::ModelError(format!("IQN bias_grad_reduce_p2 branch_{i}: {e}")))?;
}
grad_off += bs;
@@ -1267,19 +1301,40 @@ impl GpuIqnHead {
lt_ws_ptr, lt_ws_size, "iqn_bwd_dw_embed",
)?;
// Step 7e: db_embed: bias gradient reduce
// Step 7e: db_embed: 2-phase bias gradient reduce
// Phase 1: grid(ceil(BQ/256), h, 1), block(256) — partial sums per block
// Phase 2: grid(h, 1, 1), block(256) — final reduce + scale
let db_embed_ptr = self.grad_buf.raw_ptr() + (h * d) as u64 * f32_sz;
let bq_i32 = bq as i32;
let p1_blocks = self.bias_grad_num_blocks as u32;
let num_blocks_i32 = self.bias_grad_num_blocks as i32;
let partials_ptr = self.bias_grad_partials_buf.raw_ptr();
unsafe {
effective_stream
.launch_builder(&self.bias_grad_reduce_kernel)
.launch_builder(&self.bias_grad_reduce_p1_kernel)
.arg(&self.d_embed_pre_buf.raw_ptr())
.arg(&db_embed_ptr)
.arg(&partials_ptr)
.arg(&h_i32)
.arg(&bq_i32)
.launch(LaunchConfig {
grid_dim: (p1_blocks, h as u32, 1),
block_dim: (256, 1, 1),
shared_mem_bytes: 256 * 4,
})
.map_err(|e| MLError::ModelError(format!("IQN bias_grad_reduce_p1 embed: {e}")))?;
effective_stream
.launch_builder(&self.bias_grad_reduce_p2_kernel)
.arg(&partials_ptr)
.arg(&db_embed_ptr)
.arg(&num_blocks_i32)
.arg(&h_i32)
.arg(&inv_bq)
.launch(elem_config(h))
.map_err(|e| MLError::ModelError(format!("IQN bias_grad_reduce embed: {e}")))?;
.launch(LaunchConfig {
grid_dim: (h as u32, 1, 1),
block_dim: (256, 1, 1),
shared_mem_bytes: 0,
})
.map_err(|e| MLError::ModelError(format!("IQN bias_grad_reduce_p2 embed: {e}")))?;
}
// Step 7f: Reduce d_h_s2_tiled [H, B*Q] → d_h_s2 [H, B]
@@ -1854,7 +1909,8 @@ struct IqnKernels {
hadamard_sigmoid_bwd: CudaFunction,
quantile_huber_loss: CudaFunction,
bias_add: CudaFunction,
bias_grad_reduce: CudaFunction,
bias_grad_reduce_p1: CudaFunction,
bias_grad_reduce_p2: CudaFunction,
h_s2_tile: CudaFunction,
d_h_s2_reduce: CudaFunction,
cos_tile: CudaFunction,
@@ -1905,17 +1961,18 @@ fn load_iqn_kernels(
let hadamard_sigmoid_bwd = load("iqn_hadamard_sigmoid_bwd")?;
let quantile_huber_loss = load("iqn_quantile_huber_loss")?;
let bias_add = load("iqn_bias_add")?;
let bias_grad_reduce = load("iqn_bias_grad_reduce")?;
let bias_grad_reduce_p1 = load("iqn_bias_grad_reduce_p1")?;
let bias_grad_reduce_p2 = load("iqn_bias_grad_reduce_p2")?;
let h_s2_tile = load("iqn_h_s2_tile")?;
let d_h_s2_reduce = load("iqn_d_h_s2_reduce")?;
let cos_tile = load("iqn_cos_tile")?;
info!("GpuIqnHead: 19 kernels loaded (9 retained + 10 new element-wise)");
info!("GpuIqnHead: 20 kernels loaded (9 retained + 11 new element-wise)");
Ok(IqnKernels {
trunk_fwd, loss_reduce, gnorm_p1, gnorm_p2, adam, fwd_only,
ema, decode_act, sample_taus,
relu_fwd, relu_bwd, hadamard_sigmoid, hadamard_sigmoid_bwd,
quantile_huber_loss, bias_add, bias_grad_reduce,
quantile_huber_loss, bias_add, bias_grad_reduce_p1, bias_grad_reduce_p2,
h_s2_tile, d_h_s2_reduce, cos_tile,
})
}

61
crates/ml/src/cuda_pipeline/iqn_dual_head_kernel.cu
View File

@@ -1299,22 +1299,61 @@ void iqn_bias_add(
}
}
/* Bias gradient reduce: d_bias[j] = inv_bq * sum_i d_pre[j + i*out_dim]
Column-major: d_pre is [out_dim, BQ], sum along BQ axis. */
/* Bias gradient reduce — Phase 1: partial block sums.
Grid: (ceil(BQ/256), out_dim, 1), Block: (256, 1, 1), Shared: 256*sizeof(float).
d_pre is [out_dim, BQ] column-major.
partials is [num_blocks_x, out_dim] — one partial sum per (block, neuron). */
extern "C" __global__
void iqn_bias_grad_reduce(
void iqn_bias_grad_reduce_p1(
const float* __restrict__ d_pre,
float* __restrict__ d_bias,
float* __restrict__ partials,
int out_dim,
int BQ
) {
extern __shared__ float sdata[];
int j = blockIdx.y; /* which output neuron */
int tid = threadIdx.x;
int gid = blockIdx.x * blockDim.x + tid;
float val = (gid < BQ) ? 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];
}
}
/* Bias gradient reduce — Phase 2: final reduction + scale.
Grid: (out_dim, 1, 1), Block: (256, 1, 1).
partials is [num_blocks, out_dim]. */
extern "C" __global__
void iqn_bias_grad_reduce_p2(
const float* __restrict__ partials,
float* __restrict__ d_bias,
int num_blocks,
int out_dim,
int BQ,
float inv_bq
) {
int j = blockIdx.x * blockDim.x + threadIdx.x;
if (j < out_dim) {
float sum = 0.0f;
for (int i = 0; i < BQ; i++) {
sum += d_pre[j + i * out_dim];
}
int j = blockIdx.x;
int tid = threadIdx.x;
float sum = 0.0f;
for (int i = tid; i < num_blocks; i += blockDim.x) {
sum += partials[i * out_dim + j];
}
/* Warp-level reduction */
for (int offset = 16; offset > 0; offset >>= 1) {
sum += __shfl_down_sync(0xFFFFFFFF, sum, offset);
}
if (tid == 0) {
d_bias[j] = sum * inv_bq;
}
}