From 049d07a441003999f0d9dd63fec6bdfd3ceb3523 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Mon, 6 Apr 2026 13:26:30 +0200 Subject: [PATCH] fix: use self.stream for cast kernel module load, ext_stream for launches MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit get_cast_kernels uses OnceLock + cuModuleLoadData which requires the CUDA context to be bound to the current thread. When ext_stream (the trainer's stream) was passed, its context wasn't bound — causing a hang in cuModuleLoadData on H100. Fix: always use self.stream (replay buffer's stream, with bound context) for kernel compilation. Only use ext_stream for the actual kernel launches. Co-Authored-By: Claude Opus 4.6 (1M context) --- crates/ml-dqn/src/gpu_replay_buffer.rs | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/crates/ml-dqn/src/gpu_replay_buffer.rs b/crates/ml-dqn/src/gpu_replay_buffer.rs index 852e7ce8c..607caf604 100644 --- a/crates/ml-dqn/src/gpu_replay_buffer.rs +++ b/crates/ml-dqn/src/gpu_replay_buffer.rs @@ -753,9 +753,11 @@ impl GpuReplayBuffer { let stream_owned = ext_stream.cloned().unwrap_or_else(|| Arc::clone(&self.stream)); let stream = &stream_owned; // Convert bf16 td_errors to f32 using pre-allocated scratch buffer - eprintln!("PER_DIAG: getting cast kernels"); - let kernels = get_cast_kernels(stream)?; - eprintln!("PER_DIAG: launching bf16_to_f32 cast"); + eprintln!("PER_DIAG: getting cast kernels (using self.stream for module load)"); + // Use self.stream for kernel compilation (needs bound context). + // ext_stream is only used for kernel launches. + let kernels = get_cast_kernels(&self.stream)?; + eprintln!("PER_DIAG: launching bf16_to_f32 cast on ext_stream"); let ni = bs as i32; // SAFETY: update_td_f32, td_errors are valid device allocations of at least bs elements. unsafe {