foxhunt/docs/superpowers/plans/2026-03-27-full-bf16-rewrite.md

Full BF16 Rewrite — Zero F32 On GPU

For agentic workers: REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (- [ ]) syntax for tracking.

Goal: Convert every CUDA kernel, every GPU buffer, and every cuBLAS call from F32 to BF16. Zero F32 on GPU. No mixed precision.

Architecture: 4-phase conversion. Phase 1 lays the foundation (BF16 type helpers, common header already has cuda_bf16.h). Phase 2 converts cuBLAS GEMMs + activation/weight buffers. Phase 3 converts optimizer and loss kernels. Phase 4 converts environment and auxiliary kernels. Each phase compiles and passes tests independently.

Tech Stack: Rust, CUDA (__nv_bfloat16, cuda_bf16.h), cudarc 0.19.3 (CudaSlice<u16> for BF16), cuBLAS (cublasGemmEx, CUDA_R_16BF), half crate (Rust-side BF16 conversion)

Spec: docs/superpowers/specs/2026-03-27-full-bf16-rewrite-design.md

INVARIANT: After every task, SQLX_OFFLINE=true cargo check -p ml && cargo test -p ml --lib -- dqn && cargo test -p ml-dqn --lib must pass.

---

BF16 Conversion Pattern

Every kernel follows the same pattern:

CUDA side:

// Before:
extern "C" __global__ void my_kernel(float* data, int n) {
    float val = data[i];
    val = val * 2.0f;
    data[i] = val;
}

// After:
extern "C" __global__ void my_kernel(__nv_bfloat16* data, int n) {
    __nv_bfloat16 val = data[i];
    float vf = __bfloat162float(val);  // promote for arithmetic
    vf = vf * 2.0f;
    data[i] = __float2bfloat16(vf);    // demote for storage
}

Key rule: BF16 for STORAGE (all pointers, all buffers). F32 for ARITHMETIC inside kernels (cast on load, cast on store). This is how H100 tensor cores work — BF16 inputs, F32 accumulate, BF16 output. The casts are free on SM90 (native BF16↔F32 conversion instructions).

Rust side:

// Before:
let buf: CudaSlice<f32> = stream.alloc_zeros::<f32>(n)?;
stream.memcpy_htod(&host_f32, &mut buf)?;

// After:
let buf: CudaSlice<u16> = stream.alloc_zeros::<u16>(n)?;
let host_bf16: Vec<u16> = host_f32.iter().map(|&x| half::bf16::from_f32(x).to_bits()).collect();
stream.memcpy_htod(&host_bf16, &mut buf)?;

---

Phase 1: Foundation

Task 1: BF16 Rust helpers + weight set type conversion

Convert the core type infrastructure. DuelingWeightSet and BranchingWeightSet change from CudaSlice<f32> to CudaSlice<u16>. Add BF16 conversion helpers.

Files:

• Modify: crates/ml/src/cuda_pipeline/gpu_weights.rs

• Modify: crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs

• Modify: crates/ml/src/cuda_pipeline/mod.rs

• Step 1: Add BF16 helper functions to gpu_weights.rs

use half::bf16;

/// Convert F32 host vector to BF16 (u16) for GPU upload.
pub fn f32_to_bf16_host(data: &[f32]) -> Vec<u16> {
    data.iter().map(|&x| bf16::from_f32(x).to_bits()).collect()
}

/// Convert BF16 (u16) GPU download to F32 host vector.
pub fn bf16_to_f32_host(data: &[u16]) -> Vec<f32> {
    data.iter().map(|&x| bf16::from_bits(x).to_f32()).collect()
}

/// Allocate BF16 buffer (u16) on GPU, zero-initialized.
pub fn alloc_bf16(stream: &Arc<CudaStream>, n: usize, label: &str) -> Result<CudaSlice<u16>, MLError> {
    stream.alloc_zeros::<u16>(n).map_err(|e| {
        MLError::ModelError(format!("alloc bf16 {label}: {e}"))
    })
}

