Spec: every CUDA kernel, every GPU buffer, every cuBLAS call → BF16. 37 HOT kernels + 11 WARM + 7 COLD = 55 total kernel files. 12 implementation tasks across 4 phases. Phase 1: Foundation (buffer types, weight sets) Phase 2: cuBLAS GemmEx BF16×BF16→BF16 (forward + backward) Phase 3: Optimizer + loss kernels (Adam, C51, MSE, spectral norm) Phase 4: Environment + auxiliary (experience, backtest, IQN, attention) Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
578 lines
21 KiB
Markdown
578 lines
21 KiB
Markdown
# 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:**
|
||
```cuda
|
||
// 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:**
|
||
```rust
|
||
// 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**
|
||
|
||
```rust
|
||
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<f32> to CudaSlice<u16>**
|
||
|
||
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<f32>
|
||
2. Convert to Vec<u16> via `f32_to_bf16_host`
|
||
3. Upload to CudaSlice<u16>
|
||
|
||
- [ ] **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:
|
||
```rust
|
||
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**
|
||
|
||
```bash
|
||
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:
|
||
```rust
|
||
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`:
|
||
```rust
|
||
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:
|
||
```cuda
|
||
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`:
|
||
```cuda
|
||
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):
|
||
|
||
```cuda
|
||
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**
|
||
|
||
```cuda
|
||
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**
|
||
|
||
```cuda
|
||
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**
|
||
|
||
```bash
|
||
SQLX_OFFLINE=true cargo test -p ml --lib
|
||
SQLX_OFFLINE=true cargo test -p ml-dqn --lib
|
||
```
|
||
|
||
- [ ] **Step 5: Final commit**
|
||
|
||
```bash
|
||
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.
|