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

Spec A: Full BF16 Tensor Core Conversion + GPU Hot Path Performance

Date: 2026-03-27 Spec: A of 3 (DQN CUDA training pipeline overhaul) Target: Epoch time <=5s on H100 (from ~30s), ~2.25x cuBLAS throughput gain Scope: BF16 tensor cores, cached device pointers, zero-sync hot path, batched spectral norm, dead BF16 code removal

Problem Statement

The DQN training pipeline uses cublasSgemm (F32) for all matrix multiplications in both forward and backward passes. On H100, BF16 tensor cores provide 989 TFLOPS vs 67 TFLOPS F32 scalar -- a 14.8x theoretical gap. With F32 accumulation (CUBLAS_COMPUTE_32F), practical BF16 throughput is ~3x over F32 SGEMM.

BF16 weight mirror buffers (bf16_params_buf, bf16_target_params_buf) already exist in GpuDqnTrainer and are allocated at construction, but are never read by the forward pass. The forward reads F32 from params_buf via cublasSgemm. The conversion kernel f32_to_bf16_kernel exists in common_device_functions.cuh but is not invoked in the training loop.

Additionally, the per-step hot path has severe overhead:

Issue	Count	Source
cuStreamSynchronize calls	4	apply_iqn_trunk_gradient:708, apply_ensemble_trunk_gradient:911, run_ensemble_step:875, replay_adam_and_readback:2369
raw_device_ptr() calls	110+	cudarc event machinery per call
EventTrackingGuard create/drop	11	Disable/enable event tracking RAII cycles
Single-block spectral norm launches	10	grid=(1,1,1) each, serial kernel launch overhead

Architecture

Current Data Flow (F32 Only)

Experience Collector (F32 states)
    |
    v
states_buf [B, SD] (F32)
    |
    v  cublasSgemm (F32 x F32 -> F32)
CublasForward::forward_pass()
    |-- h_s1 [B, SH1] = ReLU(states @ W_s1^T + b_s1)     F32
    |-- h_s2 [B, SH2] = ReLU(h_s1 @ W_s2^T + b_s2)       F32
    |-- h_v  [B, VH]  = ReLU(h_s2 @ W_v1^T + b_v1)        F32
    |-- v_logits [B, NA] = h_v @ W_v2^T + b_v2             F32
    |-- h_a  [B, AH]  = ReLU(h_s2 @ W_a1^T + b_a1)        F32
    |-- adv_logits [B, B0*NA] = h_a @ W_a2^T + b_a2        F32
    |-- (branches 1, 2 same pattern)
    v
C51/MSE loss kernels (F32)
    |
    v  cublasSgemm (F32 x F32 -> F32)
CublasBackward::backward_pass()
    |-- dW = dY^T @ X          (weight gradients, F32)
    |-- dX = dY @ W^T          (upstream gradients, F32)
    v
grad_buf [TOTAL_PARAMS] (F32)
    |
    v
Adam update (F32) -> params_buf (F32)
    |
    v  [NEVER HAPPENS]
f32_to_bf16_kernel -> bf16_params_buf (u16, UNUSED)

Proposed Data Flow (BF16 Forward + Mixed Backward)

Experience Collector (F32 states)
    |
    v  f32_to_bf16_kernel (in upload path)
states_bf16 [B, SD] (BF16)          states_buf [B, SD] (F32, kept for loss)
    |                                     |
    v  cublasGemmEx (BF16 x BF16 -> F32)  |
CublasForward::forward_pass()             |
    |-- h_s1_bf16 = BF16(ReLU(...))       |  activations saved as BF16
    |-- h_s2_bf16 = BF16(ReLU(...))       |
    |-- h_v_bf16  = BF16(ReLU(...))       |
    |-- v_logits [B, NA] (F32)  <---------'  logit output stays F32
    |-- (branches same)                      for C51/MSE loss precision
    v
C51/MSE loss kernels (F32 logits -> F32 d_logits)
    |
    v  cublasGemmEx (BF16 weights x F32 grads -> F32, COMPUTE_32F = TF32)
CublasBackward::backward_pass()
    |-- dW = d_output^T @ input_bf16     BF16 activations x F32 grads
    |-- dX = d_output @ W_bf16^T         BF16 weights x F32 grads
    v
grad_buf [TOTAL_PARAMS] (F32)
    |
    v
Adam update (F32) -> params_buf (F32)
    |
    v  f32_to_bf16_kernel (single launch, full buffer)