• Step 2: Convert DuelingWeightSet fields from CudaSlice to CudaSlice

Change all 12 fields (w_s1, b_s1, w_s2, b_s2, w_v1, b_v1, w_v2, b_v2, w_a1, b_a1, w_a2, b_a2) from CudaSlice<f32> to CudaSlice<u16>.

• Step 3: Convert BranchingWeightSet fields similarly

Change all 8 fields from CudaSlice<f32> to CudaSlice<u16>.

• Step 4: Update extract_dueling_weights_branching and extract_branching_weights

These functions extract weights from Candle VarMaps (F32) into GPU buffers. They must now:

1. Download F32 from Candle → host Vec
2. Convert to Vec via f32_to_bf16_host
3. Upload to CudaSlice

• Step 5: Fix all compilation errors from the type change

This will cascade through MANY files:

• gpu_dqn_trainer.rs: params_buf stays CudaSlice<f32> for now (Adam master weights) — BUT we're going full BF16, so change it to CudaSlice<u16> too

• Actually: EVERY CudaSlice<f32> in GpuDqnTrainer becomes CudaSlice<u16>. ALL of them. This includes: states_buf, next_states_buf, rewards_buf, dones_buf, is_weights_buf, save_h_s1..save_h_b2, grad_buf, params_buf, target_params_buf, m_buf, v_buf, grad_norm_buf, total_loss_buf, td_errors_buf, cql_grad_scratch, and all spectral norm u/v buffers.

• The alloc_f32 helper calls throughout the constructor become alloc_bf16 (or stream.alloc_zeros::<u16>(n))

• The CachedPtrs struct: pointers are u64, type-agnostic — no change needed

• Kernel launch .arg() calls: u64 pointers are type-agnostic — no change needed

• memcpy_htod calls that upload F32 host data: convert to BF16 first

• memcpy_dtoh calls that download to F32 host: download as u16, convert back

• Step 6: Update all raw_device_ptr functions

Change raw_device_ptr(slice: &CudaSlice<f32>, ...) to accept CudaSlice<u16>. Or better: make it generic:

fn raw_device_ptr<T>(slice: &CudaSlice<T>, stream: &CudaStream) -> u64 {
    let (ptr, guard) = slice.device_ptr(stream);
    let _no_drop = std::mem::ManuallyDrop::new(guard);
    ptr
}

• Step 7: Fix all remaining compilation errors

This is the big cascading fix. Work through every error methodically. Every CudaSlice<f32> → CudaSlice<u16>.

EXCEPTION: The actions_buf and t_buf stay as CudaSlice<i32> — they hold integer action indices and Adam step counter, not floating point.

• Step 8: Verify compilation + tests

Run: SQLX_OFFLINE=true cargo check -p ml This WILL have many errors at first. Fix them all before proceeding.

• Step 9: Commit

git commit -m "feat(bf16): phase 1 — all GPU buffers CudaSlice<u16>, weight sets BF16"

---

Phase 2: cuBLAS BF16

Task 2: Convert cuBLAS forward to cublasGemmEx BF16

Files:

• Modify: crates/ml/src/cuda_pipeline/batched_forward.rs

• Step 1: Set TF32 math mode on cuBLAS handle

In CublasForward::new(), after workspace allocation:

unsafe {
    cublas_sys::cublasSetMathMode(
        handle.0,
        cublas_sys::cublasMath_t::CUBLAS_TF32_TENSOR_OP_MATH,
    );
}

• Step 2: Replace sgemm_layer with gemmex_bf16

Replace the body of sgemm_layer:

