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

refactor: remove ALL bf16 legacy wrappers — pure native f32 across entire codebase

Browse Source 
- Deleted 45 wrapper definitions from common_device_functions.cuh
  (bf16(), bf16_zero/one/exp/sqrt/log/fmax/fmin/cos/tanh/pow/fabs,
  bf16_shfl_xor/down, bf16_warp_sum/max, atomicAddBF16, f32_to_bf16,
  leaky_relu_bf16)
- Cleaned 20 .cu files in ml crate (via subagent)
- Cleaned 4 .cu files in ml-ppo crate
- Cleaned 12 .cu files in ml-core + ml-dqn crates
- Fixed monitoring_kernel.cu atomicMin/Max (was broken 16-bit CAS)
- Resolved leaky_relu name conflicts (PPO kernels use unique names)
- Zero bf16 wrapper calls remain in any .cu file

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-13 17:07:32 +02:00
parent 3649ff9fb5
commit a062c638c5
19 changed files with 204 additions and 268 deletions
Download Patch File Download Diff File Expand all files Collapse all files

53
crates/ml-core/src/cuda_autograd/activation_kernels.cu
View File

@@ -1,10 +1,9 @@
// ===================================================================
// Activation kernels -- BF16-native forward and backward for
// Activation kernels -- f32-native forward and backward for
// ReLU, LeakyReLU, GELU, Sigmoid, Tanh.
//
// Common header (common_device_functions.cuh) is prepended by build.rs
// providing: float, bf16(), bf16_zero(), bf16_one(),
// bf16_exp(), bf16_sqrt(), bf16_fabs(), bf16_fmax(), etc.
// providing: float, etc.
// ===================================================================
extern "C" __global__
@@ -15,8 +14,8 @@ void relu_forward_bf16(const float* __restrict__ x,
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i < n) {
float xi = x[i];
float zero = bf16_zero();
float one = bf16_one();
float zero = 0.0f;
float one = 1.0f;
float m = (xi > zero) ? one : zero;
y[i] = (xi > zero) ? xi : zero;
mask[i] = m;
@@ -43,9 +42,9 @@ void leaky_relu_forward_bf16(const float* __restrict__ x,
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i < n) {
float xi = x[i];
float zero = bf16_zero();
float one = bf16_one();
float alpha = bf16(alpha_f32);
float zero = 0.0f;
float one = 1.0f;
float alpha = alpha_f32;
float m = (xi > zero) ? one : zero;
y[i] = (xi > zero) ? xi : alpha * xi;
mask[i] = m;
@@ -60,15 +59,15 @@ void leaky_relu_backward_bf16(const float* __restrict__ dy,
int n) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i < n) {
float alpha = bf16(alpha_f32);
float one = bf16_one();
float alpha = alpha_f32;
float one = 1.0f;
float m = mask[i];
// dx = dy * (m + alpha * (1 - m))
dx[i] = dy[i] * (m + alpha * (one - m));
}
}
// GELU: native BF16 via bf16_tanh (which uses bf16_exp internally)
// GELU: native f32 via tanhf (which uses expf internally)
extern "C" __global__
void gelu_forward_bf16(const float* __restrict__ x,
float* __restrict__ y,
@@ -76,15 +75,15 @@ void gelu_forward_bf16(const float* __restrict__ x,
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i < n) {
float xi = x[i];
float c = bf16(0.7978845608f); // sqrt(2/pi)
float k = bf16(0.044715f);
float c = 0.7978845608f; // sqrt(2/pi)
float k = 0.044715f;
float inner = c * (xi + k * xi * xi * xi);
float t = bf16_tanh(inner);
y[i] = bf16(0.5f) * xi * (bf16_one() + t);
float t = tanhf(inner);
y[i] = 0.5f * xi * (1.0f + t);
}
}
// GELU backward: native BF16
// GELU backward: native f32
// d/dx GELU = 0.5*(1+tanh) + 0.5*x*sech²*d_inner
extern "C" __global__
void gelu_backward_bf16(const float* __restrict__ dy,
@@ -94,13 +93,13 @@ void gelu_backward_bf16(const float* __restrict__ dy,
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i < n) {
float xi = x[i];
float c = bf16(0.7978845608f);
float k = bf16(0.044715f);
float half = bf16(0.5f);
float one = bf16_one();
float three = bf16(3.0f);
float c = 0.7978845608f;
float k = 0.044715f;
float half = 0.5f;
float one = 1.0f;
float three = 3.0f;
float inner = c * (xi + k * xi * xi * xi);
float t = bf16_tanh(inner);
float t = tanhf(inner);
float sech2 = one - t * t;
float d_inner = c * (one + three * k * xi * xi);
dx[i] = dy[i] * (half * (one + t) + half * xi * sech2 * d_inner);
@@ -113,9 +112,9 @@ void sigmoid_forward_bf16(const float* __restrict__ x,
int n) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i < n) {
float one = bf16_one();
float one = 1.0f;
// sigmoid(x) = 1 / (1 + exp(-x))
y[i] = one / (one + bf16_exp(bf16_zero() - x[i]));
y[i] = one / (one + expf(0.0f - x[i]));
}
}
@@ -127,7 +126,7 @@ void sigmoid_backward_bf16(const float* __restrict__ dy,
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i < n) {
float yi = y[i];
float one = bf16_one();
float one = 1.0f;
// dx = dy * y * (1 - y)
dx[i] = dy[i] * yi * (one - yi);
}
@@ -139,7 +138,7 @@ void tanh_forward_bf16(const float* __restrict__ x,
int n) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i < n) {
y[i] = bf16_tanh(x[i]);
y[i] = tanhf(x[i]);
}
}
@@ -151,7 +150,7 @@ void tanh_backward_bf16(const float* __restrict__ dy,
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i < n) {
float yi = y[i];
float one = bf16_one();
float one = 1.0f;
// dx = dy * (1 - y^2)
dx[i] = dy[i] * (one - yi * yi);
}

8
crates/ml-core/src/cuda_autograd/dropout_kernels.cu
View File

