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

fix: wire atom_stats in replay_forward_ungraphed — second NULL call site

Browse Source 
The first fix only patched populate_q_out(). The graph-captured path
uses replay_forward_ungraphed() which had its own null_atom_stats=0.
This is the call site that actually runs during training.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-19 16:32:55 +02:00
parent 64d0eb8bfe
commit f359905fe6
1 changed files with 8 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
View File

@@ -8071,9 +8071,13 @@ impl GpuDqnTrainer {
let block_dim = 256_u32;
let grid_dim = ((batch_size as u32 + block_dim - 1) / block_dim).max(1);
let q_out_ptr = self.q_out_buf.raw_ptr();
let null_atom_stats = 0u64;
let null_q_var = 0u64;
let atom_stats_ptr = self.atom_stats_buf.raw_ptr();
let q_var_ptr = self.q_var_buf_trainer.raw_ptr();
let atom_positions_buf_ptr = self.atom_positions_buf.raw_ptr();
// Zero atom_stats before accumulation (kernel uses atomicAdd)
unsafe {
cudarc::driver::sys::cuMemsetD32Async(atom_stats_ptr, 0, 2, self.stream.cu_stream());
}
unsafe {
self.stream
.launch_builder(&self.expected_q_kernel)
@@ -8087,8 +8091,8 @@ impl GpuDqnTrainer {
.arg(&b2)
.arg(&b3)
.arg(&self.per_sample_support_ptr)
.arg(&null_atom_stats)
.arg(&null_q_var)
.arg(&atom_stats_ptr)
.arg(&q_var_ptr)
.arg(&atom_positions_buf_ptr)
.launch(LaunchConfig {
grid_dim: (grid_dim, 1, 1),