fn gemmex_bf16(
    &self, w_ptr: u64, a_ptr: u64, c_ptr: u64,
    n: usize, b: usize, k: usize, _label: &str,
) -> Result<(), MLError> {
    let alpha = 1.0_f32;
    let beta = 0.0_f32;
    unsafe {
        let status = cublas_sys::cublasGemmEx(
            self.handle.0,
            cublas_sys::cublasOperation_t::CUBLAS_OP_T,
            cublas_sys::cublasOperation_t::CUBLAS_OP_N,
            n as i32, b as i32, k as i32,
            &alpha as *const f32 as *const std::ffi::c_void,
            w_ptr as *const std::ffi::c_void,
            cublas_sys::cudaDataType_t::CUDA_R_16BF, k as i32,
            a_ptr as *const std::ffi::c_void,
            cublas_sys::cudaDataType_t::CUDA_R_16BF, k as i32,
            &beta as *const f32 as *const std::ffi::c_void,
            c_ptr as *mut std::ffi::c_void,
            cublas_sys::cudaDataType_t::CUDA_R_16BF, n as i32,
            cublas_sys::cublasComputeType_t::CUBLAS_COMPUTE_32F,
            cublas_sys::cublasGemmAlgo_t::CUBLAS_GEMM_DEFAULT_TENSOR_OP,
        );
        if status != cublas_sys::cublasStatus_t::CUBLAS_STATUS_SUCCESS {
            return Err(MLError::ModelError(format!("cublasGemmEx {_label}: {status:?}")));
        }
    }
    Ok(())
}

Note: Ctype=CUDA_R_16BF with computeType=CUBLAS_COMPUTE_32F → BF16 inputs, F32 internal accumulate, BF16 output. Full tensor core throughput.

• Step 3: Update forward_online to use BF16 pointers

Change w_ptrs: &[u64; 20] to use BF16 weight buffer pointers. Change all activation buffer refs from &CudaSlice<f32> to &CudaSlice<u16>.

• Step 4: Convert bias kernels to BF16

In bias_kernels.cu, change both kernels:

extern "C" __global__ void add_bias_relu_kernel(
    __nv_bfloat16* __restrict__ output,
    const __nv_bfloat16* __restrict__ bias,
    int out_dim, int total_elements)
{
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= total_elements) return;
    float val = __bfloat162float(output[i]) + __bfloat162float(bias[i % out_dim]);
    output[i] = __float2bfloat16((val > 0.0f) ? val : 0.0f);
}

Same for add_bias_kernel.

• Step 5: Update forward_target and forward_online_next similarly

• Step 6: Verify compilation + tests

• Step 7: Commit

---

Task 3: Convert cuBLAS backward to cublasGemmEx BF16

Files:

• Modify: crates/ml/src/cuda_pipeline/batched_backward.rs

• Step 1: Replace cublasSgemm with cublasGemmEx in backward_fc_layer

Both dW and dX computations use cublasGemmEx with BF16 inputs:

• dW = d_output^T × input: both BF16, output to BF16 grad_buf

• dX = d_output × W^T: both BF16, output to BF16 upstream gradient

• Step 2: Update backward_full to pass BF16 pointers

All activation pointers, weight pointers, and gradient pointers are now BF16 (u64 addresses into u16 buffers).

• Step 3: Convert relu_mask_kernel to BF16

In relu_mask_kernel.cu:

extern "C" __global__ void relu_mask_kernel(
    __nv_bfloat16* __restrict__ dx,
    const __nv_bfloat16* __restrict__ activation,
    int n)
{
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= n) return;
    if (__bfloat162float(activation[i]) <= 0.0f)
        dx[i] = __float2bfloat16(0.0f);
}

• Step 4: Verify + commit

---

Phase 3: Optimizer + Loss Kernels

Task 4: Adam kernel BF16

Files:

• Modify: crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu

• Step 1: Convert dqn_adam_update_kernel

All pointers (params, grads, m, v, grad_norm_sq) → __nv_bfloat16*. Internal arithmetic stays F32 (cast on load, cast on store):