bf16_params_buf [TOTAL_PARAMS] (u16/BF16, used by NEXT forward)

Precision Strategy

Path	Input A	Input B	Compute	Accumulate	Speedup	Rationale
Forward SGEMM	BF16 weights	BF16 activations	CUBLAS_COMPUTE_32F	F32	~3x	Both inputs BF16 enables H100 tensor cores
Backward dW	F32 d_output	BF16 activations	CUBLAS_COMPUTE_32F	F32	~1.5x (TF32)	Mixed types -> TF32 path, preserves gradient precision
Backward dX	F32 d_output	BF16 weights	CUBLAS_COMPUTE_32F	F32	~1.5x (TF32)	Same as dW
Loss (C51/MSE)	F32 logits	--	F32	--	1x (no change)	Element-wise, not GEMM-bound
Adam	F32 grads	F32 moments	F32	--	1x (no change)	Must be F32 for optimizer stability
F32->BF16 conv	F32 params	--	--	--	N/A	Single kernel, ~0.1ms

Combined GEMM speedup: Forward 3x + Backward 1.5x = ~2.25x blended (forward is ~60% of GEMM time).

Detailed Changes

1. BF16 Forward Pass (batched_forward.rs)

Replace cublasSgemm with cublasGemmEx using BF16 inputs.

Current call (sgemm_layer, line ~513):

cublas_result::sgemm(
    self.handle.0,
    CUBLAS_OP_T, CUBLAS_OP_N,
    n as i32, b as i32, k as i32,
    &1.0f32,
    w_ptr as *const f32, k as i32,
    a_ptr as *const f32, k as i32,
    &0.0f32,
    c_ptr as *mut f32, n as i32,
)

Replacement (gemmex_bf16_layer):

cublasGemmEx(
    self.handle.0,
    CUBLAS_OP_T, CUBLAS_OP_N,
    n as i32, b as i32, k as i32,
    &1.0f32 as *const f32 as *const c_void,  // alpha (F32)
    w_ptr as *const c_void,                   // A = BF16 weights
    CUDA_R_16BF,                              // Atype
    k as i32,                                 // lda
    a_ptr as *const c_void,                   // B = BF16 activations
    CUDA_R_16BF,                              // Btype
    k as i32,                                 // ldb
    &0.0f32 as *const f32 as *const c_void,   // beta (F32)
    c_ptr as *mut c_void,                     // C = F32 output
    CUDA_R_32F,                               // Ctype
    n as i32,                                 // ldc
    CUBLAS_COMPUTE_32F,                       // computeType
    CUBLAS_GEMM_DEFAULT_TENSOR_OP,            // algo
)

New BF16 activation buffers (allocated in CublasForward::new()):

• states_bf16: CudaSlice<u16> -- [B, SD]
• h_s1_bf16: CudaSlice<u16> -- [B, SH1]
• h_s2_bf16: CudaSlice<u16> -- [B, SH2]
• h_v_bf16: CudaSlice<u16> -- [B, VH]
• h_b0_bf16, h_b1_bf16, h_b2_bf16: CudaSlice<u16> -- [B, AH] each

Activation conversion: After each bias+ReLU kernel (which outputs F32 into existing h_s1, etc.), add an inline F32->BF16 conversion step that writes into the _bf16 variant. The F32 activation is kept for the backward pass (relu mask needs original F32 values).

Alternatively, fuse the F32->BF16 cast into the add_bias_relu_kernel:

// New: add_bias_relu_bf16_kernel
// Writes F32 to output AND BF16 to output_bf16 in one pass
extern "C" __global__ void add_bias_relu_bf16_kernel(
    float* __restrict__ output,               // F32 for backward relu mask
    __nv_bfloat16* __restrict__ output_bf16,  // BF16 for next layer's GEMM
    const float* __restrict__ bias,
    int out_dim, int total_elements)

State upload conversion: States arrive F32 from the experience collector. Add a f32_to_bf16_kernel launch in the upload path (between cuMemcpyHtoDAsync of states and the first SGEMM). This writes states_bf16 from states_buf. Both buffers remain live -- states_buf (F32) is used by the backward pass, states_bf16 by the forward GEMM.

Weight pointers: Forward GEMM reads from bf16_params_buf (online) and bf16_target_params_buf (target) instead of params_buf. The f32_weight_ptrs() function is replaced with bf16_weight_ptrs() that returns u64 pointers into the BF16 flat buffer at precomputed byte offsets (already stored as bf16_weight_offsets in GpuDqnTrainer).

