diff --git a/crates/ml/src/cuda_pipeline/training_guard_kernel.cu b/crates/ml/src/cuda_pipeline/training_guard_kernel.cu index 9d04c03a6..84a419f9b 100644 --- a/crates/ml/src/cuda_pipeline/training_guard_kernel.cu +++ b/crates/ml/src/cuda_pipeline/training_guard_kernel.cu @@ -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(); }