extern "C" __global__ void dqn_adam_update_kernel(
    __nv_bfloat16* __restrict__ params,
    const __nv_bfloat16* __restrict__ grads,
    __nv_bfloat16* __restrict__ m,
    __nv_bfloat16* __restrict__ v,
    const __nv_bfloat16* __restrict__ grad_norm_sq,
    float lr, float beta1, float beta2, float epsilon,
    float weight_decay, float max_grad_norm,
    const int* __restrict__ t_ptr, int total_params
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total_params) return;
    int t = *t_ptr;
    float g = __bfloat162float(grads[idx]);
    float norm = sqrtf(__bfloat162float(*grad_norm_sq) + 1e-12f);
    float clip_scale = (norm > max_grad_norm) ? (max_grad_norm / norm) : 1.0f;
    float clipped_g = g * clip_scale;
    float beta1_t = 1.0f - powf(beta1, (float)t);
    float beta2_t = 1.0f - powf(beta2, (float)t);
    float m_i = beta1 * __bfloat162float(m[idx]) + (1.0f - beta1) * clipped_g;
    float v_i = beta2 * __bfloat162float(v[idx]) + (1.0f - beta2) * clipped_g * clipped_g;
    m[idx] = __float2bfloat16(m_i);
    v[idx] = __float2bfloat16(v_i);
    float m_hat = m_i / beta1_t;
    float v_hat = v_i / beta2_t;
    float p = __bfloat162float(params[idx]);
    p -= lr * (m_hat / (sqrtf(v_hat) + epsilon) + weight_decay * p);
    params[idx] = __float2bfloat16(p);
}

• Step 2: Convert dqn_grad_norm_kernel to BF16 input

extern "C" __global__ void dqn_grad_norm_kernel(
    const __nv_bfloat16* __restrict__ grads,
    __nv_bfloat16* __restrict__ out_grad_norm,  // BF16 output
    int total_params)

Internal sum stays F32 (warp reduction), final output converted to BF16.

• Step 3: Convert SAXPY, clipped SAXPY, clip grad kernels

All float* → __nv_bfloat16*. Same pattern: BF16 storage, F32 arithmetic.

• Step 4: Convert spectral norm kernel

Weight matrix, u/v vectors: __nv_bfloat16*. Internal matvec and reduction: F32.

• Step 5: Convert shrink-perturb, regime-scale kernels

Same pattern.

• Step 6: Convert EMA kernel

extern "C" __global__ void ema_kernel(
    __nv_bfloat16* __restrict__ target,
    const __nv_bfloat16* __restrict__ online,
    float tau, int n)

• Step 7: Verify + commit

---

Task 5: C51 loss + gradient kernels BF16

Files:

• Modify: crates/ml/src/cuda_pipeline/c51_loss_kernel.cu

• Modify: crates/ml/src/cuda_pipeline/c51_grad_kernel.cu

• Step 1: Convert c51_loss_batched kernel

All logit inputs, saved tensors, IS weights, rewards, dones → __nv_bfloat16*. Shared memory arrays: __nv_bfloat16. Internal arithmetic (softmax, log, cross-entropy): cast to F32 for transcendentals. Output (per_sample_loss, td_errors, total_loss): BF16.

• Step 2: Convert c51_grad kernel

d_value_logits, d_adv_logits outputs: __nv_bfloat16*.

• Step 3: Convert MSE loss + gradient kernels

Same pattern for mse_loss_kernel.cu and mse_grad_kernel.cu.

• Step 4: Convert expected_q_kernel and q_stats_kernel

BF16 logit inputs, BF16 outputs.

• Step 5: Verify + commit

---

Task 6: CQL + IQN + ensemble gradient kernels BF16

Files:

• Modify: crates/ml/src/cuda_pipeline/cql_grad_kernel.cu

• Modify: crates/ml/src/cuda_pipeline/iqn_dual_head_kernel.cu

• Modify: crates/ml/src/cuda_pipeline/iqn_cvar_kernel.cu

• Modify: crates/ml/src/cuda_pipeline/ensemble_kernels.cu

• Modify: crates/ml/src/cuda_pipeline/iql_value_kernel.cu

• Step 1: Convert CQL logit gradient kernel

