jgrusewski 907b5f3355 fix: cross-stream race in PER update + pre-allocate scratch buffers ...

The replay buffer used its own CUDA stream for seg_tree_update, but
the td_errors buffer was computed on the trainer's stream. Without
synchronization, the replay stream could read incomplete data — causing
a hang on H100 where streams execute truly concurrently.

Fix: update_priorities_gpu now accepts an optional ext_stream parameter.
The fused training path passes the trainer's main stream, ensuring all
GPU work (td_errors computation + seg_tree_update) runs on the same
stream with implicit ordering.

Also pre-allocate update_td_f32, update_batch_max, update_max_merge
scratch buffers at construction — eliminates 3 cuMemAlloc calls per
training step that could trigger implicit device synchronization.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

2026-04-06 12:36:15 +02:00

13 KiB

Raw Blame History

View Raw