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

Spec A (Revised): Full BF16 Rewrite — Zero F32 On GPU

Date: 2026-03-27 Scope: Convert EVERY CUDA kernel, EVERY GPU buffer, and EVERY cuBLAS call from F32 to BF16. Zero F32 on GPU. No mixed precision. No compromises.

Principle

__nv_bfloat16 everywhere. Every float in every CUDA kernel becomes __nv_bfloat16. Every CudaSlice<f32> in Rust becomes CudaSlice<u16> (BF16 stored as u16 in cudarc). Every cublasSgemm becomes cublasGemmEx with CUDA_R_16BF for ALL three matrix types (A, B, C). Every accumulator, every intermediate, every buffer — BF16.

The only F32 that touches GPU is the alpha/beta scalar arguments to cublasGemmEx (required by the API to be host-side F32 pointers even when compute type is BF16).

Why Full BF16

1. Performance: H100 BF16 tensor cores = 989 TFLOPS vs 67 TFLOPS F32. 14.8× theoretical, ~3× practical.
2. Memory: Every buffer halved. 500K params × 4 bytes = 2MB F32. In BF16 = 1MB. Adam moments halved. Activations halved. More room for larger batches.
3. Bandwidth: HBM3 bandwidth is the bottleneck for element-wise kernels (Adam, loss, spectral norm). Half the data = 2× throughput.
4. Consistency: No precision boundaries between kernels. No F32→BF16 conversion kernels needed between stages. The output of one kernel IS the input of the next — both BF16.
5. Simplicity: One data type everywhere. No mixed-precision bookkeeping.

BF16 Numeric Properties

• Exponent: 8 bits (same as F32) → range ±3.4×10³⁸, no underflow/overflow risk
• Mantissa: 7 bits → ~1/128 relative precision → 0.78% relative error
• Adam convergence: Weight updates of magnitude lr × m_hat / sqrt(v_hat) ≈ 3e-4 × 1.0 = 3e-4. BF16 can represent this relative to weights of magnitude 0.01-10.0. No stalling.
• Loss precision: C51 cross-entropy ≈ 0.1-50.0. BF16 resolution at 50.0 is 0.25. Adequate for training signal.
• Gradient precision: Per-element gradients ≈ 1e-4 to 1e+2. BF16 handles this entire range.

Architecture

Data Flow (All BF16)

Experience Collector
    |
    v
states [B, SD] (BF16)
    |
    v  cublasGemmEx (BF16 × BF16 → BF16)
Forward Pass
    |-- h_s1 [B, SH1] = BF16(ReLU(states @ W_s1^T + b_s1))
    |-- h_s2 [B, SH2] = BF16(ReLU(h_s1 @ W_s2^T + b_s2))
    |-- v_logits [B, NA] = BF16(h_v @ W_v2^T + b_v2)
    v
C51/MSE Loss (BF16 logits → BF16 d_logits)
    |
    v  cublasGemmEx (BF16 × BF16 → BF16)
Backward Pass
    |-- dW = d_output^T × input    → BF16 (accumulated into grad_buf)
    |-- dX = d_output × W^T        → BF16 (upstream gradient)
    v
grad_buf [TOTAL_PARAMS] (BF16)
    |
    v
Adam (BF16 params, BF16 grads, BF16 m, BF16 v)
    |
    v
params_buf [TOTAL_PARAMS] (BF16) → next forward pass

No conversion kernels between stages. Output of one is input of next. All BF16.

What Changes

CUDA Shared Header (common_device_functions.cuh)

• Add #include <cuda_bf16.h> at the top
• Add BF16 helper macros: __BF16_TO_FLOAT(x), __FLOAT_TO_BF16(x)
• All float state/param typedefs → __nv_bfloat16

CUDA Trade Physics (trade_physics.cuh)

• All 11 device functions: float args/locals → __nv_bfloat16
• Intermediate arithmetic: use __nv_bfloat16 native ops (H100 has BF16 ALU)
• For complex math (sqrt, exp, log): cast to float, compute, cast back

cuBLAS Forward (batched_forward.rs)

• cublasSgemm → cublasGemmEx with ALL types CUDA_R_16BF
• cublasComputeType_t::CUBLAS_COMPUTE_32F (internal F32 accumulate, BF16 output)
• Wait — for full BF16 output: use CUBLAS_COMPUTE_16BF to get BF16 C output
  • Actually: cublasGemmEx with Ctype=CUDA_R_16BF and computeType=CUBLAS_COMPUTE_32F gives BF16 output with F32 internal accumulation. Best of both worlds.
• All activation buffers: CudaSlice<u16> (BF16)
• Weight pointers: from bf16_params_buf (already exists)
• Bias kernels: BF16 input/output

cuBLAS Backward (batched_backward.rs)

• Same cublasGemmEx pattern for dW and dX
• grad_buf: BF16 (was F32) — Adam reads BF16 gradients
• All intermediate activation gradients: BF16

Adam Kernel (dqn_utility_kernels.cu)

• params, grads, m, v: all __nv_bfloat16*
• Bias correction: BF16 arithmetic (cast to float for powf if needed)
• Weight decay: BF16
• Grad norm: BF16 input, BF16 sum → final sqrt

C51 Loss (c51_loss_kernel.cu)