All logit/gradient buffers → BF16.

• Step 2: Convert IQN dual-head kernel

Tau sampling, cosine embedding, quantile loss — all BF16 storage.

• Step 3: Convert IQN CVaR kernel

• Step 4: Convert ensemble diversity kernels

• Step 5: Convert IQL value kernel

• Step 6: Update Rust launchers (gpu_iqn_head.rs, gpu_iql_trainer.rs)

Change all CudaSlice<f32> buffer types in the launcher structs.

• Step 7: Verify + commit

---

Phase 4: Environment + Auxiliary Kernels

Task 7: Experience collector + trade physics BF16

Files:

• Modify: crates/ml/src/cuda_pipeline/trade_physics.cuh

• Modify: crates/ml/src/cuda_pipeline/experience_kernels.cu

• Modify: crates/ml/src/cuda_pipeline/gpu_experience_collector.rs

• Step 1: Convert trade_physics.cuh — all 11 device functions

All function parameters and locals → BF16 storage. Use __bfloat162float / __float2bfloat16 for arithmetic. For sqrt, exp, log: promote to F32.

• Step 2: Convert experience_env_step kernel

All 30 parameters that are float* → __nv_bfloat16*. Portfolio state, rewards, features, targets.

• Step 3: Update ExperienceCollectorConfig and launch args

All CudaSlice<f32> → CudaSlice<u16> in the Rust launcher.

• Step 4: Verify + commit

---

Task 8: Backtest evaluator BF16

Files:

• Modify: crates/ml/src/cuda_pipeline/backtest_env_kernel.cu

• Modify: crates/ml/src/cuda_pipeline/backtest_metrics_kernel.cu

• Modify: crates/ml/src/cuda_pipeline/backtest_gather_kernel.cu

• Modify: crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs

• Modify: crates/ml/src/cuda_pipeline/trade_stats_kernel.cu

• Step 1: Convert backtest_env_step kernel (already uses trade_physics.cuh)

• Step 2: Convert backtest_metrics_kernel (Sharpe, Calmar, CVaR, etc.)

• Step 3: Convert backtest_gather_kernel

• Step 4: Convert trade_stats_kernel

• Step 5: Update Rust launcher

• Step 6: Verify + commit

---

Task 9: Action selector + attention + monitoring BF16

Files:

• Modify: crates/ml/src/cuda_pipeline/epsilon_greedy_kernel.cu

• Modify: crates/ml/src/cuda_pipeline/gpu_action_selector.rs

• Modify: crates/ml/src/cuda_pipeline/attention_kernel.cu

• Modify: crates/ml/src/cuda_pipeline/attention_backward_kernel.cu

• Modify: crates/ml/src/cuda_pipeline/gpu_attention.rs

• Modify: crates/ml/src/cuda_pipeline/monitoring_kernel.cu

• Modify: crates/ml/src/cuda_pipeline/gpu_monitoring.rs

• Modify: crates/ml/src/cuda_pipeline/per_update_kernel.cu

• Step 1: Convert epsilon_greedy_kernel — BF16 Q-values

• Step 2: Convert attention forward + backward kernels

• Step 3: Convert monitoring kernel

• Step 4: Convert PER priority update kernel

• Step 5: Update all Rust launchers

• Step 6: Verify + commit

---

Task 10: Remaining kernels BF16

Files:

• Modify: crates/ml/src/cuda_pipeline/nstep_kernel.cu

• Modify: crates/ml/src/cuda_pipeline/her_episode_kernel.cu

• Modify: crates/ml/src/cuda_pipeline/her_relabel_kernel.cu

• Modify: crates/ml/src/cuda_pipeline/curiosity_training_kernel.cu

• Modify: crates/ml/src/cuda_pipeline/training_guard_kernel.cu

• Modify: crates/ml/src/cuda_pipeline/statistics_kernel.cu

• Modify: crates/ml/src/cuda_pipeline/dt_kernels.cu

