gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix: use self.stream for cast kernel module load, ext_stream for launches

Browse Source 
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) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-06 13:26:30 +02:00
parent 5770ac5f62
commit 049d07a441
1 changed files with 5 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
crates/ml-dqn/src/gpu_replay_buffer.rs
View File

@@ -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 {