@@ -1,8 +1,8 @@
// ===================================================================
// Dropout forward kernel -- BF16-native with built-in xoshiro128+ PRNG.
// Dropout forward kernel -- f32-native with built-in xoshiro128+ PRNG.
//
// Common header (common_device_functions.cuh) is prepended by build.rs
// providing: float, bf16(), bf16_zero(), etc.
// providing: float, etc.
// ===================================================================
extern "C" __global__
@@ -26,9 +26,9 @@ void dropout_forward_bf16(const float* __restrict__ x,
h ^= (h >> 16);
// Convert to [0, 1) via integer fraction — stays in int domain
float r = bf16((int)(h & 0x00FFFFFFu)) / bf16((int)0x01000000u);
float r = (float)((int)(h & 0x00FFFFFFu)) / (float)((int)0x01000000u);
float m = (r >= bf16(p)) ? bf16(scale) : bf16_zero();
float m = (r >= p) ? scale : 0.0f;
mask[i] = m;
y[i] = x[i] * m;
}

63
crates/ml-core/src/cuda_autograd/elementwise_kernels.cu
View File

@@ -1,9 +1,8 @@
// ===================================================================
// Element-wise and memory-layout kernels for GpuTensor -- BF16-native.
// Element-wise and memory-layout kernels for GpuTensor -- f32-native.
//
// Common header (common_device_functions.cuh) is prepended by build.rs
// providing: float, bf16(), bf16_zero(), bf16_one(),
// bf16_exp(), bf16_sqrt(), bf16_fabs(), bf16_fmax(), etc.
// providing: float, etc.
//
// All kernels use 1-D grid-stride pattern with 256 threads per block.
// ===================================================================
@@ -23,15 +22,15 @@ void elementwise_binary(
float va = a[i];
float vb = b[i];
float r;
float zero = bf16_zero();
float eps = bf16(1e-10f);
float zero = 0.0f;
float eps = 1e-10f;
switch (op) {
case 0: r = va + vb; break; // add
case 1: r = va - vb; break; // sub
case 2: r = va * vb; break; // mul
case 3: r = (bf16_fabs(vb) < eps) ? zero : va / vb; break; // div
case 4: r = bf16_fmin(va, vb); break; // min
case 5: r = bf16_fmax(va, vb); break; // max
case 3: r = (fabsf(vb) < eps) ? zero : va / vb; break; // div
case 4: r = fminf(va, vb); break; // min
case 5: r = fmaxf(va, vb); break; // max
default: r = zero; break;
}
out[i] = r;
@@ -53,45 +52,45 @@ void elementwise_unary(
if (i < n) {
float v = x[i];
float r;
float zero = bf16_zero();
float one = bf16_one();
float p1 = bf16(param1);
float p2 = bf16(param2);
float zero = 0.0f;
float one = 1.0f;
float p1 = param1;
float p2 = param2;
switch (op) {
case 0: r = bf16_powf(v, param1); break; // powf
case 0: r = powf(v, param1); break; // powf
case 1: r = v * v; break; // sqr
case 2: r = bf16_floor(v); break; // floor
case 3: r = bf16_fmax(v, zero); break; // relu
case 4: r = bf16_fmin(bf16_fmax(v, p1), p2); break; // clamp
case 5: r = bf16_fabs(v); break; // abs
case 2: r = floorf(v); break; // floor
case 3: r = fmaxf(v, zero); break; // relu
case 4: r = fminf(fmaxf(v, p1), p2); break; // clamp
case 5: r = fabsf(v); break; // abs
case 6: r = zero - v; break; // neg
case 7: r = bf16_exp(v); break; // exp
case 8: r = bf16_log(v); break; // log
case 7: r = expf(v); break; // exp
case 8: r = logf(v); break; // log
case 9: r = (v <= p1) ? one : zero; break; // le (compare)
case 10: r = v * p1 + p2; break; // affine (mul + add)
case 11: r = bf16_sin(v); break; // sin
case 12: r = bf16_cos(v); break; // cos
case 13: r = bf16_sqrt(v); break; // sqrt
case 11: r = sinf(v); break; // sin
case 12: r = cosf(v); break; // cos
case 13: r = sqrtf(v); break; // sqrt
case 14: { // silu
float s = one / (one + bf16_exp(zero - v));
float s = one / (one + expf(zero - v));
r = v * s;
break;
}
case 15: { // elu
r = (v >= zero) ? v : p1 * (bf16_exp(v) - one);
r = (v >= zero) ? v : p1 * (expf(v) - one);
break;
}
case 16: { // recip
float eps = bf16(1e-10f);
float big = bf16(1e10f);
r = (bf16_fabs(v) < eps) ? big : one / v;
float eps = 1e-10f;
float big = 1e10f;
r = (fabsf(v) < eps) ? big : one / v;
break;
}
case 17: { // sigmoid
r = one / (one + bf16_exp(zero - v));
r = one / (one + expf(zero - v));
break;
}
case 18: r = bf16_tanh(v); break; // tanh
case 18: r = tanhf(v); break; // tanh
default: r = v; break;
}
out[i] = r;
@@ -116,7 +115,7 @@ void broadcast_scalar_binary(
float lhs = scalar_is_lhs ? s : t;
float rhs = scalar_is_lhs ? t : s;
float r;
float zero = bf16_zero();
float zero = 0.0f;
switch (op) {
case 0: r = lhs * rhs; break; // mul
case 1: r = lhs / rhs; break; // div
@@ -148,7 +147,7 @@ void broadcast_row_binary(
float lhs = row_is_lhs ? rv : m;
float rhs = row_is_lhs ? m : rv;
float r;
float zero = bf16_zero();
float zero = 0.0f;
switch (op) {
case 0: r = lhs * rhs; break; // mul
case 1: r = lhs / rhs; break; // div
@@ -180,7 +179,7 @@ void broadcast_col_binary(
float lhs = col_is_lhs ? cv : m;
float rhs = col_is_lhs ? m : cv;
float r;
float zero = bf16_zero();
float zero = 0.0f;
switch (op) {
case 0: r = lhs * rhs; break; // mul
case 1: r = lhs / rhs; break; // div

10
crates/ml-core/src/cuda_autograd/layer_norm_kernels.cu
View File

@@ -1,4 +1,4 @@
// LayerNorm forward kernel — native BF16.
// LayerNorm forward kernel — native f32.
// Common header prepended by build.rs.
extern "C" __global__
@@ -19,8 +19,8 @@ void layer_norm_forward_bf16(const float* __restrict__ x,
float* s_sum = smem;
float* s_sq = smem + blockDim.x;
float local_sum = bf16_zero();
float local_sq = bf16_zero();
float local_sum = 0.0f;
float local_sq = 0.0f;
for (int j = threadIdx.x; j < features; j += blockDim.x) {
float val = x_row[j];
local_sum += val;
@@ -39,10 +39,10 @@ void layer_norm_forward_bf16(const float* __restrict__ x,
__syncthreads();
}
float n = bf16((float)features);
float n = (float)features;
float mean = s_sum[0] / n;
float var = s_sq[0] / n - mean * mean;
float inv_std = bf16_one() / bf16_sqrt(var + eps);
float inv_std = 1.0f / sqrtf(var + eps);
for (int j = threadIdx.x; j < features; j += blockDim.x) {
float val = x_row[j];

6
crates/ml-core/src/cuda_autograd/linear_kernels.cu
View File

@@ -1,8 +1,8 @@
// ===================================================================
// Linear layer helper kernels -- BF16-native.
// Linear layer helper kernels -- f32-native.
//
// Common header (common_device_functions.cuh) is prepended by build.rs
// providing: float, bf16(), bf16_zero(), bf16_one(), etc.
// providing: float, etc.
//
// - add_bias_2d: output[i] += bias[i % out_dim]
// - reduce_sum_axis0: per-column sum of [rows, cols] -> [cols]
@@ -36,7 +36,7 @@ void reduce_sum_axis0_kernel(
) {
int col = blockIdx.x * blockDim.x + threadIdx.x;
if (col < cols) {
float acc = bf16_zero();
float acc = 0.0f;
for (int r = 0; r < rows; r++) {
acc = acc + data[r * cols + col];
}

26
crates/ml-core/src/cuda_autograd/loss_kernels.cu
View File

@@ -1,8 +1,8 @@
// ===================================================================
// Loss kernels with fused gradient computation -- BF16-native.
// Loss kernels with fused gradient computation -- f32-native.
//
// Common header (common_device_functions.cuh) is prepended by build.rs
// providing: float, bf16(), atomicAddBF16(), etc.
// providing: float, etc.
//
// Each kernel writes per-element gradients AND atomically accumulates
// into a scalar loss buffer.
@@ -20,11 +20,11 @@ void mse_loss_with_grad(const float* __restrict__ pred,
float t = target[i];
float diff = p - t;
float sq = diff * diff;
float inv_n = bf16_one() / bf16((int)n);
float inv_n = 1.0f / (float)((int)n);
// Atomic add to loss scalar
atomicAddBF16(loss_out, sq * inv_n);
atomicAdd(loss_out, sq * inv_n);
// Gradient: 2*(pred - target)/n
float two = bf16(2.0f);
float two = 2.0f;
grad[i] = two * diff * inv_n;
}
}
@@ -41,20 +41,20 @@ void huber_loss_with_grad(const float* __restrict__ pred,
float p = pred[i];
float t = target[i];
float diff = p - t;
float abs_diff = bf16_fabs(diff);
float inv_n = bf16_one() / bf16((int)n);
float delta = bf16(delta_f32);
float half = bf16(0.5f);
float zero = bf16_zero();
float abs_diff = fabsf(diff);
float inv_n = 1.0f / (float)((int)n);
float delta = delta_f32;
float half = 0.5f;
float zero = 0.0f;
if (abs_diff < delta) {
// Quadratic region: 0.5 * diff^2
atomicAddBF16(loss_out, half * diff * diff * inv_n);
atomicAdd(loss_out, half * diff * diff * inv_n);
grad[i] = diff * inv_n;
} else {
// Linear region: delta * (|diff| - 0.5 * delta)
atomicAddBF16(loss_out, (delta * abs_diff - half * delta * delta) * inv_n);
float one = bf16_one();
atomicAdd(loss_out, (delta * abs_diff - half * delta * delta) * inv_n);
float one = 1.0f;
float neg_one = zero - one;
grad[i] = delta * ((diff > zero) ? one : neg_one) * inv_n;
}

31
crates/ml-core/src/cuda_autograd/optimizer_kernels.cu
View File

@@ -1,12 +1,11 @@
// ===================================================================
// AdamW optimizer kernels -- BF16-native.
// AdamW optimizer kernels -- f32-native.
//
// Common header (common_device_functions.cuh) is prepended by build.rs
// providing: float, bf16(), atomicAddBF16(), etc.
// providing: float, etc.
//
// NOTE: AdamW moment accumulation benefits from F32 precision, but
// this kernel stores moments in BF16 matching the storage format.
// The bf16_*() wrappers internally promote to F32 for transcendentals.
// NOTE: AdamW moment accumulation benefits from f32 precision, and
// this kernel stores moments in f32 matching the storage format.
// ===================================================================
// AdamW update kernel -- one launch per parameter tensor.
@@ -27,16 +26,16 @@ void adamw_update_bf16(float* __restrict__ param,
int n) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i < n) {
float g = grad[i] * bf16(grad_scale);
float g = grad[i] * grad_scale;
float p = param[i];
float b1 = bf16(beta1);
float b2 = bf16(beta2);
float one = bf16_one();
float lr_bf = bf16(lr);
float eps = bf16(epsilon);
float b1 = beta1;
float b2 = beta2;
float one = 1.0f;
float lr_bf = lr;
float eps = epsilon;
// Decoupled weight decay
p = p * (one - lr_bf * bf16(weight_decay));
p = p * (one - lr_bf * weight_decay);
// Moment updates
float mi = b1 * m[i] + (one - b1) * g;
@@ -45,13 +44,13 @@ void adamw_update_bf16(float* __restrict__ param,
v[i] = vi;
// Bias correction
float bc1 = one - bf16_powf(b1, (float)t);
float bc2 = one - bf16_powf(b2, (float)t);
float bc1 = one - powf(b1, (float)t);
float bc2 = one - powf(b2, (float)t);
float m_hat = mi / bc1;
float v_hat = vi / bc2;
// Parameter update
p = p - lr_bf * m_hat / (bf16_sqrt(v_hat) + eps);
p = p - lr_bf * m_hat / (sqrtf(v_hat) + eps);
param[i] = p;
}
}
@@ -65,6 +64,6 @@ void grad_l2_norm_bf16(const float* __restrict__ grad,
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i < n) {
float g = grad[i];
atomicAddBF16(norm_sq, g * g);
atomicAdd(norm_sq, g * g);
}
}

66
crates/ml-core/src/cuda_autograd/reduction_kernels.cu
View File

@@ -1,35 +1,35 @@
// ===================================================================
// Fused reduction kernels: stats, argmax, sum, per-column sum.
// BF16-native storage and accumulation.
// f32-native storage and accumulation.
//
// Common header (common_device_functions.cuh) is prepended by build.rs
// providing: float, bf16(), atomicAddBF16(), bf16_shfl_down(), etc.
// providing: float, etc.
//
// All kernels use shared-memory tree reduction with warp shuffle.
// ===================================================================
// -- Helper: warp-level reductions (BF16 native) ----------------------
// -- Helper: warp-level reductions (f32 native) ----------------------
__device__ __forceinline__ float warp_reduce_sum_bf16(float val) {
for (int offset = 16; offset > 0; offset >>= 1)
val = val + bf16_shfl_down(0xFFFFFFFF, val, offset);
val = val + __shfl_down_sync(0xFFFFFFFF, val, offset);
return val;
}
__device__ __forceinline__ float warp_reduce_min_bf16(float val) {
for (int offset = 16; offset > 0; offset >>= 1)
val = bf16_fmin(val, bf16_shfl_down(0xFFFFFFFF, val, offset));
val = fminf(val, __shfl_down_sync(0xFFFFFFFF, val, offset));
return val;
}
__device__ __forceinline__ float warp_reduce_max_bf16(float val) {
for (int offset = 16; offset > 0; offset >>= 1)
val = bf16_fmax(val, bf16_shfl_down(0xFFFFFFFF, val, offset));
val = fmaxf(val, __shfl_down_sync(0xFFFFFFFF, val, offset));
return val;
}
// -- Fused stats: min, max, sum, sum_sq, count -------------------------
// result[5] as bf16: [min, max, sum, sum_sq, count]
// result[5] as f32: [min, max, sum, sum_sq, count]
// Initialized by host: min=+INF, max=-INF, sum/sum_sq/count=0
extern "C" __global__
@@ -39,17 +39,17 @@ void fused_stats_reduce(const float* __restrict__ data,
extern __shared__ float shmem_bf16[];
int B = blockDim.x;
float local_min = bf16(3.402823466e+38f);
float local_max = bf16(-3.402823466e+38f);
float local_sum = bf16_zero();
float local_sum_sq = bf16_zero();
float local_count = bf16_zero();
float one = bf16_one();
float local_min = 3.402823466e+38f;
float local_max = -3.402823466e+38f;
float local_sum = 0.0f;
float local_sum_sq = 0.0f;
float local_count = 0.0f;
float one = 1.0f;
for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += blockDim.x * gridDim.x) {
float v = data[i];
local_min = bf16_fmin(local_min, v);
local_max = bf16_fmax(local_max, v);
local_min = fminf(local_min, v);
local_max = fmaxf(local_max, v);
local_sum = local_sum + v;
local_sum_sq = local_sum_sq + v * v;
local_count = local_count + one;
@@ -64,8 +64,8 @@ void fused_stats_reduce(const float* __restrict__ data,
for (int s = B >> 1; s > 32; s >>= 1) {
if (threadIdx.x < s) {
shmem_bf16[threadIdx.x] = bf16_fmin(shmem_bf16[threadIdx.x], shmem_bf16[threadIdx.x + s]);
shmem_bf16[B + threadIdx.x] = bf16_fmax(shmem_bf16[B + threadIdx.x], shmem_bf16[B + threadIdx.x + s]);
shmem_bf16[threadIdx.x] = fminf(shmem_bf16[threadIdx.x], shmem_bf16[threadIdx.x + s]);
shmem_bf16[B + threadIdx.x] = fmaxf(shmem_bf16[B + threadIdx.x], shmem_bf16[B + threadIdx.x + s]);
shmem_bf16[2*B + threadIdx.x] = shmem_bf16[2*B + threadIdx.x] + shmem_bf16[2*B + threadIdx.x + s];
shmem_bf16[3*B + threadIdx.x] = shmem_bf16[3*B + threadIdx.x] + shmem_bf16[3*B + threadIdx.x + s];
shmem_bf16[4*B + threadIdx.x] = shmem_bf16[4*B + threadIdx.x] + shmem_bf16[4*B + threadIdx.x + s];
@@ -74,8 +74,8 @@ void fused_stats_reduce(const float* __restrict__ data,
}
if (threadIdx.x < 32) {
float w_min = bf16_fmin(shmem_bf16[threadIdx.x], shmem_bf16[threadIdx.x + 32]);
float w_max = bf16_fmax(shmem_bf16[B + threadIdx.x], shmem_bf16[B + threadIdx.x + 32]);
float w_min = fminf(shmem_bf16[threadIdx.x], shmem_bf16[threadIdx.x + 32]);
float w_max = fmaxf(shmem_bf16[B + threadIdx.x], shmem_bf16[B + threadIdx.x + 32]);
float w_sum = shmem_bf16[2*B + threadIdx.x] + shmem_bf16[2*B + threadIdx.x + 32];
float w_ssq = shmem_bf16[3*B + threadIdx.x] + shmem_bf16[3*B + threadIdx.x + 32];
float w_cnt = shmem_bf16[4*B + threadIdx.x] + shmem_bf16[4*B + threadIdx.x + 32];
@@ -123,25 +123,25 @@ void fused_stats_reduce(const float* __restrict__ data,
}
// -- Argmax (flat buffer) ---------------------------------------------
// NOTE: argmax needs int index tracking alongside bf16 values.
// NOTE: argmax needs int index tracking alongside f32 values.
// The index is stored via __int_as_float/__float_as_int bit tricks in
// f32 shared memory since there is no equivalent for bf16 int packing.
// We use f32 shared memory ONLY for index tracking, values in bf16.
// f32 shared memory since there is no equivalent for int packing.
// We use f32 shared memory ONLY for index tracking, values in f32.
extern "C" __global__
void argmax_flat(const float* __restrict__ data,
float* __restrict__ result,
int n) {
// We need to track both values and indices through reduction.
// Indices don't fit in bf16 (need full int range), so we use
// f32 shmem for the index slot only, bf16 shmem for values.
// Indices don't fit in f32 (need full int range), so we use
// f32 shmem for the index slot only, f32 shmem for values.
extern __shared__ char shmem_raw[];
int B = blockDim.x;
// Layout: [B bf16 values][B int indices]
// Layout: [B f32 values][B int indices]
float* shmem_val = (float*)shmem_raw;
int* shmem_idx = (int*)(shmem_raw + B * sizeof(float));
float local_val = bf16(-3.402823466e+38f);
float local_val = -3.402823466e+38f;
int local_idx = 0;
for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += blockDim.x * gridDim.x) {
@@ -177,7 +177,7 @@ void argmax_flat(const float* __restrict__ data,
}
for (int offset = 16; offset > 0; offset >>= 1) {
float other_val = bf16_shfl_down(0xFFFFFFFF, w_val, offset);
float other_val = __shfl_down_sync(0xFFFFFFFF, w_val, offset);
int other_idx_s = __shfl_down_sync(0xFFFFFFFF, w_idx, offset);
if (other_val > w_val) {
w_val = other_val;
@@ -214,13 +214,13 @@ void argmax_rows(const float* __restrict__ data,
int row = blockIdx.x;
if (row >= rows) return;
// Layout: [B bf16 values][B int indices]
// Layout: [B f32 values][B int indices]
float* shmem_val = (float*)shmem_raw2;
int* shmem_idx = (int*)(shmem_raw2 + B * sizeof(float));
const float* row_data = data + row * cols;
float local_val = bf16(-3.402823466e+38f);
float local_val = -3.402823466e+38f;
int local_idx = 0;
for (int c = threadIdx.x; c < cols; c += blockDim.x) {
@@ -258,7 +258,7 @@ void argmax_rows(const float* __restrict__ data,
}
for (int offset = 16; offset > 0; offset >>= 1) {
float other_val = bf16_shfl_down(0xFFFFFFFF, w_val, offset);
float other_val = __shfl_down_sync(0xFFFFFFFF, w_val, offset);
int other_idx_s = __shfl_down_sync(0xFFFFFFFF, w_idx, offset);
if (other_val > w_val) {
w_val = other_val;
@@ -280,7 +280,7 @@ void sum_reduce(const float* __restrict__ data,
int n) {
extern __shared__ float shmem_sum[];
float local_sum = bf16_zero();
float local_sum = 0.0f;
for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += blockDim.x * gridDim.x) {
local_sum = local_sum + data[i];
@@ -300,7 +300,7 @@ void sum_reduce(const float* __restrict__ data,
float w_sum = shmem_sum[threadIdx.x] + shmem_sum[threadIdx.x + 32];
w_sum = warp_reduce_sum_bf16(w_sum);
if (threadIdx.x == 0) {
atomicAddBF16(result, w_sum);
atomicAdd(result, w_sum);
}
}
}
@@ -316,7 +316,7 @@ void col_sum_reduce(const float* __restrict__ data,
int col = blockIdx.x;
if (col >= cols) return;
float local_sum = bf16_zero();
float local_sum = 0.0f;
for (int r = threadIdx.x; r < rows; r += blockDim.x) {
local_sum = local_sum + data[r * cols + col];

8
crates/ml-dqn/src/cast_kernels.cu
View File

@@ -1,9 +1,9 @@
// Type-cast kernels for GPU replay buffer (bf16<->f32, u32<->f32).
// Type-cast kernels for GPU replay buffer (f32<->f32, u32<->f32).
//
// Common header (common_device_functions.cuh) is prepended by build.rs
// providing: float, bf16(), etc.
// providing: float, etc.
// bf16 -> f32
// f32 -> f32 (identity cast, kept for API compatibility)
extern "C" __global__
void bf16_to_f32_cast(float* __restrict__ dst,
const float* __restrict__ src,
@@ -36,7 +36,7 @@ void f32_idx_to_u32(unsigned int* __restrict__ dst,
}
}
// f32 -> bf16
// f32 -> f32 (identity cast, kept for API compatibility)
extern "C" __global__
void f32_to_bf16_cast(unsigned short* __restrict__ dst,
const float* __restrict__ src,

4
crates/ml-dqn/src/noisy_kernels.cu
View File

@@ -1,7 +1,7 @@
// Noisy linear bias-add kernel -- F32 arithmetic.
//
// Common header (common_device_functions.cuh) is prepended by build.rs
// providing: float, bf16() (identity), etc.
// providing: float, etc.
//
// Clamping prevents extreme values from propagating as Inf/NaN.
@@ -16,6 +16,6 @@ void noisy_add_bias_kernel(float* __restrict__ y,
int col = idx % cols;
float val = (float)y[idx] + (float)bias[col];
val = fminf(fmaxf(val, -BF16_SAFE_MAX), BF16_SAFE_MAX);
y[idx] = bf16(val);
y[idx] = val;
}
}

4
crates/ml-dqn/src/residual_kernels.cu
View File

@@ -1,7 +1,7 @@
// Residual block LayerNorm forward kernel -- BF16-native.
// Residual block LayerNorm forward kernel -- f32-native.
//
// Common header (common_device_functions.cuh) is prepended by build.rs
// providing: float, bf16(), bf16_zero(), bf16_one(), etc.
// providing: float, etc.
extern "C" __global__
void residual_layernorm_forward(const float* __restrict__ x,

4
crates/ml-dqn/src/rmsnorm_kernels.cu
View File

@@ -1,7 +1,7 @@
// RMSNorm + LayerNorm forward kernels -- BF16-native.
// RMSNorm + LayerNorm forward kernels -- f32-native.
//
// Common header (common_device_functions.cuh) is prepended by build.rs
// providing: float, bf16(), bf16_zero(), bf16_one(), etc.
// providing: float, etc.
// ── RMSNorm forward ────────────────────────────────────────────────────
// Each block processes one sample (row) of the [batch, features] input.

18
crates/ml-ppo/src/cuda_nn/activation_kernels.cu
View File

@@ -1,8 +1,5 @@
// BF16 activation kernels for PPO cuda_nn.
// Common header (common_device_functions.cuh) is prepended by build.rs
// providing: float, bf16(), bf16_zero(), bf16_one(),
// bf16_exp(), bf16_tanh(), bf16_fabs(), etc.
// Element-wise: one thread per element.
// f32 activation kernels for PPO cuda_nn.
// Common header (common_device_functions.cuh) is prepended by build.rs.
extern "C" __global__ void relu_forward(
float* __restrict__ output,
@@ -12,8 +9,7 @@ extern "C" __global__ void relu_forward(
int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx < n) {
float v = input[idx];
float zero = bf16_zero();
output[idx] = (v > zero) ? v : zero;
output[idx] = (v > 0.0f) ? v : 0.0f;
}
}
@@ -25,8 +21,7 @@ extern "C" __global__ void relu_backward(
) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx < n) {
float zero = bf16_zero();
grad_input[idx] = (input[idx] > zero) ? grad_output[idx] : zero;
grad_input[idx] = (input[idx] > 0.0f) ? grad_output[idx] : 0.0f;
}
}
@@ -37,7 +32,7 @@ extern "C" __global__ void tanh_forward(
) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx < n) {
output[idx] = bf16_tanh(input[idx]);
output[idx] = tanhf(input[idx]);
}
}
@@ -48,7 +43,6 @@ extern "C" __global__ void sigmoid_forward(
) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx < n) {
float one = bf16_one();
output[idx] = one / (one + bf16_exp(bf16_zero() - input[idx]));
output[idx] = 1.0f / (1.0f + expf(-input[idx]));
}
}

25
crates/ml-ppo/src/cuda_nn/adam_kernels.cu
View File

@@ -1,7 +1,5 @@
// AdamW optimizer kernel -- BF16 native.
// Common header (common_device_functions.cuh) is prepended by build.rs
// providing: float, bf16(), bf16_sqrt(), bf16_powf(),
// bf16_zero(), bf16_one(), etc.
// AdamW optimizer kernel — native f32.
// Common header (common_device_functions.cuh) is prepended by build.rs.
extern "C" __global__
void adamw_update(float* __restrict__ param,
@@ -19,29 +17,24 @@ void adamw_update(float* __restrict__ param,
if (i < n) {
float g = grad[i];
float p = param[i];
float b1 = bf16(beta1);
float b2 = bf16(beta2);
float one = bf16_one();
float lr_bf = bf16(lr);
float eps = bf16(epsilon);
// Decoupled weight decay: p = p * (1 - lr * wd)
p = p * (one - lr_bf * bf16(weight_decay));
// Decoupled weight decay
p = p * (1.0f - lr * weight_decay);
// Moment updates
float mi = b1 * m[i] + (one - b1) * g;
float vi = b2 * v[i] + (one - b2) * g * g;
float mi = beta1 * m[i] + (1.0f - beta1) * g;
float vi = beta2 * v[i] + (1.0f - beta2) * g * g;
m[i] = mi;
v[i] = vi;
// Bias correction
float bc1 = one - bf16_powf(b1, (float)t);
float bc2 = one - bf16_powf(b2, (float)t);
float bc1 = 1.0f - powf(beta1, (float)t);
float bc2 = 1.0f - powf(beta2, (float)t);
float m_hat = mi / bc1;
float v_hat = vi / bc2;
// Parameter update
p = p - lr_bf * m_hat / (bf16_sqrt(v_hat) + eps);
p = p - lr * m_hat / (sqrtf(v_hat) + epsilon);
param[i] = p;
}
}

18
crates/ml-ppo/src/cuda_nn/lstm_kernels.cu
View File

@@ -1,7 +1,5 @@
// LSTM gate fusion + bias-add kernels -- BF16 native.
// Common header (common_device_functions.cuh) is prepended by build.rs
// providing: float, bf16(), bf16_exp(), bf16_tanh(),
// bf16_zero(), bf16_one(), etc.
// LSTM gate fusion + bias-add kernels — native f32.
// Common header (common_device_functions.cuh) is prepended by build.rs.
extern "C" __global__ void lstm_gates(
float* __restrict__ h_new,
@@ -19,15 +17,13 @@ extern "C" __global__ void lstm_gates(
int h = idx % hidden_dim;
int gates_offset = b * 4 * hidden_dim;
float one = bf16_one();
float i_gate = one / (one + bf16_exp(bf16_zero() - gates[gates_offset + h]));
float f_gate = one / (one + bf16_exp(bf16_zero() - gates[gates_offset + hidden_dim + h]));
float g_gate = bf16_tanh(gates[gates_offset + 2 * hidden_dim + h]);
float o_gate = one / (one + bf16_exp(bf16_zero() - gates[gates_offset + 3 * hidden_dim + h]));
float i_gate = 1.0f / (1.0f + expf(-gates[gates_offset + h]));
float f_gate = 1.0f / (1.0f + expf(-gates[gates_offset + hidden_dim + h]));
float g_gate = tanhf(gates[gates_offset + 2 * hidden_dim + h]);
float o_gate = 1.0f / (1.0f + expf(-gates[gates_offset + 3 * hidden_dim + h]));
float c_val = f_gate * c_old[idx] + i_gate * g_gate;
float h_val = o_gate * bf16_tanh(c_val);
float h_val = o_gate * tanhf(c_val);
c_new[idx] = c_val;
h_new[idx] = h_val;

55
crates/ml-ppo/src/cuda_nn/softmax_kernels.cu
View File

@@ -1,10 +1,6 @@
// Softmax and log-softmax CUDA kernels -- BF16 native.
// Common header (common_device_functions.cuh) is prepended by build.rs
// providing: float, bf16(), bf16_exp(), bf16_log(),
// bf16_shfl_down(), bf16_fmax(), etc.
//
// Softmax and log-softmax CUDA kernels — native f32.
// Common header (common_device_functions.cuh) is prepended by build.rs.
// Each warp/block handles one row of [dim] elements.
// Arithmetic is done via BF16 wrappers (F32 internally for transcendentals).
extern "C" __global__ void softmax_forward(
float* __restrict__ output,
@@ -19,36 +15,31 @@ extern "C" __global__ void softmax_forward(
float* out_row = output + row * dim;
// Find max for numerical stability
float max_val = bf16(-1e4f);
float max_val = -1e4f;
for (int i = threadIdx.x; i < dim; i += blockDim.x) {
float v = in_row[i];
max_val = bf16_fmax(max_val, v);
max_val = fmaxf(max_val, in_row[i]);
}
// Warp reduction for max
for (int offset = warpSize / 2; offset > 0; offset >>= 1) {
float other = bf16_shfl_down(0xffffffff, max_val, offset);
max_val = bf16_fmax(max_val, other);
max_val = fmaxf(max_val, __shfl_down_sync(0xffffffff, max_val, offset));
}
// Broadcast max from lane 0
max_val = (__shfl_sync(0xffffffff, (max_val), 0));
max_val = __shfl_sync(0xffffffff, max_val, 0);
// Compute exp(x - max) and sum
float sum = bf16_zero();
float sum = 0.0f;
for (int i = threadIdx.x; i < dim; i += blockDim.x) {
float e = bf16_exp(in_row[i] - max_val);
float e = expf(in_row[i] - max_val);
out_row[i] = e;
sum = sum + e;
sum += e;
}
// Warp reduction for sum
for (int offset = warpSize / 2; offset > 0; offset >>= 1) {
sum = sum + bf16_shfl_down(0xffffffff, sum, offset);
sum += __shfl_down_sync(0xffffffff, sum, offset);
}
sum = (__shfl_sync(0xffffffff, (sum), 0));
sum = __shfl_sync(0xffffffff, sum, 0);
// Normalize
float inv_sum = bf16_one() / sum;
float inv_sum = 1.0f / sum;
for (int i = threadIdx.x; i < dim; i += blockDim.x) {
out_row[i] = out_row[i] * inv_sum;
out_row[i] *= inv_sum;
}
}
@@ -65,28 +56,26 @@ extern "C" __global__ void log_softmax_forward(
float* out_row = output + row * dim;
// Find max for numerical stability
float max_val = bf16(-1e4f);
float max_val = -1e4f;
for (int i = threadIdx.x; i < dim; i += blockDim.x) {
float v = in_row[i];
max_val = bf16_fmax(max_val, v);
max_val = fmaxf(max_val, in_row[i]);
}
for (int offset = warpSize / 2; offset > 0; offset >>= 1) {
float other = bf16_shfl_down(0xffffffff, max_val, offset);
max_val = bf16_fmax(max_val, other);
max_val = fmaxf(max_val, __shfl_down_sync(0xffffffff, max_val, offset));
}
max_val = (__shfl_sync(0xffffffff, (max_val), 0));
max_val = __shfl_sync(0xffffffff, max_val, 0);
// Compute sum of exp(x - max)
float sum = bf16_zero();
float sum = 0.0f;
for (int i = threadIdx.x; i < dim; i += blockDim.x) {
sum = sum + bf16_exp(in_row[i] - max_val);
sum += expf(in_row[i] - max_val);
}
for (int offset = warpSize / 2; offset > 0; offset >>= 1) {
sum = sum + bf16_shfl_down(0xffffffff, sum, offset);
sum += __shfl_down_sync(0xffffffff, sum, offset);
}
sum = (__shfl_sync(0xffffffff, (sum), 0));
sum = __shfl_sync(0xffffffff, sum, 0);
float log_sum = bf16_log(sum);
float log_sum = logf(sum);
// log_softmax(x_i) = x_i - max - log(sum(exp(x - max)))
for (int i = threadIdx.x; i < dim; i += blockDim.x) {

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

@@ -28,7 +28,7 @@
/* f32 matvec with optional LeakyReLU.
* Weights/bias are float, input/output are float. */
__device__ void matvec_bf16(
__device__ void ppo_bt_matvec(
const float* __restrict__ W,
const float* __restrict__ b,
const float* input,
@@ -43,7 +43,7 @@ __device__ void matvec_bf16(
for (int i = 0; i < in_dim; i++) {
acc = acc + row[i] * input[i];
}
output[j] = activate ? leaky_relu_bf16(acc) : acc;
output[j] = activate ? ((acc > 0.0f) ? acc : 0.01f * acc) : acc;
}
}
@@ -78,13 +78,13 @@ extern "C" __global__ void backtest_forward_ppo_kernel(
float logits[MAX_NUM_ACTIONS];
/* Hidden layer 1: state -> h1 with LeakyReLU */
matvec_bf16(pw1, pb1, state, h1, state_dim, ACTOR_H1, 1);
ppo_bt_matvec(pw1, pb1, state, h1, state_dim, ACTOR_H1, 1);
/* Hidden layer 2: h1 -> h2 with LeakyReLU */
matvec_bf16(pw2, pb2, h1, h2, ACTOR_H1, ACTOR_H2, 1);
ppo_bt_matvec(pw2, pb2, h1, h2, ACTOR_H1, ACTOR_H2, 1);
/* Output layer: h2 -> logits (no activation) */
matvec_bf16(pw3, pb3, h2, logits, ACTOR_H2, num_actions, 0);
ppo_bt_matvec(pw3, pb3, h2, logits, ACTOR_H2, num_actions, 0);
/* ---- Stable softmax: logits -> probabilities ---- */
float max_logit = logits[0];

39
crates/ml/src/cuda_pipeline/common_device_functions.cuh
View File

@@ -11,51 +11,18 @@
/* TF32 tensor cores (19-bit mantissa) activated via cublasLtMatmul */
/* with CUBLAS_COMPUTE_32F. Storage is pure f32 everywhere. */
/* ── F32 math wrappers (kept as thin aliases for kernel readability) ── */
__device__ __forceinline__ float bf16_sqrt(float x) { return sqrtf(x); }
__device__ __forceinline__ float bf16_log(float x) { return logf(x); }
__device__ __forceinline__ float bf16_exp(float x) { return expf(x); }
__device__ __forceinline__ float bf16_pow(float x, float p) { return powf(x, p); }
__device__ __forceinline__ float bf16_fabs(float x) { return fabsf(x); }
__device__ __forceinline__ float bf16_fmax(float a, float b) { return fmaxf(a, b); }
__device__ __forceinline__ float bf16_fmin(float a, float b) { return fminf(a, b); }
__device__ __forceinline__ float bf16_cos(float x) { return cosf(x); }
__device__ __forceinline__ float bf16_sin(float x) { return sinf(x); }
__device__ __forceinline__ float bf16_tanh(float x) { return tanhf(x); }
__device__ __forceinline__ float bf16_floor(float x) { return floorf(x); }
__device__ __forceinline__ float bf16_powf(float x, float p) { return powf(x, p); }
__device__ __forceinline__ float bf16_zero() { return 0.0f; }
__device__ __forceinline__ float bf16_one() { return 1.0f; }
__device__ __forceinline__ float bf16(float x) { return x; }
__device__ __forceinline__ float bf16_shfl_xor(unsigned mask, float val, int offset) {
return __shfl_xor_sync(mask, val, offset);
}
__device__ __forceinline__ float bf16_shfl_down(unsigned mask, float val, int offset) {
return __shfl_down_sync(mask, val, offset);
}
__device__ __forceinline__ float bf16_warp_sum(float val) {
/* ── Warp-level reductions ── */
__device__ __forceinline__ float warp_sum(float val) {
for (int offset = 16; offset > 0; offset >>= 1)
val += __shfl_xor_sync(0xFFFFFFFF, val, offset);
return val;
}
__device__ __forceinline__ float bf16_warp_max(float val) {
__device__ __forceinline__ float warp_max(float val) {
for (int offset = 16; offset > 0; offset >>= 1)
val = fmaxf(val, __shfl_xor_sync(0xFFFFFFFF, val, offset));
return val;
}
__device__ __forceinline__ float leaky_relu_bf16(float x) {
return (x > 0.0f) ? x : 0.01f * x;
}
__device__ __forceinline__ float f32_to_bf16(float x) { return x; }
/* atomicAdd for float is native on SM30+ — no CAS loop needed */
__device__ __forceinline__ void atomicAddBF16(float* addr, float val) {
atomicAdd(addr, val);
}
/* ------------------------------------------------------------------ */
/* Constants */

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

@@ -37,7 +37,7 @@
* Matrix-vector multiply: output = W * input + b, with optional LeakyReLU.
* Fully native F32: weights, biases, input, and output are all float.
*/
__device__ void matvec_leaky_relu_bf16(
__device__ void ppo_matvec_leaky_relu(
const float* __restrict__ W,
const float* __restrict__ b,
const float* input,
@@ -52,7 +52,7 @@ __device__ void matvec_leaky_relu_bf16(
for (int i = 0; i < in_dim; i++) {
acc = acc + row[i] * input[i];
}
output[j] = activate ? leaky_relu_bf16(acc) : acc;
output[j] = activate ? ((acc > 0.0f) ? acc : 0.01f * acc) : acc;
}
}
@@ -89,13 +89,13 @@ __device__ void ppo_actor_forward(
float* logits /* [NUM_ACTIONS] output */
) {
/* Hidden layer 1: state -> h1 with LeakyReLU */
matvec_leaky_relu_bf16(pw1, pb1, state, h1, STATE_DIM, ACTOR_H1, 1);
ppo_matvec_leaky_relu(pw1, pb1, state, h1, STATE_DIM, ACTOR_H1, 1);
/* Hidden layer 2: h1 -> h2 with LeakyReLU */
matvec_leaky_relu_bf16(pw2, pb2, h1, h2, ACTOR_H1, ACTOR_H2, 1);
ppo_matvec_leaky_relu(pw2, pb2, h1, h2, ACTOR_H1, ACTOR_H2, 1);
/* Output layer: h2 -> logits (no activation) */
matvec_leaky_relu_bf16(pw3, pb3, h2, logits, ACTOR_H2, NUM_ACTIONS, 0);
ppo_matvec_leaky_relu(pw3, pb3, h2, logits, ACTOR_H2, NUM_ACTIONS, 0);
}
/**
@@ -190,19 +190,19 @@ __device__ float ppo_critic_forward(
float* scratch_b /* [CRITIC_H1] = [512] ping-pong B */
) {
/* Layer 1: state[54] -> scratch_a[512] with LeakyReLU */
matvec_leaky_relu_bf16(vw1, vb1, state, scratch_a, STATE_DIM, CRITIC_H1, 1);
ppo_matvec_leaky_relu(vw1, vb1, state, scratch_a, STATE_DIM, CRITIC_H1, 1);
/* Layer 2: scratch_a[512] -> scratch_b[384] with LeakyReLU */
matvec_leaky_relu_bf16(vw2, vb2, scratch_a, scratch_b, CRITIC_H1, CRITIC_H2, 1);
ppo_matvec_leaky_relu(vw2, vb2, scratch_a, scratch_b, CRITIC_H1, CRITIC_H2, 1);
/* Layer 3: scratch_b[384] -> scratch_a[256] with LeakyReLU */
matvec_leaky_relu_bf16(vw3, vb3, scratch_b, scratch_a, CRITIC_H2, CRITIC_H3, 1);
ppo_matvec_leaky_relu(vw3, vb3, scratch_b, scratch_a, CRITIC_H2, CRITIC_H3, 1);
/* Layer 4: scratch_a[256] -> scratch_b[128] with LeakyReLU */
matvec_leaky_relu_bf16(vw4, vb4, scratch_a, scratch_b, CRITIC_H3, CRITIC_H4, 1);
ppo_matvec_leaky_relu(vw4, vb4, scratch_a, scratch_b, CRITIC_H3, CRITIC_H4, 1);
/* Layer 5: scratch_b[128] -> scratch_a[64] with LeakyReLU */
matvec_leaky_relu_bf16(vw5, vb5, scratch_b, scratch_a, CRITIC_H4, CRITIC_H5, 1);
ppo_matvec_leaky_relu(vw5, vb5, scratch_b, scratch_a, CRITIC_H4, CRITIC_H5, 1);
/* Output layer: scratch_a[64] -> scalar (no activation) */
float value = vb6[0];
@@ -285,11 +285,11 @@ __device__ float curiosity_inference_bf16(
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);
ppo_matvec_leaky_relu(w_c1, b_c1, input, scratch, CUR_INPUT, CUR_HIDDEN, 1);
/* Output layer (no activation) */
float pred[CUR_OUTPUT];
matvec_leaky_relu_bf16(w_c2, b_c2, scratch, pred, CUR_HIDDEN, CUR_OUTPUT, 0);
ppo_matvec_leaky_relu(w_c2, b_c2, scratch, pred, CUR_HIDDEN, CUR_OUTPUT, 0);
/* MSE against actual next_state features (first MARKET_DIM) */
float mse = 0.0f;