Files
foxhunt/crates
jgrusewski 29b6ec4dd9 feat(bf16): C51/MSE loss + grad kernels accept BF16 logits
All 6 loss/gradient CUDA kernels converted from float* to __nv_bfloat16*:
- c51_loss_batched: BF16 logits (12 inputs), rewards, dones, IS weights, outputs
- c51_grad_kernel: BF16 d_logits output, atomicAddBF16 for gradient accumulation
- mse_loss_batched: BF16 logits, rewards, dones, IS weights, outputs
- mse_grad_kernel: BF16 d_logits output, atomicAddBF16
- expected_q_kernel: BF16 logits in, BF16 Q-values out
- q_stats_kernel: BF16 Q-values in (monitoring scalars stay float)

Pattern: BF16 storage, F32 arithmetic (cast on load/store).
Shared memory stays float for softmax/log/exp precision.
total_loss stays float* (atomicAdd doesn't support BF16).

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-03-27 22:09:03 +01:00
..