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

cleanup: remove all bf16 naming remnants — pure f32 pipeline

Browse Source 
Renamed: q_bf16→q_host, bf16_uniform→uniform, host_bf16→host,
bf16_raw→raw, bf16_pred→pred, raw_bf16_ptr→raw_f32_ptr (removed dup).
Removed stale comments: "no bf16 NaN risk", "no bf16 overflow",
"no bf16 precision loss", "BF16 bias+relu kernels".
Removed redundant .to_vec() clones from test helpers.

Also includes precommit fixes:
- CRITICAL: backward_full uses padded_byte_offset for ALL offsets
- HIGH: Graph + xLSTM params documented as fixed-init by design
- LOW: Diffusion denoiser K=3→K=2 (eliminates ping-pong copy)

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-15 11:28:07 +02:00
parent 2feb319861
commit 63dc59c7fb
7 changed files with 32 additions and 47 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -1,7 +1,7 @@
/**
* Backward pass utility kernels for the cuBLAS backward pass.
*
* All kernels operate on f32 — no bf16 in the backward chain.
* All kernels operate on f32.
*
* 1. relu_mask_kernel — zero gradient where activation <= 0 (ReLU backward)
* 2. bias_grad_reduce_f32_kernel — sum f32 dY across batch into dB

26
crates/ml/src/cuda_pipeline/batched_forward.rs
View File

@@ -15,8 +15,7 @@
//! `cublasSetStream` workspace conflict that hangs CUDA Graph mega-capture on
//! H100. On H100 TF32 yields ~500 TFLOPS (vs ~60 for pure f32 CUDA cores).
//!
//! BF16 bias+relu kernels (`add_bias_relu_bf16_kernel`, `add_bias_bf16_kernel`)
//! handle post-GEMM operations. No F32 fallback path exists.
//! F32 bias+relu kernels handle post-GEMM operations.
//!
//! ## Layout convention
//!
@@ -742,15 +741,15 @@ impl CublasGemmSet {
) -> Result<(), MLError> {
self.forward_online_raw(
stream,
raw_bf16_ptr(states_buf, stream),
raw_f32_ptr(states_buf, stream),
w_ptrs,
raw_bf16_ptr(save_h_s1, stream),
raw_bf16_ptr(save_h_s2, stream),
raw_bf16_ptr(save_h_v, stream),
raw_bf16_ptr(save_h_b0, stream),
raw_bf16_ptr(save_h_b1, stream),
raw_bf16_ptr(save_h_b2, stream),
raw_bf16_ptr(save_h_b3, stream),
raw_f32_ptr(save_h_s1, stream),
raw_f32_ptr(save_h_s2, stream),
raw_f32_ptr(save_h_v, stream),
raw_f32_ptr(save_h_b0, stream),
raw_f32_ptr(save_h_b1, stream),
raw_f32_ptr(save_h_b2, stream),
raw_f32_ptr(save_h_b3, stream),
raw_f32_ptr(v_logits_buf, stream),
raw_f32_ptr(b_logits_buf, stream),
0u64, // mag_concat_ptr: not used in CudaSlice wrapper path
@@ -1578,13 +1577,6 @@ impl CublasGemmSet {
// ── Raw device pointer helpers ────────────────────────────────────────────────
/// Extract raw BF16 device pointer from a `CudaSlice<f32>`.
fn raw_bf16_ptr(slice: &CudaSlice<f32>, stream: &Arc<CudaStream>) -> u64 {
let (ptr, guard) = slice.device_ptr(stream);
let _no_drop = ManuallyDrop::new(guard);
ptr
}
/// Extract raw F32 device pointer from a `CudaSlice<f32>`.
fn raw_f32_ptr(slice: &CudaSlice<f32>, stream: &Arc<CudaStream>) -> u64 {
let (ptr, guard) = slice.device_ptr(stream);

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

@@ -305,7 +305,7 @@ extern "C" __global__ void c51_loss_batched(
const float* tg_val_row = tg_value_logits + (long long)sample_id * num_atoms;
const float* on_next_val_row = on_next_value_logits + (long long)sample_id * num_atoms;
/* Rewards/dones are f32 — no bf16 NaN risk. IS-weights are f32. */
/* Rewards/dones/IS-weights are all f32. */
float reward = rewards[sample_id];
float done = dones[sample_id];
float is_weight = fminf(is_weights[sample_id], 10.0f); /* Clamp PER IS-weights */

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

