- total_loss_buf: CudaSlice<half::bf16> → CudaSlice<f32> (native atomicAdd,
eliminates atomicAddBF16 CAS loop as potential NaN source)
- Loss kernels: float* total_loss + atomicAdd (was atomicAddBF16)
- Training guard: const float* loss_scalar (reads f32 directly)
- Guard check_and_accumulate: takes u64 raw ptrs (type-agnostic)
- All callers pass .raw_ptr() — works for both fused (f32) and non-fused (bf16) paths
- Readback: reads 4 bytes f32 for loss (was 2 bytes bf16)
NaN persists: the per-sample loss computation in the loss kernel produces NaN
for specific samples despite float arithmetic and ±500 activation clamping.
The NaN is within the softmax/expected-Q/TD-error chain, not from the
accumulator. Next step: add in-kernel NaN detection to pinpoint the exact
computation step.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
5328f0e33b