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

perf: attention + IQL block 32→256 (occupancy 3.1%→25%)

Browse Source 
Same optimization pattern as IQN: increase block size from 32 (1 warp)
to 256 (8 warps) for all attention and IQL kernels.

Attention (forward + backward):
- Loop strides 32→256, shared_concat dynamic shmem
- Added bf16_block_max/bf16_block_sum for 8-warp reduction
- Replaces warp-only shuffle with shared memory cross-warp reduce

IQL (forward+loss, backward, forward-only):
- Loop strides 32→256, added iql_bf16_block_sum
- grad_norm + adam kernels already at 256, unchanged

Expected: ~15-20ms/step savings from attention+IQL combined.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-06 23:37:36 +02:00
parent 6ff45466f9
commit 74f63f80b2
5 changed files with 231 additions and 127 deletions
Download Patch File Download Diff File Expand all files Collapse all files

134
crates/ml/src/cuda_pipeline/attention_backward_kernel.cu
View File

@@ -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;

92
crates/ml/src/cuda_pipeline/attention_kernel.cu
View File

@@ -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];
}
}

14
crates/ml/src/cuda_pipeline/gpu_attention.rs
View File

@@ -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::<half::bf16>();
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::<half::bf16>() * 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::<half::bf16>()) 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;

16
crates/ml/src/cuda_pipeline/gpu_iql_trainer.rs
View File

@@ -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::<half::bf16>();
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,
};

102
crates/ml/src/cuda_pipeline/iql_value_kernel.cu
View File

@@ -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];
}
}