• All logit inputs: BF16
• Shared memory: BF16
• Softmax, log, exp: cast to float for transcendentals, result back to BF16
• Cross-entropy accumulation: BF16
• Loss output: BF16

MSE Loss (mse_loss_kernel.cu)

• Same pattern — BF16 logits, BF16 loss

C51/MSE Gradient Kernels

• d_logits output: BF16

Experience Kernels (experience_kernels.cu)

• Portfolio state: BF16
• Rewards: BF16
• Trade physics: BF16

Backtest Kernel (backtest_env_kernel.cu)

• Same as experience — all BF16

Spectral Norm (dqn_utility_kernels.cu)

• Weight matrix: BF16
• u/v vectors: BF16
• sigma computation: BF16 (cast to float for sqrt)

All Other Kernels

• EMA kernel: BF16
• PER update: BF16
• Epsilon greedy: BF16 Q-values
• IQN head: BF16
• Attention: BF16
• CQL gradient: BF16
• Ensemble: BF16
• Every. Single. One.

Rust-Side Changes

Every CudaSlice<f32> in GpuDqnTrainer, CublasForward, CublasBackward, GpuExperienceCollector, GpuBacktestEvaluator, GpuActionSelector, GpuIqnHead, GpuAttention, etc. becomes CudaSlice<u16>.

The alloc_f32 helper becomes alloc_bf16. All memcpy_htod of &[f32] becomes conversion to Vec<u16> (BF16 bit pattern) first.

The DuelingWeightSet and BranchingWeightSet fields change from CudaSlice<f32> to CudaSlice<u16>.

Conversion at Boundaries

Input boundary (CPU → GPU):

• Feature vectors arrive as Vec<f32> from Rust
• Convert to Vec<u16> (BF16 via f32::to_bits() >> 16 approximation or half::bf16::from_f32)
• Upload CudaSlice<u16>

Output boundary (GPU → CPU):

• Loss scalar, grad_norm scalar: download as BF16 u16, convert to f32 on CPU for logging
• Q-value stats: same
• Checkpoint save: download BF16 params, convert to f32 for serialization (or serialize as BF16)

Implementation Phases

This is too large for one commit. Split into 4 phases, each producing compilable + testable code:

Phase 1: BF16 type foundation

• Add half crate dependency for bf16 type
• Create alloc_bf16, bf16_from_f32_vec, f32_from_bf16_vec helpers
• Add #include <cuda_bf16.h> to common_device_functions.cuh
• All tests still pass (no behavioral change yet)

Phase 2: cuBLAS BF16 (forward + backward)

• Convert cublasSgemm → cublasGemmEx BF16×BF16→BF16
• Convert activation buffers to BF16
• Convert weight buffers to BF16
• Bias kernels to BF16
• Forward pass fully BF16
• Backward pass fully BF16

Phase 3: Optimizer + loss kernels BF16

• Adam kernel: all BF16
• Gradient norm: BF16
• Clipping kernels: BF16
• C51 loss + grad: BF16
• MSE loss + grad: BF16
• Spectral norm: BF16
• EMA: BF16

Phase 4: Environment + auxiliary kernels BF16

• Experience collector: BF16
• Backtest evaluator: BF16
• Trade physics: BF16
• Action selector: BF16
• IQN head: BF16
• Attention: BF16
• CQL: BF16
• Ensemble: BF16
• PER: BF16
• All remaining kernels

Files Changed (Complete List)

CUDA Kernels (37 files + 2 headers): Every .cu and .cuh file in crates/ml/src/cuda_pipeline/

Rust Launchers (18 files): Every gpu_*.rs and batched_*.rs in crates/ml/src/cuda_pipeline/

Rust Weight Types: crates/ml/src/cuda_pipeline/gpu_weights.rs — all weight set structs

Training Orchestration: crates/ml/src/trainers/dqn/fused_training.rs

Tests: All smoke tests, gradient budget tests

Success Criteria

1. Zero float in CUDA hot path: grep -rn "float\b" *.cu in hot-path kernels returns 0 hits (except for cublasGemmEx alpha/beta host scalars and transcendental casts)
2. All 488+ tests pass
3. Epoch time ≤3s on H100 (was ~30s F32, target: BF16 tensor cores + halved bandwidth)
4. Training converges: Sharpe ratio within 10% of F32 baseline over 50 epochs
5. Peak VRAM reduced ~40%: all buffers halved except scalar counters

Risks

Risk	Mitigation
BF16 Adam stalling (updates too small)	Monitor weight update magnitude; BF16 range covers lr=3e-4 updates on weights 0.01-10
C51 log-softmax overflow	Max logit ~50 in BF16 → exp(50) = 5e21, fits BF16 range (3.4e38)
Spectral norm sigma precision	sigma_max=3.0, BF16 resolution at 3.0 is 0.015 — adequate
cuBLAS BF16 GEMM not capturable in CUDA Graph	Test on CUDA 12.x; fallback: launch outside graph (adds ~2ms)
Backtest metrics precision	Sharpe/Calmar use division — BF16 division is fine for 2-decimal metrics

Non-Goals

• F16 support (BF16 has superior dynamic range)
• Multi-GPU (single H100 target)
• CPU fallback path (GPU-only)
• Backward compatibility with F32 checkpoints (convert on load)