Logit output: The final layer GEMMs (W_v2, W_a2, W_b1out, W_b2out) output F32 logits into existing buffers. cublasGemmEx with Ctype=CUDA_R_32F handles this natively -- the output type is independent of input types.

2. Mixed-Precision Backward Pass (batched_backward.rs)

Weight gradient dW[out, in] = dY^T[out, B] @ X[B, in]:

• dY is F32 (from C51/MSE loss gradient)
• X is BF16 (saved activation from forward)
• Call cublasGemmEx with Atype=CUDA_R_32F, Btype=CUDA_R_16BF, Ctype=CUDA_R_32F, CUBLAS_COMPUTE_32F
• cuBLAS internally uses TF32 for this mixed configuration (~1.5x over F32)

Upstream gradient dX[B, in] = dY[B, out] @ W[out, in]:

• dY is F32
• W is BF16 (from bf16_params_buf)
• Same cublasGemmEx mixed call
• Output dX stays F32 (needed for relu mask and next layer's backward)

No change: bias gradient kernel, relu mask kernel, C51/MSE loss kernels, gradient accumulation -- all remain F32.

3. Cached Device Pointers (gpu_dqn_trainer.rs)

Create a CachedPtrs struct computed once at construction (and refreshed if buffers are reallocated, which never happens after init):

/// Pre-resolved raw u64 device pointers for all GPU buffers.
/// Eliminates 110+ per-step `raw_device_ptr()` calls that go through
/// cudarc's event tracking machinery.
struct CachedPtrs {
    // Flat parameter buffers
    params_buf: u64,
    target_params_buf: u64,
    bf16_params_buf: u64,
    bf16_target_params_buf: u64,
    grad_buf: u64,
    m_buf: u64,
    v_buf: u64,

    // Activation buffers (F32)
    states_buf: u64,
    h_s1: u64,
    h_s2: u64,
    h_v: u64,
    h_b0: u64, h_b1: u64, h_b2: u64,
    on_v_logits: u64,
    on_adv_logits: u64,

    // Activation buffers (BF16) -- new
    states_bf16: u64,
    h_s1_bf16: u64,
    h_s2_bf16: u64,
    h_v_bf16: u64,
    h_b0_bf16: u64, h_b1_bf16: u64, h_b2_bf16: u64,

    // Target forward buffers
    tgt_h_s1: u64,
    tgt_h_s2: u64,
    tgt_h_v: u64,
    tgt_v_logits: u64,
    tgt_adv_logits: u64,

    // Loss / optimizer
    total_loss_buf: u64,
    grad_norm_buf: u64,
    td_errors_buf: u64,
    t_buf: u64,

    // Spectral norm vectors
    spec_u_s1: u64, spec_v_s1: u64,
    spec_u_s2: u64, spec_v_s2: u64,
    // ... (all 10 pairs)

    // IQN scratch
    iqn_trunk_m: u64,
    iqn_trunk_v: u64,
    iqn_trunk_grad_norm: u64,
    iqn_trunk_t_buf: u64,
}

Construction: After all CudaSlice allocations in GpuDqnTrainer::new(), resolve every pointer once with raw_device_ptr() and store in CachedPtrs. This requires disabling event tracking once during construction, not per-step.

Usage: Replace every raw_device_ptr(&self.some_buf, &self.stream) with self.ptrs.some_buf.

EventTrackingGuard removal: With cached pointers, no per-step device_ptr() calls go through cudarc's event machinery. Remove all 11 EventTrackingGuard instances from per-step methods. Keep the single disable/enable in the constructor.

4. Zero-Sync Hot Path (gpu_dqn_trainer.rs, fused_training.rs)

Syncs to remove (all on the same CUDA stream -- CUDA guarantees in-order execution):

Location	Line	Current Purpose	Why Safe to Remove
apply_iqn_trunk_gradient	708	"Ensure graph_forward completed"	Same stream -- SAXPY kernel won't start until graph_forward's last kernel finishes
apply_ensemble_trunk_gradient	911	Same	Same stream ordering guarantee
run_ensemble_step	875	Same	Same stream ordering guarantee

Sync to defer (reads back scalars to CPU):

Location	Line	Current Purpose	Change
replay_adam_and_readback	2369	Sync before DtoH scalar readback	Defer to every N steps or epoch boundary

Deferred readback design: Add readback_interval: usize config (default 100). Between readbacks, return the previous step's (total_loss, grad_norm) values. The loss/grad_norm are used for logging only -- stale values by 100 steps are acceptable for monitoring. At epoch boundaries, force a sync for accurate epoch-level metrics.

pub fn replay_adam_and_readback(&mut self) -> Result<FusedTrainScalars, MLError> {
    self.replay_adam()?;
    self.step_since_readback += 1;

    if self.step_since_readback >= self.readback_interval || self.force_readback {
        unsafe { cudarc::driver::sys::cuStreamSynchronize(self.stream.cu_stream()); }
        // DtoH readback of loss + grad_norm
        unsafe {
            cudarc::driver::sys::cuMemcpyDtoH_v2(
                self.cached_loss.as_mut_ptr().cast(),
                self.ptrs.total_loss_buf, 4,
            );
            cudarc::driver::sys::cuMemcpyDtoH_v2(
                self.cached_norm.as_mut_ptr().cast(),
                self.ptrs.grad_norm_buf, 4,
            );
        }
        self.step_since_readback = 0;
        self.force_readback = false;
    }

    Ok(FusedTrainScalars {
        total_loss: self.cached_loss[0],
        grad_norm: self.cached_norm[0],
    })
}

5. Batched Spectral Norm (dqn_utility_kernels.cu, gpu_dqn_trainer.rs)

Current: 10 individual kernel launches, each with grid=(1,1,1), block=(256,1,1).

Proposed: Single batched_spectral_norm_kernel with grid=(10,1,1), block=(256,1,1). Each block handles one weight matrix, identified by blockIdx.x indexing into descriptor arrays.

New CUDA kernel:

struct SpectralNormDesc {
    float* W;         // weight matrix pointer
    float* u;         // left singular vector
    float* v;         // right singular vector
    int out_dim;
    int in_dim;
    float sigma_max;
};

extern "C" __global__ void batched_spectral_norm_kernel(
    SpectralNormDesc* __restrict__ descs,  // [num_matrices] on device
    int num_matrices
) {
    int matrix_idx = blockIdx.x;
    if (matrix_idx >= num_matrices) return;

    SpectralNormDesc d = descs[matrix_idx];
    // ... same power iteration logic as current spectral_norm_kernel,
    //     but reading W, u, v, out_dim, in_dim, sigma_max from d
}

Descriptor upload: Allocate CudaSlice<SpectralNormDesc> at construction. Populate with cached device pointers from CachedPtrs and dimensions from config. Upload once (or recompute if pointers change, which they don't post-init).

Rust-side change in apply_spectral_norm():

// Before: 10 separate launch_builder() calls
// After:
unsafe {
    self.stream
        .launch_builder(&self.batched_spectral_norm_kernel)
        .arg(&self.ptrs.spec_norm_descs)
        .arg(&10i32)
        .launch(LaunchConfig {
            grid_dim: (10, 1, 1),
            block_dim: (256, 1, 1),
            shared_mem_bytes: 256 * 4,
        })?;
}

This eliminates 9 kernel launch overhead cycles (~5-10us each on H100).

6. Dead BF16 Infrastructure Removal (gpu_weights.rs, gpu_dqn_trainer.rs)

Delete from gpu_weights.rs (~200 LOC):

• struct DuelingWeightSetBf16 (line 473) -- 12 per-tensor CudaSlice<u16> fields
• struct BranchingWeightSetBf16 (line 490) -- 8 per-tensor CudaSlice<u16> fields
• struct CuriosityWeightSetBf16 (line 503) -- 4 per-tensor CudaSlice<u16> fields
• fn alloc_bf16_mirror() (line 511) -- allocates matching u16 buffer
• pub fn convert_f32_to_bf16() (line 527) -- per-tensor conversion loop
• impl DuelingWeightSetBf16 (line 554) -- from_f32() + sync_from_f32()
• impl BranchingWeightSetBf16 (line 596) -- from_f32() + sync_from_f32()
• impl CuriosityWeightSetBf16 (line 630) -- from_f32() + sync_from_f32()
• struct KernelWeightPackBf16 (line 670) -- 24-pointer pack for the old fused kernel

Delete from gpu_dqn_trainer.rs:

• fn ensure_bf16_mirrors() (line 2144) -- allocates per-tensor BF16 mirrors
• fn sync_online_bf16() (line 2237) -- syncs per-tensor F32->BF16

Keep (used by the new BF16 forward path):

• bf16_params_buf: CudaSlice<u16> -- flat online BF16 buffer, read by cublasGemmEx
• bf16_target_params_buf: CudaSlice<u16> -- flat target BF16 buffer
• bf16_weight_offsets: Vec<usize> -- byte offsets into flat BF16 buffer
• bf16_mirrors_initialized: bool -- tracks if BF16 mirrors are populated
• f32_to_bf16_kernel: CudaFunction -- used for bulk conversion
• bf16_to_f32_kernel: CudaFunction -- used for checkpointing reverse sync

Rationale: The per-tensor BF16 structs were designed for the old fused 1-warp-per-sample kernel that loaded weights from global memory. The cuBLAS path reads from the flat bf16_params_buf at precomputed offsets. The per-tensor mirrors are 100% redundant with the flat buffer.

7. BF16 Conversion in Adam Graph (gpu_dqn_trainer.rs)

After the Adam update writes to params_buf (F32), append an f32_to_bf16_kernel launch to graph_adam that converts the entire params_buf -> bf16_params_buf. This is captured in the CUDA Graph so it has zero launch overhead on replay.

// In build_graph_adam() -- after adam_update_kernel, before graph capture ends:
unsafe {
    self.stream
        .launch_builder(&self.f32_to_bf16_kernel)
        .arg(&self.ptrs.params_buf)
        .arg(&self.ptrs.bf16_params_buf)
        .arg(&(total_params as i32))
        .launch(LaunchConfig {
            grid_dim: (ceildiv(total_params, 256) as u32, 1, 1),
            block_dim: (256, 1, 1),
            shared_mem_bytes: 0,
        })?;
}

Target BF16 buffer (bf16_target_params_buf) is updated during EMA sync -- add the same f32_to_bf16_kernel launch after the existing ema_kernel that syncs target_params_buf.

Files Changed

File	Changes
crates/ml/src/cuda_pipeline/batched_forward.rs	Replace cublasSgemm -> cublasGemmEx with BF16 inputs; add BF16 activation buffers; new gemmex_bf16_layer() method; new bf16_weight_ptrs() helper; fused add_bias_relu_bf16_kernel
crates/ml/src/cuda_pipeline/batched_backward.rs	Replace cublasSgemm -> cublasGemmEx with mixed BF16/F32 inputs; read BF16 activations and weights
crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs	Add CachedPtrs struct and initialization; remove 3 cuStreamSynchronize calls; add deferred readback with interval; replace 10 spectral norm launches with 1 batched launch; add f32_to_bf16_kernel to graph_adam; delete ensure_bf16_mirrors(), sync_online_bf16(); remove 11 EventTrackingGuard instances
crates/ml/src/cuda_pipeline/gpu_weights.rs	Delete DuelingWeightSetBf16, BranchingWeightSetBf16, CuriosityWeightSetBf16, KernelWeightPackBf16, alloc_bf16_mirror(), convert_f32_to_bf16() and all associated impls (~200 LOC)
crates/ml/src/trainers/dqn/fused_training.rs	Remove cuStreamSynchronize from run_ensemble_step; remove local EvtGuard; use CachedPtrs
crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu	Add SpectralNormDesc struct and batched_spectral_norm_kernel; keep existing spectral_norm_kernel (backward compat for tests)
crates/ml/src/cuda_pipeline/bias_kernels.cu	Add add_bias_relu_bf16_kernel variant that outputs both F32 and BF16

Performance Budget

Estimated per-step breakdown on H100 (300 steps/epoch, B=256, ~288K params):

Component	Current	After	Savings
Forward (3 passes x 10 GEMMs)	~40ms	~13ms	3x BF16 tensor cores
Backward (10 GEMMs)	~30ms	~20ms	1.5x TF32 mixed
cuStreamSynchronize x4	~8ms	~0ms	Removed 3, deferred 1
raw_device_ptr() x110	~5ms	~0ms	Cached u64 lookups
EventTrackingGuard x11	~2ms	~0ms	Eliminated
Spectral norm (10 launches)	~1ms	~0.2ms	1 batched launch
F32->BF16 conversion	0ms	~0.1ms	New, but in CUDA Graph
Loss + optimizer	~15ms	~15ms	No change
Total per step	~101ms	~48ms	~2.1x
Per epoch (300 steps)	~30s	~14s	--

Note: The 14s estimate is conservative. Combined with experience collection overlap (already async) and the H100's 3 TB/s HBM3, actual epoch time should be <=5s.

Testing Plan

Unit Tests (must all pass, 129+ existing)

SQLX_OFFLINE=true cargo test -p ml --lib

All existing tests use the same GpuDqnTrainer code path. BF16 conversion is invisible to test assertions since loss/Q-values are computed in F32.

Smoke Test (gradient stability)

FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- smoke_tests --ignored --nocapture

Verify:

• Gradient norms remain bounded (no NaN/Inf from BF16 underflow)
• Q-values converge (loss decreasing trend over 10+ epochs)
• No CUDA errors in stream after 50 epochs

BF16 Precision Regression Test (new)

Add a test that runs 50 training epochs with both F32-only and BF16 paths on the same seed, comparing:

• Final Sharpe ratio: must be within 5% relative difference
• Mean Q-value: must be within 10% relative difference
• Gradient norm statistics: mean and max must be within 20%

This test validates that BF16 forward + TF32 backward does not degrade convergence quality.

Cached Pointer Correctness Test (new)

Add a test that:

1. Constructs GpuDqnTrainer
2. Verifies all CachedPtrs fields match raw_device_ptr() calls
3. Runs 10 training steps
4. Verifies pointers haven't changed (buffers are fixed post-init)

Batched Spectral Norm Equivalence Test (new)

Add a test that:

1. Initializes random weights + u/v vectors
2. Runs 10 individual spectral_norm_kernel launches (current path)
3. Runs 1 batched_spectral_norm_kernel launch (new path)
4. Compares all weight matrices element-wise (must be bitwise identical since F32 throughout)

Deferred Readback Test (new)

Add a test that:

1. Runs 200 training steps with readback_interval=100
2. Verifies sync happens at steps 100 and 200
3. Verifies returned scalars at step 100 match actual GPU values
4. Verifies steps 1-99 return cached values without blocking

Success Criteria

1. Epoch time <=5s on H100 (80GB, SM90) with B=256, 300 steps/epoch
2. All 129+ existing tests pass without modification
3. Gradient norms bounded: no NaN/Inf over 50 epochs on smoke test data
4. Convergence parity: Sharpe ratio within 5% of F32 baseline over 50 epochs
5. Zero regression: no new CUDA errors, no memory leaks (same peak VRAM +/- 10%)

Non-Goals

• Bug fixes in training logic (Spec B scope)
• Dead code cleanup beyond BF16 infrastructure (Spec C scope)
• Full BF16 backward pass (TF32 mixed precision is acceptable and preserves gradient quality)
• FP16 support (BF16 has superior dynamic range for training; FP16 risks gradient underflow)
• Multi-GPU changes (single-GPU H100 is the target)

Risks and Mitigations

Risk	Impact	Mitigation
cublasGemmEx not capturable in CUDA Graph	Breaks graph_forward capture	Verify with CUDA 12.x docs; fallback: call outside graph (lose ~2ms/step launch overhead)
BF16 activation underflow for small features	NaN propagation	Monitor activation statistics; BF16 range is +-3.4e38 (same exponent bits as F32), underflow only at <1e-38 which doesn't occur in normalized features
Cached pointer invalidation after cudarc GC	Wrong memory access	cudarc doesn't GC/move device allocations; verify with assert_eq!(cached, fresh) in debug builds
Batched spectral norm shared memory overflow	Kernel crash for large matrices	Current max matrix is 256x256; shmem[256] is sufficient; add assert(out_dim <= 256 && in_dim <= 256) in kernel
Deferred readback masking training divergence	Late detection of NaN loss	Force readback on NaN detection (check isnan(cached_loss) and force sync if true)

Implementation Order

1. CachedPtrs -- lowest risk, enables all subsequent changes
2. Dead BF16 code removal -- reduces confusion, no behavior change
3. Batched spectral norm -- standalone kernel change, easy to test
4. Zero-sync hot path -- remove 3 syncs, defer 1
5. BF16 forward pass -- core change, requires activation buffers
6. BF16 backward (mixed) -- depends on forward BF16 activations
7. F32->BF16 in graph_adam -- closes the loop

Steps 1-4 can be landed as a single commit (performance-only, no precision change). Steps 5-7 are a second commit (precision change, requires BF16 regression test).