Files
foxhunt/crates/ml/tests
jgrusewski 5d63762ab3 fix(sp15-wave4.1b-OOB-followup): pre-load bn_tanh_concat_dd_kernel — fixes forward-capture SEGV
Root cause: SP15 Wave 4.1b (eb9515e41) migrated the trainer's 3 production
bottleneck-concat call sites (online forward, target forward, DDQN argmax)
from the pre-loaded `bn_tanh_concat_kernel: CudaFunction` field to a
`launch_sp15_bn_concat_dd` free function that resolved the symbol via
`load_cubin` + `load_function` ON EVERY CALL. That pattern is incompatible
with CUDA Graph capture: `cuModuleLoadData` and `cuModuleGetFunction` are
HOST-side driver API calls that allocate memory and mutate driver state,
and they are NOT capturable inside a `cuStreamBeginCapture` region. Calling
the launcher from `submit_forward_ops_main` (graph-captured forward child)
caused a SEGV at exit 139 — the loader raced with the capture-mode driver
state, the host-side corruption surfaced as a segfault before
`CAPTURE_PHASE_FORWARD_DONE` could print.

Diagnostic evidence (L40S smoke `train-vg5f7` on commit `bfc3ffa9d`):
ALL 16 step-0 ungraphed checkpoints printed clean. Capture begins:
  CAPTURE_PHASE_BEGIN
  CAPTURE_PHASE_PER_SAMPLE_DONE
  CAPTURE_PHASE_COUNTERS_DONE
  CAPTURE_PHASE_SPECTRAL_DONE
Then: SEGV. The next checkpoint that didn't print was
CAPTURE_PHASE_FORWARD_DONE -> SEGV is inside `submit_forward_ops_main`'s
`forward` child capture. The same function ran cleanly ungraphed in
step 0 because the loader-host-branch was harmless without active
capture.

The experience collector's equivalent caller (gpu_experience_collector.rs:4127)
was already doing this correctly — pre-loaded `bn_tanh_concat_fn` field
populated at construction. Wave 4.1b's mistake was asymmetric: it kept the
collector's pre-load pattern but introduced an on-demand loader for the
trainer's call sites.

Fix (atomic, restores graph-safe contract):
- Add `bn_tanh_concat_dd_kernel: CudaFunction` field on `GpuDqnTrainer`
  (back-fills what Wave 4.1b removed, with updated docstring naming the
  new dd_pct-aware kernel).
- Pre-load the symbol in `compile_training_kernels` from the same `module`
  as `bn_tanh_bw` / `bn_bias_grad` (forward-child captured replay group).
  Tuple grows 43 -> 44 CudaFunctions; struct ctor wires the field.
- Change `launch_sp15_bn_concat_dd` signature to take `&CudaFunction` as
  parameter; remove the per-call `load_cubin` + `load_function`. All 3
  trainer call sites pass `&self.bn_tanh_concat_dd_kernel`.
- Promote `DQN_UTILITY_CUBIN` from `pub(crate)` to `pub` so the oracle
  tests in `crates/ml/tests/sp15_phase1_oracle_tests.rs` can pre-load
  the kernel handle (the launcher no longer hides this for them).
- Update the 2 test call sites (kernel-level oracle + Wave 4.1c behavioral
  KL test) to load the kernel handle once up-front and pass it through.

Verification:
- `cargo check -p ml --tests` clean (release + dev profile).
- `cargo test -p ml --test sp15_phase1_oracle_tests -- --ignored`:
  36/36 pass, including the Wave 4.1c behavioral KL test that exercises
  two end-to-end forward passes through the modified launcher.

Refs: SP15 Wave 4.1b (eb9515e41), pearl_no_host_branches_in_captured_graph,
feedback_no_partial_refactor, feedback_wire_everything_up.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-05-07 13:18:14 +02:00
..