diff --git a/crates/ml/src/cuda_pipeline/iql_value_kernel.cu b/crates/ml/src/cuda_pipeline/iql_value_kernel.cu index f79e32ee8..ae2a21eac 100644 --- a/crates/ml/src/cuda_pipeline/iql_value_kernel.cu +++ b/crates/ml/src/cuda_pipeline/iql_value_kernel.cu @@ -8,31 +8,27 @@ * Architecture: 2-hidden-layer MLP with SiLU activation. * 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=(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) + * Launch configs: + * Forward + loss: grid=(B, 1, 1), block=(256, 1, 1) + * Backward (per-sample): grid=(B, 1, 1), block=(256, 1, 1) + * Weight grad reduce: grid=(ceil(P/256), 1, 1), block=(256, 1, 1) + * Grad norm phase 1: grid=(ceil(P/256), 1, 1), block=(256, 1, 1) + * Grad norm phase 2: grid=(1, 1, 1), block=(1, 1, 1) + * Adam: grid=(ceil(P/256), 1, 1), block=(256, 1, 1) + * Forward-only: grid=(B, 1, 1), block=(256, 1, 1) + * Advantage weights: grid=(ceil(B/256), 1, 1), block=(256, 1, 1) * - * Compile-time defines (injected via NVRTC dim_overrides): + * Compile-time defines (injected via build.rs dim_overrides): * STATE_DIM -- tensor-core-aligned state width (e.g. 48 or 56) * VALUE_HIDDEN_DIM -- hidden layer width (default 128) - * IQL_EXPECTILE_TAU -- expectile parameter (default 0.7) * - * ALL storage and computation in native float. - * Common header (common_device_functions.cuh) provides: - * bf16(), bf16_zero(), bf16_one(), bf16_sqrt(), bf16_exp(), bf16_log(), - * bf16_fabs(), bf16_fmax(), bf16_fmin(), bf16_tanh(), bf16_pow(), - * bf16_shfl_xor(), bf16_shfl_down(), bf16_warp_sum(), atomicAddBF16(). + * ALL storage and computation in native f32. Zero atomicAdd — fully deterministic. */ #ifndef VALUE_HIDDEN_DIM #define VALUE_HIDDEN_DIM 128 #endif -#ifndef IQL_EXPECTILE_TAU -#define IQL_EXPECTILE_TAU 0.7f -#endif - /* Parameter sizes that don't depend on state_dim */ #define V_B1_SIZE (VALUE_HIDDEN_DIM) #define V_W2_SIZE (VALUE_HIDDEN_DIM * VALUE_HIDDEN_DIM) @@ -56,20 +52,20 @@ __device__ __forceinline__ void iql_compute_offsets( } /* ------------------------------------------------------------------ */ -/* SiLU activation: x * sigmoid(x) -- native BF16 */ +/* SiLU activation: x * sigmoid(x) */ /* ------------------------------------------------------------------ */ __device__ __forceinline__ float silu(float x) { - return x / (bf16_one() + bf16_exp(-x)); + return x / (1.0f + expf(-x)); } __device__ __forceinline__ float silu_grad(float x) { - float s = bf16_one() / (bf16_one() + bf16_exp(-x)); - return s * (bf16_one() + x * (bf16_one() - s)); + float s = 1.0f / (1.0f + expf(-x)); + return s * (1.0f + x * (1.0f - s)); } -/* Block-level bf16 sum reduction: 256 threads -> single sum. +/* Block-level f32 sum reduction: 256 threads -> single sum. * Returns the result broadcast to ALL threads in the block. */ -__device__ __forceinline__ float iql_bf16_block_sum( +__device__ __forceinline__ float iql_block_sum( float val, float* warp_sums /* [8] in shared memory */ ) { int tid = threadIdx.x; @@ -77,13 +73,13 @@ __device__ __forceinline__ float iql_bf16_block_sum( 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); + val += __shfl_xor_sync(0xFFFFFFFF, val, offset); if (warp_lane == 0) warp_sums[warp_id] = val; __syncthreads(); if (warp_id == 0) { - float v = (warp_lane < (blockDim.x / 32)) ? warp_sums[warp_lane] : bf16_zero(); + float v = (warp_lane < (blockDim.x / 32)) ? warp_sums[warp_lane] : 0.0f; for (int off = 16; off > 0; off >>= 1) - v = v + bf16_shfl_xor(0xFFFFFFFF, v, off); + v += __shfl_xor_sync(0xFFFFFFFF, v, off); warp_sums[0] = v; } __syncthreads(); @@ -102,12 +98,11 @@ __device__ __forceinline__ float iql_bf16_block_sum( * Inputs: * states [B, STATE_DIM] -- input states * q_values [B] -- target Q(s,a) from DQN - * params [V_TOTAL_PARAMS] -- flat weight buffer + * params [total_params] -- flat weight buffer * * Outputs: * v_out [B] -- V(s) predictions * loss_out [B] -- per-sample expectile loss - * total_loss [1] -- batch-mean loss (atomicAdd) * save_pre1 [B, H] -- pre-activation layer 1 (for backward) * save_pre2 [B, H] -- pre-activation layer 2 (for backward) * save_h1 [B, H] -- post-activation layer 1 (for backward) @@ -120,13 +115,13 @@ void iql_forward_loss_kernel( const float* __restrict__ params, float* __restrict__ v_out, float* __restrict__ loss_out, - float* __restrict__ total_loss, float* __restrict__ save_pre1, float* __restrict__ save_pre2, float* __restrict__ save_h1, float* __restrict__ save_h2, int batch_size, - int state_dim /* runtime state dimension */ + int state_dim, + float expectile_tau ) { int sample = blockIdx.x; @@ -156,7 +151,7 @@ void iql_forward_loss_kernel( for (int j = tid; j < VALUE_HIDDEN_DIM; j += 256) { float acc = b1[j]; for (int k = 0; k < state_dim; k++) { - acc = acc + w1[j * state_dim + k] * x[k]; + acc += w1[j * state_dim + k] * x[k]; } pre1_ptr[j] = acc; h1_ptr[j] = silu(acc); @@ -170,7 +165,7 @@ void iql_forward_loss_kernel( for (int j = tid; j < VALUE_HIDDEN_DIM; j += 256) { float acc = b2[j]; for (int k = 0; k < VALUE_HIDDEN_DIM; k++) { - acc = acc + w2[j * VALUE_HIDDEN_DIM + k] * h1_ptr[k]; + acc += w2[j * VALUE_HIDDEN_DIM + k] * h1_ptr[k]; } pre2_ptr[j] = acc; h2_ptr[j] = silu(acc); @@ -178,50 +173,37 @@ void iql_forward_loss_kernel( __syncthreads(); /* ---- Output layer: Linear(H -> 1) ---- */ - /* Block-reduce the dot product across all 256 threads */ - float v_acc = bf16_zero(); + float v_acc = 0.0f; for (int k = tid; k < VALUE_HIDDEN_DIM; k += 256) { - v_acc = v_acc + w3[k] * h2_ptr[k]; + v_acc += w3[k] * h2_ptr[k]; } - /* Block-level sum reduction */ - v_acc = iql_bf16_block_sum(v_acc, warp_sums); + v_acc = iql_block_sum(v_acc, warp_sums); - /* Thread 0 has the full sum */ if (tid == 0) { float v_val = v_acc + b3[0]; v_out[sample] = v_val; /* Expectile loss: L_tau(u) = |tau - 1(u<0)| * u^2 */ - float u = bf16(q_values[sample]) - v_val; - float bf_tau = bf16(IQL_EXPECTILE_TAU); - float weight = (u >= bf16_zero()) ? bf_tau : (bf16_one() - bf_tau); - float sample_loss = weight * u * u; - loss_out[sample] = sample_loss; - - /* Accumulate batch-mean loss */ - atomicAddBF16(total_loss, sample_loss / bf16((float)batch_size)); + float u = q_values[sample] - v_val; + float weight = (u >= 0.0f) ? expectile_tau : (1.0f - expectile_tau); + loss_out[sample] = weight * u * u; } } /* ------------------------------------------------------------------ */ -/* Backward Kernel */ +/* Per-Sample Backward Kernel (zero atomicAdd) */ /* ------------------------------------------------------------------ */ /** * Per-sample backward pass through V(s) MLP. * - * 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). + * Each block (one sample) writes to its own slice in grads_per_sample, + * eliminating all cross-sample atomicAdd. Within a block, each thread + * handles non-overlapping parameter indices — no write conflicts. * - * Gradient flow: - * dL/dV = -2 * weight * u (where u = Q - V, weight = tau or 1-tau) - * dV/dw3[k] = h2[k] - * dV/db3 = 1 - * ... backprop through SiLU and linear layers ... + * A separate iql_weight_grad_reduce kernel sums across samples. */ extern "C" __global__ -void iql_backward_kernel( +void iql_backward_per_sample( const float* __restrict__ states, const float* __restrict__ q_values, const float* __restrict__ v_out, @@ -230,9 +212,11 @@ void iql_backward_kernel( const float* __restrict__ save_pre2, const float* __restrict__ save_h1, const float* __restrict__ save_h2, - float* __restrict__ grads, + float* __restrict__ grads_per_sample, /* [B, total_params] */ int batch_size, - int state_dim /* runtime state dimension */ + int state_dim, + int total_params, + float expectile_tau ) { int sample = blockIdx.x; @@ -243,16 +227,17 @@ void iql_backward_kernel( 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); - const float* w1 = params + off_w1; const float* w2 = params + off_w2; const float* w3 = params + off_w3; - float* gw1 = grads + off_w1; - float* gb1 = grads + off_b1; - float* gw2 = grads + off_w2; - float* gb2 = grads + off_b2; - float* gw3 = grads + off_w3; - float* gb3 = grads + off_b3; + /* Per-sample gradient slice — no overlap with other samples */ + float* g = grads_per_sample + sample * total_params; + float* gw1 = g + off_w1; + float* gb1 = g + off_b1; + float* gw2 = g + off_w2; + float* gb2 = g + off_b2; + float* gw3 = g + off_w3; + float* gb3 = g + off_b3; const float* x = states + sample * state_dim; const float* h1 = save_h1 + sample * VALUE_HIDDEN_DIM; @@ -261,72 +246,175 @@ void iql_backward_kernel( const float* pre2 = save_pre2 + sample * VALUE_HIDDEN_DIM; float v_val = v_out[sample]; - float q_val = bf16(q_values[sample]); + float q_val = q_values[sample]; /* dL/dV = -2 * weight * (Q - V) / batch_size */ float u = q_val - v_val; - float bf_tau = bf16(IQL_EXPECTILE_TAU); - float weight = (u >= bf16_zero()) ? bf_tau : (bf16_one() - bf_tau); - float dldv = bf16(-2.0f) * weight * u / bf16((float)batch_size); + float weight = (u >= 0.0f) ? expectile_tau : (1.0f - expectile_tau); + float dldv = -2.0f * weight * u / (float)batch_size; /* ---- Output layer gradient: dL/dw3, dL/db3 ---- */ - /* dV/dw3[k] = h2[k], dV/db3 = 1 */ for (int k = tid; k < VALUE_HIDDEN_DIM; k += 256) { - atomicAddBF16(&gw3[k], dldv * h2[k]); + gw3[k] = dldv * h2[k]; } if (tid == 0) { - atomicAddBF16(&gb3[0], dldv); + 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 = tid; j < VALUE_HIDDEN_DIM; j += 256) { float dh2_j = dldv * w3[j]; float dpre2_j = dh2_j * silu_grad(pre2[j]); - /* dL/db2[j] = dpre2_j */ - atomicAddBF16(&gb2[j], dpre2_j); + gb2[j] = dpre2_j; - /* dL/dw2[j, k] = dpre2_j * h1[k] */ for (int k = 0; k < VALUE_HIDDEN_DIM; k++) { - atomicAddBF16(&gw2[j * VALUE_HIDDEN_DIM + k], dpre2_j * h1[k]); + gw2[j * VALUE_HIDDEN_DIM + k] = dpre2_j * h1[k]; } } __syncthreads(); /* ---- Backprop through layer 1 ---- */ - /* dL/dh1[k] = sum_j(dL/dpre2[j] * w2[j, k]) */ for (int k = tid; k < VALUE_HIDDEN_DIM; k += 256) { - float dh1_k = bf16_zero(); + float dh1_k = 0.0f; for (int j = 0; j < VALUE_HIDDEN_DIM; j++) { float dh2_j = dldv * w3[j]; float dpre2_j = dh2_j * silu_grad(pre2[j]); - dh1_k = dh1_k + dpre2_j * w2[j * VALUE_HIDDEN_DIM + k]; + dh1_k += dpre2_j * w2[j * VALUE_HIDDEN_DIM + k]; } float dpre1_k = dh1_k * silu_grad(pre1[k]); - /* dL/db1[k] = dpre1_k */ - atomicAddBF16(&gb1[k], dpre1_k); + gb1[k] = dpre1_k; - /* dL/dw1[k, d] = dpre1_k * x[d] */ for (int d = 0; d < state_dim; d++) { - atomicAddBF16(&gw1[k * state_dim + d], dpre1_k * x[d]); + gw1[k * state_dim + d] = dpre1_k * x[d]; } } } +/* ------------------------------------------------------------------ */ +/* Weight Gradient Reduce (deterministic cross-sample sum) */ +/* ------------------------------------------------------------------ */ +/** + * Sum per-sample gradients into the shared gradient buffer. + * + * grads[i] = sum_{b=0}^{B-1} grads_per_sample[b * total_params + i] + * + * Fixed summation order — fully deterministic. + * Launch: grid=(ceil(total_params/256)), block=256. + */ +extern "C" __global__ +void iql_weight_grad_reduce( + const float* __restrict__ grads_per_sample, /* [B, total_params] */ + float* __restrict__ grads, /* [total_params] */ + int batch_size, + int total_params +) +{ + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= total_params) return; + + float sum = 0.0f; + for (int b = 0; b < batch_size; b++) { + sum += grads_per_sample[b * total_params + i]; + } + grads[i] = sum; +} + +/* ------------------------------------------------------------------ */ +/* Loss Reduce (deterministic sequential sum) */ +/* ------------------------------------------------------------------ */ +/** + * Reduce per-sample losses to a single batch-mean scalar. + * + * total_loss[0] = (1/B) * sum(loss_out[b]) + * + * Sequential for determinism. Launch: grid=1, block=1. + */ +extern "C" __global__ +void iql_loss_reduce( + const float* __restrict__ loss_out, /* [B] per-sample loss */ + float* __restrict__ total_loss, /* [1] output */ + int batch_size +) +{ + float sum = 0.0f; + for (int b = 0; b < batch_size; b++) { + sum += loss_out[b]; + } + total_loss[0] = sum / (float)batch_size; +} + +/* ------------------------------------------------------------------ */ +/* Gradient Norm — Phase 1 (per-block partial sums, no atomicAdd) */ +/* ------------------------------------------------------------------ */ +/** + * Each block computes a partial sum-of-squares and writes to block_sums. + * No atomicAdd — one output per block. + */ +extern "C" __global__ +void iql_grad_norm_phase1( + const float* __restrict__ grads, + float* __restrict__ block_sums, /* [num_blocks] */ + int total_params +) +{ + int tid = blockIdx.x * blockDim.x + threadIdx.x; + + float partial = 0.0f; + for (int i = tid; i < total_params; i += gridDim.x * blockDim.x) { + float g = grads[i]; + partial += g * g; + } + + /* Block-level reduction via shared memory */ + __shared__ float shared[256]; + shared[threadIdx.x] = partial; + __syncthreads(); + + for (int s = blockDim.x / 2; s > 0; s >>= 1) { + if (threadIdx.x < s) + shared[threadIdx.x] += shared[threadIdx.x + s]; + __syncthreads(); + } + + if (threadIdx.x == 0) + block_sums[blockIdx.x] = shared[0]; +} + +/* ------------------------------------------------------------------ */ +/* Gradient Norm — Phase 2 (sequential sum of partials) */ +/* ------------------------------------------------------------------ */ +/** + * Sequential sum of block partial sums → sum-of-squares output. + * Adam kernel reads this as raw sum-of-squares (takes sqrt internally). + * + * Launch: grid=1, block=1. + */ +extern "C" __global__ +void iql_grad_norm_phase2( + const float* __restrict__ block_sums, + float* __restrict__ norm_out, /* [1] sum-of-squares */ + int num_blocks +) +{ + float sum = 0.0f; + for (int i = 0; i < num_blocks; i++) + sum += block_sums[i]; + norm_out[0] = sum; +} + /* ------------------------------------------------------------------ */ /* Adam Update Kernel */ /* ------------------------------------------------------------------ */ /** * AdamW optimizer step over the flat parameter buffer. * - * grid = ceil(V_TOTAL_PARAMS / 256), block = 256. + * grid = ceil(total_params / 256), block = 256. * Each thread updates one parameter. * - * Supports gradient clipping via max_grad_norm (applied before Adam step - * using the pre-computed grad_norm scalar). + * grad_norm_sq_buf contains raw sum-of-squares from phase1+phase2. + * Gradient clipping: scale = max_grad_norm / sqrt(grad_norm_sq). */ extern "C" __global__ void iql_adam_kernel( @@ -334,7 +422,7 @@ void iql_adam_kernel( float* __restrict__ grads, float* __restrict__ m, /* first moment */ float* __restrict__ v, /* second moment */ - const float* __restrict__ grad_norm_buf, + const float* __restrict__ grad_norm_sq_buf, /* [1] sum-of-squares */ float lr, float beta1, float beta2, @@ -348,81 +436,54 @@ void iql_adam_kernel( int i = blockIdx.x * blockDim.x + threadIdx.x; if (i >= total_params) return; - /* Convert host scalar params on first use */ - float bf_lr = bf16(lr); - float bf_beta1 = bf16(beta1); - float bf_beta2 = bf16(beta2); - float bf_eps = bf16(eps); - float bf_wd = bf16(weight_decay); - float bf_max_gn = bf16(max_grad_norm); - - /* Gradient clipping */ - float gn = grad_norm_buf[0]; - float clip_scale = (gn > bf_max_gn && gn > bf16_zero()) ? (bf_max_gn / gn) : bf16_one(); + /* Gradient clipping from sum-of-squares */ + float gn = sqrtf(grad_norm_sq_buf[0]); + float clip_scale = (gn > max_grad_norm && gn > 0.0f) ? (max_grad_norm / gn) : 1.0f; float g = grads[i] * clip_scale; /* AdamW: decoupled weight decay */ float p = params[i]; - p = p - bf_lr * bf_wd * p; + p -= lr * weight_decay * p; /* Adam moment updates */ - float m_i = bf_beta1 * m[i] + (bf16_one() - bf_beta1) * g; - float v_i = bf_beta2 * v[i] + (bf16_one() - bf_beta2) * g * g; + float m_i = beta1 * m[i] + (1.0f - beta1) * g; + float v_i = beta2 * v[i] + (1.0f - beta2) * g * g; /* Bias correction */ int t = t_buf[0]; - float bf_t = bf16((float)t); - float m_hat = m_i / (bf16_one() - bf16_pow(bf_beta1, bf_t)); - float v_hat = v_i / (bf16_one() - bf16_pow(bf_beta2, bf_t)); + float m_hat = m_i / (1.0f - powf(beta1, (float)t)); + float v_hat = v_i / (1.0f - powf(beta2, (float)t)); /* Parameter update */ - p = p - bf_lr * m_hat / (bf16_sqrt(v_hat) + bf_eps); + p -= lr * m_hat / (sqrtf(v_hat) + eps); params[i] = p; m[i] = m_i; v[i] = v_i; - grads[i] = bf16_zero(); /* zero gradient for next step */ + grads[i] = 0.0f; /* zero gradient for next step */ } /* ------------------------------------------------------------------ */ -/* Gradient Norm Kernel */ +/* Gather Q(s, a_taken) from q_out_buf */ /* ------------------------------------------------------------------ */ /** - * Compute L2 norm of the gradient buffer via parallel reduction. + * Extract Q-value for the taken action from the full Q-value buffer. + * q_taken[b] = q_out[b * total_actions + actions[b]] * - * grid = ceil(V_TOTAL_PARAMS / 256), block = 256. - * Each block reduces 256 elements, atomicAdds the partial sum into grad_norm_buf[0]. - * Caller must zero grad_norm_buf[0] before launch. + * Launch: grid=ceil(B/256), block=256. */ extern "C" __global__ -void iql_grad_norm_kernel( - const float* __restrict__ grads, - float* __restrict__ grad_norm_buf, - int total_params +void iql_gather_q_taken( + const float* __restrict__ q_out, /* [B, total_actions] */ + const int* __restrict__ actions, /* [B] taken action indices */ + float* __restrict__ q_taken, /* [B] output */ + int batch_size, + int total_actions ) { - __shared__ float sdata[256]; - - int tid = threadIdx.x; - int i = blockIdx.x * blockDim.x + threadIdx.x; - - /* Load and square */ - float val = (i < total_params) ? grads[i] : bf16_zero(); - sdata[tid] = val * val; - __syncthreads(); - - /* Tree reduction within block */ - for (int s = 128; s > 0; s >>= 1) { - if (tid < s) { - sdata[tid] = sdata[tid] + sdata[tid + s]; - } - __syncthreads(); - } - - /* Block 0 writes partial sum */ - if (tid == 0) { - atomicAddBF16(grad_norm_buf, bf16_sqrt(sdata[0])); - } + int b = blockIdx.x * blockDim.x + threadIdx.x; + if (b >= batch_size) return; + q_taken[b] = q_out[b * total_actions + actions[b]]; } /* ------------------------------------------------------------------ */ @@ -441,7 +502,7 @@ void iql_forward_kernel( const float* __restrict__ params, float* __restrict__ v_out, int batch_size, - int state_dim /* runtime state dimension */ + int state_dim ) { int sample = blockIdx.x; @@ -470,7 +531,7 @@ void iql_forward_kernel( for (int j = tid; j < VALUE_HIDDEN_DIM; j += 256) { float acc = b1[j]; for (int k = 0; k < state_dim; k++) { - acc = acc + w1[j * state_dim + k] * x[k]; + acc += w1[j * state_dim + k] * x[k]; } sh1[j] = silu(acc); } @@ -480,20 +541,319 @@ void iql_forward_kernel( for (int j = tid; j < VALUE_HIDDEN_DIM; j += 256) { float acc = b2[j]; for (int k = 0; k < VALUE_HIDDEN_DIM; k++) { - acc = acc + w2[j * VALUE_HIDDEN_DIM + k] * sh1[k]; + acc += w2[j * VALUE_HIDDEN_DIM + k] * sh1[k]; } sh2[j] = silu(acc); } __syncthreads(); - /* Output: dot product + bias -- block-level sum reduction */ - float v_acc = bf16_zero(); + /* Output: dot product + bias */ + float v_acc = 0.0f; for (int k = tid; k < VALUE_HIDDEN_DIM; k += 256) { - v_acc = v_acc + w3[k] * sh2[k]; + v_acc += w3[k] * sh2[k]; } - v_acc = iql_bf16_block_sum(v_acc, warp_sums); + v_acc = iql_block_sum(v_acc, warp_sums); if (tid == 0) { v_out[sample] = v_acc + b3[0]; } } + +/* ------------------------------------------------------------------ */ +/* Advantage Weight Kernel */ +/* ------------------------------------------------------------------ */ +/** + * Compute IQL advantage weights for policy modulation: + * w[b] = clamp(exp(beta * (Q_taken[b] - V(s)[b])), 0.01, 100) + * + * Q_taken is gathered from q_out_buf using the taken action index. + * Weights are clamped for numerical stability. + * + * Launch: grid=ceil(B/256), block=256. + */ +extern "C" __global__ +void iql_compute_advantage_weights( + const float* __restrict__ q_out, /* [B, total_actions] */ + const int* __restrict__ actions, /* [B] taken action indices */ + const float* __restrict__ v_out, /* [B] V(s) from IQL */ + float* __restrict__ adv_weights, /* [B] output weights */ + float beta, + int batch_size, + int total_actions +) +{ + int b = blockIdx.x * blockDim.x + threadIdx.x; + if (b >= batch_size) return; + + float q_taken = q_out[b * total_actions + actions[b]]; + float adv = q_taken - v_out[b]; + float w = expf(beta * adv); + /* Clamp for numerical stability */ + adv_weights[b] = fminf(fmaxf(w, 0.01f), 100.0f); +} + +/* ------------------------------------------------------------------ */ +/* TD-Error Modulation Kernel */ +/* ------------------------------------------------------------------ */ +/** + * Modulate TD errors by advantage weights and staleness decay: + * K = max(exp(beta * 3 * max(sigma_adv, 1e-6)), 1.1) + * w = clamp(adv_weights[b], 1/K, K) + * age = (write_pos - indices[b] + capacity) % capacity + * decay = exp(-staleness_lambda * age / max(staleness_tau, 1.0)) + * td_errors[b] *= w * decay + * + * In-place update of td_errors for PER priority recomputation. + * Launch: grid=ceil(B/256), block=256. + */ +extern "C" __global__ +void iql_modulate_td_errors( + float* __restrict__ td_errors, /* [B] in-place modulated */ + const float* __restrict__ adv_weights, /* [B] advantage weights */ + const int* __restrict__ indices, /* [B] buffer indices for staleness */ + float sigma_adv, + float beta, + float staleness_lambda, + float staleness_tau, + int write_pos, + int capacity, + int batch_size +) +{ + int b = blockIdx.x * blockDim.x + threadIdx.x; + if (b >= batch_size) return; + + float K = expf(beta * 3.0f * fmaxf(sigma_adv, 1e-6f)); + K = fmaxf(K, 1.1f); + float inv_K = 1.0f / K; + + float w = fminf(fmaxf(adv_weights[b], inv_K), K); + + int age = (write_pos - indices[b] + capacity) % capacity; + float decay = expf(-staleness_lambda * (float)age / fmaxf(staleness_tau, 1.0f)); + + td_errors[b] *= w * decay; +} + +/* ------------------------------------------------------------------ */ +/* Advantage Variance Reduction Kernel */ +/* ------------------------------------------------------------------ */ +/** + * Compute mean and variance of advantage weights over a batch. + * Single-thread kernel — launched with grid=1, block=1. + * stats[0] = mean(adv_weights) + * stats[1] = var(adv_weights) + * + * Used to adaptively modulate the TD-error clamp range K. + * Launch: grid=1, block=1. + */ +extern "C" __global__ +void iql_adv_variance_reduce( + const float* __restrict__ adv_weights, /* [B] */ + float* __restrict__ stats, /* [2]: mean, variance */ + int batch_size +) +{ + float sum = 0.0f; + for (int b = 0; b < batch_size; b++) + sum += adv_weights[b]; + float mean = sum / (float)batch_size; + + float var_sum = 0.0f; + for (int b = 0; b < batch_size; b++) { + float d = adv_weights[b] - mean; + var_sum += d * d; + } + stats[0] = mean; + stats[1] = var_sum / (float)batch_size; +} + +/* ------------------------------------------------------------------ */ +/* Per-Sample C51 Atom Support Kernel */ +/* ------------------------------------------------------------------ */ +/** + * Compute per-sample [v_min, v_max, delta_z] for C51 atom projection: + * spread = max_{a} |Q(s,a) - V(s)| + * half_w = spread * (1 + gamma) + * v_min[b] = V(s) - half_w + * v_max[b] = V(s) + half_w + * delta_z = (v_max - v_min) / (num_atoms - 1) + * + * Output: per_sample_support[b*3+0]=v_min, [b*3+1]=v_max, [b*3+2]=delta_z + * Launch: grid=ceil(B/256), block=256. + */ +extern "C" __global__ +void iql_compute_per_sample_support( + const float* __restrict__ v_out, /* [B] */ + const float* __restrict__ q_out, /* [B, total_actions] */ + float* __restrict__ per_sample_support, /* [B*3] */ + float gamma, + int batch_size, + int total_actions, + int num_atoms +) +{ + int b = blockIdx.x * blockDim.x + threadIdx.x; + if (b >= batch_size) return; + + float v = v_out[b]; + const float* q = q_out + b * total_actions; + + float spread = 0.0f; + for (int a = 0; a < total_actions; a++) { + float dist = fabsf(q[a] - v); + spread = fmaxf(spread, dist); + } + + float half_w = spread * (1.0f + gamma); + float v_min = v - half_w; + float v_max = v + half_w; + float delta_z = (v_max - v_min) / (float)(num_atoms - 1); + + per_sample_support[b * 3 + 0] = v_min; + per_sample_support[b * 3 + 1] = v_max; + per_sample_support[b * 3 + 2] = delta_z; +} + +/* ------------------------------------------------------------------ */ +/* Per-Branch Advantage Scale Kernel */ +/* ------------------------------------------------------------------ */ +/** + * Compute per-branch normalized advantage scales for 4-branch DQN: + * For each branch d, marginalise Q over all joints where sub_d == a_d. + * a_branch[d] = |Q_marginal - V(s)| normalised by max across branches. + * + * Output: branch_scales[b*4+d] in [0,1] for each sample and branch. + * Launch: grid=ceil(B/256), block=256. + */ +extern "C" __global__ +void iql_per_branch_advantage( + const float* __restrict__ q_out, /* [B, total_actions] */ + const float* __restrict__ v_out, /* [B] */ + const int* __restrict__ actions, /* [B] factored action indices */ + float* __restrict__ branch_scales, /* [B*4] */ + int batch_size, + int total_actions, + int b0_size, int b1_size, int b2_size, int b3_size +) +{ + int b = blockIdx.x * blockDim.x + threadIdx.x; + if (b >= batch_size) return; + + float v = v_out[b]; + const float* q = q_out + b * total_actions; + + int factored = actions[b]; + int max_act = b0_size * b1_size * b2_size * b3_size; + if (factored < 0 || factored >= max_act) factored = 0; + int a0 = factored / (b1_size * b2_size * b3_size); + int a1 = (factored / (b2_size * b3_size)) % b1_size; + int a2 = (factored / b3_size) % b2_size; + int a3 = factored % b3_size; + int branch_actions[4] = { a0, a1, a2, a3 }; + int branch_sizes[4] = { b0_size, b1_size, b2_size, b3_size }; + + float a_branch[4]; + float max_a = 0.0f; + + for (int d = 0; d < 4; d++) { + int a_d = branch_actions[d]; + float q_sum = 0.0f; + int count = 0; + + for (int joint = 0; joint < total_actions; joint++) { + int sub_d; + if (d == 0) sub_d = joint / (b1_size * b2_size * b3_size); + else if (d == 1) sub_d = (joint / (b2_size * b3_size)) % b1_size; + else if (d == 2) sub_d = (joint / b3_size) % b2_size; + else sub_d = joint % b3_size; + + if (sub_d == a_d) { + q_sum += q[joint]; + count++; + } + } + + float q_marginal = (count > 0) ? (q_sum / (float)count) : v; + a_branch[d] = fabsf(q_marginal - v); + max_a = fmaxf(max_a, a_branch[d]); + } + + float inv_max = (max_a > 1e-8f) ? (1.0f / max_a) : 1.0f; + for (int d = 0; d < 4; d++) { + branch_scales[b * 4 + d] = a_branch[d] * inv_max; + } +} + +/* ------------------------------------------------------------------ */ +/* Expectile Gap Kernels */ +/* ------------------------------------------------------------------ */ +/** + * Compute element-wise gap between high and low IQL expectile values: + * gap[b] = v_high[b] - v_low[b] + * + * Used for per-sample adaptive epsilon computation. + * Launch: grid=ceil(B/256), block=256. + */ +extern "C" __global__ +void iql_expectile_gap( + const float* __restrict__ v_high, /* [B] high-expectile value estimates */ + const float* __restrict__ v_low, /* [B] low-expectile value estimates */ + float* __restrict__ gap, /* [B] output gaps */ + int batch_size +) +{ + int b = blockIdx.x * blockDim.x + threadIdx.x; + if (b >= batch_size) return; + gap[b] = v_high[b] - v_low[b]; +} + +/** + * Reduce gap array to scalar mean. + * Single-thread kernel — launched with grid=1, block=1. + * gap_mean[0] = mean(gap[0..batch_size-1]) + * + * Launch: grid=1, block=1. + */ +extern "C" __global__ +void iql_gap_mean_reduce( + const float* __restrict__ gap, /* [B] */ + float* __restrict__ gap_mean, /* [1] scalar output */ + int batch_size +) +{ + float sum = 0.0f; + for (int b = 0; b < batch_size; b++) + sum += gap[b]; + gap_mean[0] = sum / (float)batch_size; +} + +/* ------------------------------------------------------------------ */ +/* Per-Sample Adaptive Epsilon Kernel */ +/* ------------------------------------------------------------------ */ +/** + * Compute per-sample exploration epsilon from normalised expectile gap: + * x = gap[b] / max(gap_mean[0], 1e-8) - 1.0 + * sig = sigmoid(x) + * per_sample_eps[b] = base_epsilon * sig + * + * Samples with wider value uncertainty get higher exploration probability. + * Launch: grid=ceil(B/256), block=256. + */ +extern "C" __global__ +void iql_compute_per_sample_epsilon( + const float* __restrict__ gap, /* [B] per-sample expectile gaps */ + const float* __restrict__ gap_mean, /* [1] scalar mean gap */ + float* __restrict__ per_sample_eps, /* [B] output epsilon values */ + float base_epsilon, + int batch_size +) +{ + int b = blockIdx.x * blockDim.x + threadIdx.x; + if (b >= batch_size) return; + + float gm = fmaxf(gap_mean[0], 1e-8f); + float x = gap[b] / gm - 1.0f; + float sig = 1.0f / (1.0f + expf(-x)); + per_sample_eps[b] = base_epsilon * sig; +}