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

cleanup: remove bf16 wrappers from training guard Q-value kernels

Browse Source 
qvalue_stats_reduce and qvalue_divergence_check used bf16() identity
wrappers on native f32 data — legacy from the bf16 era. These
wrappers can cause subtle precision differences across GPU
architectures. Replaced with native fminf/fmaxf/__shfl_xor_sync.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-14 09:35:32 +02:00
parent 9260100200
commit 2b612b478f
1 changed files with 49 additions and 59 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -123,31 +123,28 @@ 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);
local_all = local_all + qv;
if (!isnan(qv) && !isinf(qv)) {
sample_max = fmaxf(sample_max, qv);
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);
local_sum = local_sum + sample_max;
if (sample_max > -1e30f) {
local_min = fminf(local_min, sample_max);
local_max = fmaxf(local_max, sample_max);
local_sum += sample_max;
}
}
@@ -160,25 +157,25 @@ 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_sum[tid] = s_sum[tid] + s_sum[tid + stride];
s_all_sum[tid] = s_all_sum[tid] + s_all_sum[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 + stride];
s_all_sum[tid]+= s_all_sum[tid + stride];
}
__syncthreads();
}
/* Parallel reduction phase 2: warp-level reduction (no __syncthreads needed) */
/* Parallel reduction phase 2: warp-level (no __syncthreads needed) */
if (tid < 32) {
float my_min = s_min[tid];
float my_max = s_max[tid];
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 += __shfl_xor_sync(0xFFFFFFFF, my_sum, mask);
my_all_sum += __shfl_xor_sync(0xFFFFFFFF, my_all_sum, mask);
}
if (tid == 0) {
s_min[0] = my_min;
@@ -189,15 +186,14 @@ extern "C" __global__ void qvalue_stats_reduce(
}
__syncthreads();
/* 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[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 */
output[0] = (s_min[0] > 1e29f) ? 0.0f : s_min[0];
output[1] = (s_max[0] < -1e29f) ? 0.0f : s_max[0];
output[2] = (n > 0.0f) ? (s_sum[0] / n) : 0.0f;
output[3] = (total_n > 0.0f) ? (s_all_sum[0] / total_n) : 0.0f;
__threadfence_system();
}
}
@@ -220,49 +216,43 @@ 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++) {
float qv = q_values[a];
float qv_f = (float)qv;
if (isnan(qv_f) || isinf(qv_f)) continue;
if (isnan(qv) || isinf(qv)) continue;
q_min = bf16_fmin(q_min, qv);
q_max = bf16_fmax(q_max, qv);
q_sum = q_sum + qv;
q_sq = q_sq + qv * qv;
q_min = fminf(q_min, qv);
q_max = fmaxf(q_max, qv);
q_sum += qv;
q_sq += qv * qv;
valid++;
}
float q_mean_bf = bf16_zero();
float q_variance_bf = bf16_zero();
float q_mean = 0.0f;
float q_var = 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 = (float)valid;
q_mean = q_sum / n;
q_var = fmaxf((q_sq / n) - q_mean * q_mean, 0.0f);
}
/* 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;
output[3] = q_var;
output[4] = (float)diverged;
__threadfence_system();
}