# 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)