• Modify: crates/ml/src/cuda_pipeline/ppo_experience_kernel.cu

• Modify: crates/ml/src/cuda_pipeline/backtest_forward_supervised_kernel.cu

• Modify: crates/ml/src/cuda_pipeline/backtest_forward_ppo_kernel.cu

• Modify: crates/ml/src/cuda_pipeline/signal_adapter_kernel.cu

• Modify: crates/ml/src/cuda_pipeline/backward_kernels.cu

• Step 1: Convert n-step, HER, curiosity kernels

• Step 2: Convert training guard, statistics, DT kernels

• Step 3: Convert PPO, supervised, signal adapter kernels

• Step 4: Convert backward_kernels.cu (legacy backward)

• Step 5: Update all Rust launchers

• Step 6: Verify + commit

---

Task 11: Fused training context + weight sync BF16

Files:

• Modify: crates/ml/src/trainers/dqn/fused_training.rs

• Modify: crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs

• Step 1: Remove F32→BF16 conversion step from graph_adam

Since params_buf IS BF16 now, there's no conversion needed after Adam. The Adam kernel writes BF16 directly to params_buf. The bf16_params_buf separate buffer is now REDUNDANT — params_buf IS the BF16 buffer.

• Step 2: Eliminate bf16_params_buf / bf16_target_params_buf

Since all buffers are BF16, the separate BF16 mirrors are no longer needed. Forward reads directly from params_buf (which is now BF16). Delete the mirror buffers and all conversion code.

• Step 3: Update unflatten_online_weights for BF16

The unflatten copies from params_buf to individual weight set tensors. Both are now CudaSlice<u16>. The copy is u16 → u16, same as before but type changes.

• Step 4: Update checkpoint save/load

Checkpoint serialization downloads params to host. Must convert BF16 → F32 for serialization (or serialize as BF16 with a version flag).

• Step 5: Verify + commit

---

Task 12: BF16-native tests

Files:

• Modify: crates/ml/src/trainers/dqn/smoke_tests/gradient_budget.rs

• Modify: crates/ml/src/trainers/dqn/smoke_tests/training_stability.rs

• Step 1: Update test helpers to construct BF16 buffers natively

alloc_dueling and alloc_branching now produce CudaSlice<u16>. Use half::bf16::from_f32(0.1).to_bits() for test values.

• Step 2: Update gradient budget tests for BF16 precision

Tolerance assertions may need widening (BF16 has ~0.78% relative error vs F32's ~0.00001%).

• Step 3: Add BF16 roundtrip correctness test

Upload F32 → BF16, download BF16 → F32, verify within BF16 precision bounds.

• Step 4: Verify ALL tests pass

SQLX_OFFLINE=true cargo test -p ml --lib
SQLX_OFFLINE=true cargo test -p ml-dqn --lib

• Step 5: Final commit

git commit -m "feat(bf16): complete — zero F32 on GPU, all kernels BF16"

---

Execution Dependencies

Task 1 (foundation) → Task 2 (forward) → Task 3 (backward) → Task 4 (Adam)
                                                              → Task 5 (C51/MSE)
                                                              → Task 6 (CQL/IQN/ensemble)
                    → Task 7 (experience) → Task 8 (backtest) → Task 9 (action/attn/monitoring)
                                                               → Task 10 (remaining)
                    → Task 11 (fused training) — after Tasks 2-6
                    → Task 12 (tests) — after all

Task 1 is the foundation — changes types everywhere. Tasks 2-10 convert individual kernel groups. Task 11 wires up the training pipeline. Task 12 validates.

Verification After Each Phase

Phase 1 (Task 1): cargo check --workspace — may have many errors, fix all. Phase 2 (Tasks 2-3): cargo check -p ml && cargo test -p ml --lib -- dqn Phase 3 (Tasks 4-6): Same + cargo test -p ml --lib -- gradient_budget Phase 4 (Tasks 7-12): Full test suite including smoke tests if GPU available.