diff --git a/crates/ml/src/cuda_pipeline/attention_backward_kernel.cu b/crates/ml/src/cuda_pipeline/attention_backward_kernel.cu index bafcb0744..10ef924bf 100644 --- a/crates/ml/src/cuda_pipeline/attention_backward_kernel.cu +++ b/crates/ml/src/cuda_pipeline/attention_backward_kernel.cu @@ -22,7 +22,7 @@ * ln_beta: [D] offset 4*D*D + 5*D * Total: 4*D^2 + 5*D * - * Grid: (B, 1, 1), Block: (32, 1, 1) -- one warp per sample. + * Grid: (B, 1, 1), Block: (256, 1, 1) -- 256 threads (8 warps) per sample. * Weight gradients accumulated via atomicAdd across batch dimension. * * ALL storage and computation in native __nv_bfloat16. @@ -48,6 +48,50 @@ __device__ __forceinline__ __nv_bfloat16 bf16_shfl(unsigned mask, __nv_bfloat16 return *(__nv_bfloat16*)&raw; } +/* Block-level bf16 max reduction: 256 threads -> single max value. + * Returns the result broadcast to ALL threads in the block. */ +__device__ __forceinline__ __nv_bfloat16 bf16_block_max_bwd( + __nv_bfloat16 val, __nv_bfloat16* warp_sums /* [8] in shared memory */ +) { + int tid = threadIdx.x; + int warp_id = tid / 32; + int warp_lane = tid % 32; + for (int offset = 16; offset > 0; offset >>= 1) + val = bf16_fmax(val, bf16_shfl_xor(0xFFFFFFFF, val, offset)); + if (warp_lane == 0) warp_sums[warp_id] = val; + __syncthreads(); + if (warp_id == 0) { + __nv_bfloat16 v = (warp_lane < (blockDim.x / 32)) ? warp_sums[warp_lane] : bf16(-1e30f); + for (int off = 16; off > 0; off >>= 1) + v = bf16_fmax(v, bf16_shfl_xor(0xFFFFFFFF, v, off)); + warp_sums[0] = v; + } + __syncthreads(); + return warp_sums[0]; +} + +/* Block-level bf16 sum reduction: 256 threads -> single sum. + * Returns the result broadcast to ALL threads in the block. */ +__device__ __forceinline__ __nv_bfloat16 bf16_block_sum_bwd( + __nv_bfloat16 val, __nv_bfloat16* warp_sums /* [8] in shared memory */ +) { + int tid = threadIdx.x; + int warp_id = tid / 32; + int warp_lane = tid % 32; + for (int offset = 16; offset > 0; offset >>= 1) + val = val + bf16_shfl_xor(0xFFFFFFFF, val, offset); + if (warp_lane == 0) warp_sums[warp_id] = val; + __syncthreads(); + if (warp_id == 0) { + __nv_bfloat16 v = (warp_lane < (blockDim.x / 32)) ? warp_sums[warp_lane] : bf16_zero(); + for (int off = 16; off > 0; off >>= 1) + v = v + bf16_shfl_xor(0xFFFFFFFF, v, off); + warp_sums[0] = v; + } + __syncthreads(); + return warp_sums[0]; +} + /** * @param d_output [B, D] gradient from downstream (DQN trunk backward d_h_s2) * @param states_input [B, D] saved input states from forward pass @@ -72,7 +116,7 @@ extern "C" __global__ void attention_backward_kernel( int sample = blockIdx.x; if (sample >= B) return; - int tid = threadIdx.x; /* 0-31 warp lane */ + int tid = threadIdx.x; /* 0-255 */ const __nv_bfloat16* x = states_input + sample * D; const __nv_bfloat16* dy = d_output + sample * D; @@ -102,9 +146,11 @@ extern "C" __global__ void attention_backward_kernel( __nv_bfloat16* d_ln_gamma = db_O_g + D; __nv_bfloat16* d_ln_beta = d_ln_gamma + D; - /* Shared memory layout: concat[D] + d_proj[D] */ - __shared__ __nv_bfloat16 concat[128]; /* recomputed multi-head attention output */ - __shared__ __nv_bfloat16 d_proj[128]; /* d_projection = d_prenorm from LN backward */ + /* Shared memory layout: concat[D] + d_proj[D] + warp_reduce[8] */ + extern __shared__ __nv_bfloat16 shmem_bwd[]; + __nv_bfloat16* concat = shmem_bwd; /* [D] recomputed multi-head attention output */ + __nv_bfloat16* d_proj = shmem_bwd + D; /* [D] d_projection = d_prenorm from LN backward */ + __nv_bfloat16* warp_reduce = shmem_bwd + 2 * D; /* [8] for block reductions */ /* =================================================================== * Step 1: Recompute forward pass to recover concat and pre_ln @@ -118,7 +164,7 @@ extern "C" __global__ void attention_backward_kernel( int head_start = h * Dh; /* Compute Q, K for attention scores */ - for (int f = tid; f < Dh; f += 32) { + for (int f = tid; f < Dh; f += 256) { int gf = head_start + f; __nv_bfloat16 q = b_Q[gf]; __nv_bfloat16 k = b_K[gf]; @@ -133,26 +179,23 @@ extern "C" __global__ void attention_backward_kernel( /* Softmax over head */ __nv_bfloat16 max_s = bf16(-1e4f); - for (int f = tid; f < Dh; f += 32) { + for (int f = tid; f < Dh; f += 256) { __nv_bfloat16 s = concat[head_start + f]; max_s = bf16_fmax(max_s, s); } - for (int mask = 16; mask >= 1; mask >>= 1) - max_s = bf16_fmax(max_s, bf16_shfl_xor(0xFFFFFFFF, max_s, mask)); - max_s = bf16_shfl(0xFFFFFFFF, max_s, 0); + max_s = bf16_block_max_bwd(max_s, warp_reduce); __nv_bfloat16 local_sum = bf16_zero(); - for (int f = tid; f < Dh; f += 32) { + for (int f = tid; f < Dh; f += 256) { __nv_bfloat16 e = bf16_exp(concat[head_start + f] - max_s); concat[head_start + f] = e; local_sum = local_sum + e; } - for (int mask = 16; mask >= 1; mask >>= 1) - local_sum = local_sum + bf16_shfl_xor(0xFFFFFFFF, local_sum, mask); - __nv_bfloat16 inv_sum = bf16_shfl(0xFFFFFFFF, bf16_one() / (local_sum + bf16(1e-8f)), 0); + __nv_bfloat16 total_sum = bf16_block_sum_bwd(local_sum, warp_reduce); + __nv_bfloat16 inv_sum = bf16_one() / (total_sum + bf16(1e-8f)); /* Apply attention to V, store as concat */ - for (int f = tid; f < Dh; f += 32) { + for (int f = tid; f < Dh; f += 256) { __nv_bfloat16 attn = concat[head_start + f] * inv_sum; int gf = head_start + f; __nv_bfloat16 v = b_V[gf]; @@ -166,7 +209,7 @@ extern "C" __global__ void attention_backward_kernel( /* Now concat[0..D) is the recomputed multi-head attention output. * Compute pre_ln[f] = x[f] + sum_j(concat[j] * W_O[j*D+f]) + b_O[f] * Store pre_ln in d_proj temporarily (we'll overwrite with d_projection later). */ - for (int f = tid; f < D; f += 32) { + for (int f = tid; f < D; f += 256) { __nv_bfloat16 proj = b_O[f]; for (int j = 0; j < D; j++) proj = proj + concat[j] * W_O[j * D + f]; @@ -183,21 +226,19 @@ extern "C" __global__ void attention_backward_kernel( /* Compute mean of pre_ln */ __nv_bfloat16 local_mean = bf16_zero(); - for (int f = tid; f < D; f += 32) + for (int f = tid; f < D; f += 256) local_mean = local_mean + d_proj[f]; /* d_proj holds pre_ln here */ - for (int mask = 16; mask >= 1; mask >>= 1) - local_mean = local_mean + bf16_shfl_xor(0xFFFFFFFF, local_mean, mask); - __nv_bfloat16 mean = bf16_shfl(0xFFFFFFFF, local_mean / bf16((float)D), 0); + __nv_bfloat16 total_mean = bf16_block_sum_bwd(local_mean, warp_reduce); + __nv_bfloat16 mean = total_mean / bf16((float)D); /* Compute variance of pre_ln */ __nv_bfloat16 local_var = bf16_zero(); - for (int f = tid; f < D; f += 32) { + for (int f = tid; f < D; f += 256) { __nv_bfloat16 diff = d_proj[f] - mean; local_var = local_var + diff * diff; } - for (int mask = 16; mask >= 1; mask >>= 1) - local_var = local_var + bf16_shfl_xor(0xFFFFFFFF, local_var, mask); - __nv_bfloat16 var = bf16_shfl(0xFFFFFFFF, local_var / bf16((float)D), 0); + __nv_bfloat16 total_var = bf16_block_sum_bwd(local_var, warp_reduce); + __nv_bfloat16 var = total_var / bf16((float)D); __nv_bfloat16 inv_std = bf16_one() / bf16_sqrt(var + bf16(1e-5f)); /* LN backward: @@ -208,23 +249,19 @@ extern "C" __global__ void attention_backward_kernel( /* First pass: accumulate sum(dy*gamma) and sum(dy*gamma*x_hat) */ __nv_bfloat16 sum_dy_gamma = bf16_zero(); __nv_bfloat16 sum_dy_gamma_xhat = bf16_zero(); - for (int f = tid; f < D; f += 32) { + for (int f = tid; f < D; f += 256) { __nv_bfloat16 x_hat = (d_proj[f] - mean) * inv_std; /* d_proj holds pre_ln */ __nv_bfloat16 dy_g = dy[f] * ln_gamma[f]; sum_dy_gamma = sum_dy_gamma + dy_g; sum_dy_gamma_xhat = sum_dy_gamma_xhat + dy_g * x_hat; } - for (int mask = 16; mask >= 1; mask >>= 1) { - sum_dy_gamma = sum_dy_gamma + bf16_shfl_xor(0xFFFFFFFF, sum_dy_gamma, mask); - sum_dy_gamma_xhat = sum_dy_gamma_xhat + bf16_shfl_xor(0xFFFFFFFF, sum_dy_gamma_xhat, mask); - } - sum_dy_gamma = bf16_shfl(0xFFFFFFFF, sum_dy_gamma, 0); - sum_dy_gamma_xhat = bf16_shfl(0xFFFFFFFF, sum_dy_gamma_xhat, 0); + sum_dy_gamma = bf16_block_sum_bwd(sum_dy_gamma, warp_reduce); + sum_dy_gamma_xhat = bf16_block_sum_bwd(sum_dy_gamma_xhat, warp_reduce); /* Second pass: compute d_prenorm, accumulate LN param grads. * Overwrite d_proj with d_prenorm (we're done with pre_ln). */ __nv_bfloat16 bf16_D = bf16((float)D); - for (int f = tid; f < D; f += 32) { + for (int f = tid; f < D; f += 256) { __nv_bfloat16 x_hat = (d_proj[f] - mean) * inv_std; __nv_bfloat16 dy_g = dy[f] * ln_gamma[f]; @@ -250,7 +287,7 @@ extern "C" __global__ void attention_backward_kernel( * =================================================================== */ /* Initialize d_input with residual gradient */ - for (int f = tid; f < D; f += 32) + for (int f = tid; f < D; f += 256) atomicAddBF16(&dx[f], d_proj[f]); /* d_residual via atomicAdd (heads also add) */ __syncthreads(); @@ -267,7 +304,7 @@ extern "C" __global__ void attention_backward_kernel( * =================================================================== */ /* b_O gradient */ - for (int f = tid; f < D; f += 32) + for (int f = tid; f < D; f += 256) atomicAddBF16(&db_O_g[f], d_proj[f]); /* W_O gradient and d_concat computation. @@ -278,7 +315,7 @@ extern "C" __global__ void attention_backward_kernel( * into d_proj (overwriting d_projection which we no longer need). */ /* Phase 1: accumulate W_O gradients (concat[j] * d_proj[f]) */ - for (int f = tid; f < D; f += 32) { + for (int f = tid; f < D; f += 256) { __nv_bfloat16 dp = d_proj[f]; for (int j = 0; j < D; j++) atomicAddBF16(&dW_O[j * D + f], concat[j] * dp); @@ -286,7 +323,7 @@ extern "C" __global__ void attention_backward_kernel( __syncthreads(); /* Phase 2: compute d_concat and store back in d_proj */ - for (int j = tid; j < D; j += 32) { + for (int j = tid; j < D; j += 256) { __nv_bfloat16 d_c = bf16_zero(); for (int f = 0; f < D; f++) d_c = d_c + d_proj[f] * W_O[j * D + f]; @@ -296,7 +333,7 @@ extern "C" __global__ void attention_backward_kernel( __syncthreads(); /* Copy d_concat from concat back to d_proj, restore concat for head backward */ - for (int f = tid; f < D; f += 32) + for (int f = tid; f < D; f += 256) d_proj[f] = concat[f]; /* d_proj now holds d_concat */ __syncthreads(); @@ -318,7 +355,7 @@ extern "C" __global__ void attention_backward_kernel( /* Recompute Q, K, V and attention for this head. * Use concat[head_start..] as scratch for attn values. */ - for (int f = tid; f < Dh; f += 32) { + for (int f = tid; f < Dh; f += 256) { int gf = head_start + f; __nv_bfloat16 q = b_Q[gf]; __nv_bfloat16 k = b_K[gf]; @@ -333,29 +370,26 @@ extern "C" __global__ void attention_backward_kernel( /* Softmax recompute */ __nv_bfloat16 max_s = bf16(-1e4f); - for (int f = tid; f < Dh; f += 32) { + for (int f = tid; f < Dh; f += 256) { __nv_bfloat16 s = concat[head_start + f]; max_s = bf16_fmax(max_s, s); } - for (int mask = 16; mask >= 1; mask >>= 1) - max_s = bf16_fmax(max_s, bf16_shfl_xor(0xFFFFFFFF, max_s, mask)); - max_s = bf16_shfl(0xFFFFFFFF, max_s, 0); + max_s = bf16_block_max_bwd(max_s, warp_reduce); __nv_bfloat16 sm_sum = bf16_zero(); - for (int f = tid; f < Dh; f += 32) { + for (int f = tid; f < Dh; f += 256) { __nv_bfloat16 e = bf16_exp(concat[head_start + f] - max_s); concat[head_start + f] = e; sm_sum = sm_sum + e; } - for (int mask = 16; mask >= 1; mask >>= 1) - sm_sum = sm_sum + bf16_shfl_xor(0xFFFFFFFF, sm_sum, mask); - __nv_bfloat16 inv_sum = bf16_shfl(0xFFFFFFFF, bf16_one() / (sm_sum + bf16(1e-8f)), 0); + __nv_bfloat16 total_sm = bf16_block_sum_bwd(sm_sum, warp_reduce); + __nv_bfloat16 inv_sum = bf16_one() / (total_sm + bf16(1e-8f)); /* concat[gf] now has unnorm exp. attn[f] = concat[gf] * inv_sum */ /* Compute sum(d_attn * attn) for softmax backward */ __nv_bfloat16 sum_da_a = bf16_zero(); - for (int f = tid; f < Dh; f += 32) { + for (int f = tid; f < Dh; f += 256) { int gf = head_start + f; __nv_bfloat16 attn = concat[gf] * inv_sum; @@ -367,13 +401,11 @@ extern "C" __global__ void attention_backward_kernel( __nv_bfloat16 d_attn = d_proj[gf] * v; /* d_proj holds d_concat */ sum_da_a = sum_da_a + d_attn * attn; } - for (int mask = 16; mask >= 1; mask >>= 1) - sum_da_a = sum_da_a + bf16_shfl_xor(0xFFFFFFFF, sum_da_a, mask); - sum_da_a = bf16_shfl(0xFFFFFFFF, sum_da_a, 0); + sum_da_a = bf16_block_sum_bwd(sum_da_a, warp_reduce); /* Compute d_Q, d_K, d_V and accumulate weight/bias/input gradients */ __nv_bfloat16 sqrt_Dh = bf16_sqrt(bf16((float)Dh)); - for (int f = tid; f < Dh; f += 32) { + for (int f = tid; f < Dh; f += 256) { int gf = head_start + f; __nv_bfloat16 attn = concat[gf] * inv_sum; diff --git a/crates/ml/src/cuda_pipeline/attention_kernel.cu b/crates/ml/src/cuda_pipeline/attention_kernel.cu index 407a2283e..b4bd856a2 100644 --- a/crates/ml/src/cuda_pipeline/attention_kernel.cu +++ b/crates/ml/src/cuda_pipeline/attention_kernel.cu @@ -15,7 +15,7 @@ * Residual: output = state + projected * LayerNorm: output = layernorm(output) * - * Grid: (B, 1, 1), Block: (32, 1, 1). One warp per sample. + * Grid: (B, 1, 1), Block: (256, 1, 1). 256 threads (8 warps) per sample. */ /* -- Configuration constants (overridden by NVRTC injection) ---------- */ @@ -27,6 +27,52 @@ #endif #define ATTN_HEAD_DIM (ATTN_STATE_DIM / ATTN_NUM_HEADS) +/* Block-level bf16 max reduction: 256 threads -> single max value. + * Returns the result broadcast to ALL threads in the block. */ +__device__ __forceinline__ __nv_bfloat16 bf16_block_max( + __nv_bfloat16 val, __nv_bfloat16* warp_sums /* [8] in shared memory */ +) { + int tid = threadIdx.x; + int warp_id = tid / 32; + int warp_lane = tid % 32; + /* Intra-warp max */ + for (int offset = 16; offset > 0; offset >>= 1) + val = bf16_fmax(val, bf16_shfl_xor(0xFFFFFFFF, val, offset)); + if (warp_lane == 0) warp_sums[warp_id] = val; + __syncthreads(); + if (warp_id == 0) { + __nv_bfloat16 v = (warp_lane < (blockDim.x / 32)) ? warp_sums[warp_lane] : bf16(-1e30f); + for (int off = 16; off > 0; off >>= 1) + v = bf16_fmax(v, bf16_shfl_xor(0xFFFFFFFF, v, off)); + warp_sums[0] = v; + } + __syncthreads(); + return warp_sums[0]; +} + +/* Block-level bf16 sum reduction: 256 threads -> single sum. + * Returns the result broadcast to ALL threads in the block. */ +__device__ __forceinline__ __nv_bfloat16 bf16_block_sum( + __nv_bfloat16 val, __nv_bfloat16* warp_sums /* [8] in shared memory */ +) { + int tid = threadIdx.x; + int warp_id = tid / 32; + int warp_lane = tid % 32; + /* Intra-warp sum */ + for (int offset = 16; offset > 0; offset >>= 1) + val = val + bf16_shfl_xor(0xFFFFFFFF, val, offset); + if (warp_lane == 0) warp_sums[warp_id] = val; + __syncthreads(); + if (warp_id == 0) { + __nv_bfloat16 v = (warp_lane < (blockDim.x / 32)) ? warp_sums[warp_lane] : bf16_zero(); + for (int off = 16; off > 0; off >>= 1) + v = v + bf16_shfl_xor(0xFFFFFFFF, v, off); + warp_sums[0] = v; + } + __syncthreads(); + return warp_sums[0]; +} + /** * Forward pass: multi-head self-attention + residual + layer norm. * @@ -57,7 +103,7 @@ extern "C" __global__ void multihead_feature_attention( int sample = blockIdx.x; if (sample >= B) return; - int tid = threadIdx.x; /* 0-31 warp lane */ + int tid = threadIdx.x; /* 0-255 */ const __nv_bfloat16* x = states + sample * D; __nv_bfloat16* out = output + sample * D; @@ -74,15 +120,17 @@ extern "C" __global__ void multihead_feature_attention( const __nv_bfloat16* ln_gamma = b_O + D; /* [D] */ const __nv_bfloat16* ln_beta = ln_gamma + D; /* [D] */ - /* Accumulate multi-head output into a shared buffer */ - __shared__ __nv_bfloat16 shared_concat[128]; /* max D */ + /* Shared memory layout: concat[D] + reduction scratch [8 bf16] */ + extern __shared__ __nv_bfloat16 shmem_fwd[]; + __nv_bfloat16* shared_concat = shmem_fwd; /* [D] */ + __nv_bfloat16* warp_reduce = shmem_fwd + D; /* [8] for block reductions */ /* Process each head sequentially (H=4, each head has Dh=18 dims) */ for (int h = 0; h < H; h++) { int head_start = h * Dh; /* Compute Q_h, K_h, V_h for features this thread handles */ - for (int f = tid; f < Dh; f += 32) { + for (int f = tid; f < Dh; f += 256) { int global_f = head_start + f; __nv_bfloat16 q = b_Q[global_f]; __nv_bfloat16 k = b_K[global_f]; @@ -104,27 +152,23 @@ extern "C" __global__ void multihead_feature_attention( /* Softmax over this head's Dh dimensions */ __nv_bfloat16 max_s = bf16(-1e30f); - for (int f = tid; f < Dh; f += 32) { + for (int f = tid; f < Dh; f += 256) { __nv_bfloat16 s = shared_concat[head_start + f]; max_s = bf16_fmax(max_s, s); } - for (int mask = 16; mask >= 1; mask >>= 1) - max_s = bf16_fmax(max_s, bf16_shfl_xor(0xFFFFFFFF, max_s, mask)); - /* Broadcast max from lane 0 */ - max_s = bf16((float)__shfl_sync(0xFFFFFFFF, (float)max_s, 0)); + max_s = bf16_block_max(max_s, warp_reduce); __nv_bfloat16 local_sum = bf16_zero(); - for (int f = tid; f < Dh; f += 32) { + for (int f = tid; f < Dh; f += 256) { __nv_bfloat16 e = bf16_exp(shared_concat[head_start + f] - max_s); shared_concat[head_start + f] = e; local_sum = local_sum + e; } - for (int mask = 16; mask >= 1; mask >>= 1) - local_sum = local_sum + bf16_shfl_xor(0xFFFFFFFF, local_sum, mask); - __nv_bfloat16 inv_sum = bf16((float)__shfl_sync(0xFFFFFFFF, (float)(bf16_one() / (local_sum + bf16(1e-8f))), 0)); + __nv_bfloat16 total_sum = bf16_block_sum(local_sum, warp_reduce); + __nv_bfloat16 inv_sum = bf16_one() / (total_sum + bf16(1e-8f)); /* Apply attention to V, store in shared_concat */ - for (int f = tid; f < Dh; f += 32) { + for (int f = tid; f < Dh; f += 256) { __nv_bfloat16 attn = shared_concat[head_start + f] * inv_sum; /* Recompute V[f] */ @@ -140,7 +184,7 @@ extern "C" __global__ void multihead_feature_attention( /* -- Step 2: Output projection + residual -- */ /* out[f] = x[f] + sum_j(concat[j] * W_O[j * D + f]) + b_O[f] */ - for (int f = tid; f < D; f += 32) { + for (int f = tid; f < D; f += 256) { __nv_bfloat16 proj = b_O[f]; for (int j = 0; j < D; j++) { proj = proj + shared_concat[j] * W_O[j * D + f]; @@ -152,23 +196,21 @@ extern "C" __global__ void multihead_feature_attention( /* -- Step 3: Layer normalization -- */ /* mean = avg(out[f]) */ __nv_bfloat16 local_mean = bf16_zero(); - for (int f = tid; f < D; f += 32) local_mean = local_mean + out[f]; - for (int mask = 16; mask >= 1; mask >>= 1) - local_mean = local_mean + bf16_shfl_xor(0xFFFFFFFF, local_mean, mask); - __nv_bfloat16 mean = bf16((float)__shfl_sync(0xFFFFFFFF, (float)(local_mean / bf16((float)D)), 0)); + for (int f = tid; f < D; f += 256) local_mean = local_mean + out[f]; + __nv_bfloat16 total_mean = bf16_block_sum(local_mean, warp_reduce); + __nv_bfloat16 mean = total_mean / bf16((float)D); /* variance = avg((out[f] - mean)^2) */ __nv_bfloat16 local_var = bf16_zero(); - for (int f = tid; f < D; f += 32) { + for (int f = tid; f < D; f += 256) { __nv_bfloat16 diff = out[f] - mean; local_var = local_var + diff * diff; } - for (int mask = 16; mask >= 1; mask >>= 1) - local_var = local_var + bf16_shfl_xor(0xFFFFFFFF, local_var, mask); - __nv_bfloat16 inv_std = bf16((float)__shfl_sync(0xFFFFFFFF, (float)(bf16_one() / bf16_sqrt(local_var / bf16((float)D) + bf16(1e-5f))), 0)); + __nv_bfloat16 total_var = bf16_block_sum(local_var, warp_reduce); + __nv_bfloat16 inv_std = bf16_one() / bf16_sqrt(total_var / bf16((float)D) + bf16(1e-5f)); /* Normalize + affine */ - for (int f = tid; f < D; f += 32) { + for (int f = tid; f < D; f += 256) { out[f] = ln_gamma[f] * (out[f] - mean) * inv_std + ln_beta[f]; } } diff --git a/crates/ml/src/cuda_pipeline/gpu_attention.rs b/crates/ml/src/cuda_pipeline/gpu_attention.rs index df065d316..f1e1c7c78 100644 --- a/crates/ml/src/cuda_pipeline/gpu_attention.rs +++ b/crates/ml/src/cuda_pipeline/gpu_attention.rs @@ -220,10 +220,12 @@ impl GpuAttention { ).map_err(|e| MLError::ModelError(format!("attention save input DtoD: {e}")))?; } + // Shared memory: shared_concat[D] + warp_reduce[8] in bf16 + let fwd_shmem = (d + 8) * std::mem::size_of::(); let launch_cfg = LaunchConfig { grid_dim: (b as u32, 1, 1), - block_dim: (32, 1, 1), - shared_mem_bytes: (d.max(128) * 4) as u32, // shared_concat[128] hardcoded in kernel + block_dim: (256, 1, 1), + shared_mem_bytes: fwd_shmem as u32, }; let b_i32 = b as i32; let state_dim_i32 = d as i32; @@ -285,12 +287,12 @@ impl GpuAttention { ).map_err(|e| MLError::ModelError(format!("attention d_input zero: {e}")))?; } - // Launch backward kernel: one warp per sample - // Shared memory: concat[128] + d_proj[128] in bf16 = 2*128*sizeof(bf16) - let shared_mem = (d.max(128) * std::mem::size_of::() * 2) as u32; + // Launch backward kernel: 256 threads (8 warps) per sample + // Shared memory: concat[D] + d_proj[D] + warp_reduce[8] in bf16 + let shared_mem = ((2 * d + 8) * std::mem::size_of::()) as u32; let launch_cfg = LaunchConfig { grid_dim: (b as u32, 1, 1), - block_dim: (32, 1, 1), + block_dim: (256, 1, 1), shared_mem_bytes: shared_mem, }; let b_i32 = b as i32; diff --git a/crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs index 0ff5dfd09..52ccb774b 100644 --- a/crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs @@ -221,8 +221,8 @@ impl GpuIqlTrainer { /// /// Executes the full forward + expectile loss + backward + Adam cycle: /// 1. Zero total_loss and grad_norm scalars - /// 2. Forward + loss kernel (one warp per sample) - /// 3. Backward kernel (one warp per sample, atomicAdd gradients) + /// 2. Forward + loss kernel (8 warps per sample) + /// 3. Backward kernel (8 warps per sample, atomicAdd gradients) /// 4. Grad norm kernel (parallel reduction) /// 5. Adam update kernel (one thread per parameter) /// @@ -250,11 +250,13 @@ impl GpuIqlTrainer { let batch_size_i32 = b as i32; let total_params_i32 = self.total_params as i32; - // 1. Forward + loss kernel + // 1. Forward + loss kernel (256 threads: 8 warps per sample) + // Shared memory: warp_sums[8] in bf16 for block-level reduction + let fwd_shmem = 8 * std::mem::size_of::(); let fwd_config = LaunchConfig { grid_dim: (b as u32, 1, 1), - block_dim: (32, 1, 1), - shared_mem_bytes: 0, + block_dim: (256, 1, 1), + shared_mem_bytes: fwd_shmem as u32, }; // Safety: states_f32 and q_values_f32 are valid F32 CudaSlice buffers @@ -279,10 +281,10 @@ impl GpuIqlTrainer { .map_err(|e| MLError::ModelError(format!("IQL forward+loss kernel: {e}")))?; } - // 2. Backward kernel + // 2. Backward kernel (256 threads: 8 warps per sample) let bwd_config = LaunchConfig { grid_dim: (b as u32, 1, 1), - block_dim: (32, 1, 1), + block_dim: (256, 1, 1), shared_mem_bytes: 0, }; diff --git a/crates/ml/src/cuda_pipeline/iql_value_kernel.cu b/crates/ml/src/cuda_pipeline/iql_value_kernel.cu index d270c406b..10afc2490 100644 --- a/crates/ml/src/cuda_pipeline/iql_value_kernel.cu +++ b/crates/ml/src/cuda_pipeline/iql_value_kernel.cu @@ -9,8 +9,8 @@ * state -> Linear(STATE_DIM, H) -> SiLU -> Linear(H, H) -> SiLU -> Linear(H, 1) -> V(s) * * Launch config: - * Forward + loss: grid=(batch_size, 1, 1), block=(32, 1, 1) -- one warp per sample - * Backward: grid=(batch_size, 1, 1), block=(32, 1, 1) + * Forward + loss: grid=(batch_size, 1, 1), block=(256, 1, 1) -- 8 warps per sample + * Backward: grid=(batch_size, 1, 1), block=(256, 1, 1) * Adam: grid=(ceil(total_params/256), 1, 1), block=(256, 1, 1) * * Compile-time defines (injected via NVRTC dim_overrides): @@ -67,13 +67,36 @@ __device__ __forceinline__ __nv_bfloat16 silu_grad(__nv_bfloat16 x) { return s * (bf16_one() + x * (bf16_one() - s)); } +/* Block-level bf16 sum reduction: 256 threads -> single sum. + * Returns the result broadcast to ALL threads in the block. */ +__device__ __forceinline__ __nv_bfloat16 iql_bf16_block_sum( + __nv_bfloat16 val, __nv_bfloat16* warp_sums /* [8] in shared memory */ +) { + int tid = threadIdx.x; + int warp_id = tid / 32; + int warp_lane = tid % 32; + /* Intra-warp sum via shfl_xor */ + for (int offset = 16; offset > 0; offset >>= 1) + val = val + bf16_shfl_xor(0xFFFFFFFF, val, offset); + if (warp_lane == 0) warp_sums[warp_id] = val; + __syncthreads(); + if (warp_id == 0) { + __nv_bfloat16 v = (warp_lane < (blockDim.x / 32)) ? warp_sums[warp_lane] : bf16_zero(); + for (int off = 16; off > 0; off >>= 1) + v = v + bf16_shfl_xor(0xFFFFFFFF, v, off); + warp_sums[0] = v; + } + __syncthreads(); + return warp_sums[0]; +} + /* ------------------------------------------------------------------ */ /* Forward + Expectile Loss Kernel */ /* ------------------------------------------------------------------ */ /** * Per-sample forward pass through V(s) MLP + expectile loss computation. * - * One warp (32 threads) per sample. Warp-cooperative matvec for each layer. + * 256 threads (8 warps) per sample. Block-cooperative matvec for each layer. * Saves pre-activation values for backward pass. * * Inputs: @@ -109,7 +132,10 @@ void iql_forward_loss_kernel( int sample = blockIdx.x; if (sample >= batch_size) return; - int lane = threadIdx.x; /* 0..31 */ + int tid = threadIdx.x; /* 0..255 */ + + /* Block-level reduction scratch (8 warps) */ + __shared__ __nv_bfloat16 warp_sums[8]; int off_w1, off_b1, off_w2, off_b2, off_w3, off_b3; iql_compute_offsets(state_dim, &off_w1, &off_b1, &off_w2, &off_b2, &off_w3, &off_b3); @@ -124,11 +150,10 @@ void iql_forward_loss_kernel( const __nv_bfloat16* x = states + sample * state_dim; /* ---- Layer 1: Linear(state_dim -> H) + SiLU ---- */ - /* Each lane computes a subset of H outputs via warp reduction */ __nv_bfloat16* pre1_ptr = save_pre1 + sample * VALUE_HIDDEN_DIM; __nv_bfloat16* h1_ptr = save_h1 + sample * VALUE_HIDDEN_DIM; - for (int j = lane; j < VALUE_HIDDEN_DIM; j += 32) { + for (int j = tid; j < VALUE_HIDDEN_DIM; j += 256) { __nv_bfloat16 acc = b1[j]; for (int k = 0; k < state_dim; k++) { acc = acc + w1[j * state_dim + k] * x[k]; @@ -136,13 +161,13 @@ void iql_forward_loss_kernel( pre1_ptr[j] = acc; h1_ptr[j] = silu(acc); } - __syncwarp(); + __syncthreads(); /* ---- Layer 2: Linear(H -> H) + SiLU ---- */ __nv_bfloat16* pre2_ptr = save_pre2 + sample * VALUE_HIDDEN_DIM; __nv_bfloat16* h2_ptr = save_h2 + sample * VALUE_HIDDEN_DIM; - for (int j = lane; j < VALUE_HIDDEN_DIM; j += 32) { + for (int j = tid; j < VALUE_HIDDEN_DIM; j += 256) { __nv_bfloat16 acc = b2[j]; for (int k = 0; k < VALUE_HIDDEN_DIM; k++) { acc = acc + w2[j * VALUE_HIDDEN_DIM + k] * h1_ptr[k]; @@ -150,20 +175,19 @@ void iql_forward_loss_kernel( pre2_ptr[j] = acc; h2_ptr[j] = silu(acc); } - __syncwarp(); + __syncthreads(); /* ---- Output layer: Linear(H -> 1) ---- */ - /* Warp-reduce the dot product across lanes */ + /* Block-reduce the dot product across all 256 threads */ __nv_bfloat16 v_acc = bf16_zero(); - for (int k = lane; k < VALUE_HIDDEN_DIM; k += 32) { + for (int k = tid; k < VALUE_HIDDEN_DIM; k += 256) { v_acc = v_acc + w3[k] * h2_ptr[k]; } - /* Warp reduce */ - for (int offset = 16; offset > 0; offset >>= 1) { - v_acc = v_acc + bf16_shfl_down(0xFFFFFFFF, v_acc, offset); - } - /* Lane 0 has the full sum */ - if (lane == 0) { + /* Block-level sum reduction */ + v_acc = iql_bf16_block_sum(v_acc, warp_sums); + + /* Thread 0 has the full sum */ + if (tid == 0) { __nv_bfloat16 v_val = v_acc + b3[0]; v_out[sample] = v_val; @@ -185,9 +209,10 @@ void iql_forward_loss_kernel( /** * Per-sample backward pass through V(s) MLP. * - * Computes gradients of the expectile loss w.r.t. all parameters. - * Uses atomicAdd to accumulate into the flat gradient buffer (safe for - * concurrent samples since different samples write overlapping param indices). + * 256 threads (8 warps) per sample. Computes gradients of the expectile + * loss w.r.t. all parameters. Uses atomicAdd to accumulate into the flat + * gradient buffer (safe for concurrent samples since different samples + * write overlapping param indices). * * Gradient flow: * dL/dV = -2 * weight * u (where u = Q - V, weight = tau or 1-tau) @@ -213,7 +238,7 @@ void iql_backward_kernel( int sample = blockIdx.x; if (sample >= batch_size) return; - int lane = threadIdx.x; + int tid = threadIdx.x; /* 0..255 */ int off_w1, off_b1, off_w2, off_b2, off_w3, off_b3; iql_compute_offsets(state_dim, &off_w1, &off_b1, &off_w2, &off_b2, &off_w3, &off_b3); @@ -246,17 +271,17 @@ void iql_backward_kernel( /* ---- Output layer gradient: dL/dw3, dL/db3 ---- */ /* dV/dw3[k] = h2[k], dV/db3 = 1 */ - for (int k = lane; k < VALUE_HIDDEN_DIM; k += 32) { + for (int k = tid; k < VALUE_HIDDEN_DIM; k += 256) { atomicAddBF16(&gw3[k], dldv * h2[k]); } - if (lane == 0) { + if (tid == 0) { atomicAddBF16(&gb3[0], dldv); } /* ---- Backprop through layer 2 ---- */ /* dL/dh2[k] = dldv * w3[k] */ /* dL/dpre2[k] = dL/dh2[k] * silu'(pre2[k]) */ - for (int j = lane; j < VALUE_HIDDEN_DIM; j += 32) { + for (int j = tid; j < VALUE_HIDDEN_DIM; j += 256) { __nv_bfloat16 dh2_j = dldv * w3[j]; __nv_bfloat16 dpre2_j = dh2_j * silu_grad(pre2[j]); @@ -268,11 +293,11 @@ void iql_backward_kernel( atomicAddBF16(&gw2[j * VALUE_HIDDEN_DIM + k], dpre2_j * h1[k]); } } - __syncwarp(); + __syncthreads(); /* ---- Backprop through layer 1 ---- */ /* dL/dh1[k] = sum_j(dL/dpre2[j] * w2[j, k]) */ - for (int k = lane; k < VALUE_HIDDEN_DIM; k += 32) { + for (int k = tid; k < VALUE_HIDDEN_DIM; k += 256) { __nv_bfloat16 dh1_k = bf16_zero(); for (int j = 0; j < VALUE_HIDDEN_DIM; j++) { __nv_bfloat16 dh2_j = dldv * w3[j]; @@ -408,6 +433,7 @@ void iql_grad_norm_kernel( * * Used for computing advantages A(s,a) = Q(s,a) - V(s) at inference time. * Same architecture as forward_loss but skips loss and activation saves. + * 256 threads (8 warps) per sample. */ extern "C" __global__ void iql_forward_kernel( @@ -421,7 +447,7 @@ void iql_forward_kernel( int sample = blockIdx.x; if (sample >= batch_size) return; - int lane = threadIdx.x; + int tid = threadIdx.x; /* 0..255 */ int off_w1, off_b1, off_w2, off_b2, off_w3, off_b3; iql_compute_offsets(state_dim, &off_w1, &off_b1, &off_w2, &off_b2, &off_w3, &off_b3); @@ -435,39 +461,39 @@ void iql_forward_kernel( const __nv_bfloat16* x = states + sample * state_dim; - /* Shared memory for intermediate activations (reused per layer) */ + /* Shared memory for intermediate activations + block reduction */ __shared__ __nv_bfloat16 sh1[VALUE_HIDDEN_DIM]; __shared__ __nv_bfloat16 sh2[VALUE_HIDDEN_DIM]; + __shared__ __nv_bfloat16 warp_sums[8]; /* Layer 1: SiLU */ - for (int j = lane; j < VALUE_HIDDEN_DIM; j += 32) { + for (int j = tid; j < VALUE_HIDDEN_DIM; j += 256) { __nv_bfloat16 acc = b1[j]; for (int k = 0; k < state_dim; k++) { acc = acc + w1[j * state_dim + k] * x[k]; } sh1[j] = silu(acc); } - __syncwarp(); + __syncthreads(); /* Layer 2: SiLU */ - for (int j = lane; j < VALUE_HIDDEN_DIM; j += 32) { + for (int j = tid; j < VALUE_HIDDEN_DIM; j += 256) { __nv_bfloat16 acc = b2[j]; for (int k = 0; k < VALUE_HIDDEN_DIM; k++) { acc = acc + w2[j * VALUE_HIDDEN_DIM + k] * sh1[k]; } sh2[j] = silu(acc); } - __syncwarp(); + __syncthreads(); - /* Output: dot product + bias */ + /* Output: dot product + bias -- block-level sum reduction */ __nv_bfloat16 v_acc = bf16_zero(); - for (int k = lane; k < VALUE_HIDDEN_DIM; k += 32) { + for (int k = tid; k < VALUE_HIDDEN_DIM; k += 256) { v_acc = v_acc + w3[k] * sh2[k]; } - for (int offset = 16; offset > 0; offset >>= 1) { - v_acc = v_acc + bf16_shfl_down(0xFFFFFFFF, v_acc, offset); - } - if (lane == 0) { + v_acc = iql_bf16_block_sum(v_acc, warp_sums); + + if (tid == 0) { v_out[sample] = v_acc + b3[0]; } }