feat(dqn): SP2 — replace 8 check_nan_f32 calls with fused launch

run_nan_checks_post_backward now fires a single fused kernel launch
covering slots 24-35 in 12 blocks. Reduces graph-capture overhead from
8 launches to 1 per step. Slots 33-35 inline checks at backward
orchestration sites (gpu_dqn_trainer.rs:18580/:18601/:6939) unchanged
— multi-point localization preserved.

Method signature simplified — IQN pointers no longer per-step args.
populate_nan_check_meta (called once at construction in fused_training.rs)
baked Option<u64> nullity into the metadata buffer entries; the fused
kernel's null-pointer guard handles slots 27/28 when IQN inactive
without per-step branching at the Rust caller.

Both call sites in fused_training.rs (ungraphed + graph-captured paths)
updated together per feedback_no_partial_refactor.

This is the F0 regression fix — Phase B's per-step kernel-launch
overhead was the F0 cause across 3 SP1 smokes (F0 ~ 35 vs baseline
55.87). Gate 1 smoke (Task A6) validates F0 >= 53.08.
This commit is contained in:
jgrusewski
2026-04-30 09:22:52 +02:00
parent fbf48df9de
commit b5a064f6ca
3 changed files with 37 additions and 126 deletions

View File

@@ -2270,3 +2270,5 @@ SP2 Phase A1 — fused NaN-check kernel (2026-04-29): foundational kernel-only c
SP2 Phase A2 — nan-check (ptr, len) device tables (2026-04-29): added two device buffers on `GpuDqnTrainer` to feed the fused NaN-check kernel from A1: `nan_check_buf_ptrs: CudaSlice<u64>` (12 entries, one per slot 24-35) and `nan_check_buf_lens: CudaSlice<i32>` (12 entries). Both are `stream.alloc_zeros` allocations adjacent to `nan_flags_buf`. Two separate buffers chosen over a packed-stride layout for direct ABI match with the kernel signature `(const float* const* buf_ptrs, const int* buf_lens, ...)` — no struct-stride alignment concerns. Sized at 12 (slots 24-35); slots 36-47 headroom is reserved for SP3 Mech 5 extension which will resize both buffers in lockstep. Slot 31 (deferred ensemble, cross-struct on `FusedDqnTraining` per A1's null-skip pattern) entry will hold 0 — kernel skips that block on `buf == nullptr`. Allocated as zeros; populated once in A3 (slot pointers are stable across the trainer's lifetime — no per-step host→device traffic, satisfying `feedback_no_htod_htoh_only_mapped_pinned` once A3 selects the population path: mapped-pinned helper from `mapped_pinned.rs` for the one-shot construction-time write). Unused yet — A3 wires the population + Rust launch wrapper, A4 replaces the 8 individual `check_nan_f32` calls in `run_nan_checks_post_backward` with a single fused launch. Field-level only — no behavioral change.
SP2 Phase A3 — populate + launch wrapper (2026-04-29): refactored A2's `nan_check_buf_ptrs: CudaSlice<u64>` / `nan_check_buf_lens: CudaSlice<i32>` to mapped-pinned (`MappedU64Buffer` / `MappedI32Buffer` from `mapped_pinned.rs`) per `feedback_no_htod_htoh_only_mapped_pinned` — host-side write through the mapped pages is visible to the kernel via the `dev_ptr` returned by `cuMemHostGetDevicePointer_v2`, eliminating the HtoD copy entirely. Added `populate_nan_check_meta(b, sh2, iqn_d_h_s2_ptr, iqn_d_branch_logits_buf_ptr, iqn_q)` on `GpuDqnTrainer` — one-shot construction-time writer of 12 `(ptr, len)` tuples covering slots 24-35 (mirrors the per-slot accessor table at the end of `docs/dqn-backward-nan-audit.md` and the per-slot launches in `run_nan_checks_post_backward`). Slot 31 deferred ensemble entry: `(0, 0)`. Slots 27/28 IQN entries: `Option<u64>` ptrs gated on `gpu_iqn.is_some()` — when IQN is inactive the entries stay zero and the fused kernel skips them via its null-pointer guard. Slots 33-35 (`bw_d_h_s2` inline checks) hold null entries — fused kernel skips; the inline `check_nan_f32` calls at the three backward orchestration phases (`launch_cublas_backward_to` post-main / post-aux / `apply_iqn_trunk_gradient` post-iqn) continue to fire individually for entry-point localization. Added `launch_nan_check_fused_f32` Rust wrapper — single launch with `grid_dim=12, block_dim=256, BASE_FLAG_IDX=24`, mirrors the kernel signature `(buf_ptrs_dev, buf_lens_dev, base_flag_idx, nan_flags_ptr)`. Kernel registered in the precompiled-cubin loader (`compile_training_kernels` tuple 43→44, 38→39 utility kernels logged) and stored on `GpuDqnTrainer` as `nan_check_fused_f32_kernel: CudaFunction` — same `module` (forward_child / aux_child / post_aux_child captured replay group) as the per-buffer `dqn_nan_check_f32` to preserve graph-capture compatibility for A4's call-site replacement. Constructor-time wire-up lives in `FusedTrainingCtx::new` (after `gpu_iqn` is constructed) — chosen over the GpuDqnTrainer constructor because `gpu_iqn` is owned by `FusedTrainingCtx`, mirroring the same `Option<u64>` argument pattern used by `apply_iqn_trunk_gradient(iqn_d_h_s2_ptr, ...)` and `run_nan_checks_post_backward`. Wrapper unused at A3 commit time — A4 replaces the 8-launch sequence in `run_nan_checks_post_backward` with the single fused launch. Zero new HtoD copies; pre-commit Invariant 7 satisfied; no ISV slots added; sticky-flag semantics preserved at the kernel.
SP2 Phase A4 — call-site replacement (2026-04-29): replaced the 8 per-step `check_nan_f32` launches inside `GpuDqnTrainer::run_nan_checks_post_backward` with a single delegated call to `launch_nan_check_fused_f32` (12 blocks × 256 threads, one block per slot 24-35). Method signature simplified to `run_nan_checks_post_backward(&mut self) -> Result<(), MLError>` — IQN pointers (`iqn_d_h_s2_ptr`, `iqn_d_branch_logits_buf_ptr`) and `batch_size` are no longer per-step args; their values were already baked into the mapped-pinned `(ptr, len)` metadata buffer at construction time via `populate_nan_check_meta` (A3). Both call sites in `fused_training.rs` (ungraphed step at line 1556 + graph-captured `post_aux` closure at line 2374) updated atomically per `feedback_no_partial_refactor` — old `iqn_d_h_s2_ptr` / `iqn_d_branch_logits_buf_ptr` derivation blocks removed (no leftover dead code). Slots 33-35 (`bw_d_h_s2` multi-point snapshots) inline checks at backward orchestration sites (`launch_cublas_backward_to` post-main + post-aux, `apply_iqn_trunk_gradient` post-iqn) unchanged — fused kernel's null-pointer guard skips them in the metadata table; the per-phase inline calls continue to provide entry-point localization across the 3 backward phases. Sticky-flag semantics preserved (kernel only writes 1, never clears); flags accumulate within fold and reset on fold boundary via `reset_nan_flags()`. Behavioral change at the diagnostic level: graph-capture overhead reduced from 8 launches per step to 1 (per `feedback_no_atomicadd` and the kernel's grid-strided block-local `__syncthreads_or` reduce). This is the F0 regression remediation — Phase B's 8 per-step launches were the F0 cause across 3 SP1 smokes (F0 ≈ 35 vs baseline 55.87); Gate 1 smoke (Task A6) validates F0 ≥ 53.08. Zero new HtoD/DtoD/HtoH copies; pre-commit Invariant 7 satisfied; no ISV slots added.