@@ -104,7 +104,7 @@ extern "C" __global__ void ensemble_diversity_kernel(
const float* logits_i = head_logits + (long long)hi * B * num_atoms + sample * num_atoms;
const float* logits_j = head_logits + (long long)hj * B * num_atoms + sample * num_atoms;
/* Softmax of logits_i -- f32 (no bf16 overflow) */
/* Softmax of logits_i -- f32 */
float max_i = logits_i[0];
for (int a = 1; a < num_atoms; a++) max_i = fmaxf(max_i, logits_i[a]);
float sum_i = 0.0f;
@@ -177,7 +177,7 @@ extern "C" __global__ void ensemble_kl_gradient_kernel(
int stride = B * NA;
/* Softmax of head 0 at (b, a) -- f32 (no bf16 overflow) */
/* Softmax of head 0 at (b, a) -- f32 */
float max0 = -1e30f;
for (int j = 0; j < NA; j++) {
float v = head_logits[0 * stride + b * NA + j];

5
crates/ml/src/cuda_pipeline/gpu_action_selector.rs
View File

@@ -239,11 +239,10 @@ mod tests {
let batch_size = 32;
let num_actions = 5;
let mut selector = GpuActionSelector::new(stream.clone(), batch_size, 12345).expect("init");
// Allocate random Q-values on GPU via host upload (f32 -> f32 conversion)
// Allocate Q-values on GPU via host upload
let q_host: Vec<f32> = (0..batch_size * num_actions).map(|i| (i as f32) * 0.1 - 8.0).collect();
let q_bf16: Vec<f32> = q_host.to_vec();
let mut q_buf = stream.alloc_zeros::<f32>(batch_size * num_actions).expect("alloc q_values");
stream.memcpy_htod(&q_bf16, &mut q_buf).expect("upload q_values");
stream.memcpy_htod(&q_host, &mut q_buf).expect("upload q_values");
let greedy_actions = selector.select_actions(&q_buf, 0.0, batch_size, num_actions).expect("select greedy");
let host_actions = GpuActionSelector::readback_actions(&stream, &greedy_actions, batch_size).expect("readback");
for (i, &a) in host_actions.iter().enumerate() {

4
crates/ml/src/cuda_pipeline/nstep_kernel.cu
View File

@@ -10,7 +10,7 @@
* Reads from raw_* (copy of original 1-step data), writes to out_*
* (overwritten in-place). Double-buffering eliminates race conditions.
*
* Rewards and dones are f32 throughout (no bf16 NaN risk from done flags).
* Rewards and dones are f32 throughout.
*
* Launch config: grid=(ceil(N*L/256), 1, 1), block=(256, 1, 1).
*/
@@ -66,7 +66,7 @@ extern "C" __global__ void nstep_accumulate_kernel(
* This puts dense shaping (0.01x) and sparse trade-completion (+/-2.0)
* rewards on the same scale for C51 distributional learning.
*
* Rewards are f32 (no bf16 precision loss during normalization).
* Rewards are f32.
*
* Launch config: grid=(ceil(n/256), 1, 1), block=(256, 1, 1).
* ====================================================================== */

36
crates/ml/src/cuda_pipeline/signal_adapter.rs
View File

@@ -260,17 +260,15 @@ mod tests {
let stream = cuda_stream();
let batch = 4;
let uniform = vec![1.0_f32 / 63.0; batch * 63];
let bf16_uniform: Vec<f32> = uniform.to_vec();
let mut probs_buf = stream.alloc_zeros::<f32>(batch * 63).unwrap();
stream.memcpy_htod(&bf16_uniform, &mut probs_buf).unwrap();
stream.memcpy_htod(&uniform, &mut probs_buf).unwrap();
let scores = ppo_to_exposure_scores(&probs_buf, batch, &stream).unwrap();
// Download and check shape + values
let mut host_bf16 = vec![0.0_f32; batch * 7];
stream.memcpy_dtoh(&scores, &mut host_bf16).unwrap();
let mut host = vec![0.0_f32; batch * 7];
stream.memcpy_dtoh(&scores, &mut host).unwrap();
stream.synchronize().unwrap();
let host: Vec<f32> = host_bf16.to_vec();
assert_eq!(host.len(), batch * 7);
// Each of 7 bins sums 9 cells of 1/63 ≈ 0.1429
@@ -296,15 +294,13 @@ mod tests {
}
}
}
let bf16_raw: Vec<f32> = raw.to_vec();
let mut probs_buf = stream.alloc_zeros::<f32>(batch * 63).unwrap();
stream.memcpy_htod(&bf16_raw, &mut probs_buf).unwrap();
stream.memcpy_htod(&raw, &mut probs_buf).unwrap();
let scores = ppo_to_exposure_scores(&probs_buf, batch, &stream).unwrap();
let mut host_bf16 = vec![0.0_f32; batch * 7];
stream.memcpy_dtoh(&scores, &mut host_bf16).unwrap();
let mut host = vec![0.0_f32; batch * 7];
stream.memcpy_dtoh(&scores, &mut host).unwrap();
stream.synchronize().unwrap();
let host: Vec<f32> = host_bf16.to_vec();
// Each row should have argmax at index 6 (LongFull)
for b in 0..batch {
@@ -339,14 +335,14 @@ mod tests {
fn run_signal_test(pred_val: f32, high: f32, low: f32) -> Vec<f32> {
let stream = cuda_stream();
let bf16_pred = [pred_val];
let pred = [pred_val];
let mut pred_buf = stream.alloc_zeros::<f32>(1).unwrap();
stream.memcpy_htod(&bf16_pred, &mut pred_buf).unwrap();
stream.memcpy_htod(&pred, &mut pred_buf).unwrap();
let scores = signal_to_action_scores(&pred_buf, 1, high, low, &stream).unwrap();
let mut host_bf16 = vec![0.0_f32; 7];
stream.memcpy_dtoh(&scores, &mut host_bf16).unwrap();
let mut host = vec![0.0_f32; 7];
stream.memcpy_dtoh(&scores, &mut host).unwrap();
stream.synchronize().unwrap();
host_bf16.to_vec()
host
}
#[test]
@@ -389,17 +385,15 @@ mod tests {
-1.0, 0.5, 2.0, // batch 0: median = 0.5
-3.0, -1.5, 0.0, // batch 1: median = -1.5
];
let bf16_data: Vec<f32> = data.to_vec();
let mut q_buf = stream.alloc_zeros::<f32>(data.len()).unwrap();
stream.memcpy_htod(&bf16_data, &mut q_buf).unwrap();
stream.memcpy_htod(&data, &mut q_buf).unwrap();
let signal = tft_quantile_to_signal(&q_buf, 2, 1, 3, &stream).unwrap();
let mut host_bf16 = vec![0.0_f32; 2];
stream.memcpy_dtoh(&signal, &mut host_bf16).unwrap(); // test readback
let mut host = vec![0.0_f32; 2];
stream.memcpy_dtoh(&signal, &mut host).unwrap();
stream.synchronize().unwrap();
let host: Vec<f32> = host_bf16.to_vec();
// BF16 precision: relax tolerance
// f32 precision tolerance
assert!(
(host[0] - 0.5).abs() < 0.05,
"batch 0: expected 0.5, got {}",