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

feat: add 8 IQL CUDA kernels — TD modulation, adv variance, per-sample support/epsilon, branch advantage, expectile gap

Browse Source 
Appends to iql_value_kernel.cu: iql_modulate_td_errors, iql_adv_variance_reduce,
iql_compute_per_sample_support, iql_per_branch_advantage, iql_expectile_gap,
iql_gap_mean_reduce, iql_compute_per_sample_epsilon. All f32, zero atomics, fully
deterministic. Builds clean.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-13 14:59:37 +02:00
parent f0cda76229
commit a8cb75894c
1 changed files with 505 additions and 145 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

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