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

refactor: remove all bf16 legacy wrappers from CUDA kernels — pure native f32

Browse Source 
Replaced all bf16()/bf16_zero()/bf16_one()/bf16_exp/bf16_sqrt/bf16_fmax/
bf16_fmin/bf16_fabs/bf16_log/bf16_shfl_xor/bf16_shfl_down/atomicAddBF16/
f32_to_bf16 wrapper calls with their native f32 equivalents across all 25
.cu kernel files. Updated stale comments. Fixed monitoring_kernel.cu atomic
CAS helpers to use 32-bit int CAS instead of broken 16-bit. Build verified.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-13 16:44:20 +02:00
parent 25eb2f5e20
commit 3649ff9fb5
25 changed files with 548 additions and 608 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -26,10 +26,8 @@
* Weight gradients accumulated via atomicAdd across batch dimension.
*
* 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().
* Common header (common_device_functions.cuh) provides native f32 helpers
* (wrappers are now identity functions, kept for backward compatibility).
*/
/* -- Configuration constants (overridden by NVRTC injection) ---------- */
@@ -41,14 +39,12 @@
#endif
#define ATTN_HEAD_DIM (ATTN_STATE_DIM / ATTN_NUM_HEADS)
/* BF16 broadcast shuffle (lane src) — not in common header */
/* F32 broadcast shuffle (lane src) */
__device__ __forceinline__ float bf16_shfl(unsigned mask, float val, int src) {
unsigned short raw = *(unsigned short*)&val;
raw = __shfl_sync(mask, raw, src);
return *(float*)&raw;
return __shfl_sync(mask, val, src);
}
/* Block-level bf16 max reduction: 256 threads -> single max value.
/* Block-level f32 max reduction: 256 threads -> single max value.
* Returns the result broadcast to ALL threads in the block. */
__device__ __forceinline__ float bf16_block_max_bwd(
float val, float* warp_sums /* [8] in shared memory */
@@ -57,20 +53,20 @@ __device__ __forceinline__ float bf16_block_max_bwd(
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));
val = fmaxf(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(-1e30f);
float v = (warp_lane < (blockDim.x / 32)) ? warp_sums[warp_lane] : -1e30f;
for (int off = 16; off > 0; off >>= 1)
v = bf16_fmax(v, bf16_shfl_xor(0xFFFFFFFF, v, off));
v = fmaxf(v, __shfl_xor_sync(0xFFFFFFFF, v, off));
warp_sums[0] = v;
}
__syncthreads();
return warp_sums[0];
}
/* 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 bf16_block_sum_bwd(
float val, float* warp_sums /* [8] in shared memory */
@@ -79,13 +75,13 @@ __device__ __forceinline__ float bf16_block_sum_bwd(
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);
val = 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 = v + __shfl_xor_sync(0xFFFFFFFF, v, off);
warp_sums[0] = v;
}
__syncthreads();
@@ -179,26 +175,26 @@ extern "C" __global__ void attention_backward_kernel(
int gf = head_start + f;
float q = fwd_Q[gf];
float k = fwd_K[gf];
concat[gf] = ((q) * (k)) / (bf16_sqrt(bf16((float)Dh)));
concat[gf] = ((q) * (k)) / sqrtf((float)Dh);
}
__syncthreads();
/* Softmax over head */
float max_s = bf16(-1e4f);
float max_s = -1e4f;
for (int f = tid; f < Dh; f += 256) {
float s = concat[head_start + f];
max_s = bf16_fmax(max_s, s);
max_s = fmaxf(max_s, s);
}
max_s = bf16_block_max_bwd(max_s, warp_reduce);
float local_sum = bf16_zero();
float local_sum = 0.0f;
for (int f = tid; f < Dh; f += 256) {
float e = bf16_exp(concat[head_start + f] - max_s);
float e = expf(concat[head_start + f] - max_s);
concat[head_start + f] = e;
local_sum = local_sum + e;
}
float total_sum = bf16_block_sum_bwd(local_sum, warp_reduce);
float inv_sum = bf16_one() / (total_sum + bf16(1e-8f));
float inv_sum = 1.0f / (total_sum + 1e-8f);
/* Apply attention to saved V, store as concat */
for (int f = tid; f < Dh; f += 256) {
@@ -229,21 +225,21 @@ extern "C" __global__ void attention_backward_kernel(
* =================================================================== */
/* Compute mean of pre_ln */
float local_mean = bf16_zero();
float local_mean = 0.0f;
for (int f = tid; f < D; f += 256)
local_mean = local_mean + d_proj[f]; /* d_proj holds pre_ln here */
float total_mean = bf16_block_sum_bwd(local_mean, warp_reduce);
float mean = total_mean / bf16((float)D);
float mean = total_mean / (float)D;
/* Compute variance of pre_ln */
float local_var = bf16_zero();
float local_var = 0.0f;
for (int f = tid; f < D; f += 256) {
float diff = d_proj[f] - mean;
local_var = local_var + diff * diff;
}
float total_var = bf16_block_sum_bwd(local_var, warp_reduce);
float var = total_var / bf16((float)D);
float inv_std = bf16_one() / bf16_sqrt(var + bf16(1e-5f));
float var = total_var / (float)D;
float inv_std = 1.0f / sqrtf(var + 1e-5f);
/* LN backward:
* d_prenorm[f] = (1/D) * inv_std * (D * dy[f]*gamma[f]
@@ -251,8 +247,8 @@ extern "C" __global__ void attention_backward_kernel(
* where x_hat = (pre_ln - mean) * inv_std */
/* First pass: accumulate sum(dy*gamma) and sum(dy*gamma*x_hat) */
float sum_dy_gamma = bf16_zero();
float sum_dy_gamma_xhat = bf16_zero();
float sum_dy_gamma = 0.0f;
float sum_dy_gamma_xhat = 0.0f;
for (int f = tid; f < D; f += 256) {
float x_hat = (d_proj[f] - mean) * inv_std; /* d_proj holds pre_ln */
float dy_g = dy[f] * ln_gamma[f];
@@ -264,13 +260,13 @@ extern "C" __global__ void attention_backward_kernel(
/* Second pass: compute d_prenorm, accumulate LN param grads.
* Overwrite d_proj with d_prenorm (we're done with pre_ln). */
float bf16_D = bf16((float)D);
float f32_D = (float)D;
for (int f = tid; f < D; f += 256) {
float x_hat = (d_proj[f] - mean) * inv_std;
float dy_g = dy[f] * ln_gamma[f];
float d_prenorm = inv_std / bf16_D *
(bf16_D * dy_g - sum_dy_gamma - x_hat * sum_dy_gamma_xhat);
float d_prenorm = inv_std / f32_D *
(f32_D * dy_g - sum_dy_gamma - x_hat * sum_dy_gamma_xhat);
d_proj[f] = d_prenorm; /* overwrite: now d_proj holds d_prenorm */
@@ -292,7 +288,7 @@ extern "C" __global__ void attention_backward_kernel(
/* Initialize d_input with residual gradient */
for (int f = tid; f < D; f += 256)
atomicAddBF16(&dx[f], d_proj[f]); /* d_residual via atomicAdd (heads also add) */
atomicAdd(&dx[f], d_proj[f]); /* d_residual via atomicAdd (heads also add) */
__syncthreads();
/* ===================================================================
@@ -323,7 +319,7 @@ extern "C" __global__ void attention_backward_kernel(
/* Phase 2: compute d_concat and store back in d_proj */
for (int j = tid; j < D; j += 256) {
float d_c = bf16_zero();
float d_c = 0.0f;
for (int f = 0; f < D; f++)
d_c = d_c + d_proj[f] * W_O[j * D + f];
/* Temporarily store in register; we need to sync before overwriting d_proj */
@@ -358,31 +354,31 @@ extern "C" __global__ void attention_backward_kernel(
int gf = head_start + f;
float q = fwd_Q[gf];
float k = fwd_K[gf];
concat[gf] = ((q) * (k)) / (bf16_sqrt(bf16((float)Dh)));
concat[gf] = ((q) * (k)) / sqrtf((float)Dh);
}
__syncthreads();
/* Softmax over scores */
float max_s = bf16(-1e4f);
float max_s = -1e4f;
for (int f = tid; f < Dh; f += 256) {
float s = concat[head_start + f];
max_s = bf16_fmax(max_s, s);
max_s = fmaxf(max_s, s);
}
max_s = bf16_block_max_bwd(max_s, warp_reduce);
float sm_sum = bf16_zero();
float sm_sum = 0.0f;
for (int f = tid; f < Dh; f += 256) {
float e = bf16_exp(concat[head_start + f] - max_s);
float e = expf(concat[head_start + f] - max_s);
concat[head_start + f] = e;
sm_sum = sm_sum + e;
}
float total_sm = bf16_block_sum_bwd(sm_sum, warp_reduce);
float inv_sum = bf16_one() / (total_sm + bf16(1e-8f));
float inv_sum = 1.0f / (total_sm + 1e-8f);
/* concat[gf] now has unnorm exp. attn[f] = concat[gf] * inv_sum */
/* Compute sum(d_attn * attn) for softmax backward */
float sum_da_a = bf16_zero();
float sum_da_a = 0.0f;
for (int f = tid; f < Dh; f += 256) {
int gf = head_start + f;
float attn = concat[gf] * inv_sum;
@@ -394,7 +390,7 @@ extern "C" __global__ void attention_backward_kernel(
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 */
float sqrt_Dh = bf16_sqrt(bf16((float)Dh));
float sqrt_Dh = sqrtf((float)Dh);
for (int f = tid; f < Dh; f += 256) {
int gf = head_start + f;
float attn = concat[gf] * inv_sum;
@@ -438,7 +434,7 @@ extern "C" __global__ void attention_backward_kernel(
float grad_j = d_q * W_Q[j * D + gf]
+ d_k * W_K[j * D + gf]
+ d_v * W_V[j * D + gf];
atomicAddBF16(&dx[j], grad_j);
atomicAdd(&dx[j], grad_j);
}
}
__syncthreads();
@@ -537,46 +533,34 @@ extern "C" __global__ void attn_adam_kernel(
int tid = blockIdx.x * blockDim.x + threadIdx.x;
if (tid >= 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);
/* Skip NaN/Inf gradients — leave param, momentum, variance unchanged. */
{
float g_check = (float)grads[tid];
if (!isfinite(g_check)) { grads[tid] = bf16_zero(); return; }
}
float g_check = (float)grads[tid];
if (!isfinite(g_check)) { grads[tid] = 0.0f; return; }
/* Gradient clipping */
float grad_norm = norm[0];
float clip_scale = (grad_norm > bf_max_gn && grad_norm > bf16_zero())
? bf_max_gn / grad_norm : bf16_one();
float clip_scale = (grad_norm > max_grad_norm && grad_norm > 0.0f)
? max_grad_norm / grad_norm : 1.0f;
float g = grads[tid] * clip_scale;
/* AdamW: weight decay applied to params directly */
params[tid] = params[tid] * (bf16_one() - bf_lr * bf_wd);
params[tid] = params[tid] * (1.0f - lr * weight_decay);
/* Moment updates */
float m_val = bf_beta1 * m[tid] + (bf16_one() - bf_beta1) * g;
float v_val = bf_beta2 * v[tid] + (bf16_one() - bf_beta2) * g * g;
float m_val = beta1 * m[tid] + (1.0f - beta1) * g;
float v_val = beta2 * v[tid] + (1.0f - beta2) * g * g;
m[tid] = m_val;
v[tid] = v_val;
/* Bias correction */
int adam_t = adam_t_buf[0];
float bf_t = bf16((float)adam_t);
float bc1 = bf16_one() - bf16_pow(bf_beta1, bf_t);
float bc2 = bf16_one() - bf16_pow(bf_beta2, bf_t);
float m_hat = m_val / (bc1 + bf16(1e-12f));
float v_hat = v_val / (bc2 + bf16(1e-12f));
float bc1 = 1.0f - powf(beta1, (float)adam_t);
float bc2 = 1.0f - powf(beta2, (float)adam_t);
float m_hat = m_val / (bc1 + 1e-12f);
float v_hat = v_val / (bc2 + 1e-12f);
/* Parameter update */
params[tid] = params[tid] - bf_lr * m_hat / (bf16_sqrt(v_hat) + bf_eps);
params[tid] = params[tid] - lr * m_hat / (sqrtf(v_hat) + eps);
/* Zero gradient for next step */
grads[tid] = bf16_zero();
grads[tid] = 0.0f;
}

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

@@ -27,7 +27,7 @@
#endif
#define ATTN_HEAD_DIM (ATTN_STATE_DIM / ATTN_NUM_HEADS)
/* Block-level bf16 max reduction: 256 threads -> single max value.
/* Block-level max reduction: 256 threads -> single max value.
* Returns the result broadcast to ALL threads in the block. */
__device__ __forceinline__ float bf16_block_max(
float val, float* warp_sums /* [8] in shared memory */
@@ -37,20 +37,20 @@ __device__ __forceinline__ float bf16_block_max(
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));
val = fmaxf(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(-1e30f);
float v = (warp_lane < (blockDim.x / 32)) ? warp_sums[warp_lane] : -1e30f;
for (int off = 16; off > 0; off >>= 1)
v = bf16_fmax(v, bf16_shfl_xor(0xFFFFFFFF, v, off));
v = fmaxf(v, __shfl_xor_sync(0xFFFFFFFF, v, off));
warp_sums[0] = v;
}
__syncthreads();
return warp_sums[0];
}
/* Block-level bf16 sum reduction: 256 threads -> single sum.
/* Block-level sum reduction: 256 threads -> single sum.
* Returns the result broadcast to ALL threads in the block. */
__device__ __forceinline__ float bf16_block_sum(
float val, float* warp_sums /* [8] in shared memory */
@@ -60,13 +60,13 @@ __device__ __forceinline__ float bf16_block_sum(
int warp_lane = tid % 32;
/* Intra-warp sum */
for (int offset = 16; offset > 0; offset >>= 1)
val = val + bf16_shfl_xor(0xFFFFFFFF, val, offset);
val = 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 = v + __shfl_xor_sync(0xFFFFFFFF, v, off);
warp_sums[0] = v;
}
__syncthreads();
@@ -153,7 +153,7 @@ extern "C" __global__ void multihead_feature_attention(
save_V[global_f] = v;
/* Attention score: Q[f] * K[f] / sqrt(Dh) */
float score = (q * k) / bf16_sqrt(bf16((float)Dh));
float score = (q * k) / sqrtf((float)Dh);
/* Store for softmax */
shared_concat[head_start + f] = score;
@@ -161,21 +161,21 @@ extern "C" __global__ void multihead_feature_attention(
__syncthreads();
/* Softmax over this head's Dh dimensions */
float max_s = bf16(-1e30f);
float max_s = -1e30f;
for (int f = tid; f < Dh; f += 256) {
float s = shared_concat[head_start + f];
max_s = bf16_fmax(max_s, s);
max_s = fmaxf(max_s, s);
}
max_s = bf16_block_max(max_s, warp_reduce);
float local_sum = bf16_zero();
float local_sum = 0.0f;
for (int f = tid; f < Dh; f += 256) {
float e = bf16_exp(shared_concat[head_start + f] - max_s);
float e = expf(shared_concat[head_start + f] - max_s);
shared_concat[head_start + f] = e;
local_sum = local_sum + e;
}
float total_sum = bf16_block_sum(local_sum, warp_reduce);
float inv_sum = bf16_one() / (total_sum + bf16(1e-8f));
float inv_sum = 1.0f / (total_sum + 1e-8f);
/* Apply attention to V (read from saved projection), store in shared_concat */
for (int f = tid; f < Dh; f += 256) {
@@ -200,19 +200,19 @@ extern "C" __global__ void multihead_feature_attention(
/* -- Step 3: Layer normalization -- */
/* mean = avg(out[f]) */
float local_mean = bf16_zero();
float local_mean = 0.0f;
for (int f = tid; f < D; f += 256) local_mean = local_mean + out[f];
float total_mean = bf16_block_sum(local_mean, warp_reduce);
float mean = total_mean / bf16((float)D);
float mean = total_mean / (float)D;
/* variance = avg((out[f] - mean)^2) */
float local_var = bf16_zero();
float local_var = 0.0f;
for (int f = tid; f < D; f += 256) {
float diff = out[f] - mean;
local_var = local_var + diff * diff;
}
float total_var = bf16_block_sum(local_var, warp_reduce);
float inv_std = bf16_one() / bf16_sqrt(total_var / bf16((float)D) + bf16(1e-5f));
float inv_std = 1.0f / sqrtf(total_var / (float)D + 1e-5f);
/* Normalize + affine */
for (int f = tid; f < D; f += 256) {

12
crates/ml/src/cuda_pipeline/backtest_forward_ppo_kernel.cu
View File

@@ -26,7 +26,7 @@
#define N_EXPOSURE_BINS 5
#define ACTIONS_PER_BIN 9 /* 3 order x 3 urgency */
/* BF16-native matvec with optional LeakyReLU.
/* f32 matvec with optional LeakyReLU.
* Weights/bias are float, input/output are float. */
__device__ void matvec_bf16(
const float* __restrict__ W,
@@ -89,16 +89,16 @@ extern "C" __global__ void backtest_forward_ppo_kernel(
/* ---- Stable softmax: logits -> probabilities ---- */
float max_logit = logits[0];
for (int i = 1; i < num_actions; i++) {
max_logit = bf16_fmax(max_logit, logits[i]);
max_logit = fmaxf(max_logit, logits[i]);
}
float probs[MAX_NUM_ACTIONS];
float sum_exp = bf16_zero();
float sum_exp = 0.0f;
for (int i = 0; i < num_actions; i++) {
probs[i] = bf16_exp(logits[i] - max_logit);
probs[i] = expf(logits[i] - max_logit);
sum_exp = sum_exp + probs[i];
}
float inv_sum = bf16_one() / bf16_fmax(sum_exp, bf16(1e-8f));
float inv_sum = 1.0f / fmaxf(sum_exp, 1e-8f);
for (int i = 0; i < num_actions; i++) {
probs[i] = probs[i] * inv_sum;
}
@@ -108,7 +108,7 @@ extern "C" __global__ void backtest_forward_ppo_kernel(
* Exposure bin k (k=0..4) sums probs of actions k*9 .. k*9+8. */
float exposure_scores[N_EXPOSURE_BINS];
for (int k = 0; k < N_EXPOSURE_BINS; k++) {
float sum = bf16_zero();
float sum = 0.0f;
int start = k * ACTIONS_PER_BIN;
for (int j = 0; j < ACTIONS_PER_BIN; j++) {
sum = sum + probs[start + j];

8
crates/ml/src/cuda_pipeline/backtest_forward_supervised_kernel.cu
View File

@@ -28,10 +28,10 @@ extern "C" __global__ void signal_to_action_kernel(
if (idx >= N) return;
float pred = predictions[idx];
float high_bf = bf16(high_threshold_bps);
float low_bf = bf16(low_threshold_bps);
float neg_high = bf16(-high_threshold_bps);
float neg_low = bf16(-low_threshold_bps);
float high_bf = high_threshold_bps;
float low_bf = low_threshold_bps;
float neg_high = -high_threshold_bps;
float neg_low = -low_threshold_bps;
int action;
if (pred < neg_high) {

6
crates/ml/src/cuda_pipeline/backtest_gather_kernel.cu
View File

@@ -32,7 +32,7 @@
// Per window: return, volatility, volume_trend, momentum
extern "C" __global__ void gather_states(
const float* __restrict__ features, // [n_windows * max_len * feat_dim] (bf16 from data pipeline)
const float* __restrict__ features, // [n_windows * max_len * feat_dim] (f32 from data pipeline)
const float* __restrict__ portfolio, // [n_windows * 8] (f32 from env kernel)
float* states_out, // [n_windows * padded_sd] (f32 for cublasLtMatmul, padded stride)
int n_windows,
@@ -77,7 +77,7 @@ extern "C" __global__ void gather_states(
// Determine market_dim: if OFI is present, market features are feat_dim - ofi_dim
int market_dim = (ofi_dim > 0) ? (feat_dim - ofi_dim) : feat_dim;
// ── 1. Market features [0 .. market_dim) ── bf16 input → f32 output //
// ── 1. Market features [0 .. market_dim) ── f32 input → f32 output //
int i = 0;
for (; i + 3 < market_dim; i += 4) {
states_out[out_base + i] = (__ldg(&features[feat_base + i]));
@@ -208,7 +208,7 @@ extern "C" __global__ void gather_states(
}
}
// ── 4. OFI features (if present) ── bf16 input → f32 output //
// ── 4. OFI features (if present) ── f32 input → f32 output //
if (ofi_dim > 0) {
int ofi_out_base = market_dim + 8 + 16;
// OFI features are stored after market features in the feature vector

2
crates/ml/src/cuda_pipeline/c51_grad_kernel.cu
View File

@@ -14,7 +14,7 @@
extern "C" __global__ void c51_grad_kernel(
const float* __restrict__ current_lp, // [B, 4, NA]
const float* __restrict__ projected, // [B, 4, NA]
const float* __restrict__ is_weights, // [B] f32 (bf16 overflows to Inf)
const float* __restrict__ is_weights, // [B] f32
const int* __restrict__ actions, // [B] factored
float* __restrict__ d_value_logits, // [B, NA] f32
float* __restrict__ d_adv_logits, // [B, (B0+B1+B2+B3)*NA] f32

14
crates/ml/src/cuda_pipeline/c51_loss_kernel.cu
View File

@@ -424,7 +424,7 @@ extern "C" __global__ void c51_loss_batched(
/* Copy and save current log-probs */
for (int j = tid; j < num_atoms; j += BLOCK_THREADS) {
shmem_current_lp[j] = shmem_lp[j];
save_current_lp[save_off + j] = bf16(shmem_lp[j]);
save_current_lp[save_off + j] = shmem_lp[j];
}
__syncthreads();
@@ -525,9 +525,9 @@ extern "C" __global__ void c51_loss_batched(
__syncthreads();
}
/* Save projected (BF16 output) */
/* Save projected (f32 output) */
for (int j = tid; j < num_atoms; j += BLOCK_THREADS)
save_projected[save_off + j] = bf16(shmem_lp[j]);
save_projected[save_off + j] = shmem_lp[j];
__syncthreads();
/* ═══ STEP e: Cross-entropy loss ════════════════════════ */
@@ -570,8 +570,8 @@ extern "C" __global__ void c51_loss_batched(
if (tid == 0) {
float clamped_ce = fminf(avg_ce, MAX_PER_SAMPLE_CE);
float weighted_loss = clamped_ce * is_weight;
per_sample_loss[sample_id] = (weighted_loss);
td_errors[sample_id] = bf16(clamped_ce);
per_sample_loss[sample_id] = weighted_loss;
td_errors[sample_id] = clamped_ce;
/* total_loss and q_divergence reduced by separate deterministic kernel.
* No atomicAdd — fully deterministic training. */
/* q_divergence is monitoring-only (does not affect gradients).
@@ -698,8 +698,8 @@ extern "C" __global__ void c51_mixup_ce(
if (tid == 0) {
float clamped_ce = fminf(avg_ce, MAX_PER_SAMPLE_CE);
float weighted_loss = clamped_ce * is_weight;
per_sample_loss[sample_id] = bf16(weighted_loss);
td_errors[sample_id] = bf16(clamped_ce);
per_sample_loss[sample_id] = weighted_loss;
td_errors[sample_id] = clamped_ce;
/* total_loss reduced by deterministic c51_loss_reduce kernel.
* No atomicAdd — fully deterministic training. */
}

10
crates/ml/src/cuda_pipeline/curiosity_inference_kernel.cu
View File

@@ -10,7 +10,7 @@
*
* Requires common_device_functions.cuh to be prepended (provides MARKET_DIM,
* CUR_INPUT, CUR_HIDDEN, CUR_OUTPUT, DQN_ORDER_ACTIONS, DQN_URGENCY_ACTIONS,
* DQN_NUM_ACTIONS, bf16 helpers).
* DQN_NUM_ACTIONS).
*/
extern "C" __global__ void curiosity_inference_error(
@@ -35,9 +35,9 @@ extern "C" __global__ void curiosity_inference_error(
/* ── Build input: first MARKET_DIM state features + 3-class action one-hot ── */
float input[CUR_INPUT];
for (int i = 0; i < MARKET_DIM; i++) { input[i] = state_bf[i]; }
input[MARKET_DIM + 0] = bf16_zero();
input[MARKET_DIM + 1] = bf16_zero();
input[MARKET_DIM + 2] = bf16_zero();
input[MARKET_DIM + 0] = 0.0f;
input[MARKET_DIM + 1] = 0.0f;
input[MARKET_DIM + 2] = 0.0f;
/* Action to category one-hot — decode exposure from factored action.
* Factored action = exposure_idx * (b1_size * b2_size) + order_idx * b2_size + urgency_idx.
@@ -47,7 +47,7 @@ extern "C" __global__ void curiosity_inference_error(
if (exposure_idx < DQN_NUM_ACTIONS / 2) category = 0; /* Short */
else if (exposure_idx == DQN_NUM_ACTIONS / 2) category = 1; /* Flat */
else category = 2; /* Long */
input[MARKET_DIM + category] = bf16_one();
input[MARKET_DIM + category] = 1.0f;
/* ── Layer 1: hidden = LeakyReLU(w1 * input + b1, alpha=0.01) ── */
float hidden[CUR_HIDDEN];

121
crates/ml/src/cuda_pipeline/curiosity_training_kernel.cu
View File

@@ -10,7 +10,7 @@
*
* Requires common_device_functions.cuh to be prepended for CUR_INPUT/CUR_HIDDEN/CUR_OUTPUT.
*
* ALL-BF16: Native float everywhere. Zero float on GPU.
* All native f32.
*/
/* Total trainable parameters */
@@ -45,7 +45,7 @@ extern "C" __global__ void curiosity_shift_states(
extern "C" __global__ void curiosity_zero_grads(float* grads, int n) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i < n) grads[i] = bf16_zero();
if (i < n) grads[i] = 0.0f;
}
/* ------------------------------------------------------------------ */
@@ -55,15 +55,15 @@ extern "C" __global__ void curiosity_zero_grads(float* grads, int n) {
/* Warp-level sum reduction via butterfly shuffle (all lanes get result) */
__device__ __forceinline__ float warp_sum_cur(float val) {
for (int offset = 16; offset > 0; offset >>= 1)
val = val + bf16_shfl_xor(0xFFFFFFFF, val, offset);
val = val + __shfl_xor_sync(0xFFFFFFFF, val, offset);
return val;
}
/**
* One thread per sample. Computes forward pass, MSE loss, and backpropagates
* gradients via atomicAddBF16 with warp-level pre-reduction (32x fewer atomics).
* gradients via atomicAdd with warp-level pre-reduction (32x fewer atomics).
* Threads within a warp process different samples but accumulate gradients
* to the same weight indices, so warp reduction before atomicAddBF16 is valid.
* to the same weight indices, so warp reduction before atomicAdd is valid.
*/
extern "C" __global__ void curiosity_forward_backward(
const float* __restrict__ states, /* [N, state_dim] */
@@ -92,9 +92,9 @@ extern "C" __global__ void curiosity_forward_backward(
/* Build input: first MARKET_DIM state features + 3-class action one-hot */
float input[CUR_INPUT];
for (int i = 0; i < MARKET_DIM; i++) { input[i] = state_bf[i]; }
input[MARKET_DIM + 0] = bf16_zero();
input[MARKET_DIM + 1] = bf16_zero();
input[MARKET_DIM + 2] = bf16_zero();
input[MARKET_DIM + 0] = 0.0f;
input[MARKET_DIM + 1] = 0.0f;
input[MARKET_DIM + 2] = 0.0f;
/* Action to category one-hot -- decode exposure from factored action.
* Factored action = exposure_idx * (b1*b2) + order_idx * b2 + urgency_idx.
@@ -105,7 +105,7 @@ extern "C" __global__ void curiosity_forward_backward(
if (exposure_idx < DQN_NUM_ACTIONS / 2) category = 0; /* Short */
else if (exposure_idx == DQN_NUM_ACTIONS / 2) category = 1; /* Flat */
else category = 2; /* Long */
input[MARKET_DIM + category] = bf16_one();
input[MARKET_DIM + category] = 1.0f;
/* Layer 1: pre_act = w1 * input + b1, hidden = LeakyReLU(pre_act, alpha=0.01) */
float pre_act[CUR_HIDDEN];
@@ -116,7 +116,7 @@ extern "C" __global__ void curiosity_forward_backward(
sum += w1[h * CUR_INPUT + i] * input[i];
}
pre_act[h] = sum;
hidden[h] = (sum > bf16_zero()) ? sum : bf16(0.01f) * sum;
hidden[h] = (sum > 0.0f) ? sum : 0.01f * sum;
}
/* Layer 2: pred = w2 * hidden + b2 (no activation) */
@@ -132,24 +132,24 @@ extern "C" __global__ void curiosity_forward_backward(
/* ---- Loss: MSE(pred, next_state[:MARKET_DIM]) ---- */
/* d_loss/d_pred[o] = 2 * (pred[o] - next_state_bf[o]) / CUR_OUTPUT */
float d_pred[CUR_OUTPUT];
float inv_out = bf16(2.0f / (float)CUR_OUTPUT);
float inv_out = 2.0f / (float)CUR_OUTPUT;
for (int o = 0; o < CUR_OUTPUT; o++) {
d_pred[o] = (pred[o] - next_state_bf[o]) * inv_out;
}
/* ---- Backward: Layer 2 (warp-reduced atomics) ---- */
float d_hidden[CUR_HIDDEN];
for (int h = 0; h < CUR_HIDDEN; h++) d_hidden[h] = bf16_zero();
for (int h = 0; h < CUR_HIDDEN; h++) d_hidden[h] = 0.0f;
for (int o = 0; o < CUR_OUTPUT; o++) {
{
float val = warp_sum_cur(d_pred[o]);
if ((threadIdx.x & 31) == 0) atomicAddBF16(&grad_b2[o], val);
if ((threadIdx.x & 31) == 0) atomicAdd(&grad_b2[o], val);
}
for (int h = 0; h < CUR_HIDDEN; h++) {
{
float val = warp_sum_cur(d_pred[o] * hidden[h]);
if ((threadIdx.x & 31) == 0) atomicAddBF16(&grad_w2[o * CUR_HIDDEN + h], val);
if ((threadIdx.x & 31) == 0) atomicAdd(&grad_w2[o * CUR_HIDDEN + h], val);
}
d_hidden[h] += w2[o * CUR_HIDDEN + h] * d_pred[o];
}
@@ -157,19 +157,19 @@ extern "C" __global__ void curiosity_forward_backward(
/* ---- Backward: LeakyReLU ---- */
for (int h = 0; h < CUR_HIDDEN; h++) {
if (pre_act[h] <= bf16_zero()) d_hidden[h] = d_hidden[h] * bf16(0.01f);
if (pre_act[h] <= 0.0f) d_hidden[h] = d_hidden[h] * 0.01f;
}
/* ---- Backward: Layer 1 (warp-reduced atomics) ---- */
for (int h = 0; h < CUR_HIDDEN; h++) {
{
float val = warp_sum_cur(d_hidden[h]);
if ((threadIdx.x & 31) == 0) atomicAddBF16(&grad_b1[h], val);
if ((threadIdx.x & 31) == 0) atomicAdd(&grad_b1[h], val);
}
for (int i = 0; i < CUR_INPUT; i++) {
{
float val = warp_sum_cur(d_hidden[h] * input[i]);
if ((threadIdx.x & 31) == 0) atomicAddBF16(&grad_w1[h * CUR_INPUT + i], val);
if ((threadIdx.x & 31) == 0) atomicAdd(&grad_w1[h * CUR_INPUT + i], val);
}
}
}
@@ -200,23 +200,16 @@ extern "C" __global__ void curiosity_adam_step(
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i >= num_params) return;
/* Convert host scalars to bf16 */
float lr_bf = bf16(lr);
float beta1_bf = bf16(beta1);
float beta2_bf = bf16(beta2);
float eps_bf = bf16(eps);
float one_bf = bf16_one();
float g = grads[i] / bf16((float)batch_size);
float g = grads[i] / (float)batch_size;
/* Skip NaN/Inf gradients */
if (!isfinite((float)g)) return;
if (!isfinite(g)) return;
m[i] = beta1_bf * m[i] + (one_bf - beta1_bf) * g;
v[i] = beta2_bf * v[i] + (one_bf - beta2_bf) * g * g;
float m_hat = m[i] / bf16(1.0f - powf(beta1, (float)step));
float v_hat = v[i] / bf16(1.0f - powf(beta2, (float)step));
params[i] = params[i] - lr_bf * m_hat / (bf16_sqrt(v_hat) + eps_bf);
m[i] = beta1 * m[i] + (1.0f - beta1) * g;
v[i] = beta2 * v[i] + (1.0f - beta2) * g * g;
float m_hat = m[i] / (1.0f - powf(beta1, (float)step));
float v_hat = v[i] / (1.0f - powf(beta2, (float)step));
params[i] = params[i] - lr * m_hat / (sqrtf(v_hat) + eps);
}
/* ------------------------------------------------------------------ */
@@ -253,23 +246,16 @@ extern "C" __global__ void curiosity_adam_step_fused(
p = p3; g = g3; mm = m3; vv = v3; local_i = idx - n0 - n1 - n2;
}
/* Convert host scalars to bf16 */
float lr_bf = bf16(lr);
float beta1_bf = bf16(beta1);
float beta2_bf = bf16(beta2);
float eps_bf = bf16(eps);
float one_bf = bf16_one();
float grad = g[local_i] / bf16((float)batch_size);
float grad = g[local_i] / (float)batch_size;
/* Skip NaN/Inf gradients */
if (!isfinite((float)grad)) return;
if (!isfinite(grad)) return;
mm[local_i] = beta1_bf * mm[local_i] + (one_bf - beta1_bf) * grad;
vv[local_i] = beta2_bf * vv[local_i] + (one_bf - beta2_bf) * grad * grad;
float m_hat = mm[local_i] / bf16(1.0f - powf(beta1, (float)step));
float v_hat = vv[local_i] / bf16(1.0f - powf(beta2, (float)step));
p[local_i] = p[local_i] - lr_bf * m_hat / (bf16_sqrt(v_hat) + eps_bf);
mm[local_i] = beta1 * mm[local_i] + (1.0f - beta1) * grad;
vv[local_i] = beta2 * vv[local_i] + (1.0f - beta2) * grad * grad;
float m_hat = mm[local_i] / (1.0f - powf(beta1, (float)step));
float v_hat = vv[local_i] / (1.0f - powf(beta2, (float)step));
p[local_i] = p[local_i] - lr * m_hat / (sqrtf(v_hat) + eps);
}
/* ------------------------------------------------------------------ */
@@ -285,7 +271,7 @@ extern "C" __global__ void curiosity_adam_step_fused(
* Phase 1 (all threads): Zero gradient buffers via grid-stride loop,
* then __syncthreads() within each block.
* Phase 2 (all threads): Forward + backward pass (one sample per thread),
* accumulate gradients via warp-reduced atomicAddBF16.
* accumulate gradients via warp-reduced atomicAdd.
* Phase 3 (last-arriving block only): After all blocks complete phase 2,
* the last block to arrive (detected via atomic counter) applies
* Adam update across all parameters in a grid-stride loop.
@@ -354,16 +340,16 @@ extern "C" __global__ void curiosity_fused_zero_fwd_bwd_adam(
/* Build input: first MARKET_DIM state features + 3-class action one-hot */
float input[CUR_INPUT];
for (int i = 0; i < MARKET_DIM; i++) { input[i] = state_bf2[i]; }
input[MARKET_DIM + 0] = bf16_zero();
input[MARKET_DIM + 1] = bf16_zero();
input[MARKET_DIM + 2] = bf16_zero();
input[MARKET_DIM + 0] = 0.0f;
input[MARKET_DIM + 1] = 0.0f;
input[MARKET_DIM + 2] = 0.0f;
/* Action to category one-hot */
int category;
if (action_idx <= 1) category = 0; /* Short100/Short50 */
else if (action_idx == 2) category = 1; /* Flat */
else category = 2; /* Long50/Long100 */
input[MARKET_DIM + category] = bf16_one();
input[MARKET_DIM + category] = 1.0f;
/* Layer 1: pre_act = w1 * input + b1, hidden = LeakyReLU(pre_act, 0.01) */
float pre_act[CUR_HIDDEN];
@@ -374,7 +360,7 @@ extern "C" __global__ void curiosity_fused_zero_fwd_bwd_adam(
sum += w1[h * CUR_INPUT + i] * input[i];
}
pre_act[h] = sum;
hidden[h] = (sum > bf16_zero()) ? sum : bf16(0.01f) * sum;
hidden[h] = (sum > 0.0f) ? sum : 0.01f * sum;
}
/* Layer 2: pred = w2 * hidden + b2 (no activation) */
@@ -389,24 +375,24 @@ extern "C" __global__ void curiosity_fused_zero_fwd_bwd_adam(
/* ---- Loss: MSE(pred, next_state[:MARKET_DIM]) ---- */
float d_pred[CUR_OUTPUT];
float inv_out = bf16(2.0f / (float)CUR_OUTPUT);
float inv_out = 2.0f / (float)CUR_OUTPUT;
for (int o = 0; o < CUR_OUTPUT; o++) {
d_pred[o] = (pred[o] - next_state_bf2[o]) * inv_out;
}
/* ---- Backward: Layer 2 (warp-reduced atomics) ---- */
float d_hidden[CUR_HIDDEN];
for (int h = 0; h < CUR_HIDDEN; h++) d_hidden[h] = bf16_zero();
for (int h = 0; h < CUR_HIDDEN; h++) d_hidden[h] = 0.0f;
for (int o = 0; o < CUR_OUTPUT; o++) {
{
float val = warp_sum_cur(d_pred[o]);
if ((threadIdx.x & 31) == 0) atomicAddBF16(&grad_b2[o], val);
if ((threadIdx.x & 31) == 0) atomicAdd(&grad_b2[o], val);
}
for (int h = 0; h < CUR_HIDDEN; h++) {
{
float val = warp_sum_cur(d_pred[o] * hidden[h]);
if ((threadIdx.x & 31) == 0) atomicAddBF16(&grad_w2[o * CUR_HIDDEN + h], val);
if ((threadIdx.x & 31) == 0) atomicAdd(&grad_w2[o * CUR_HIDDEN + h], val);
}
d_hidden[h] += w2[o * CUR_HIDDEN + h] * d_pred[o];
}
@@ -414,19 +400,19 @@ extern "C" __global__ void curiosity_fused_zero_fwd_bwd_adam(
/* ---- Backward: LeakyReLU ---- */
for (int h = 0; h < CUR_HIDDEN; h++) {
if (pre_act[h] <= bf16_zero()) d_hidden[h] = d_hidden[h] * bf16(0.01f);
if (pre_act[h] <= 0.0f) d_hidden[h] = d_hidden[h] * 0.01f;
}
/* ---- Backward: Layer 1 (warp-reduced atomics) ---- */
for (int h = 0; h < CUR_HIDDEN; h++) {
{
float val = warp_sum_cur(d_hidden[h]);
if ((threadIdx.x & 31) == 0) atomicAddBF16(&grad_b1[h], val);
if ((threadIdx.x & 31) == 0) atomicAdd(&grad_b1[h], val);
}
for (int i = 0; i < CUR_INPUT; i++) {
{
float val = warp_sum_cur(d_hidden[h] * input[i]);
if ((threadIdx.x & 31) == 0) atomicAddBF16(&grad_w1[h * CUR_INPUT + i], val);
if ((threadIdx.x & 31) == 0) atomicAdd(&grad_w1[h * CUR_INPUT + i], val);
}
}
}
@@ -460,13 +446,6 @@ extern "C" __global__ void curiosity_fused_zero_fwd_bwd_adam(
int block_tid = threadIdx.x;
int block_size = blockDim.x;
/* Convert host scalars to bf16 once for the Adam loop */
float lr_bf = bf16(lr);
float beta1_bf = bf16(beta1);
float beta2_bf = bf16(beta2);
float eps_bf = bf16(eps);
float one_bf = bf16_one();
for (int i = block_tid; i < total_grad_elems; i += block_size) {
/* Determine which param group and local offset */
float* p; float* g; float* mmv; float* vvv;
@@ -481,11 +460,11 @@ extern "C" __global__ void curiosity_fused_zero_fwd_bwd_adam(
p = b2; g = grad_b2; mmv = adam_m_b2; vvv = adam_v_b2; local_i = i - w1_len - b1_len - w2_len;
}
float grad = g[local_i] / bf16((float)batch_size);
mmv[local_i] = beta1_bf * mmv[local_i] + (one_bf - beta1_bf) * grad;
vvv[local_i] = beta2_bf * vvv[local_i] + (one_bf - beta2_bf) * grad * grad;
float m_hat = mmv[local_i] / bf16(1.0f - powf(beta1, (float)adam_step));
float v_hat = vvv[local_i] / bf16(1.0f - powf(beta2, (float)adam_step));
p[local_i] = p[local_i] - lr_bf * m_hat / (bf16_sqrt(v_hat) + eps_bf);
float grad = g[local_i] / (float)batch_size;
mmv[local_i] = beta1 * mmv[local_i] + (1.0f - beta1) * grad;
vvv[local_i] = beta2 * vvv[local_i] + (1.0f - beta2) * grad * grad;
float m_hat = mmv[local_i] / (1.0f - powf(beta1, (float)adam_step));
float v_hat = vvv[local_i] / (1.0f - powf(beta2, (float)adam_step));
p[local_i] = p[local_i] - lr * m_hat / (sqrtf(v_hat) + eps);
}
}

58
crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu
View File

@@ -18,7 +18,7 @@
* GRADIENT NORM KERNEL (Phase 3a)
*
* Computes gradient L2 norm (sum of squares) via warp + block reduction
* and atomicAddBF16 into a single output BF16. Must be followed by
* and atomicAdd into a single output f32. Must be followed by
* dqn_adam_update_kernel which reads the completed norm.
*
* Launch config: grid=(ceil(TOTAL_PARAMS/256), 1, 1), block=(256, 1, 1).
@@ -56,7 +56,7 @@ extern "C" __global__ void dqn_grad_norm_kernel(
}
}
/* Phase 2: Reduce block_sums[0..num_blocks] → sqrt → f32 + bf16 output.
/* Phase 2: Reduce block_sums[0..num_blocks] → sqrt → f32 output.
* Single block, 256 threads. Each thread sums a stripe, then warp+block reduce.
* Launch: grid=(1), block=(256). */
extern "C" __global__ void dqn_grad_norm_finalize(
@@ -169,7 +169,7 @@ extern "C" __global__ void dqn_saxpy_kernel(
int n
) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i < n) y[i] = y[i] + bf16(alpha) * x[i];
if (i < n) y[i] = y[i] + alpha * x[i];
}
/* ══════════════════════════════════════════════════════════════════════
@@ -264,14 +264,14 @@ extern "C" __global__ void dqn_clipped_saxpy_kernel(
extern "C" __global__ void dqn_clip_grad_kernel(
float* __restrict__ grads, /* [TOTAL_PARAMS] f32 gradient accumulator */
const float* __restrict__ grad_norm_ptr, /* bf16 L2 norm from finalize kernel */
const float* __restrict__ grad_norm_ptr, /* f32 L2 norm from finalize kernel */
float max_norm,
int total_params
) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx >= total_params) return;
float norm_f = (float)(*grad_norm_ptr); /* bf16 L2 norm, already sqrt'd */
float norm_f = (float)(*grad_norm_ptr); /* f32 L2 norm, already sqrt'd */
if (norm_f > max_norm) {
float scale = max_norm / norm_f;
grads[idx] = grads[idx] * scale;
@@ -334,7 +334,7 @@ extern "C" __global__ void dqn_zero_kernel(
int n
) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i < n) buf[i] = bf16_zero();
if (i < n) buf[i] = 0.0f;
}
/* ══════════════════════════════════════════════════════════════════════
@@ -369,10 +369,10 @@ extern "C" __global__ void dqn_regime_scale_kernel(
float cusum_diff = sample_cusum - target_cusum;
float dist_sq = adx_diff * adx_diff + cusum_diff * cusum_diff;
float temperature = bf16(0.3f);
float sim = bf16_exp(-(dist_sq) / (temperature * temperature + bf16(1e-8f)));
float temperature = 0.3f;
float sim = expf(-(dist_sq) / (temperature * temperature + 1e-8f));
/* Clamp to [0.5, 2.0]: never fully suppress, max 2x upweight */
float scale = bf16(0.5f) + bf16(1.5f) * sim;
float scale = 0.5f + 1.5f * sim;
td_errors[i] = td_errors[i] * scale;
}
@@ -413,16 +413,16 @@ extern "C" __global__ void dqn_shrink_perturb_kernel(
/* Per-element LCG PRNG: hash(seed, index) → uniform random */
unsigned int state = seed ^ (unsigned int)(i * 2654435761u);
state = state * 1664525u + 1013904223u;
float u1 = bf16((float)(state >> 8) / 16777216.0f); /* (0, 1) */
float u1 = (float)(state >> 8) / 16777216.0f; /* (0, 1) */
state = state * 1664525u + 1013904223u;
float u2 = bf16((float)(state >> 8) / 16777216.0f);
float u2 = (float)(state >> 8) / 16777216.0f;
/* Box-Muller: uniform → Gaussian N(0, sigma) */
u1 = bf16_fmax(u1, bf16(1e-7f)); /* prevent log(0) */
float noise = bf16(sigma) * bf16_sqrt(bf16(-2.0f) * bf16_log(u1)) * bf16_cos(u2 * bf16(6.283185307f));
u1 = fmaxf(u1, 1e-7f); /* prevent log(0) */
float noise = sigma * sqrtf(-2.0f * logf(u1)) * cosf(u2 * 6.283185307f);
/* Shrink-and-Perturb: blend old weights with noise */
params[i] = bf16(alpha) * params[i] + (bf16_one() - bf16(alpha)) * noise;
params[i] = alpha * params[i] + (1.0f - alpha) * noise;
}
/* ══════════════════════════════════════════════════════════════════════
@@ -440,7 +440,7 @@ extern "C" __global__ void dqn_relu_mask_kernel(
) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i >= n) return;
if (activation[i] <= bf16_zero()) dx[i] = bf16_zero();
if (activation[i] <= 0.0f) dx[i] = 0.0f;
}
/* ══════════════════════════════════════════════════════════════════════
@@ -631,7 +631,7 @@ extern "C" __global__ void bn_tanh_concat_kernel(
extern "C" __global__ void feature_importance_kernel(
float* __restrict__ importance, /* [market_dim] output */
const float* __restrict__ grad_w_s1, /* [shared_h1, state_dim] f32 gradient */
const float* __restrict__ states, /* [batch_size, state_dim_padded] bf16 */
const float* __restrict__ states, /* [batch_size, state_dim_padded] f32 */
int batch_size,
int shared_h1,
int state_dim,
@@ -672,7 +672,7 @@ extern "C" __global__ void feature_importance_kernel(
* ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void lottery_ticket_mask(
float* __restrict__ params_f32, /* [total_params] f32 master weights */
float* __restrict__ params_bf16, /* [total_params] bf16 shadow weights */
float* __restrict__ params_bf16, /* [total_params] f32 shadow weights */
const float* __restrict__ mask, /* [total_params] binary mask (0 or 1) */
int n
) {
@@ -722,7 +722,7 @@ extern "C" __global__ void causal_intervene_feature(
long long base = (long long)b * state_dim_padded;
for (int j = 0; j < state_dim_padded; j++) {
scratch[base + j] = (j == feature_k) ? bf16_zero() : states[base + j];
scratch[base + j] = (j == feature_k) ? 0.0f : states[base + j];
}
}
@@ -1088,14 +1088,14 @@ extern "C" __global__ void gradient_project(
* Only the first bn_dim columns of d_concat are used (portfolio columns skipped).
* d_concat has stride concat_dim, d_bn has stride bn_dim.
*
* tanh_val is read from bn_hidden (bf16, already has tanh applied from forward).
* tanh_val is read from bn_hidden (f32, already has tanh applied from forward).
*
* Grid: ceil(B * bn_dim / 256), Block: 256.
* ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void bn_tanh_backward_kernel(
float* __restrict__ d_bn, /* [B, bn_dim] f32 output */
const float* __restrict__ d_concat, /* [B, concat_dim] f32 input */
const float* __restrict__ bn_hidden, /* [B, bn_dim] bf16 (tanh values from forward) */
const float* __restrict__ bn_hidden, /* [B, bn_dim] f32 (tanh values from forward) */
int batch_size,
int bn_dim,
int concat_dim
@@ -1204,7 +1204,7 @@ extern "C" __global__ void stochastic_depth_rng(
/* ══════════════════════════════════════════════════════════════════════
* STOCHASTIC DEPTH SCALE KERNEL (#21)
*
* Scales a bf16 hidden activation buffer by a per-layer scalar.
* Scales a f32 hidden activation buffer by a per-layer scalar.
* Used for stochastic depth: each hidden layer is either kept (scaled
* by 1/(1-p) for expected-value correction) or dropped (scaled by 0).
*
@@ -1232,7 +1232,7 @@ extern "C" __global__ void stochastic_depth_scale(
/* ══════════════════════════════════════════════════════════════════════
* BF16 → I32 CAST KERNEL
*
* Converts action indices stored as bf16 (from GpuTensor) to i32
* Converts action indices stored as f32 (from GpuTensor) to i32
* (required by loss/grad kernels). Eliminates the synchronous
* GPU→CPU→GPU round-trip in upload_batch_gpu that was a per-step
* serialization bottleneck.
@@ -1267,15 +1267,15 @@ extern "C" __global__ void bf16_to_i32_kernel(
* One warp per relabeled sample — coalesced column writes.
* ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void her_inplace_relabel(
float* __restrict__ dst_states, /* [batch_size, dst_stride] padded bf16 */
float* __restrict__ dst_next_states, /* [batch_size, dst_stride] padded bf16 */
float* __restrict__ dst_states, /* [batch_size, dst_stride] padded f32 */
float* __restrict__ dst_next_states, /* [batch_size, dst_stride] padded f32 */
float* __restrict__ dst_rewards, /* [batch_size] f32 */
const float* __restrict__ src_next_states, /* [batch_size, src_stride] f32 (PER stores f32) */
const int* __restrict__ donor_indices, /* [her_batch_size] i32 */
int offset, /* normal_count: first HER row in dst */
int goal_dim, /* number of goal columns to replace */
int src_stride, /* unpadded state_dim in src (f32 elements) */
int dst_stride, /* padded state_dim in dst (bf16 elements) */
int dst_stride, /* padded state_dim in dst (f32 elements) */
int state_dim /* min(src_stride, dst_stride) for full copy */
) {
int i = blockIdx.x; /* HER sample index */
@@ -1284,13 +1284,11 @@ extern "C" __global__ void her_inplace_relabel(
int dst_row = offset + i;
/* Replace goal columns (first goal_dim elements) with donor's achieved goal.
* Source is f32 (PER buffer stores full-precision states). Convert to bf16
* for the padded staging buffer. */
* Source is f32 (PER buffer stores full-precision states). */
for (int d = lane; d < goal_dim; d += 32) {
float achieved = src_next_states[donor * src_stride + d];
float achieved_bf16 = (achieved);
dst_states[dst_row * dst_stride + d] = achieved_bf16;
dst_next_states[dst_row * dst_stride + d] = achieved_bf16;
dst_states[dst_row * dst_stride + d] = achieved;
dst_next_states[dst_row * dst_stride + d] = achieved;
}
/* Set reward = 1.0 (only lane 0) */

53
crates/ml/src/cuda_pipeline/dt_kernels.cu
View File

@@ -93,7 +93,7 @@ extern "C" __global__ void dt_qkv_projection_kernel(
const float* x = input + bt * E;
for (int dd = d; dd < E; dd += blockDim.x) {
float q = bf16_zero(), k = bf16_zero(), v = bf16_zero();
float q = 0.0f, k = 0.0f, v = 0.0f;
for (int j = 0; j < E; j++) {
float xj = x[j];
q = q + xj * W_Q[j * E + dd];
@@ -123,7 +123,7 @@ extern "C" __global__ void dt_qkv_projection_kernel(
* final[b][t][d] = sum_h( out_h[t][dh] * W_O[h*Dh + dh, d] )
*
* NOTE: Attention score computation uses F32 accumulation for numerical
* stability (softmax is sensitive to precision). Results stored as BF16.
* stability. All storage and computation in native f32.
* ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void dt_causal_attention_kernel(
const float* __restrict__ Q, /* [B, T, E] */
@@ -177,7 +177,7 @@ extern "C" __global__ void dt_causal_attention_kernel(
/* Compute softmax weights and weighted sum of V */
float attn_out[128]; /* max Dh — static array to avoid dynamic alloc */
for (int dd = 0; dd < Dh; dd++) attn_out[dd] = bf16_zero();
for (int dd = 0; dd < Dh; dd++) attn_out[dd] = 0.0f;
float sum_exp = 0.0f;
for (int j = 0; j <= t_idx; j++) {
@@ -185,7 +185,7 @@ extern "C" __global__ void dt_causal_attention_kernel(
for (int dd = 0; dd < Dh; dd++) {
score += (float)Q_b[t_idx * E + head_off + dd] * (float)sh_K[j * Dh + dd];
}
float weight = bf16(expf(score * inv_sqrt_dh - max_score));
float weight = expf(score * inv_sqrt_dh - max_score);
sum_exp += expf(score * inv_sqrt_dh - max_score);
for (int dd = 0; dd < Dh; dd++) {
attn_out[dd] = attn_out[dd] + weight * sh_V[j * Dh + dd];
@@ -193,7 +193,7 @@ extern "C" __global__ void dt_causal_attention_kernel(
}
/* Normalize by softmax denominator */
float inv_sum = bf16(1.0f / (sum_exp + 1e-8f));
float inv_sum = 1.0f / (sum_exp + 1e-8f);
for (int dd = 0; dd < Dh; dd++) {
attn_out[dd] = attn_out[dd] * inv_sum;
}
@@ -226,11 +226,11 @@ extern "C" __global__ void dt_causal_attention_kernel(
/* Each head's projection contribution via atomicAdd */
for (int dd = 0; dd < E; dd++) {
float proj = bf16_zero();
float proj = 0.0f;
for (int dh = 0; dh < Dh; dh++) {
proj = proj + sh_head_out[t_idx * Dh + dh] * W_O[(head_off + dh) * E + dd];
}
atomicAddBF16(&out[dd], proj);
atomicAdd(&out[dd], proj);
}
}
@@ -240,7 +240,7 @@ extern "C" __global__ void dt_causal_attention_kernel(
* output = gamma * (input - mean) / sqrt(var + eps) + beta
*
* NOTE: Mean/variance computation uses F32 accumulation for numerical
* stability (reduction over large vectors). Final output stored as BF16.
* stability. All computation and storage in native f32.
*
* Grid: (B*T, 1, 1)
* Block: (min(E, 256), 1, 1)
@@ -317,9 +317,9 @@ extern "C" __global__ void dt_layernorm_kernel(
__syncthreads();
float inv_std = sh_inv_std;
/* Normalize + affine — final store as BF16 */
/* Normalize + affine */
for (int d = tid; d < E; d += blockDim.x) {
out[d] = bf16((float)gamma[d] * ((float)x[d] - mean) * inv_std + (float)beta[d]);
out[d] = (float)gamma[d] * ((float)x[d] - mean) * inv_std + (float)beta[d];
}
}
@@ -331,8 +331,8 @@ extern "C" __global__ void dt_layernorm_kernel(
*
* GELU approximation: 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3)))
*
* NOTE: GELU computation uses F32 for the intermediate hidden activations
* (tanh precision matters). Final output stored as BF16.
* NOTE: GELU computation uses F32 for the intermediate hidden activations.
* All computation and storage in native f32.
*
* Grid: (B*T, 1, 1)
* Block: (min(E, 256), 1, 1)
@@ -381,7 +381,7 @@ extern "C" __global__ void dt_ffn_kernel(
for (int ff = 0; ff < FFN; ff++) {
val += sh_hidden[ff] * (float)W2[ff * E + d];
}
out[d] = bf16(val);
out[d] = val;
}
}
@@ -466,8 +466,8 @@ extern "C" __global__ void dt_cross_entropy_kernel(
/* Clamp to prevent NaN propagation */
loss = fminf(fmaxf(loss, 0.0f), 100.0f);
per_sample_loss[n] = bf16(loss);
atomicAddBF16(total_loss, bf16(loss / (float)N));
per_sample_loss[n] = loss;
atomicAdd(total_loss, loss / (float)N);
}
/* ══════════════════════════════════════════════════════════════════════
@@ -476,7 +476,7 @@ extern "C" __global__ void dt_cross_entropy_kernel(
* dL/d_logits = softmax(logits) - one_hot(target)
*
* NOTE: Softmax computation uses F32 for numerical stability.
* Gradient output stored as BF16.
* Gradient output stored as F32.
*
* Grid: (B*T, 1, 1)
* Block: (min(A, 256), 1, 1)
@@ -514,7 +514,7 @@ extern "C" __global__ void dt_ce_backward_kernel(
for (int a = tid; a < A; a += blockDim.x) {
float softmax_a = expf((float)lg[a] - max_val) * inv_sum;
float one_hot = (a == target) ? 1.0f : 0.0f;
d_lg[a] = bf16((softmax_a - one_hot) * scale);
d_lg[a] = (softmax_a - one_hot) * scale;
}
}
@@ -552,7 +552,7 @@ extern "C" __global__ void dt_linear_backward_kernel(
/* Compute d_input[n][i] = sum_o(d_output[n][o] * W[i][o]) */
for (int i = tid; i < I; i += blockDim.x) {
float val = bf16_zero();
float val = 0.0f;
for (int o = 0; o < O; o++) {
val = val + dout[o] * W[i * O + o];
}
@@ -563,13 +563,13 @@ extern "C" __global__ void dt_linear_backward_kernel(
for (int i = tid; i < I; i += blockDim.x) {
float xi = x[i];
for (int o = 0; o < O; o++) {
atomicAddBF16(&dW[i * O + o], xi * dout[o]);
atomicAdd(&dW[i * O + o], xi * dout[o]);
}
}
/* Bias gradient: only one thread per sample contributes */
for (int o = tid; o < O; o += blockDim.x) {
atomicAddBF16(&db[o], dout[o]);
atomicAdd(&db[o], dout[o]);
}
}
@@ -586,7 +586,7 @@ extern "C" __global__ void dt_zero_kernel(
int n
) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i < n) buf[i] = bf16_zero();
if (i < n) buf[i] = 0.0f;
}
/* ══════════════════════════════════════════════════════════════════════
@@ -631,10 +631,9 @@ extern "C" __global__ void dt_return_to_go_kernel(
float* rtg = returns_to_go + ep * episode_len;
/* Backward scan: R_{T-1} = r_{T-1}, R_t = r_t + gamma * R_{t+1} */
float bf16_gamma = bf16(gamma);
float cumul = bf16_zero();
float cumul = 0.0f;
for (int t = episode_len - 1; t >= 0; t--) {
cumul = r[t] + bf16_gamma * cumul;
cumul = r[t] + gamma * cumul;
rtg[t] = cumul;
}
}
@@ -689,9 +688,9 @@ extern "C" __global__ void dt_build_trajectories_kernel(
out[1 + f] = feat[f];
}
/* Slot state_dim+1: action as bf16 (state_dim = feat_dim) */
/* Slot state_dim+1: action as f32 (state_dim = feat_dim) */
int act = actions[bar_idx]; /* index into global bar actions */
out[feat_dim + 1] = bf16((float)act);
out[feat_dim + 1] = (float)act;
/* Target action for cross-entropy loss */
target_actions[ep * T + t] = act;
@@ -730,7 +729,7 @@ extern "C" __global__ void dt_compute_rewards_actions_kernel(
} else {
reward = 0.0f; /* Last bar has no future */
}
rewards[i] = bf16(reward);
rewards[i] = reward;
/* Expert action: Long=8 (max exposure), Short=0 (min exposure), Flat=4 (neutral) */
int action;

26
crates/ml/src/cuda_pipeline/ensemble_kernels.cu
View File

@@ -38,19 +38,19 @@ extern "C" __global__ void ensemble_aggregate_kernel(
int total = B * num_actions;
if (idx >= total) return;
float sum = bf16_zero();
float sum_sq = bf16_zero();
float sum = 0.0f;
float sum_sq = 0.0f;
for (int k = 0; k < K; k++) {
float q = head_q_values[(long long)k * total + idx];
sum = sum + q;
sum_sq = sum_sq + q * q;
}
float k_bf = bf16((float)K);
float mean = sum / k_bf;
float k_f = (float)K;
float mean = sum / k_f;
mean_q[idx] = mean;
/* Var = E[X^2] - E[X]^2, clamped to 0 for numerical stability */
float v = sum_sq / k_bf - mean * mean;
var_q[idx] = bf16_fmax(v, bf16_zero());
float v = sum_sq / k_f - mean * mean;
var_q[idx] = fmaxf(v, 0.0f);
}
/* ======================================================================
@@ -83,7 +83,7 @@ extern "C" __global__ void ensemble_diversity_kernel(
int total_work = B * num_pairs;
int idx = blockIdx.x * blockDim.x + threadIdx.x;
float kl_sum = bf16_zero();
float kl_sum = 0.0f;
if (idx < total_work) {
/* Decode sample index and pair index */
@@ -125,13 +125,13 @@ extern "C" __global__ void ensemble_diversity_kernel(
float log_ratio_ji = logf(pj / (pi + 1e-8f) + 1e-8f);
kl += 0.5f * (pi * log_ratio_ij + pj * log_ratio_ji);
}
kl_sum = bf16(kl);
kl_sum = kl;
}
/* -- Hierarchical reduction: warp -> block -> atomicAdd ------------ */
/* Warp-level reduction via shuffle (no shared memory) */
for (int offset = 16; offset > 0; offset >>= 1)
kl_sum = kl_sum + bf16_shfl_xor(0xFFFFFFFF, kl_sum, offset);
kl_sum = kl_sum + __shfl_xor_sync(0xFFFFFFFF, kl_sum, offset);
/* Block-level cross-warp reduction via shared memory (padded, same as grad_norm) */
__shared__ float warp_sums[16]; /* 8 warps x 2 stride (padded) */
@@ -142,11 +142,11 @@ extern "C" __global__ void ensemble_diversity_kernel(
/* First warp reduces across warps */
if (warp_id == 0) {
float val = (warp_lane < blockDim.x / 32) ? warp_sums[warp_lane * 2] : bf16_zero();
float val = (warp_lane < blockDim.x / 32) ? warp_sums[warp_lane * 2] : 0.0f;
for (int offset = 16; offset > 0; offset >>= 1)
val = val + bf16_shfl_xor(0xFFFFFFFF, val, offset);
val = val + __shfl_xor_sync(0xFFFFFFFF, val, offset);
if (warp_lane == 0)
atomicAddBF16(diversity_loss, val);
atomicAdd(diversity_loss, val);
}
}
@@ -206,5 +206,5 @@ extern "C" __global__ void ensemble_kl_gradient_kernel(
grad /= (float)(K - 1);
/* NEGATIVE: maximize divergence (subtract from loss) */
d_logits_0[b * NA + a] = bf16(-diversity_weight * grad);
d_logits_0[b * NA + a] = -diversity_weight * grad;
}

30
crates/ml/src/cuda_pipeline/epsilon_greedy_kernel.cu
View File

@@ -30,13 +30,13 @@ extern "C" __global__ void epsilon_greedy_select(
if (idx >= batch_size) return;
unsigned int rng = rng_states[idx];
float r = bf16(gpu_random(&rng));
float eps_bf = bf16(epsilon);
float r = gpu_random(&rng);
float eps_bf = epsilon;
unsigned int action;
if (r < eps_bf) {
/* Random exploration */
float rand_val = bf16(gpu_random(&rng));
float rand_val = gpu_random(&rng);
action = (unsigned int)((float)rand_val * (float)num_actions);
if (action >= (unsigned int)num_actions) action = (unsigned int)(num_actions - 1);
} else {
@@ -116,8 +116,8 @@ extern "C" __global__ void epsilon_greedy_routed(
if (idx >= batch_size) return;
unsigned int rng = rng_states[idx];
float r = bf16(gpu_random(&rng));
float eps_bf = bf16(epsilon);
float r = gpu_random(&rng);
float eps_bf = epsilon;
/* Step 1: Epsilon-greedy action selection (exposure index 0-4) */
int exposure_idx;
@@ -156,7 +156,7 @@ extern "C" __global__ void epsilon_greedy_routed(
/* Step 3: Fill simulation (deterministic splitmix64 hash) */
/* Normalized vol for fill probability calculation */
float norm_vol = (median_vol > 0.0f) ? bf16(volatility / median_vol) : bf16_one();
float norm_vol = (median_vol > 0.0f) ? (volatility / median_vol) : 1.0f;
float cost_adj;
int filled = simulate_fill_check(
order_type, urgency, (float)norm_vol, spread_bps,
@@ -209,16 +209,16 @@ extern "C" __global__ void branching_action_select(
/* Head 1: Exposure (5 actions) -- with UCB bonus + Q-gap conviction filter */
int exposure;
float r1 = bf16(gpu_random(&rng));
float r1 = gpu_random(&rng);
if (r1 < eps) {
exposure = (int)(gpu_random(&rng) * 5.0f);
if (exposure >= 5) exposure = 4;
} else {
const float* qe = q_exposure + idx * 5;
float best = qe[0] + ((bonus_exposure != NULL) ? bonus_exposure[0] : bf16_zero());
float best = qe[0] + ((bonus_exposure != NULL) ? bonus_exposure[0] : 0.0f);
exposure = 0;
for (int a = 1; a < 5; a++) {
float q_plus_bonus = qe[a] + ((bonus_exposure != NULL) ? bonus_exposure[a] : bf16_zero());
float q_plus_bonus = qe[a] + ((bonus_exposure != NULL) ? bonus_exposure[a] : 0.0f);
if (q_plus_bonus > best) { best = q_plus_bonus; exposure = a; }
}
/* Adaptive Q-gap filter: fraction of exposure Q-range. */
@@ -239,7 +239,7 @@ extern "C" __global__ void branching_action_select(
/* Head 2: Order type (3 actions) -- Boltzmann softmax with UCB bonus */
int order;
float r2 = bf16(gpu_random(&rng));
float r2 = gpu_random(&rng);
if (r2 < eps) {
/* Random exploration */
order = (int)(gpu_random(&rng) * 3.0f);
@@ -247,12 +247,12 @@ extern "C" __global__ void branching_action_select(
} else {
/* Boltzmann softmax over order Q-values */
const float* qo = q_order + idx * 3;
float q_max_o = qo[0] + ((bonus_order != NULL) ? bonus_order[0] : bf16_zero());
float q_max_o = qo[0] + ((bonus_order != NULL) ? bonus_order[0] : 0.0f);
float q_min_o = q_max_o;
float qo_adj[3];
qo_adj[0] = q_max_o;
for (int a = 1; a < 3; a++) {
qo_adj[a] = qo[a] + ((bonus_order != NULL) ? bonus_order[a] : bf16_zero());
qo_adj[a] = qo[a] + ((bonus_order != NULL) ? bonus_order[a] : 0.0f);
q_max_o = fmaxf(q_max_o, qo_adj[a]);
q_min_o = fminf(q_min_o, qo_adj[a]);
}
@@ -274,7 +274,7 @@ extern "C" __global__ void branching_action_select(
/* Head 3: Urgency (3 actions) -- Boltzmann softmax with UCB bonus */
int urgency;
float r3 = bf16(gpu_random(&rng));
float r3 = gpu_random(&rng);
if (r3 < eps) {
/* Random exploration */
urgency = (int)(gpu_random(&rng) * 3.0f);
@@ -282,12 +282,12 @@ extern "C" __global__ void branching_action_select(
} else {
/* Boltzmann softmax over urgency Q-values */
const float* qu = q_urgency + idx * 3;
float q_max_u = qu[0] + ((bonus_urgency != NULL) ? bonus_urgency[0] : bf16_zero());
float q_max_u = qu[0] + ((bonus_urgency != NULL) ? bonus_urgency[0] : 0.0f);
float q_min_u = q_max_u;
float qu_adj[3];
qu_adj[0] = q_max_u;
for (int a = 1; a < 3; a++) {
qu_adj[a] = qu[a] + ((bonus_urgency != NULL) ? bonus_urgency[a] : bf16_zero());
qu_adj[a] = qu[a] + ((bonus_urgency != NULL) ? bonus_urgency[a] : 0.0f);
q_max_u = fmaxf(q_max_u, qu_adj[a]);
q_min_u = fminf(q_min_u, qu_adj[a]);
}

103
crates/ml/src/cuda_pipeline/experience_kernels.cu
View File

@@ -52,7 +52,7 @@
*/
/* ------------------------------------------------------------------ */
/* Portfolio stride for experience kernels (23 bf16 per episode). */
/* Portfolio stride for experience kernels (23 f32 per episode). */
/* portfolio_sim_kernel uses its own stride of 8 — do NOT change it. */
/* ------------------------------------------------------------------ */
#define PORTFOLIO_STRIDE 23
@@ -366,11 +366,11 @@ extern "C" __global__ void saboteur_select_best(
*
* Grid: ceil(N / 256), Block: 256. One thread per episode.
*
* @param market_features [total_bars, market_dim] raw market feature matrix (bf16)
* @param market_features [total_bars, market_dim] raw market feature matrix (f32)
* @param episode_starts [N] global bar index of episode start
* @param current_timesteps [N] timestep offset within episode (read-only)
* @param portfolio_states [N, 3] {position, cash, portfolio_value} (bf16)
* @param batch_states [N, state_dim] output: assembled state batch (bf16)
* @param portfolio_states [N, 3] {position, cash, portfolio_value} (f32)
* @param batch_states [N, state_dim] output: assembled state batch (f32)
* @param N number of episodes
* @param total_bars number of bars in market_features
* @param state_dim full state dimension (tensor-core aligned)
@@ -441,7 +441,7 @@ extern "C" __global__ void experience_state_gather(
return;
}
/* -- Market features: [0 .. market_dim) -- bf16 input → f32 output */
/* -- Market features: [0 .. market_dim) -- f32 input → f32 output */
const float* mf_row = market_features + (long long)bar_idx * market_dim;
for (int k = 0; k < market_dim; k++)
out[k] = (float)mf_row[k];
@@ -546,10 +546,7 @@ extern "C" __global__ void experience_state_gather(
float entry_price = ps[12];
float trade_start_pnl = ps[13];
/* ── Portfolio features computed in FLOAT to prevent bf16 overflow ──
* BF16 max ~65504. ES position × price_diff can exceed this (2.0 × 5000 = 10000,
* but accumulated P&L or subtraction of similar prices overflow easily).
* Float arithmetic eliminates the root cause of NaN states in the replay buffer. */
/* ── Portfolio features computed in FLOAT ── */
float f_position = (float)position;
float f_cash = (float)cash;
float f_portfolio_value = (float)portfolio_value;
@@ -618,8 +615,7 @@ extern "C" __global__ void experience_state_gather(
if (past_idx >= 0 && slot + 3 < state_dim) {
const float* now_row = market_features + (long long)bar_idx * market_dim;
const float* past_row = market_features + (long long)past_idx * market_dim;
/* Float arithmetic — bf16 close prices (~5000) subtracted produce
* tiny differences that lose all precision in bf16 (7-bit mantissa). */
/* Float arithmetic for return computation. */
float f_close_now = (float)now_row[0];
float f_close_past = (float)past_row[0];
@@ -702,10 +698,10 @@ extern "C" __global__ void experience_state_gather(
*
* Grid: ceil(N / 256), Block: 256. One thread per episode.
*
* @param q_values [N, q_stride] Q-values from cuBLAS (bf16)
* @param q_values [N, q_stride] Q-values from cuBLAS (f32)
* @param out_actions [N] selected factored action index (output)
* @param rng_states [N] per-episode LCG RNG counter (updated in place)
* @param out_q_gaps [N] output: Q-gap per episode for conviction sizing (bf16)
* @param out_q_gaps [N] output: Q-gap per episode for conviction sizing (f32)
* @param epsilon exploration probability in [0, 1]
* @param N number of episodes
* @param b0_size direction branch size (3)
@@ -713,7 +709,7 @@ extern "C" __global__ void experience_state_gather(
* @param b2_size order branch size (3)
* @param b3_size urgency branch size (3)
* @param q_gap_threshold min Q-gap for trade entry (0.0 = disabled)
* @param portfolio_states [N, 20] read-only for hold enforcement (bf16)
* @param portfolio_states [N, 20] read-only for hold enforcement (f32)
* @param min_hold_bars minimum bars to hold before switching
* @param max_position max allowed position (host scalar)
*/
@@ -789,12 +785,12 @@ extern "C" __global__ void experience_action_select(
* Portfolio state: ps[0] = position, ps[10] = hold_time.
* When portfolio_states is NULL (backtest evaluator), skip hold enforcement. */
int in_hold = 0;
float cur_position = bf16_zero();
float cur_position = 0.0f;
if (min_hold_bars > 0 && portfolio_states != NULL) {
int ps_base = i * 23; /* PORTFOLIO_STRIDE = 23 */
float hold_time_val = portfolio_states[ps_base + 10];
cur_position = portfolio_states[ps_base + 0];
in_hold = (hold_time_val > bf16_zero() && hold_time_val < bf16((float)min_hold_bars));
in_hold = (hold_time_val > 0.0f && hold_time_val < (float)min_hold_bars);
}
if (in_hold) {
@@ -820,7 +816,7 @@ extern "C" __global__ void experience_action_select(
}
/* Zero Q-gap during hold — conviction sizing is meaningless when forced */
if (out_q_gaps != NULL) {
out_q_gaps[i] = bf16_zero();
out_q_gaps[i] = 0.0f;
}
} else {
/* Branch 0: direction — BOLTZMANN SOFTMAX selection.
@@ -1028,7 +1024,7 @@ extern "C" __global__ void experience_action_select(
float q_best = q_b0[argmax_n(q_b0, b0_size)];
float q_flat = q_b0[flat_idx];
float gap = q_best - q_flat;
out_q_gaps[i] = (gap > bf16_zero()) ? gap : bf16_zero();
out_q_gaps[i] = (gap > 0.0f) ? gap : 0.0f;
}
}
@@ -1171,11 +1167,7 @@ extern "C" __global__ void experience_env_step(
* [0:1] = preprocessed (log-return normalized) — for Q-network
* [2:3] = raw dollar prices — for portfolio simulation P&L + tx costs */
const float* tgt = targets + (long long)bar_idx * 4;
/* Trade physics (execute_trade, compute_tx_cost, etc.) use float internally
* because they involve multi-step accumulation where bf16 precision is
* insufficient (e.g. cash -= delta * price; cash -= cost). We convert
* prices to float here for the physics section, then store results back
* as bf16 at the end. */
/* Trade physics (execute_trade, compute_tx_cost, etc.) use float internally. */
float raw_close = (tgt[2]);
float raw_next = (tgt[3]);
@@ -1184,11 +1176,7 @@ extern "C" __global__ void experience_env_step(
if (raw_next <= 0.0f) raw_next = raw_close;
/* ---- Read full portfolio state (PORTFOLIO_STRIDE=23) ---- */
/* Portfolio arithmetic uses float accumulators because trade physics
* functions (execute_trade, apply_margin_cap, etc.) in trade_physics.cuh
* are all float. Converting the entire physics engine to bf16 would
* require rewriting the shared header used by backtest_env_kernel too.
* We load bf16 → float here, compute, then store float → bf16 at end. */
/* Portfolio arithmetic uses float accumulators throughout. */
float* ps = portfolio_states + (long long)i * PORTFOLIO_STRIDE;
float position = (ps[0]);
float cash = (ps[1]);
@@ -1220,22 +1208,22 @@ extern "C" __global__ void experience_env_step(
out_dones[out_off] = 1.0f;
/* Reset portfolio for next episode (fresh capital) */
float init_cap = peak_equity; /* use peak as initial for next episode */
ps[0] = bf16_zero(); /* position = flat */
ps[0] = 0.0f; /* position = flat */
ps[1] = (init_cap); /* cash = initial_capital */
ps[2] = (init_cap); /* portfolio_value */
ps[7] = (init_cap); /* peak_equity */
ps[8] = bf16_zero(); /* flat_counter */
ps[8] = 0.0f; /* flat_counter */
ps[9] = (init_cap); /* prev_equity */
ps[10] = bf16_zero(); /* hold_time */
ps[11] = bf16_zero(); /* realized_pnl */
ps[12] = bf16_zero(); /* entry_price */
ps[13] = bf16_zero(); /* trade_start_pnl */
ps[14] = bf16_zero(); /* win_count (reset Kelly) */
ps[15] = bf16_zero(); /* loss_count */
ps[16] = bf16_zero(); /* sum_wins */
ps[17] = bf16_zero(); /* sum_losses */
ps[18] = bf16_zero(); /* sum_returns */
ps[19] = bf16_zero(); /* sum_sq_returns */
ps[10] = 0.0f; /* hold_time */
ps[11] = 0.0f; /* realized_pnl */
ps[12] = 0.0f; /* entry_price */
ps[13] = 0.0f; /* trade_start_pnl */
ps[14] = 0.0f; /* win_count (reset Kelly) */
ps[15] = 0.0f; /* loss_count */
ps[16] = 0.0f; /* sum_wins */
ps[17] = 0.0f; /* sum_losses */
ps[18] = 0.0f; /* sum_returns */
ps[19] = 0.0f; /* sum_sq_returns */
current_timesteps[i] = 0; /* reset episode timer */
return;
}
@@ -1964,22 +1952,22 @@ extern "C" __global__ void experience_env_step(
* hits the pre-trade floor check, producing stuck episodes with fake MaxDD. */
if (done) {
float init_cap = (peak_equity); /* use peak as starting capital for next ep */
ps[0] = bf16_zero(); /* position = flat */
ps[0] = 0.0f; /* position = flat */
ps[1] = init_cap; /* cash */
ps[2] = init_cap; /* portfolio_value */
ps[7] = init_cap; /* peak_equity */
ps[8] = bf16_zero(); /* flat_counter */
ps[8] = 0.0f; /* flat_counter */
ps[9] = init_cap; /* prev_equity */
ps[10] = bf16_zero(); /* hold_time */
ps[11] = bf16_zero(); /* realized_pnl */
ps[12] = bf16_zero(); /* entry_price */
ps[13] = bf16_zero(); /* trade_start_pnl */
ps[14] = bf16_zero(); /* win/loss/Kelly counters */
ps[15] = bf16_zero();
ps[16] = bf16_zero();
ps[17] = bf16_zero();
ps[18] = bf16_zero();
ps[19] = bf16_zero();
ps[10] = 0.0f; /* hold_time */
ps[11] = 0.0f; /* realized_pnl */
ps[12] = 0.0f; /* entry_price */
ps[13] = 0.0f; /* trade_start_pnl */
ps[14] = 0.0f; /* win/loss/Kelly counters */
ps[15] = 0.0f;
ps[16] = 0.0f;
ps[17] = 0.0f;
ps[18] = 0.0f;
ps[19] = 0.0f;
current_timesteps[i] = 0;
}
}
@@ -1991,9 +1979,8 @@ extern "C" __global__ void experience_env_step(
* NOT part of the timestep-loop experience collection — kept for
* GpuPortfolioSimulator compatibility in gpu_portfolio.rs.
*
* NOTE: portfolio_state, portfolio_out, rewards_out are bf16 on Rust side.
* Internal arithmetic uses float accumulators for precision, converts at
* boundaries.
* NOTE: portfolio_state, portfolio_out, rewards_out are f32 on Rust side.
* Internal arithmetic uses float accumulators for precision.
* ====================================================================== */
/* action_to_exposure and action_to_tx_cost are provided by common_device_functions.cuh */
@@ -2161,9 +2148,9 @@ extern "C" __global__ void portfolio_sim_kernel(
*
* Grid: ceil(N / 256), Block: 256. One thread per episode.
*
* @param v_logits [N, NA] value head logits (bf16)
* @param b_logits [N, (B0+B1+B2+B3)*NA] branch advantage logits (bf16)
* @param q_values [N, B0+B1+B2+B3] output: expected Q per action (bf16)
* @param v_logits [N, NA] value head logits (f32)
* @param b_logits [N, (B0+B1+B2+B3)*NA] branch advantage logits (f32)
* @param q_values [N, B0+B1+B2+B3] output: expected Q per action (f32)
* @param N number of samples
* @param num_atoms C51 atom count (NA)
* @param b0_size direction branch size

8
crates/ml/src/cuda_pipeline/her_relabel_kernel.cu
View File

@@ -68,21 +68,21 @@ void her_relabel_kernel(
/* 3. Scalar work: copy action/done, recompute sparse reward (lane 0 only). */
if (lane_id == 0) {
out_actions[sample_idx] = actions[src_idx];
out_dones[sample_idx] = dones[src_idx]; /* BF16 copy - OK */
out_dones[sample_idx] = dones[src_idx];
/* Goal distance: L2 between source's achieved goal and donor's achieved goal.
* Source achieved = next_states[src_idx, 0..GOAL_DIM]
* Donor achieved = next_states[donor_idx, 0..GOAL_DIM] */
float dist_sq = bf16_zero();
float dist_sq = 0.0f;
for (int d = 0; d < GOAL_DIM; d++) {
float src_g = next_states[src_idx * state_dim + d];
float donor_g = next_states[donor_idx * state_dim + d];
float diff = src_g - donor_g;
dist_sq = dist_sq + diff * diff;
}
float dist = bf16_sqrt(dist_sq);
float dist = sqrtf(dist_sq);
/* Sparse reward: +1.0 if goal achieved, -0.01 otherwise */
out_rewards[sample_idx] = (dist < bf16(GOAL_THRESHOLD)) ? bf16_one() : bf16(-0.01f);
out_rewards[sample_idx] = (dist < GOAL_THRESHOLD) ? 1.0f : -0.01f;
}
}

14
crates/ml/src/cuda_pipeline/iqn_cvar_kernel.cu
View File

@@ -29,13 +29,13 @@ extern "C" __global__ void iqn_cvar_kernel(
if (alpha_count < 1) alpha_count = 1;
// Find the alpha_count smallest quantile values using partial sort.
float cvar_sum = bf16_zero();
float cvar_sum = 0.0f;
if (alpha_count == 1) {
// Fast path: just find the minimum
float min_val = bf16(1e30f);
float min_val = 1e30f;
for (int t = 0; t < N_TAU; t++) {
float v = q_values[i * N_TAU * TBA + t * TBA + exposure_idx];
min_val = bf16_fmin(min_val, v);
min_val = fminf(min_val, v);
}
cvar_sum = min_val;
} else {
@@ -66,11 +66,11 @@ extern "C" __global__ void iqn_cvar_kernel(
for (int j = 0; j < alpha_count; j++) cvar_sum = cvar_sum + sorted[j];
}
float cvar = cvar_sum / bf16((float)alpha_count);
float cvar = cvar_sum / (float)alpha_count;
// Scale: [0.25, 1.0]. CVaR >= 0 -> full size. CVaR < 0 -> reduce.
float scale = (cvar >= bf16_zero())
? bf16_one()
: bf16_fmax(bf16(0.25f), bf16_one() + cvar * bf16(5.0f));
float scale = (cvar >= 0.0f)
? 1.0f
: fmaxf(0.25f, 1.0f + cvar * 5.0f);
scales_out[i] = scale;
}

46
crates/ml/src/cuda_pipeline/iqn_dual_head_kernel.cu
View File

@@ -175,8 +175,8 @@ __device__ __forceinline__ void iqn_compute_offsets(
/**
* Per-sample IQN forward pass + quantile Huber loss.
*
* Reads h_s2 (shared trunk activation, bf16 from DQN trunk), converts to f32
* at boundary. All internal computation in f32.
* Reads h_s2 (shared trunk activation, f32 from DQN trunk).
* All internal computation in f32.
*
* τ values are pre-sampled and uploaded via the `taus` buffer (f32).
*
@@ -187,9 +187,9 @@ __device__ __forceinline__ void iqn_compute_offsets(
*/
extern "C" __global__
void iqn_forward_loss_kernel(
/* Inputs (bf16 from DQN trunk) */
const float* __restrict__ h_s2, /* [B, hidden_dim] trunk activations (bf16) */
const float* __restrict__ target_h_s2, /* [B, hidden_dim] target trunk activations (bf16) */
/* Inputs (f32 from DQN trunk) */
const float* __restrict__ h_s2, /* [B, hidden_dim] trunk activations (f32) */
const float* __restrict__ target_h_s2, /* [B, hidden_dim] target trunk activations (f32) */
/* Inputs (f32) */
const float* __restrict__ taus, /* [B, N] online τ samples ∈ (0,1) */
const float* __restrict__ target_taus, /* [B, N] target τ samples */
@@ -233,8 +233,8 @@ void iqn_forward_loss_kernel(
iqn_compute_offsets(hidden_dim, embed_dim, b0_size, b1_size, b2_size, b3_size, off);
/* ── Pointers into this sample's data ── */
const float* my_h_s2_bf16 = h_s2 + sample * hidden_dim;
const float* my_target_h_s2_bf16 = target_h_s2 + sample * hidden_dim;
const float* my_h_s2 = h_s2 + sample * hidden_dim;
const float* my_target_h_s2 = target_h_s2 + sample * hidden_dim;
const float* my_taus = taus + sample * IQN_NUM_QUANTILES;
const float* my_target_taus = target_taus + sample * IQN_NUM_QUANTILES;
@@ -268,14 +268,14 @@ void iqn_forward_loss_kernel(
const float* tw_b3 = target_params + off[8];
const float* tb_b3 = target_params + off[9];
/* ── Load h_s2 into distributed registers (bf16 → f32 at boundary) ── */
/* ── Load h_s2 into distributed registers ── */
float h_dist[IQN_DIST_MAX];
for (int d = tid; d < hidden_dim; d += IQN_BLOCK_SIZE)
h_dist[d / IQN_BLOCK_SIZE] = (float)my_h_s2_bf16[d];
h_dist[d / IQN_BLOCK_SIZE] = my_h_s2[d];
float h_target_dist[IQN_DIST_MAX];
for (int d = tid; d < hidden_dim; d += IQN_BLOCK_SIZE)
h_target_dist[d / IQN_BLOCK_SIZE] = (float)my_target_h_s2_bf16[d];
h_target_dist[d / IQN_BLOCK_SIZE] = my_target_h_s2[d];
/* ── Process each ONLINE quantile ── */
/* For each τ_i, compute quantile Q-values for the taken actions */
@@ -898,7 +898,7 @@ void iqn_adam_kernel(
/**
* IQN forward pass for inference — no loss, no activation saves.
*
* Reads bf16 h_s2 from DQN trunk, converts to f32 at boundary.
* Reads f32 h_s2 from DQN trunk.
* All params and taus are f32. Output is f32.
*
* Computes expected Q-values per action per branch by averaging over
@@ -908,7 +908,7 @@ void iqn_adam_kernel(
*/
extern "C" __global__
void iqn_forward_kernel(
const float* __restrict__ h_s2, /* [B, hidden_dim] (bf16 from DQN trunk) */
const float* __restrict__ h_s2, /* [B, hidden_dim] (f32 from DQN trunk) */
const float* __restrict__ taus, /* [B, N] (f32) */
const float* __restrict__ params, /* IQN weights (f32) */
const float* __restrict__ cos_features, /* [N, embed_dim] precomputed cosine features */
@@ -935,7 +935,7 @@ void iqn_forward_kernel(
int off[10];
iqn_compute_offsets(hidden_dim, embed_dim, b0_size, b1_size, b2_size, b3_size, off);
const float* my_h_s2_bf16 = h_s2 + sample * hidden_dim;
const float* my_h_s2 = h_s2 + sample * hidden_dim;
const float* my_taus = taus + sample * IQN_NUM_QUANTILES;
const float* w_embed = params + off[0];
const float* b_embed = params + off[1];
@@ -948,10 +948,10 @@ void iqn_forward_kernel(
const float* w_b3 = params + off[8];
const float* b_b3 = params + off[9];
/* Load h_s2 (bf16 → f32 at boundary) */
/* Load h_s2 into distributed registers */
float h_dist[IQN_DIST_MAX];
for (int h = tid; h < hidden_dim; h += IQN_BLOCK_SIZE)
h_dist[h / IQN_BLOCK_SIZE] = (float)my_h_s2_bf16[h];
h_dist[h / IQN_BLOCK_SIZE] = my_h_s2[h];
/* Accumulate Q-values across quantiles (for mean).
* Max tba = 3+3+3+3 = 12 — fits easily in registers. */
@@ -1035,18 +1035,18 @@ void iqn_forward_kernel(
* Used to compute target_h_s2 from next_states + target DQN weights.
* One block (256 threads) per sample. Uses shared memory for h1 intermediate.
*
* Inputs: bf16 states and DQN trunk weights (from external DQN system).
* Inputs: f32 states and DQN trunk weights (from external DQN system).
* Output: f32 h_s2_out (consumed by IQN forward/backward kernels as f32).
*
* Runtime params: state_dim, shared_h1, hidden_dim (=SHARED_H2)
*/
extern "C" __global__
void iqn_trunk_forward_kernel(
const float* __restrict__ states, /* [B, state_dim] (bf16) */
const float* __restrict__ w_s1, /* [shared_h1, state_dim] (bf16 DQN trunk weights) */
const float* __restrict__ b_s1, /* [shared_h1] (bf16) */
const float* __restrict__ w_s2, /* [hidden_dim, shared_h1] (bf16) */
const float* __restrict__ b_s2, /* [hidden_dim] (bf16) */
const float* __restrict__ states, /* [B, state_dim] (f32) */
const float* __restrict__ w_s1, /* [shared_h1, state_dim] (f32 DQN trunk weights) */
const float* __restrict__ b_s1, /* [shared_h1] (f32) */
const float* __restrict__ w_s2, /* [hidden_dim, shared_h1] (f32) */
const float* __restrict__ b_s2, /* [hidden_dim] (f32) */
float* __restrict__ h_s2_out, /* [B, hidden_dim] (f32 output) */
int batch_size,
int state_dim, /* runtime state dimension (replaces IQN_STATE_DIM) */
@@ -1072,7 +1072,7 @@ void iqn_trunk_forward_kernel(
int padded_sd = (state_dim + 127) & ~127;
const float* my_state = states + sample * padded_sd;
/* Layer 1: h1 = leaky_relu(W_s1 @ state + b_s1) — bf16 inputs, f32 compute */
/* Layer 1: h1 = leaky_relu(W_s1 @ state + b_s1) — f32 inputs and compute */
for (int h = tid; h < shared_h1; h += IQN_BLOCK_SIZE) {
float acc = (float)b_s1[h];
const float* w_row = w_s1 + h * state_dim;
@@ -1082,7 +1082,7 @@ void iqn_trunk_forward_kernel(
}
__syncthreads();
/* Layer 2: h2 = leaky_relu(W_s2 @ h1 + b_s2) — bf16 weights, f32 h1 and output */
/* Layer 2: h2 = leaky_relu(W_s2 @ h1 + b_s2) — f32 weights, h1 and output */
for (int h = tid; h < hidden_dim; h += IQN_BLOCK_SIZE) {
float acc = (float)b_s2[h];
const float* w_row = w_s2 + h * shared_h1;

80
crates/ml/src/cuda_pipeline/monitoring_kernel.cu
View File

@@ -1,25 +1,25 @@
// Custom atomic min/max for float via 16-bit CAS
// Custom atomic min/max for float via 32-bit CAS
__device__ void atomicMinBF16(float* addr, float val) {
unsigned short* addr_us = (unsigned short*)addr;
unsigned short old_bits = *addr_us;
unsigned short assumed;
int* addr_i = (int*)addr;
int old_bits = *addr_i;
int assumed;
do {
assumed = old_bits;
float old_bf = *(float*)&assumed;
float new_bf = bf16_fmin(old_bf, val);
old_bits = atomicCAS(addr_us, assumed, *(unsigned short*)&new_bf);
float old_f = __int_as_float(assumed);
float new_f = fminf(old_f, val);
old_bits = atomicCAS(addr_i, assumed, __float_as_int(new_f));
} while (assumed != old_bits);
}
__device__ void atomicMaxBF16(float* addr, float val) {
unsigned short* addr_us = (unsigned short*)addr;
unsigned short old_bits = *addr_us;
unsigned short assumed;
int* addr_i = (int*)addr;
int old_bits = *addr_i;
int assumed;
do {
assumed = old_bits;
float old_bf = *(float*)&assumed;
float new_bf = bf16_fmax(old_bf, val);
old_bits = atomicCAS(addr_us, assumed, *(unsigned short*)&new_bf);
float old_f = __int_as_float(assumed);
float new_f = fmaxf(old_f, val);
old_bits = atomicCAS(addr_i, assumed, __float_as_int(new_f));
} while (assumed != old_bits);
}
@@ -50,18 +50,18 @@ extern "C" __global__ void monitoring_reduce(
__shared__ int s_nonzero;
if (tid == 0) {
s_sum = bf16_zero(); s_sq_sum = bf16_zero();
s_min = bf16(1e30f); s_max = bf16(-1e30f);
s_sum = 0.0f; s_sq_sum = 0.0f;
s_min = 1e30f; s_max = -1e30f;
s_nonzero = 0;
for (int i = 0; i < 9; i++) s_exp[i] = 0;
for (int i = 0; i < 3; i++) { s_ord[i] = 0; s_urg[i] = 0; }
}
__syncthreads();
float local_sum = bf16_zero();
float local_sq = bf16_zero();
float local_min = bf16(1e30f);
float local_max = bf16(-1e30f);
float local_sum = 0.0f;
float local_sq = 0.0f;
float local_min = 1e30f;
float local_max = -1e30f;
int local_nonzero = 0;
int lc_exp[9] = {0,0,0,0,0,0,0,0,0};
int lc_ord[3] = {0,0,0};
@@ -69,15 +69,15 @@ extern "C" __global__ void monitoring_reduce(
for (int i = tid; i < N; i += stride) {
float rf = rewards[i];
float r = bf16(rf);
float r = rf;
/* Only accumulate non-zero rewards (trade completions).
* Sparse reward design: reward=0.0 during hold/flat, non-zero at trade exit.
* Computing mean/std/sharpe over all bars masks the signal in 98% zeros. */
if (rf != 0.0f) {
local_sum = local_sum + r;
local_sq = local_sq + r * r;
local_min = bf16_fmin(local_min, r);
local_max = bf16_fmax(local_max, r);
local_min = fminf(local_min, r);
local_max = fmaxf(local_max, r);
local_nonzero++;
}
/* Decode dir*mag + order + urgency from factored action:
@@ -94,10 +94,10 @@ extern "C" __global__ void monitoring_reduce(
// Warp-level reduction
for (int mask = 16; mask >= 1; mask >>= 1) {
local_sum = local_sum + bf16_shfl_xor(0xFFFFFFFF, local_sum, mask);
local_sq = local_sq + bf16_shfl_xor(0xFFFFFFFF, local_sq, mask);
local_min = bf16_fmin(local_min, bf16_shfl_xor(0xFFFFFFFF, local_min, mask));
local_max = bf16_fmax(local_max, bf16_shfl_xor(0xFFFFFFFF, local_max, mask));
local_sum = local_sum + __shfl_xor_sync(0xFFFFFFFF, local_sum, mask);
local_sq = local_sq + __shfl_xor_sync(0xFFFFFFFF, local_sq, mask);
local_min = fminf(local_min, __shfl_xor_sync(0xFFFFFFFF, local_min, mask));
local_max = fmaxf(local_max, __shfl_xor_sync(0xFFFFFFFF, local_max, mask));
local_nonzero += __shfl_xor_sync(0xFFFFFFFF, local_nonzero, mask);
for (int i = 0; i < 9; i++)
lc_exp[i] += __shfl_xor_sync(0xFFFFFFFF, lc_exp[i], mask);
@@ -107,8 +107,8 @@ extern "C" __global__ void monitoring_reduce(
}
}
if ((tid & 31) == 0) {
atomicAddBF16(&s_sum, local_sum);
atomicAddBF16(&s_sq_sum, local_sq);
atomicAdd(&s_sum, local_sum);
atomicAdd(&s_sq_sum, local_sq);
atomicMinBF16(&s_min, local_min);
atomicMaxBF16(&s_max, local_max);
atomicAdd(&s_nonzero, local_nonzero);
@@ -122,22 +122,22 @@ extern "C" __global__ void monitoring_reduce(
* not total bars. With sparse rewards (98% zero), per-bar mean ≈ 0 always.
* Per-trade mean gives a meaningful signal for monitoring. */
int n_trades = s_nonzero;
float denom = (n_trades > 0) ? bf16((float)n_trades) : bf16(1.0f);
float denom = (n_trades > 0) ? (float)n_trades : 1.0f;
float mean = s_sum / denom;
float var = (n_trades > 1)
? s_sq_sum / denom - mean * mean
: bf16_zero();
float std_val = bf16_sqrt(bf16_fmax(var, bf16_zero()));
: 0.0f;
float std_val = sqrtf(fmaxf(var, 0.0f));
summary[0] = mean;
summary[1] = std_val;
summary[2] = (n_trades > 0) ? s_min : bf16_zero();
summary[3] = (n_trades > 0) ? s_max : bf16_zero();
summary[4] = (std_val > bf16(1e-8f)) ? mean / std_val : bf16_zero();
for (int i = 0; i < 9; i++) summary[5 + i] = bf16((float)s_exp[i]);
for (int i = 0; i < 3; i++) summary[14 + i] = bf16((float)s_ord[i]);
for (int i = 0; i < 3; i++) summary[17 + i] = bf16((float)s_urg[i]);
summary[20] = bf16((float)N);
summary[21] = bf16((float)n_trades);
for (int i = 22; i < 24; i++) summary[i] = bf16_zero();
summary[2] = (n_trades > 0) ? s_min : 0.0f;
summary[3] = (n_trades > 0) ? s_max : 0.0f;
summary[4] = (std_val > 1e-8f) ? mean / std_val : 0.0f;
for (int i = 0; i < 9; i++) summary[5 + i] = (float)s_exp[i];
for (int i = 0; i < 3; i++) summary[14 + i] = (float)s_ord[i];
for (int i = 0; i < 3; i++) summary[17 + i] = (float)s_urg[i];
summary[20] = (float)N;
summary[21] = (float)n_trades;
for (int i = 22; i < 24; i++) summary[i] = 0.0f;
}
}

2
crates/ml/src/cuda_pipeline/mse_grad_kernel.cu
View File

@@ -13,7 +13,7 @@
extern "C" __global__ void mse_grad_kernel(
const float* __restrict__ save_probs, // [B, 4, NA] softmax probs
const float* __restrict__ save_eq_td, // [B, 4, NA] layout: [td_error, E_Q, 0, ...]
const float* __restrict__ is_weights, // [B] f32 (bf16 overflows to Inf)
const float* __restrict__ is_weights, // [B] f32
const int* __restrict__ actions, // [B] factored
float* __restrict__ d_value_logits, // [B, NA] f32
float* __restrict__ d_adv_logits, // [B, (B0+B1+B2+B3)*NA] f32

23
crates/ml/src/cuda_pipeline/mse_loss_kernel.cu
View File

@@ -1,8 +1,7 @@
/**
* MSE loss kernel on expected Q-values for DQN warmup.
*
* Mixed-precision: reads/writes BF16 global memory, ALL arithmetic in float.
* BF16 softmax exp() overflows at logit > 11.1 → NaN. Float avoids this.
* Pure f32: reads/writes f32 global memory, all arithmetic in float.
*
* Grid = (batch_size, 1, 1)
* Block = (BLOCK_THREADS=256, 1, 1)
@@ -95,7 +94,7 @@ __device__ float block_softmax_expected_q_f(
float lp = shmem_logits[i] - bmax - log_sum_exp;
float prob = expf(lp);
shmem_logits[i] = prob; /* overwrite logits with probs */
if (save_probs_out) save_probs_out[i] = bf16(prob);
if (save_probs_out) save_probs_out[i] = prob;
}
__syncthreads();
@@ -108,7 +107,7 @@ __device__ float block_softmax_expected_q_f(
/* ══════════════════════════════════════════════════════════════════════
* MAIN KERNEL: mse_loss_batched (float arithmetic, BF16 I/O)
* MAIN KERNEL: mse_loss_batched (native f32)
* ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void mse_loss_batched(
@@ -137,7 +136,7 @@ extern "C" __global__ void mse_loss_batched(
const int* __restrict__ actions, /* [B] factored action indices 0-44 */
const float* __restrict__ rewards, /* [B] f32 */
const float* __restrict__ dones, /* [B] f32 (0.0/1.0) */
const float* __restrict__ is_weights, /* [B] PER importance-sampling weights (f32 — bf16 overflows to Inf) */
const float* __restrict__ is_weights, /* [B] PER importance-sampling weights (f32) */
/* ── Outputs ──────────────────────────────────────────────────── */
float* __restrict__ per_sample_loss, /* [B] IS-weighted loss per sample */
@@ -173,7 +172,7 @@ extern "C" __global__ void mse_loss_batched(
/* ── Spectral decoupling: L2 penalty on Q-value logit magnitudes ── */
float spectral_decoupling_lambda /* Pezeshki et al. 2021 (0.0=disabled, 0.01=default) */
) {
/* Shared memory is now float (4 bytes/elem) — doubled from BF16 version */
/* Shared memory: float (4 bytes/elem) */
extern __shared__ float shmem_f[];
/* Read adaptive v_range from device buffer (graph-safe, L1 cached) */
@@ -477,13 +476,13 @@ extern "C" __global__ void mse_loss_batched(
total_mse += mse;
total_abs_td += fabsf(td);
/* Save td_error and E[Q] for gradient kernel (BF16 output) */
/* Save td_error and E[Q] for gradient kernel */
if (tid == 0) {
save_projected[save_off + 0] = bf16(td);
save_projected[save_off + 1] = bf16(online_eq);
save_projected[save_off + 0] = td;
save_projected[save_off + 1] = online_eq;
}
for (int j = tid + 2; j < num_atoms; j += BLOCK_THREADS)
save_projected[save_off + j] = bf16_zero();
save_projected[save_off + j] = 0.0f;
__syncthreads();
} /* end branch loop */
@@ -506,8 +505,8 @@ extern "C" __global__ void mse_loss_batched(
if (tid == 0) {
float weighted_loss = avg_mse * is_weight;
per_sample_loss[sample_id] = bf16(weighted_loss);
td_errors[sample_id] = bf16(avg_td);
per_sample_loss[sample_id] = weighted_loss;
td_errors[sample_id] = avg_td;
/* total_loss reduced by deterministic c51_loss_reduce kernel.
* No atomicAdd — fully deterministic training. */
}

240
crates/ml/src/cuda_pipeline/ppo_experience_kernel.cu
View File

@@ -1,5 +1,5 @@
/**
* Zero-Roundtrip PPO Experience Collection Kernel — Native BF16
* Zero-Roundtrip PPO Experience Collection Kernel — Native F32
*
* Requires common_device_functions.cuh prepended via NVRTC source concatenation.
* Launch config: grid=(ceil(N/32),1,1), block=(32,1,1).
@@ -8,9 +8,7 @@
* Phase A: Forward rollout (L timesteps) — actor, critic, portfolio, rewards
* Phase B: Backward GAE scan — advantages and returns
*
* ALL internal computation uses float. Host scalar parameters
* (gamma, gae_lambda, etc.) arrive as float and are converted once at
* kernel entry via bf16().
* ALL internal computation uses float.
*/
/* Override NUM_ACTIONS for PPO: use full 45-action factored space */
@@ -32,13 +30,12 @@
#define MAX_GAE_LEN 500
/* ------------------------------------------------------------------ */
/* BF16 matvec overload: bf16 weights, bf16 input, bf16 output */
/* F32 matvec: f32 weights, f32 input, f32 output */
/* ------------------------------------------------------------------ */
/**
* Matrix-vector multiply: output = W * input + b, with optional LeakyReLU.
* Fully native BF16: weights, biases, input, and output are all float.
* Accumulation in BF16 (matching tensor-core semantics on SM80+).
* Fully native F32: weights, biases, input, and output are all float.
*/
__device__ void matvec_leaky_relu_bf16(
const float* __restrict__ W,
@@ -60,11 +57,11 @@ __device__ void matvec_leaky_relu_bf16(
}
/* ------------------------------------------------------------------ */
/* PPO-Specific Device Functions (Native BF16) */
/* PPO-Specific Device Functions (Native F32) */
/* ------------------------------------------------------------------ */
/**
* PPO Actor MLP forward pass — native BF16.
* PPO Actor MLP forward pass — native F32.
*
* state[54] -> h1[128] (LeakyReLU) -> h2[64] (LeakyReLU) -> logits[45]
*
@@ -102,7 +99,7 @@ __device__ void ppo_actor_forward(
}
/**
* Stable softmax + categorical sampling — native BF16.
* Stable softmax + categorical sampling — native F32.
*
* 1. Find max logit for numerical stability
* 2. exp(logit - max) and accumulate sum
@@ -126,25 +123,25 @@ __device__ void softmax_sample(
/* Step 1: Find max logit for numerical stability */
float max_logit = logits[0];
for (int i = 1; i < NUM_ACTIONS; i++) {
max_logit = bf16_fmax(max_logit, logits[i]);
max_logit = fmaxf(max_logit, logits[i]);
}
/* Step 2: exp(logit - max) and accumulate sum */
float sum_exp = bf16_zero();
float sum_exp = 0.0f;
for (int i = 0; i < NUM_ACTIONS; i++) {
probs[i] = bf16_exp(logits[i] - max_logit);
probs[i] = expf(logits[i] - max_logit);
sum_exp = sum_exp + probs[i];
}
/* Step 3: Normalize to probabilities */
float inv_sum = bf16_one() / bf16_fmax(sum_exp, bf16(1e-8f));
float inv_sum = 1.0f / fmaxf(sum_exp, 1e-8f);
for (int i = 0; i < NUM_ACTIONS; i++) {
probs[i] = probs[i] * inv_sum;
}
/* Step 4: CDF scan + random draw -> action index */
float u = bf16(gpu_random(rng));
float cdf = bf16_zero();
float u = gpu_random(rng);
float cdf = 0.0f;
int action = NUM_ACTIONS - 1; /* default to last action */
for (int i = 0; i < NUM_ACTIONS; i++) {
cdf = cdf + probs[i];
@@ -155,9 +152,9 @@ __device__ void softmax_sample(
}
/* Step 5: log probability for PPO loss */
float p = bf16_fmax(probs[action], bf16(1e-8f)); /* clamp for log safety */
float p = fmaxf(probs[action], 1e-8f); /* clamp for log safety */
*out_action = action;
*out_log_prob = bf16_log(p);
*out_log_prob = logf(p);
}
/**
@@ -234,8 +231,8 @@ __device__ float ppo_critic_forward(
* @param advantages Output advantage estimates [L] (float)
* @param returns Output return targets [L] (float)
* @param L Number of timesteps
* @param gm Discount factor (bf16)
* @param lm GAE lambda parameter (bf16)
* @param gm Discount factor
* @param lm GAE lambda parameter
*/
__device__ void compute_gae_backward(
const float* rewards,
@@ -247,10 +244,10 @@ __device__ void compute_gae_backward(
float gm,
float lm
) {
float gae = bf16_zero();
float gae = 0.0f;
for (int t = L - 1; t >= 0; t--) {
float not_done = bf16_one() - dones[t];
float not_done = 1.0f - dones[t];
float delta = rewards[t] + gm * values[t + 1] * not_done - values[t];
gae = delta + gm * lm * not_done * gae;
advantages[t] = gae;
@@ -259,8 +256,7 @@ __device__ void compute_gae_backward(
}
/**
* BF16 curiosity inference: builds input in bf16, calls through bf16 matvec,
* returns bf16 MSE clamped to max_reward.
* Curiosity inference: builds input, runs MLP, returns MSE clamped to max_reward.
*/
__device__ float curiosity_inference_bf16(
const float* state,
@@ -284,9 +280,9 @@ __device__ float curiosity_inference_bf16(
else if (action_idx == 2) category = 1; /* Flat */
else category = 2; /* Long50/Long100 */
input[MARKET_DIM + 0] = (category == 0) ? bf16_one() : bf16_zero();
input[MARKET_DIM + 1] = (category == 1) ? bf16_one() : bf16_zero();
input[MARKET_DIM + 2] = (category == 2) ? bf16_one() : bf16_zero();
input[MARKET_DIM + 0] = (category == 0) ? 1.0f : 0.0f;
input[MARKET_DIM + 1] = (category == 1) ? 1.0f : 0.0f;
input[MARKET_DIM + 2] = (category == 2) ? 1.0f : 0.0f;
/* Hidden layer */
matvec_leaky_relu_bf16(w_c1, b_c1, input, scratch, CUR_INPUT, CUR_HIDDEN, 1);
@@ -296,18 +292,18 @@ __device__ float curiosity_inference_bf16(
matvec_leaky_relu_bf16(w_c2, b_c2, scratch, pred, CUR_HIDDEN, CUR_OUTPUT, 0);
/* MSE against actual next_state features (first MARKET_DIM) */
float mse = bf16_zero();
float mse = 0.0f;
for (int i = 0; i < CUR_OUTPUT; i++) {
float diff = pred[i] - next_state[i];
mse = mse + diff * diff;
}
mse = mse / bf16((float)CUR_OUTPUT);
mse = mse / (float)CUR_OUTPUT;
return bf16_fmin(mse, max_reward);
return fminf(mse, max_reward);
}
/**
* BF16 barrier_init: initializes triple-barrier on position open.
* barrier_init: initializes triple-barrier on position open.
* barrier_state[5] and barrier_config[3] are both float.
*/
__device__ __forceinline__ void barrier_init_bf16(
@@ -323,12 +319,12 @@ __device__ __forceinline__ void barrier_init_bf16(
barrier_state[0] = entry_price;
barrier_state[1] = entry_price * profit_mult;
barrier_state[2] = entry_price * loss_mult;
barrier_state[3] = bf16((float)(current_step + max_hold));
barrier_state[4] = bf16_zero();
barrier_state[3] = (float)(current_step + max_hold);
barrier_state[4] = 0.0f;
}
/**
* BF16 barrier_check: check triple-barrier conditions.
* barrier_check: check triple-barrier conditions.
* Returns label: +1 = profit, -1 = loss, 0 = pending/expired.
*/
__device__ __forceinline__ int barrier_check_bf16(
@@ -337,15 +333,15 @@ __device__ __forceinline__ int barrier_check_bf16(
int current_step,
float position
) {
if (barrier_state[0] <= bf16_zero()) return 0;
if (barrier_state[0] <= 0.0f) return 0;
float upper = barrier_state[1];
float lower = barrier_state[2];
int expiry = (int)(barrier_state[3]);
float sign = (position >= bf16_zero()) ? bf16_one() : bf16(-1.0f);
float sign = (position >= 0.0f) ? 1.0f : -1.0f;
int label = 0;
if (sign > bf16_zero()) {
if (sign > 0.0f) {
if (price >= upper) label = 1;
else if (price <= lower) label = -1;
} else {
@@ -356,27 +352,27 @@ __device__ __forceinline__ int barrier_check_bf16(
if (label == 0 && current_step >= expiry) {
float entry = barrier_state[0];
float pnl = (price - entry) * sign;
label = (pnl > bf16_zero()) ? 1 : -1;
label = (pnl > 0.0f) ? 1 : -1;
}
if (label != 0) {
barrier_state[4] = bf16((float)label);
barrier_state[4] = (float)label;
}
return label;
}
/** BF16 barrier reset. */
/** barrier reset. */
__device__ __forceinline__ void barrier_reset_bf16(float* barrier_state) {
barrier_state[0] = bf16_zero();
barrier_state[1] = bf16_zero();
barrier_state[2] = bf16_zero();
barrier_state[3] = bf16_zero();
barrier_state[4] = bf16_zero();
barrier_state[0] = 0.0f;
barrier_state[1] = 0.0f;
barrier_state[2] = 0.0f;
barrier_state[3] = 0.0f;
barrier_state[4] = 0.0f;
}
/**
* BF16 diversity entropy penalty.
* Returns bf16 penalty value. Window/meta are int arrays (unchanged).
* Diversity entropy penalty.
* Returns f32 penalty value. Window/meta are int arrays (unchanged).
*/
__device__ float diversity_entropy_bf16(
int* diversity_window,
@@ -397,7 +393,7 @@ __device__ float diversity_entropy_bf16(
diversity_meta[1] = count;
}
if (count < 2) return bf16_zero();
if (count < 2) return 0.0f;
int counts[DQN_NUM_ACTIONS];
for (int i = 0; i < DQN_NUM_ACTIONS; i++) counts[i] = 0;
@@ -406,22 +402,22 @@ __device__ float diversity_entropy_bf16(
if (c >= 0 && c < DQN_NUM_ACTIONS) counts[c]++;
}
/* Shannon entropy in bf16 */
float entropy = bf16_zero();
float inv_n = bf16_one() / bf16((float)count);
/* Shannon entropy in f32 */
float entropy = 0.0f;
float inv_n = 1.0f / (float)count;
for (int c = 0; c < DQN_NUM_ACTIONS; c++) {
if (counts[c] > 0) {
float p = bf16((float)counts[c]) * inv_n;
float p = (float)counts[c] * inv_n;
/* log2(p) = log(p) / log(2) */
entropy = entropy - p * bf16_log(p) / bf16(0.6931472f);
entropy = entropy - p * logf(p) / 0.6931472f;
}
}
float max_entropy = bf16_log(bf16((float)DQN_NUM_ACTIONS)) / bf16(0.6931472f);
float max_entropy = logf((float)DQN_NUM_ACTIONS) / 0.6931472f;
if (entropy < max_entropy) {
return bf16_zero() - ((max_entropy - entropy) / max_entropy) * bf16(0.1f);
return -((max_entropy - entropy) / max_entropy) * 0.1f;
}
return bf16_zero();
return 0.0f;
}
/* ------------------------------------------------------------------ */
@@ -429,7 +425,7 @@ __device__ float diversity_entropy_bf16(
/* ------------------------------------------------------------------ */
/**
* Full PPO experience collection kernel — native BF16.
* Full PPO experience collection kernel — native F32.
*
* Each thread runs one independent episode of L timesteps (Phase A),
* then performs a backward GAE scan (Phase B).
@@ -470,7 +466,7 @@ extern "C" __global__ void ppo_full_experience_kernel(
const float* __restrict__ cur_w2, /* [CUR_OUTPUT, CUR_HIDDEN] */
const float* __restrict__ cur_b2, /* [CUR_OUTPUT] */
/* ---- Per-episode mutable state arrays (BF16) ---- */
/* ---- Per-episode mutable state arrays (f32) ---- */
float* portfolio_states, /* [N, PORTFOLIO_STATE_SIZE] */
float* barrier_states, /* [N, BARRIER_STATE_SIZE] */
int* diversity_windows, /* [N, DIVERSITY_WINDOW] */
@@ -479,7 +475,7 @@ extern "C" __global__ void ppo_full_experience_kernel(
/* ---- Barrier config (shared) ---- */
const float* __restrict__ barrier_config, /* [3]: profit_mult, loss_mult, max_bars */
/* ---- Scalar configs (host floats, converted to bf16 at entry) ---- */
/* ---- Scalar configs (host floats) ---- */
float max_position,
int episode_length,
int total_bars,
@@ -496,7 +492,7 @@ extern "C" __global__ void ppo_full_experience_kernel(
/* ---- RNG states [N] ---- */
unsigned int* rng_states,
/* ---- Output arrays (BF16) ---- */
/* ---- Output arrays (f32) ---- */
float* out_states, /* [N, L, STATE_DIM] */
int* out_actions, /* [N, L] */
float* out_log_probs, /* [N, L] */
@@ -507,15 +503,15 @@ extern "C" __global__ void ppo_full_experience_kernel(
int tid = blockIdx.x * blockDim.x + threadIdx.x;
if (tid >= N) return;
/* ---- Convert host scalar params to bf16 once ---- */
float bf_max_position = bf16(max_position);
float bf_gamma = bf16(gamma);
float bf_gae_lambda = bf16(gae_lambda);
float bf_curiosity_max_rew = bf16(curiosity_max_reward);
float bf_barrier_scale = bf16(barrier_scale);
float bf_diversity_scale = bf16(diversity_scale);
float bf_curiosity_scale = bf16(curiosity_scale);
float bf_risk_weight = bf16(risk_weight);
/* ---- Local aliases for scalar params ---- */
float bf_max_position = max_position;
float bf_gamma = gamma;
float bf_gae_lambda = gae_lambda;
float bf_curiosity_max_rew = curiosity_max_reward;
float bf_barrier_scale = barrier_scale;
float bf_diversity_scale = diversity_scale;
float bf_curiosity_scale = curiosity_scale;
float bf_risk_weight = risk_weight;
/* ---- Load per-thread portfolio state ---- */
int ps_off = tid * PORTFOLIO_STATE_SIZE;
@@ -580,13 +576,13 @@ extern "C" __global__ void ppo_full_experience_kernel(
/* Handle out-of-data */
if (global_bar >= total_bars - 1) {
for (int i = 0; i < STATE_DIM; i++)
out_states[out_off * STATE_DIM + i] = bf16_zero();
out_states[out_off * STATE_DIM + i] = 0.0f;
out_actions[out_off] = 0;
out_log_probs[out_off] = bf16_zero();
out_log_probs[out_off] = 0.0f;
out_dones[out_off] = 1;
gae_values[t] = bf16_zero();
gae_rewards[t] = bf16_zero();
gae_dones[t] = bf16_one();
gae_values[t] = 0.0f;
gae_rewards[t] = 0.0f;
gae_dones[t] = 1.0f;
continue;
}
@@ -602,12 +598,12 @@ extern "C" __global__ void ppo_full_experience_kernel(
float current_close_raw = targets[t_off + 2];
float next_close_raw = targets[t_off + 3];
float price = (current_close_raw != bf16_zero()) ? current_close_raw : current_close;
if (price <= bf16_zero()) price = bf16_one();
float price = (current_close_raw != 0.0f) ? current_close_raw : current_close;
if (price <= 0.0f) price = 1.0f;
float current_value = cash + position * price;
float current_norm = current_value / initial_cap;
float max_pos_norm = (price > bf16_zero()) ? initial_cap / price : bf16_one();
float max_pos_norm = (price > 0.0f) ? initial_cap / price : 1.0f;
float pos_norm = position / max_pos_norm;
state[MARKET_DIM + 0] = current_norm; /* normalized value */
@@ -637,62 +633,62 @@ extern "C" __global__ void ppo_full_experience_kernel(
gae_values[t] = value;
/* ---- Step 5: Portfolio simulation ---- */
float target_exposure = bf16(factored_action_to_exposure(action_idx));
float target_exposure = factored_action_to_exposure(action_idx);
float target_position = target_exposure * bf_max_position;
float tx_rate = bf16(ppo_action_to_tx_cost(action_idx));
float tx_rate = ppo_action_to_tx_cost(action_idx);
/* Detect reversal (sign change) */
int is_reversal = (position > bf16_zero() && target_position < bf16_zero()) ||
(position < bf16_zero() && target_position > bf16_zero());
int is_reversal = (position > 0.0f && target_position < 0.0f) ||
(position < 0.0f && target_position > 0.0f);
if (is_reversal) {
/* Phase 1: Close current position */
float close_cash = position * price;
float close_cost = bf16_fabs(position) * price * tx_rate;
float close_cost = fabsf(position) * price * tx_rate;
cash = cash + close_cash - close_cost;
cum_costs = cum_costs + close_cost;
/* Phase 2: Open opposite position */
float reserve = (reserve_pct > bf16_zero()) ? current_value * (reserve_pct / bf16(100.0f)) : bf16_zero();
float affordable = bf16_fmax(cash - reserve, bf16_zero());
float max_contracts = (price > bf16_zero()) ? affordable / (price * (bf16_one() + tx_rate)) : bf16_zero();
max_contracts = bf16_floor(max_contracts);
float actual = bf16_fmin(max_contracts, bf16_fabs(target_position));
float reserve = (reserve_pct > 0.0f) ? current_value * (reserve_pct / 100.0f) : 0.0f;
float affordable = fmaxf(cash - reserve, 0.0f);
float max_contracts = (price > 0.0f) ? affordable / (price * (1.0f + tx_rate)) : 0.0f;
max_contracts = floorf(max_contracts);
float actual = fminf(max_contracts, fabsf(target_position));
if (actual > bf16_zero()) {
float new_pos = (target_position > bf16_zero()) ? actual : bf16_zero() - actual;
if (actual > 0.0f) {
float new_pos = (target_position > 0.0f) ? actual : -actual;
float open_cost = actual * price * tx_rate;
cash = cash - new_pos * price - open_cost;
cum_costs = cum_costs + open_cost;
position = new_pos;
entry_price = price;
} else {
position = bf16_zero();
entry_price = bf16_zero();
position = 0.0f;
entry_price = 0.0f;
}
} else {
/* Non-reversal: adjust position directly */
float delta = target_position - position;
if (bf16_fabs(delta) > bf16_zero()) {
float trade_cost = bf16_fabs(delta) * price * tx_rate;
if (fabsf(delta) > 0.0f) {
float trade_cost = fabsf(delta) * price * tx_rate;
cum_costs = cum_costs + trade_cost;
cash = cash - trade_cost;
/* Cash reserve check for buys */
if (delta > bf16_zero() && reserve_pct > bf16_zero()) {
if (delta > 0.0f && reserve_pct > 0.0f) {
float pv = cash + position * price;
float reserve = pv * (reserve_pct / bf16(100.0f));
float reserve = pv * (reserve_pct / 100.0f);
float buy_cost = delta * price;
if (cash - buy_cost < reserve) {
float affordable = bf16_fmax(cash - reserve, bf16_zero());
delta = bf16_fmin(delta, (price > bf16_zero()) ? bf16_floor(affordable / price) : bf16_zero());
float affordable = fmaxf(cash - reserve, 0.0f);
delta = fminf(delta, (price > 0.0f) ? floorf(affordable / price) : 0.0f);
}
}
if (delta > bf16_zero()) {
if (delta > 0.0f) {
entry_price = price;
} else if (target_position == bf16_zero()) {
entry_price = bf16_zero();
} else if (target_position == 0.0f) {
entry_price = 0.0f;
}
cash = cash - delta * price;
position = position + delta;
@@ -702,7 +698,7 @@ extern "C" __global__ void ppo_full_experience_kernel(
/* ---- Step 6: Barrier tracking ---- */
float old_barrier_entry = barrier_st[0];
if (entry_price > bf16_zero() && old_barrier_entry <= bf16_zero()) {
if (entry_price > 0.0f && old_barrier_entry <= 0.0f) {
barrier_init_bf16(barrier_st, barrier_config, entry_price, global_bar);
}
@@ -715,8 +711,8 @@ extern "C" __global__ void ppo_full_experience_kernel(
float div_penalty = diversity_entropy_bf16(div_window, div_meta, action_idx);
/* ---- Step 8: Mark-to-market + build next_state for curiosity ---- */
float next_price = (next_close_raw != bf16_zero()) ? next_close_raw : next_close;
if (next_price <= bf16_zero()) next_price = price;
float next_price = (next_close_raw != 0.0f) ? next_close_raw : next_close;
if (next_price <= 0.0f) next_price = price;
float next_value = cash + position * next_price;
float next_norm = next_value / initial_cap;
@@ -729,13 +725,13 @@ extern "C" __global__ void ppo_full_experience_kernel(
for (int i = 0; i < MARKET_DIM; i++)
next_state[i] = state[i];
}
float next_max_pos_norm = (next_price > bf16_zero()) ? initial_cap / next_price : bf16_one();
float next_max_pos_norm = (next_price > 0.0f) ? initial_cap / next_price : 1.0f;
next_state[MARKET_DIM + 0] = next_norm;
next_state[MARKET_DIM + 1] = position / next_max_pos_norm;
next_state[MARKET_DIM + 2] = spread;
/* ---- Step 9: Curiosity inference ---- */
float curiosity_reward = bf16_zero();
float curiosity_reward = 0.0f;
if (cur_w1 != 0) {
curiosity_reward = curiosity_inference_bf16(
state, next_state, action_idx,
@@ -745,20 +741,20 @@ extern "C" __global__ void ppo_full_experience_kernel(
}
/* ---- Step 10: Risk penalty (drawdown) ---- */
float pnl_reward = bf16_zero();
if (current_norm > bf16_zero()) {
float pnl_reward = 0.0f;
if (current_norm > 0.0f) {
pnl_reward = (next_norm - current_norm) / current_norm;
}
float abs_pos = bf16_fabs(pos_norm);
if (abs_pos > bf16(0.8f)) {
pnl_reward = pnl_reward - (abs_pos - bf16(0.8f)) * bf16(5.0f) * bf_risk_weight;
float abs_pos = fabsf(pos_norm);
if (abs_pos > 0.8f) {
pnl_reward = pnl_reward - (abs_pos - 0.8f) * 5.0f * bf_risk_weight;
}
/* ---- Step 11: Reward combination ---- */
float barrier_mult = bf16_one();
float barrier_mult = 1.0f;
if (barrier_label != 0) {
barrier_mult = bf16_one() + bf_barrier_scale * bf16((float)barrier_label);
barrier_mult = 1.0f + bf_barrier_scale * (float)barrier_label;
}
float combined_reward = pnl_reward * barrier_mult
@@ -781,15 +777,15 @@ extern "C" __global__ void ppo_full_experience_kernel(
/* Store for GAE backward scan */
gae_rewards[t] = combined_reward;
gae_dones[t] = bf16((float)done);
gae_dones[t] = (float)done;
/* ---- Step 14: Episode reset on done ---- */
if (done) {
cash = initial_cap;
position = bf16_zero();
entry_price = bf16_zero();
cum_costs = bf16_zero();
last_price = bf16_zero();
position = 0.0f;
entry_price = 0.0f;
cum_costs = 0.0f;
last_price = 0.0f;
step_in_episode = 0;
barrier_reset_bf16(barrier_st);
/* Reset diversity window */
@@ -809,7 +805,7 @@ extern "C" __global__ void ppo_full_experience_kernel(
int last_out_off = tid * actual_L + last_t;
if (out_dones[last_out_off] == 1) {
/* Last step was terminal — bootstrap value is 0 */
gae_values[actual_L] = bf16_zero();
gae_values[actual_L] = 0.0f;
} else {
/* Last step was not terminal — run critic on next_state for bootstrap */
int last_global_bar = ep_start + last_t;
@@ -828,17 +824,17 @@ extern "C" __global__ void ppo_full_experience_kernel(
}
/* Approximate portfolio features from current thread state */
float boot_price_raw = bf16_zero();
float boot_price_raw = 0.0f;
if (next_bar_boot < total_bars) {
int t_off_boot = next_bar_boot * 4;
boot_price_raw = targets[t_off_boot + 2];
if (boot_price_raw <= bf16_zero()) boot_price_raw = targets[t_off_boot + 0];
if (boot_price_raw <= 0.0f) boot_price_raw = targets[t_off_boot + 0];
}
if (boot_price_raw <= bf16_zero()) boot_price_raw = bf16_one();
if (boot_price_raw <= 0.0f) boot_price_raw = 1.0f;
float boot_value = cash + position * boot_price_raw;
float boot_norm = boot_value / initial_cap;
float boot_max_pos = (boot_price_raw > bf16_zero()) ? initial_cap / boot_price_raw : bf16_one();
float boot_max_pos = (boot_price_raw > 0.0f) ? initial_cap / boot_price_raw : 1.0f;
boot_state[MARKET_DIM + 0] = boot_norm;
boot_state[MARKET_DIM + 1] = position / boot_max_pos;

2
crates/ml/src/cuda_pipeline/relu_mask_kernel.cu
View File

@@ -3,7 +3,7 @@
*
* dx[i] = (activation[i] > 0.0f) ? dx[i] : 0.0f
*
* dx is f32 (backward dX scratch), activation is bf16 (saved forward output).
* dx is f32 (backward dX scratch), activation is f32 (saved forward output).
*
* Launch config: grid=(ceil(n/256), 1, 1), block=(256, 1, 1).
*/

44
crates/ml/src/cuda_pipeline/statistics_kernel.cu
View File

@@ -5,7 +5,7 @@
* derived statistics (mean, variance, Sharpe, win rate) from these.
* This avoids the multi-block atomicAdd-of-means correctness bug.
*
* Output layout: 10 bf16 values
* Output layout: 10 f32 values
* [0] reward_sum
* [1] reward_sq_sum
* [2] global_min_pv (min portfolio value -- for drawdown)
@@ -38,25 +38,25 @@ extern "C" __global__ void batch_statistics(
int tid = threadIdx.x;
/* Load phase -- grid-stride loop */
float local_sum = bf16_zero();
float local_sq_sum = bf16_zero();
float local_min_val = bf16(1e30f);
float local_max_val = bf16(-1e30f);
float local_sum = 0.0f;
float local_sq_sum = 0.0f;
float local_min_val = 1e30f;
float local_max_val = -1e30f;
int local_wins = 0;
float local_pos_sum = bf16_zero();
float local_pos_sq_sum = bf16_zero();
float local_pos_sum = 0.0f;
float local_pos_sq_sum = 0.0f;
for (int i = tid; i < N; i += blockDim.x) {
float r = rewards[i];
local_sum = local_sum + r;
local_sq_sum = local_sq_sum + r * r;
if (r > bf16_zero()) local_wins++;
if (r > 0.0f) local_wins++;
float pv = portfolio_values[i];
local_min_val = bf16_fmin(local_min_val, pv);
local_max_val = bf16_fmax(local_max_val, pv);
local_min_val = fminf(local_min_val, pv);
local_max_val = fmaxf(local_max_val, pv);
float pos = bf16_fabs(positions[i]);
float pos = fabsf(positions[i]);
local_pos_sum = local_pos_sum + pos;
local_pos_sq_sum = local_pos_sq_sum + pos * pos;
}
@@ -75,8 +75,8 @@ extern "C" __global__ void batch_statistics(
if (tid < s) {
s_sum[tid] = s_sum[tid] + s_sum[tid + s];
s_sq_sum[tid] = s_sq_sum[tid] + s_sq_sum[tid + s];
s_min[tid] = bf16_fmin(s_min[tid], s_min[tid + s]);
s_max[tid] = bf16_fmax(s_max[tid], s_max[tid + s]);
s_min[tid] = fminf(s_min[tid], s_min[tid + s]);
s_max[tid] = fmaxf(s_max[tid], s_max[tid + s]);
s_win[tid] += s_win[tid + s];
s_pos_sum[tid] = s_pos_sum[tid] + s_pos_sum[tid + s];
s_pos_sq_sum[tid] = s_pos_sq_sum[tid] + s_pos_sq_sum[tid + s];
@@ -94,13 +94,13 @@ extern "C" __global__ void batch_statistics(
float my_pos_sum = s_pos_sum[tid];
float my_pos_sq_sum = s_pos_sq_sum[tid];
for (int mask = 16; mask >= 1; mask >>= 1) {
my_sum = my_sum + bf16_shfl_xor(0xFFFFFFFF, my_sum, mask);
my_sq_sum = my_sq_sum + bf16_shfl_xor(0xFFFFFFFF, my_sq_sum, mask);
my_min = bf16_fmin(my_min, bf16_shfl_xor(0xFFFFFFFF, my_min, mask));
my_max = bf16_fmax(my_max, bf16_shfl_xor(0xFFFFFFFF, my_max, mask));
my_sum = my_sum + __shfl_xor_sync(0xFFFFFFFF, my_sum, mask);
my_sq_sum = my_sq_sum + __shfl_xor_sync(0xFFFFFFFF, my_sq_sum, mask);
my_min = fminf(my_min, __shfl_xor_sync(0xFFFFFFFF, my_min, mask));
my_max = fmaxf(my_max, __shfl_xor_sync(0xFFFFFFFF, my_max, mask));
my_win += __shfl_xor_sync(0xFFFFFFFF, my_win, mask);
my_pos_sum = my_pos_sum + bf16_shfl_xor(0xFFFFFFFF, my_pos_sum, mask);
my_pos_sq_sum = my_pos_sq_sum + bf16_shfl_xor(0xFFFFFFFF, my_pos_sq_sum, mask);
my_pos_sum = my_pos_sum + __shfl_xor_sync(0xFFFFFFFF, my_pos_sum, mask);
my_pos_sq_sum = my_pos_sq_sum + __shfl_xor_sync(0xFFFFFFFF, my_pos_sq_sum, mask);
}
if (tid == 0) {
s_sum[0] = my_sum;
@@ -120,11 +120,11 @@ extern "C" __global__ void batch_statistics(
output_stats[1] = s_sq_sum[0]; /* reward_sq_sum */
output_stats[2] = s_min[0]; /* global_min_pv */
output_stats[3] = s_max[0]; /* global_max_pv */
output_stats[4] = bf16((float)s_win[0]); /* win_count */
output_stats[5] = bf16((float)N); /* total_count */
output_stats[4] = (float)s_win[0]; /* win_count */
output_stats[5] = (float)N; /* total_count */
output_stats[6] = s_sum[0]; /* total_pnl (alias) */
output_stats[7] = s_pos_sum[0]; /* position_sum */
output_stats[8] = s_pos_sq_sum[0]; /* position_sq_sum */
output_stats[9] = bf16_zero(); /* reserved */
output_stats[9] = 0.0f; /* reserved */
}
}

10
crates/ml/src/cuda_pipeline/trade_stats_kernel.cu
View File

@@ -23,13 +23,13 @@ extern "C" __global__ void trade_stats_reduce(
// Init shared memory
if (tid == 0) {
for (int k = 0; k < 6; k++) s_sums[k] = bf16_zero();
for (int k = 0; k < 6; k++) s_sums[k] = 0.0f;
}
__syncthreads();
// Thread-local accumulators — native BF16
// Thread-local accumulators — native f32
float local_sums[6];
for (int k = 0; k < 6; k++) local_sums[k] = bf16_zero();
for (int k = 0; k < 6; k++) local_sums[k] = 0.0f;
for (int i = tid; i < N; i += stride) {
int base = i * PORTFOLIO_STRIDE;
@@ -41,13 +41,13 @@ extern "C" __global__ void trade_stats_reduce(
// Warp-level reduction before atomic to shared (reduces contention 32x)
for (int mask = 16; mask >= 1; mask >>= 1) {
for (int k = 0; k < 6; k++) {
local_sums[k] = local_sums[k] + bf16_shfl_xor(0xFFFFFFFF, local_sums[k], mask);
local_sums[k] = local_sums[k] + __shfl_xor_sync(0xFFFFFFFF, local_sums[k], mask);
}
}
// Only lane 0 of each warp atomicAdd
if ((tid & 31) == 0) {
for (int k = 0; k < 6; k++) {
atomicAddBF16(&s_sums[k], (float)local_sums[k]);
atomicAdd(&s_sums[k], local_sums[k]);
}
}
__syncthreads();

88
crates/ml/src/cuda_pipeline/training_guard_kernel.cu
View File

@@ -39,7 +39,7 @@
/* ------------------------------------------------------------------ */
extern "C" __global__ void training_guard_check_and_accumulate(
const float* __restrict__ loss_scalar, /* GPU-resident f32 scalar */
const float* __restrict__ grad_norm_scalar, /* GPU-resident bf16 scalar */
const float* __restrict__ grad_norm_scalar, /* GPU-resident f32 scalar */
float* output, /* pinned host buffer (7 floats) */
float* acc_buf, /* device accumulator (3 f32) */
float clip_threshold,
@@ -115,30 +115,29 @@ extern "C" __global__ void qvalue_stats_reduce(
int tid = threadIdx.x;
float local_min = bf16(1e30f);
float local_max = bf16(-1e30f);
float local_sum = bf16_zero();
float local_all = bf16_zero();
float local_min = 1e30f;
float local_max = -1e30f;
float local_sum = 0.0f;
float local_all = 0.0f;
/* Grid-stride loop: each thread processes multiple samples */
for (int i = tid; i < batch_size; i += blockDim.x) {
const float* row = q_values + i * num_actions;
/* Find max Q across actions for this sample */
float sample_max = bf16(-1e30f);
float sample_max = -1e30f;
for (int a = 0; a < num_actions; a++) {
float qv = row[a];
/* Skip NaN/Inf: cast to F32 for isnan/isinf check */
float qv_f = (float)qv;
if (!isnan(qv_f) && !isinf(qv_f)) {
sample_max = bf16_fmax(sample_max, qv);
/* Skip NaN/Inf */
if (!isnan(qv) && !isinf(qv)) {
sample_max = fmaxf(sample_max, qv);
local_all = local_all + qv;
}
}
if (sample_max > bf16(-1e30f)) { /* at least one valid Q-value */
local_min = bf16_fmin(local_min, sample_max);
local_max = bf16_fmax(local_max, sample_max);
if (sample_max > -1e30f) { /* at least one valid Q-value */
local_min = fminf(local_min, sample_max);
local_max = fmaxf(local_max, sample_max);
local_sum = local_sum + sample_max;
}
}
@@ -152,8 +151,8 @@ extern "C" __global__ void qvalue_stats_reduce(
/* Parallel reduction phase 1: shared memory down to 32 threads */
for (int stride = blockDim.x / 2; stride >= 32; stride >>= 1) {
if (tid < stride) {
s_min[tid] = bf16_fmin(s_min[tid], s_min[tid + stride]);
s_max[tid] = bf16_fmax(s_max[tid], s_max[tid + stride]);
s_min[tid] = fminf(s_min[tid], s_min[tid + stride]);
s_max[tid] = fmaxf(s_max[tid], s_max[tid + stride]);
s_sum[tid] = s_sum[tid] + s_sum[tid + stride];
s_all_sum[tid] = s_all_sum[tid] + s_all_sum[tid + stride];
}
@@ -167,10 +166,10 @@ extern "C" __global__ void qvalue_stats_reduce(
float my_sum = s_sum[tid];
float my_all_sum = s_all_sum[tid];
for (int mask = 16; mask >= 1; mask >>= 1) {
my_min = bf16_fmin(my_min, bf16_shfl_xor(0xFFFFFFFF, my_min, mask));
my_max = bf16_fmax(my_max, bf16_shfl_xor(0xFFFFFFFF, my_max, mask));
my_sum = my_sum + bf16_shfl_xor(0xFFFFFFFF, my_sum, mask);
my_all_sum = my_all_sum + bf16_shfl_xor(0xFFFFFFFF, my_all_sum, mask);
my_min = fminf(my_min, __shfl_xor_sync(0xFFFFFFFF, my_min, mask));
my_max = fmaxf(my_max, __shfl_xor_sync(0xFFFFFFFF, my_max, mask));
my_sum = my_sum + __shfl_xor_sync(0xFFFFFFFF, my_sum, mask);
my_all_sum = my_all_sum + __shfl_xor_sync(0xFFFFFFFF, my_all_sum, mask);
}
if (tid == 0) {
s_min[0] = my_min;
@@ -183,11 +182,11 @@ extern "C" __global__ void qvalue_stats_reduce(
/* Thread 0 writes final results to host-mapped float buffer */
if (tid == 0) {
float n = bf16((float)batch_size);
float total_n = n * bf16((float)num_actions);
float n = (float)batch_size;
float total_n = n * (float)num_actions;
output[0] = (s_min[0] > bf16(1e29f)) ? 0.0f : (float)s_min[0]; /* q_min */
output[1] = (s_max[0] < bf16(-1e29f)) ? 0.0f : (float)s_max[0]; /* q_max */
output[0] = (s_min[0] > 1e29f) ? 0.0f : s_min[0]; /* q_min */
output[1] = (s_max[0] < -1e29f) ? 0.0f : s_max[0]; /* q_max */
output[2] = ((float)n > 0.0f) ? (float)(s_sum[0] / n) : 0.0f; /* q_mean */
output[3] = ((float)total_n > 0.0f) ? (float)(s_all_sum[0] / total_n) : 0.0f; /* mean_of_all */
__threadfence_system();
@@ -212,10 +211,10 @@ extern "C" __global__ void qvalue_divergence_check(
int num_actions,
float divergence_threshold
) {
float q_min = bf16(1e30f);
float q_max = bf16(-1e30f);
float q_sum = bf16_zero();
float q_sq = bf16_zero();
float q_min = 1e30f;
float q_max = -1e30f;
float q_sum = 0.0f;
float q_sq = 0.0f;
int valid = 0;
for (int a = 0; a < num_actions; a++) {
@@ -223,38 +222,37 @@ extern "C" __global__ void qvalue_divergence_check(
float qv_f = (float)qv;
if (isnan(qv_f) || isinf(qv_f)) continue;
q_min = bf16_fmin(q_min, qv);
q_max = bf16_fmax(q_max, qv);
q_min = fminf(q_min, qv);
q_max = fmaxf(q_max, qv);
q_sum = q_sum + qv;
q_sq = q_sq + qv * qv;
valid++;
}
float q_mean_bf = bf16_zero();
float q_variance_bf = bf16_zero();
float q_mean_f = 0.0f;
float q_variance_f = 0.0f;
if (valid > 0) {
float n_bf = bf16((float)valid);
q_mean_bf = q_sum / n_bf;
float mean_sq = q_sq / n_bf;
q_variance_bf = mean_sq - q_mean_bf * q_mean_bf;
q_variance_bf = bf16_fmax(q_variance_bf, bf16_zero()); /* numerical floor */
float n_f = (float)valid;
q_mean_f = q_sum / n_f;
float mean_sq = q_sq / n_f;
q_variance_f = mean_sq - q_mean_f * q_mean_f;
q_variance_f = fmaxf(q_variance_f, 0.0f); /* numerical floor */
}
/* Sentinel: if no valid values, zero out extremes */
if (valid == 0) {
q_min = bf16_zero();
q_max = bf16_zero();
q_min = 0.0f;
q_max = 0.0f;
}
float div_thresh = bf16(divergence_threshold);
int diverged = (bf16_fabs(q_min) > div_thresh ||
bf16_fabs(q_max) > div_thresh) ? 1 : 0;
int diverged = (fabsf(q_min) > divergence_threshold ||
fabsf(q_max) > divergence_threshold) ? 1 : 0;
/* Write to host-mapped float buffer */
output[0] = (float)q_min;
output[1] = (float)q_max;
output[2] = (float)q_mean_bf;
output[3] = (float)q_variance_bf;
output[0] = q_min;
output[1] = q_max;
output[2] = q_mean_f;
output[3] = q_variance_f;
output[4] = (float)diverged;
__threadfence_system();
}