From 5d63762ab3c9d9bbbe52aff3e700d13de710ff1c Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Thu, 7 May 2026 13:18:14 +0200 Subject: [PATCH] =?UTF-8?q?fix(sp15-wave4.1b-OOB-followup):=20pre-load=20b?= =?UTF-8?q?n=5Ftanh=5Fconcat=5Fdd=5Fkernel=20=E2=80=94=20fixes=20forward-c?= =?UTF-8?q?apture=20SEGV?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 98 ++++++++++++------- crates/ml/tests/sp15_phase1_oracle_tests.rs | 24 ++++- docs/dqn-wire-up-audit.md | 2 + 3 files changed, 87 insertions(+), 37 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 7e25f26c4..cbefa5100 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -76,7 +76,7 @@ use super::sp4_isv_slots::{ }; // ── Precompiled cubins (build.rs → include_bytes! → ZERO runtime nvcc) ────── -pub(crate) static DQN_UTILITY_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/dqn_utility_kernels.cubin")); +pub static DQN_UTILITY_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/dqn_utility_kernels.cubin")); static EMA_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/ema_kernel.cubin")); static RELU_MASK_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/relu_mask_kernel.cubin")); static C51_LOSS_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/c51_loss_kernel.cubin")); @@ -1070,8 +1070,26 @@ pub fn launch_sp15_dd_state_reduce( /// buffer and computing the dim. /// /// Grid: ceil(B * concat_dd_dim / 256), Block: 256. +/// +/// SP15 Wave 4.1b OOB fix (2026-05-07): the kernel handle is now passed in +/// by the caller as a pre-loaded `&CudaFunction`, NOT resolved via +/// `load_cubin` + `load_function` on every call. The on-demand resolution +/// 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 them from `submit_forward_ops_main` +/// (which runs inside the `forward` child capture) caused a SEGV at exit +/// 139 — the captured-step path entered the loader while the stream was +/// already in capture mode, the cudarc wrapper raced with the driver's +/// internal capture tracking, and the host-side state corruption surfaced +/// as a segfault before `CAPTURE_PHASE_FORWARD_DONE` could print. Pre- +/// loading the handle in `compile_training_kernels` (mirroring the legacy +/// `bn_tanh_concat_kernel` field that Wave 4.1b removed) restores the +/// graph-safe contract: each child capture sees only kernel launches, not +/// driver-state mutations. See `pearl_no_host_branches_in_captured_graph`. pub fn launch_sp15_bn_concat_dd( stream: &Arc, + kernel: &CudaFunction, bn_hidden: cudarc::driver::sys::CUdeviceptr, states: cudarc::driver::sys::CUdeviceptr, isv: cudarc::driver::sys::CUdeviceptr, @@ -1082,22 +1100,11 @@ pub fn launch_sp15_bn_concat_dd( state_dim_padded: i32, concat_dd_dim: i32, ) -> Result<(), MLError> { - let module = stream - .context() - .load_cubin(DQN_UTILITY_CUBIN.to_vec()) - .map_err(|e| MLError::ModelError(format!( - "load dqn_utility_kernels cubin for bn_tanh_concat_dd: {e}" - )))?; - let kernel = module - .load_function("bn_tanh_concat_dd_kernel") - .map_err(|e| MLError::ModelError(format!( - "load bn_tanh_concat_dd_kernel function: {e}" - )))?; let total_elems = (batch_size as i64).saturating_mul(concat_dd_dim as i64); let blocks = (((total_elems as u64) + 255) / 256).max(1) as u32; unsafe { stream - .launch_builder(&kernel) + .launch_builder(kernel) .arg(&bn_hidden) .arg(&states) .arg(&isv) @@ -5633,17 +5640,30 @@ pub struct GpuDqnTrainer { on_next_bn_hidden_buf: CudaSlice, /// DDQN argmax-pass mirror of `bn_concat_buf` (online-on-next_states). on_next_bn_concat_buf: CudaSlice, - // SP15 Wave 4.1b removed the legacy `bn_tanh_concat_kernel` field — - // production callers (3 forward sites here + 1 in `gpu_experience_collector`) - // migrated to `launch_sp15_bn_concat_dd` (the bottleneck-aware - // dd_pct-appending launcher), which loads the - // `bn_tanh_concat_dd_kernel` symbol on-demand from the shared - // `dqn_utility_kernels.cubin`. The legacy `bn_tanh_concat_kernel` - // and `bn_tanh_concat_f32_kernel` symbols stay in the cubin source - // for the SP15 oracle tests in - // `crates/ml/tests/sp15_phase1_oracle_tests.rs` (parity check - // between old and new kernels), but no production code path resolves - // them anymore. + /// SP15 Wave 4.1b OOB fix (2026-05-07): pre-loaded handle for the + /// bottleneck-aware `bn_tanh_concat_dd_kernel` (replaces the legacy + /// `bn_tanh_concat_kernel` + `bn_tanh_concat_f32_kernel` fields with a + /// single dd_pct-appending kernel). Wave 4.1b originally migrated the 3 + /// production trainer call sites (online forward, target forward, DDQN + /// argmax) 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: those CUDA driver + /// API calls allocate memory and mutate driver state, which is NOT + /// capturable inside `cuStreamBeginCapture`. 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 and the host-side corruption surfaced as a segfault before + /// `CAPTURE_PHASE_FORWARD_DONE` could print. The experience collector's + /// equivalent caller (line 4127 of gpu_experience_collector.rs) was + /// already pre-loading correctly via `bn_tanh_concat_fn` field — this + /// field restores the same graph-safe pattern on the trainer side. All + /// 3 trainer call sites consume this single handle (Hopper's + /// "no-CUfunction-sharing-across-children" comment elsewhere refers to + /// `dqn_utility_kernels` workhorse kernels like `dqn_saxpy_f32_kernel` + /// that are invoked from MULTIPLE child captures — this kernel only + /// runs in the forward + ddqn children, mirroring `bn_tanh_backward_kernel`'s + /// single-handle pattern for the forward+backward chain). + bn_tanh_concat_dd_kernel: CudaFunction, /// #31 Backward scratch: d_loss/d_bn_concat [B, concat_dim] f32. bn_d_concat_buf: CudaSlice, /// #31 Backward scratch: d_loss/d_bn [B, bn_dim] f32 (after tanh derivative). @@ -17743,7 +17763,7 @@ impl GpuDqnTrainer { // per array. Stack is set once in DQNTrainer::new() (64KB for all kernels). // ── Compile utility kernels (grad_norm, adam_update, etc.) ─ - let (grad_norm_kernel, grad_norm_finalize_kernel, grad_norm_finalize_adam, grad_norm_finalize_post_aux, adam_update_kernel, adam_update_post_aux, scale_f32_post_aux, saxpy_kernel, zero_kernel, regime_scale_kernel, shrink_perturb, _relu_mask_in_module, spectral_norm_kernel, clip_grad_kernel, pad_states_kernel, saxpy_f32_kernel, scale_f32_kernel, stochastic_depth_kernel, stochastic_depth_rng_kernel, bn_tanh_backward_kernel, bn_bias_grad_kernel, vaccine_dot_kernel, vaccine_project_kernel, vaccine_dot_finalize, causal_intervene_kernel_fn, causal_reduce_kernel_fn, causal_mean_reduce_kernel_fn, pruning_mask_kernel, pruning_compute_kernel, her_inplace_kernel, nan_check_f32_kernel, nan_check_f32_kernel_b, nan_check_fused_f32_kernel, clamp_finite_f32_kernel, popart_normalize_kernel, saxpy_f32_adam_grad, saxpy_f32_aux, scale_f32_aux, distill_saxpy_aux, grad_norm_standalone_post_aux, shrink_perturb_ungraphed, scale_f32_ungraphed, popart_robust_kernel) = + let (grad_norm_kernel, grad_norm_finalize_kernel, grad_norm_finalize_adam, grad_norm_finalize_post_aux, adam_update_kernel, adam_update_post_aux, scale_f32_post_aux, saxpy_kernel, zero_kernel, regime_scale_kernel, shrink_perturb, _relu_mask_in_module, spectral_norm_kernel, clip_grad_kernel, pad_states_kernel, saxpy_f32_kernel, scale_f32_kernel, stochastic_depth_kernel, stochastic_depth_rng_kernel, bn_tanh_backward_kernel, bn_bias_grad_kernel, vaccine_dot_kernel, vaccine_project_kernel, vaccine_dot_finalize, causal_intervene_kernel_fn, causal_reduce_kernel_fn, causal_mean_reduce_kernel_fn, pruning_mask_kernel, pruning_compute_kernel, her_inplace_kernel, nan_check_f32_kernel, nan_check_f32_kernel_b, nan_check_fused_f32_kernel, clamp_finite_f32_kernel, popart_normalize_kernel, saxpy_f32_adam_grad, saxpy_f32_aux, scale_f32_aux, distill_saxpy_aux, grad_norm_standalone_post_aux, shrink_perturb_ungraphed, scale_f32_ungraphed, popart_robust_kernel, bn_tanh_concat_dd_kernel) = compile_training_kernels(&stream, &config)?; // Separate grad_norm instance for non-graph launches (outside CUDA graph). @@ -22584,8 +22604,9 @@ impl GpuDqnTrainer { on_next_bn_concat_buf, bn_d_concat_buf, bn_d_hidden_buf, - // bn_tanh_concat_kernel field removed in SP15 Wave 4.1b — see - // the field-removal comment block in the struct definition. + // SP15 Wave 4.1b OOB fix (2026-05-07): pre-loaded handle, see + // the field docstring in the struct definition. + bn_tanh_concat_dd_kernel, bn_tanh_backward_kernel, bn_bias_grad_kernel, // ── Plan 4 Task 1B-iii: VSN feature-selection buffers ────── @@ -27417,6 +27438,7 @@ impl GpuDqnTrainer { // so by the time this kernel reads slot 406 the value is fresh. launch_sp15_bn_concat_dd( &self.stream, + &self.bn_tanh_concat_dd_kernel, bn_hidden_ptr, on_next_states_for_bn, self.isv_signals_dev_ptr, @@ -27754,6 +27776,7 @@ impl GpuDqnTrainer { launch_sp15_bn_concat_dd( &self.stream, + &self.bn_tanh_concat_dd_kernel, bn_hidden_ptr, states_for_bn, self.isv_signals_dev_ptr, @@ -27948,6 +27971,7 @@ impl GpuDqnTrainer { let b_i32 = b as i32; launch_sp15_bn_concat_dd( &self.stream, + &self.bn_tanh_concat_dd_kernel, tg_bn_hidden_ptr, tg_states_for_bn, self.isv_signals_dev_ptr, @@ -31678,7 +31702,7 @@ fn deflate_rank_one(mat: &[f32], n: usize, lambda_1: f32) -> Vec { fn compile_training_kernels( stream: &Arc, config: &GpuDqnTrainConfig, -) -> Result<(CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction), MLError> { +) -> Result<(CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction), MLError> { info!( state_dim = ml_core::state_layout::STATE_DIM, total_params = compute_total_params(config), @@ -31802,11 +31826,15 @@ fn compile_training_kernels( .map_err(|e| MLError::ModelError(format!("bn_tanh_backward_kernel load: {e}")))?; let bn_bias_grad = module.load_function("bn_bias_grad_kernel") .map_err(|e| MLError::ModelError(format!("bn_bias_grad_kernel load: {e}")))?; - // SP15 Wave 4.1b: legacy `bn_tanh_concat_kernel` no longer loaded — - // production callers migrated to `launch_sp15_bn_concat_dd`, which - // resolves the `bn_tanh_concat_dd_kernel` symbol from the same cubin. - // The legacy symbol remains in the source for the SP15 oracle parity - // tests but is not surfaced as a `CudaFunction` field on the trainer. + // SP15 Wave 4.1b OOB fix (2026-05-07): pre-load the dd_pct-appending + // bottleneck concat kernel from the same `module` as `bn_tanh_bw` / + // `bn_bias_grad` (forward-child captured replay group). The 3 trainer + // production callers (online forward, target forward, DDQN argmax) + // consume this handle via `&self.bn_tanh_concat_dd_kernel`; resolving + // the symbol on-demand inside the captured forward child caused a + // SEGV (cf. the field-level docstring in the struct definition). + let bn_tanh_concat_dd = module.load_function("bn_tanh_concat_dd_kernel") + .map_err(|e| MLError::ModelError(format!("bn_tanh_concat_dd_kernel load: {e}")))?; let her_inplace = module.load_function("her_inplace_relabel") .map_err(|e| MLError::ModelError(format!("her_inplace_relabel load: {e}")))?; @@ -31835,8 +31863,8 @@ fn compile_training_kernels( let popart_robust = ungraphed_module.load_function("popart_normalize_robust") .map_err(|e| MLError::ModelError(format!("popart_normalize_robust load: {e}")))?; - info!("GpuDqnTrainer: 39 utility kernels loaded from precompiled cubin (5 CUmodules)"); - Ok((grad_norm, grad_norm_finalize, grad_norm_finalize_adam, grad_norm_finalize_post_aux, adam_update, adam_update_post_aux, scale_f32_post_aux, saxpy, zero, regime_scale, shrink_perturb, _relu_mask_from_module, spectral_norm, clip_grad, pad_states, saxpy_f32, scale_f32, stochastic_depth, stochastic_depth_rng, bn_tanh_bw, bn_bias_grad, vaccine_dot, vaccine_project, vaccine_dot_finalize, causal_intervene, causal_reduce, causal_mean_reduce, pruning_mask_fn, pruning_compute_fn, her_inplace, nan_check_f32, nan_check_f32_b, nan_check_fused_f32, clamp_finite_f32, popart_normalize, saxpy_f32_adam_grad, saxpy_f32_aux, scale_f32_aux, distill_saxpy_aux, grad_norm_standalone_post_aux, shrink_perturb_ungraphed, scale_f32_ungraphed, popart_robust)) + info!("GpuDqnTrainer: 40 utility kernels loaded from precompiled cubin (5 CUmodules)"); + Ok((grad_norm, grad_norm_finalize, grad_norm_finalize_adam, grad_norm_finalize_post_aux, adam_update, adam_update_post_aux, scale_f32_post_aux, saxpy, zero, regime_scale, shrink_perturb, _relu_mask_from_module, spectral_norm, clip_grad, pad_states, saxpy_f32, scale_f32, stochastic_depth, stochastic_depth_rng, bn_tanh_bw, bn_bias_grad, vaccine_dot, vaccine_project, vaccine_dot_finalize, causal_intervene, causal_reduce, causal_mean_reduce, pruning_mask_fn, pruning_compute_fn, her_inplace, nan_check_f32, nan_check_f32_b, nan_check_fused_f32, clamp_finite_f32, popart_normalize, saxpy_f32_adam_grad, saxpy_f32_aux, scale_f32_aux, distill_saxpy_aux, grad_norm_standalone_post_aux, shrink_perturb_ungraphed, scale_f32_ungraphed, popart_robust, bn_tanh_concat_dd)) } /// Load the standalone Polyak EMA kernel from precompiled cubin. diff --git a/crates/ml/tests/sp15_phase1_oracle_tests.rs b/crates/ml/tests/sp15_phase1_oracle_tests.rs index 2cf00f2d3..9fe0fe299 100644 --- a/crates/ml/tests/sp15_phase1_oracle_tests.rs +++ b/crates/ml/tests/sp15_phase1_oracle_tests.rs @@ -880,9 +880,16 @@ mod gpu { #[test] #[ignore = "requires GPU"] fn bn_tanh_concat_dd_kernel_writes_dd_pct_column() { - use ml::cuda_pipeline::gpu_dqn_trainer::launch_sp15_bn_concat_dd; + use ml::cuda_pipeline::gpu_dqn_trainer::{launch_sp15_bn_concat_dd, DQN_UTILITY_CUBIN}; let stream = make_test_stream(); + // SP15 Wave 4.1b OOB fix (2026-05-07): the launcher now consumes a + // pre-loaded `&CudaFunction` (graph-safe), so tests resolve the + // symbol once up-front rather than relying on the launcher itself. + let module = stream.context().load_cubin(DQN_UTILITY_CUBIN.to_vec()) + .expect("load dqn_utility_kernels cubin"); + let kernel = module.load_function("bn_tanh_concat_dd_kernel") + .expect("load bn_tanh_concat_dd_kernel symbol"); // Synthetic batch matches production layout: bn_dim=16, // market_dim=42, STATE_DIM=128 → state_dim_padded=128 (no @@ -940,6 +947,7 @@ mod gpu { launch_sp15_bn_concat_dd( &stream, + &kernel, bn_hidden_buf.dev_ptr, states_buf.dev_ptr, isv_buf.dev_ptr, @@ -3670,7 +3678,8 @@ mod sp15_wave_4_1c_behavioral { use cudarc::cublas::sys as cublas_sys; use cudarc::driver::{CudaContext, CudaStream}; - use ml::cuda_pipeline::gpu_dqn_trainer::launch_sp15_bn_concat_dd; + use cudarc::driver::CudaFunction; + use ml::cuda_pipeline::gpu_dqn_trainer::{launch_sp15_bn_concat_dd, DQN_UTILITY_CUBIN}; use ml::cuda_pipeline::mapped_pinned::MappedF32Buffer; use ml::cuda_pipeline::shared_cublas_handle::PerStreamCublasHandles; use ml::cuda_pipeline::sp15_isv_slots::{DD_PCT_INDEX, SP15_SLOT_END}; @@ -3729,6 +3738,7 @@ mod sp15_wave_4_1c_behavioral { fn synthetic_forward( stream: &Arc, cublas_handle: cublas_sys::cublasHandle_t, + bn_concat_kernel: &CudaFunction, bn_hidden_dev: u64, states_dev: u64, isv_buf: &MappedF32Buffer, @@ -3764,6 +3774,7 @@ mod sp15_wave_4_1c_behavioral { // write returns). launch_sp15_bn_concat_dd( stream, + bn_concat_kernel, bn_hidden_dev, states_dev, isv_buf.dev_ptr, @@ -3940,10 +3951,18 @@ mod sp15_wave_4_1c_behavioral { .classic_handle_for(&stream) .expect("classic_handle_for default stream"); + // ── Pre-load the bn_tanh_concat_dd_kernel handle (graph-safe pattern; + // SP15 Wave 4.1b OOB fix 2026-05-07: launcher consumes &CudaFunction). ── + let util_module = stream.context().load_cubin(DQN_UTILITY_CUBIN.to_vec()) + .expect("load dqn_utility_kernels cubin"); + let bn_concat_kernel = util_module.load_function("bn_tanh_concat_dd_kernel") + .expect("load bn_tanh_concat_dd_kernel symbol"); + // ── At-ATH pass (dd_pct = 0.0) ── let logits_ath = synthetic_forward( &stream, cublas_handle, + &bn_concat_kernel, bn_hidden_buf.dev_ptr, states_buf.dev_ptr, &isv_buf, @@ -3962,6 +3981,7 @@ mod sp15_wave_4_1c_behavioral { let logits_dd = synthetic_forward( &stream, cublas_handle, + &bn_concat_kernel, bn_hidden_buf.dev_ptr, states_buf.dev_ptr, &isv_buf, diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index 9c52e51bd..7a8571c28 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -2,6 +2,8 @@ **Status:** Populated during Plan 1 Task 6 (A.5 orphan audit). Updated on every commit per Invariant 7. +SP15 Wave 4.1b OOB-followup — pre-load bn_tanh_concat_dd_kernel (2026-05-07): root-caused the L40S smoke `train-vg5f7` (commit `bfc3ffa9d`) capture-mode SEGV to a hot-path `cuModuleLoadData` + `cuModuleGetFunction` inside `submit_forward_ops_main`'s captured `forward` child. Wave 4.1b (eb9515e41) deleted the legacy `bn_tanh_concat_kernel: CudaFunction` field from `GpuDqnTrainer` and migrated the 3 production trainer call sites (online forward at `launch_cublas_forward` ~27755, target forward at ~27949, DDQN argmax at `submit_forward_ops_ddqn` ~27418) to a `launch_sp15_bn_concat_dd` free function that resolved `bn_tanh_concat_dd_kernel` via `stream.context().load_cubin(...)` + `module.load_function(...)` ON EVERY CALL. The collector's matching call site (gpu_experience_collector.rs:4127) already pre-loaded into `bn_tanh_concat_fn` field 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 three captured-graph call sites. Inside a `cuStreamBeginCapture` region, those CUDA driver API calls are NOT capturable (they allocate memory and mutate driver state), and the resulting host-side state corruption surfaced as a SEGV at exit 139 before `CAPTURE_PHASE_FORWARD_DONE` could print. Per `pearl_no_host_branches_in_captured_graph`. Fix: re-add `bn_tanh_concat_dd_kernel: CudaFunction` field on `GpuDqnTrainer`, pre-load it in `compile_training_kernels` from the same `module` as `bn_tanh_bw` / `bn_bias_grad` (forward-child captured replay group, same `DQN_UTILITY_CUBIN`), wire through ctor + the 43->44 utility-kernel return tuple, and change `launch_sp15_bn_concat_dd` to take `&CudaFunction` as a parameter (no more per-call resolution). All 3 trainer call sites now pass `&self.bn_tanh_concat_dd_kernel`. Promoted `DQN_UTILITY_CUBIN` from `pub(crate)` -> `pub` so the SP15 oracle tests + Wave 4.1c behavioral KL test (which used the launcher's old free-function pattern) can pre-load the symbol once in their setup blocks. Verification: `cargo check -p ml --tests` clean (dev + release); `cargo test -p ml --test sp15_phase1_oracle_tests -- --ignored` 36/36 pass including the dd_pct trunk-input KL test that exercises two end-to-end forwards through the modified launcher. + SP15 Wave 5 SEGV diagnostic — checkpoints inside capture_training_graph (2026-05-07): smoke `train-fp7xx` (commit `a8d6c3304`) printed all 16 Phase-6/7/8 checkpoints cleanly through `STEP0_PHASE_PER_PRIORITY_DONE`. The `CAPTURE_DONE` checkpoint never printed, exit changed from 139 (SEGV) to 143 (SIGTERM) — process was killed externally ~40s after PER_PRIORITY_DONE, suggesting `capture_training_graph()` either hangs or runs too long. Added 12 sub-checkpoints across capture_training_graph: `CAPTURE_PHASE_BEGIN`, `CAPTURE_PHASE_PER_SAMPLE_DONE`, `CAPTURE_PHASE_COUNTERS_DONE`, `CAPTURE_PHASE_SPECTRAL_DONE`, `CAPTURE_PHASE_FORWARD_DONE`, `CAPTURE_PHASE_DDQN_DONE`, `CAPTURE_PHASE_AUX_DONE`, `CAPTURE_PHASE_POST_AUX_DONE`, `CAPTURE_PHASE_ADAM_GRAD_DONE`, `CAPTURE_PHASE_ADAM_UPDATE_DONE`, `CAPTURE_PHASE_MAINTENANCE_DONE`, `CAPTURE_PHASE_IQL_MODULATE_DONE`, `CAPTURE_PHASE_PER_PRIORITY_DONE`, `CAPTURE_PHASE_CHILDREN_STORED`, `CAPTURE_PHASE_PARENT_COMPOSED`. Will pinpoint which child capture or compose-step hangs/crashes. Diagnostic-only — no behavior change. SP15 Wave 5 SEGV diagnostic — extend checkpoints into Phase 6/7/8 (2026-05-07): smoke `train-xq9hg` (commit `e9d9afbd6`) got past all 13 original checkpoints (`BEGIN` → `NAN_CHECKS_DONE`) cleanly, confirming the SEGV is downstream in Phase 6 (TLOB bwd + Adam, Mamba2 bwd + Adam, OFI embed bwd + Adam, pruning_mask, branch_grad_balance, grad_norm, submit_adam_ops, q_mag/dir_bin_means_reduce, update_isv_signals), Phase 7 (submit_maintenance_ops), or Phase 8 (submit_iql_modulate_ops, submit_per_priority_ops, capture_training_graph). 14 more `STEP0_PHASE_*` checkpoints added across these phases — the next L40S smoke will localize the SEGV to a single operation. Same diagnostic-only contract; will be reverted with the root-cause fix.