fix: restore c51_grad 19th param + HER (32,1,1) — match working SHA graph_mega

graph_mega on Hopper requires EXACT node structure match with working
SHA 11b1a1ca9. Two remaining differences caused the hang:
- c51_grad: 18→19 params (restored unused q_mean_ema_ptr)
- HER relabel: (1,1,1)→(32,1,1) (restored warp-parallel launch)

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-17 14:58:58 +02:00
parent 709fffe923
commit 924aaa775e
2 changed files with 4 additions and 2 deletions

View File

@@ -26,7 +26,8 @@ extern "C" __global__ void c51_grad_kernel(
const float* __restrict__ branch_scales, /* [B, 4] per-sample per-branch gradient scale */
const float* __restrict__ per_sample_support, /* [B, 3] per-sample [v_min, v_max, delta_z] */
const float* __restrict__ liquid_mod, /* [4] pinned device-mapped per-branch modulators */
const float* __restrict__ atom_positions) /* [4, num_atoms] adaptive positions. NULL = linear. */
const float* __restrict__ atom_positions, /* [4, num_atoms] adaptive positions. NULL = linear. */
const float* __restrict__ q_mean_ema_ptr) /* [1] pinned — unused but preserves 19-param graph node on Hopper */
{
int tid = blockIdx.x * blockDim.x + threadIdx.x;
int total_elems = batch_size * num_atoms;

View File

@@ -3415,7 +3415,7 @@ impl GpuDqnTrainer {
let launch_cfg = cudarc::driver::LaunchConfig {
grid_dim: (her_batch_size as u32, 1, 1),
block_dim: (1, 1, 1),
block_dim: (32, 1, 1), // must match working SHA — graph_mega node structure on Hopper
shared_mem_bytes: 0,
};
@@ -8690,6 +8690,7 @@ impl GpuDqnTrainer {
.arg(&self.per_sample_support_ptr)
.arg(&self.liquid_mod_buf.raw_ptr())
.arg(&self.atom_positions_buf)
.arg(&self.q_mean_ema_dev_ptr) // unused but preserves 19-param graph node structure on Hopper
.launch(LaunchConfig {
grid_dim: (blocks, 1, 1),
block_dim: (256, 1, 1),