From 3f94802d8f8baea9a044f1828d9de7fa51028146 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sun, 19 Apr 2026 12:33:35 +0200 Subject: [PATCH] =?UTF-8?q?perf:=20iqn=5Fbias=5Fgrad=5Freduce=20=E2=80=94?= =?UTF-8?q?=202-phase=20warp-parallel=20(16ms=C3=97130=20=E2=86=92=20<0.2m?= =?UTF-8?q?s=C3=97130)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- crates/ml/src/cuda_pipeline/gpu_iqn_head.rs | 89 +++++++++++++++---- .../src/cuda_pipeline/iqn_dual_head_kernel.cu | 61 ++++++++++--- 2 files changed, 123 insertions(+), 27 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/gpu_iqn_head.rs b/crates/ml/src/cuda_pipeline/gpu_iqn_head.rs index 2df0bc6c7..b13407fc7 100644 --- a/crates/ml/src/cuda_pipeline/gpu_iqn_head.rs +++ b/crates/ml/src/cuda_pipeline/gpu_iqn_head.rs @@ -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, /// d_h_s2_tiled from hadamard_sigmoid_bwd [H, B*Q] d_h_s2_tiled_buf: CudaSlice, + /// Scratch for 2-phase bias_grad_reduce: [ceil(BQ/256), max_out_dim] + bias_grad_partials_buf: CudaSlice, + /// 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, }) } diff --git a/crates/ml/src/cuda_pipeline/iqn_dual_head_kernel.cu b/crates/ml/src/cuda_pipeline/iqn_dual_head_kernel.cu index 5edc8d808..9d3a267c0 100644 --- a/crates/ml/src/cuda_pipeline/iqn_dual_head_kernel.cu +++ b/crates/ml/src/cuda_pipeline/iqn_dual_head_kernel.cu @@ -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; } }