diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 842e4dd84..9e14814d2 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -7386,6 +7386,22 @@ pub struct GpuDqnTrainer { #[allow(dead_code)] aux_trunk_b3_grad: CudaSlice, + /// SP14 Layer C Phase C.5a-fixup (2026-05-08): pre-activation gradient + /// scratch for aux trunk Layer-1 (`dh_aux1_pre [B_max, AUX_TRUNK_H1=256]`). + /// `aux_trunk_backward.launch` writes this scratch in its dh_pre kernel + /// (overwrites; not accumulated) and reads it back in the dW_reduce + /// kernel for `dW1 = x_in.T @ dh_aux1_pre` + `db1 = sum_b dh_aux1_pre`. + /// Required by `gpu_aux_trunk.rs:266-267` — C.5a missed it. DEAD until + /// C.5b atomically wires the aux trunk backward into the gradient chain. + #[allow(dead_code)] + dh_aux1_pre_scratch: CudaSlice, + /// SP14 Layer C Phase C.5a-fixup: pre-activation gradient scratch for + /// aux trunk Layer-2 (`dh_aux2_pre [B_max, AUX_TRUNK_H2=128]`). Same + /// role + lifecycle as `dh_aux1_pre_scratch`; consumed by Layer-2/3 + /// matmul reductions in `aux_trunk_backward.launch`. + #[allow(dead_code)] + dh_aux2_pre_scratch: CudaSlice, + /// `[num_blocks_total]` per-block partial sum-of-squares written by /// `dqn_grad_norm_kernel` across all 6 grad tensors (offsets matching /// `aux_trunk_grad_block_offsets`). Sized to the maximum total block @@ -10617,6 +10633,32 @@ impl GpuDqnTrainer { self.target_params_buf.raw_ptr() } + /// SP14 Layer C Phase C.5a-fixup (2026-05-08): raw device pointers to + /// the aux trunk's 6 parameter tensors (`w1, b1, w2, b2, w3, b3`). + /// + /// Returned in the same order as the Adam launcher's per-tensor tuple + /// list (`launch_aux_trunk_adam_update`). The collector calls the + /// matching `set_trainer_aux_trunk_param_ptrs` setter once after the + /// fused training context is created (cold-path, mirrors the existing + /// `set_trainer_params_ptr` zero-copy contract — pointers are stable + /// for the lifetime of the trainer's CudaSlice fields). + /// + /// C.5b consumes these ptrs to launch `aux_trunk_forward.launch(...)` + /// from the collector's per-rollout-step body, reading the trainer's + /// authoritative aux trunk weights without a DtoD copy. The accessor + /// itself is dead in C.5a-fixup — populated but unread until C.5b + /// flips the contract. + pub fn aux_trunk_param_ptrs(&self) -> (u64, u64, u64, u64, u64, u64) { + ( + self.aux_trunk_w1.raw_ptr(), + self.aux_trunk_b1.raw_ptr(), + self.aux_trunk_w2.raw_ptr(), + self.aux_trunk_b2.raw_ptr(), + self.aux_trunk_w3.raw_ptr(), + self.aux_trunk_b3.raw_ptr(), + ) + } + /// SP15 Phase 3.5.4.c (2026-05-07): return the /// `(w_b0out_dev_ptr, n_weights, fan_in)` triple for the /// directional advantage-head last-Linear weight tensor (param- @@ -22657,6 +22699,15 @@ impl GpuDqnTrainer { let aux_trunk_w3_grad = alloc_f32(&stream, aux_trunk_w3_count, "aux_trunk_w3_grad")?; let aux_trunk_b3_grad = alloc_f32(&stream, aux_trunk_out_dim, "aux_trunk_b3_grad")?; + // SP14 Layer C Phase C.5a-fixup (2026-05-08): pre-activation gradient + // scratch buffers required by `aux_trunk_backward.launch` + // (`gpu_aux_trunk.rs:266-267`). Sized to `B_max × layer_h` so the + // backward kernel can write `dh_pre` once and read it back for + // dW/db reductions without aliasing. C.5a missed these. DEAD until + // C.5b atomically wires the backward chain. + let dh_aux1_pre_scratch = alloc_f32(&stream, b * aux_trunk_h1, "dh_aux1_pre_scratch")?; + let dh_aux2_pre_scratch = alloc_f32(&stream, b * aux_trunk_h2, "dh_aux2_pre_scratch")?; + // Per-tensor block counts for the dqn_grad_norm_kernel launches // (block_dim=256). One partial slot per (block × tensor) pair — // we lay them out in a contiguous buffer with fixed offsets so @@ -23593,6 +23644,9 @@ impl GpuDqnTrainer { aux_trunk_b2_grad, aux_trunk_w3_grad, aux_trunk_b3_grad, + // SP14 Layer C Phase C.5a-fixup (2026-05-08): dh_pre scratch buffers (DEAD). + dh_aux1_pre_scratch, + dh_aux2_pre_scratch, aux_trunk_grad_norm_partials, aux_trunk_grad_block_offsets, aux_trunk_grad_norm_buf, diff --git a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs index 7de206c31..d37648aad 100644 --- a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs +++ b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs @@ -1015,6 +1015,45 @@ pub struct GpuExperienceCollector { #[allow(dead_code)] exp_h_aux2: cudarc::driver::CudaSlice, + /// SP14 Layer C Phase C.5a-fixup (2026-05-08): zero-copy device pointers + /// to the trainer's aux trunk parameter tensors. Mirrors the existing + /// `trainer_params_ptr` field's zero-copy contract: the trainer owns + /// the authoritative parameter buffers; the collector reads them + /// directly via these stable raw pointers (no DtoD copy). Set to 0 + /// initially; updated by `set_trainer_aux_trunk_param_ptrs` once the + /// fused training context is created. + /// + /// Order matches `GpuDqnTrainer::aux_trunk_param_ptrs()`: + /// w1 [SH1, 256] → b1 [256] → w2 [256, 128] → b2 [128] + /// → w3 [128, SH2] → b3 [SH2]. + /// + /// C.5b consumes these via `exp_aux_trunk_forward_ops.launch(...)` + /// inserted post-`forward_online_f32`. DEAD until C.5b. + #[allow(dead_code)] + aux_trunk_w1_ptr: u64, + #[allow(dead_code)] + aux_trunk_b1_ptr: u64, + #[allow(dead_code)] + aux_trunk_w2_ptr: u64, + #[allow(dead_code)] + aux_trunk_b2_ptr: u64, + #[allow(dead_code)] + aux_trunk_w3_ptr: u64, + #[allow(dead_code)] + aux_trunk_b3_ptr: u64, + + /// SP14 Layer C Phase C.5a-fixup (2026-05-08): collector-side aux trunk + /// forward orchestrator. Owns the pre-loaded `aux_trunk_forward` + /// CudaFunction handle on the collector's stream + /// (per `pearl_no_host_branches_in_captured_graph` — no `load_cubin` + /// inside per-step body, no cross-stream sync). Mirrors the trainer's + /// `aux_trunk_forward_ops` field. C.5b's collector-side wire-up + /// invokes `exp_aux_trunk_forward_ops.launch(...)` post- + /// `forward_online_f32`, reading the trainer's parameter pointers via + /// `aux_trunk_{w,b}{1,2,3}_ptr`. DEAD until C.5b. + #[allow(dead_code)] + exp_aux_trunk_forward_ops: super::gpu_aux_trunk::AuxTrunkForwardOps, + /// SP14 β-migration step 3 (2026-05-07): per-rollout-step label /// producer kernel handle. Sibling of the trajectory-wide /// `aux_sign_label_kernel` (line 3776 above) — same per-thread O(1) @@ -2073,6 +2112,17 @@ impl GpuExperienceCollector { alloc_episodes * AUX_TRUNK_H2 )))?; + // SP14 Layer C Phase C.5a-fixup (2026-05-08): pre-load the + // collector-side `aux_trunk_forward` CudaFunction handle on the + // collector's stream. Mirrors the trainer-side + // `AuxTrunkForwardOps::new(&stream)` instantiation per + // `pearl_no_host_branches_in_captured_graph` — every per-step + // launcher uses cached handles instead of `load_cubin` inside the + // hot path. DEAD until C.5b atomically inserts + // `exp_aux_trunk_forward_ops.launch(...)` post-`forward_online_f32`. + let exp_aux_trunk_forward_ops = + super::gpu_aux_trunk::AuxTrunkForwardOps::new(&stream)?; + // SP14 β-migration step 3 (2026-05-07): per-rollout-step label // producer kernel. Loaded on the collector's stream so its launch // chains directly after `forward_online_f32` (which populates @@ -2528,6 +2578,17 @@ impl GpuExperienceCollector { exp_h_s2_aux, exp_h_aux1, exp_h_aux2, + // SP14 Layer C Phase C.5a-fixup (2026-05-08): zero-copy ptrs to + // trainer's aux trunk params + collector-side forward orchestrator. + // All DEAD until C.5b atomically inserts the forward launch + + // wires the setter call from training_loop.rs. + aux_trunk_w1_ptr: 0, + aux_trunk_b1_ptr: 0, + aux_trunk_w2_ptr: 0, + aux_trunk_b2_ptr: 0, + aux_trunk_w3_ptr: 0, + aux_trunk_b3_ptr: 0, + exp_aux_trunk_forward_ops, exp_aux_sign_label_per_step_kernel, // SP15 Wave 5 follow-up (2026-05-07): pre-loaded SP15 kernel handles. exp_sp15_dd_state_kernel, @@ -5788,6 +5849,43 @@ impl GpuExperienceCollector { info!(trainer_params_ptr = ptr, "Experience collector: zero-copy trainer params pointer set"); } + /// SP14 Layer C Phase C.5a-fixup (2026-05-08): set the raw device + /// pointers to the trainer's 6 aux trunk parameter tensors. + /// + /// Mirrors `set_trainer_params_ptr`'s zero-copy contract for the + /// SP14 aux trunk (3-layer Linear→ELU→Linear→ELU→Linear MLP allocated + /// in `GpuDqnTrainer::new` per Phase C.2). Pointers are stable for + /// the lifetime of the trainer's `aux_trunk_{w,b}{1,2,3}` `CudaSlice` + /// fields; called once from `training_loop.rs` adjacent to the + /// existing `set_trainer_params_ptr` call site. + /// + /// Argument order matches `GpuDqnTrainer::aux_trunk_param_ptrs()`: + /// `(w1, b1, w2, b2, w3, b3)`. + /// + /// DEAD in C.5a-fixup — the populated ptrs are not yet read by any + /// collector-side launch (the `aux_trunk_forward.launch(...)` call + /// is inserted atomically in C.5b alongside the + /// `aux_heads_fwd input → h_s2_aux` kernel-rename + zero-fill revert). + pub fn set_trainer_aux_trunk_param_ptrs( + &mut self, + w1: u64, b1: u64, + w2: u64, b2: u64, + w3: u64, b3: u64, + ) { + self.aux_trunk_w1_ptr = w1; + self.aux_trunk_b1_ptr = b1; + self.aux_trunk_w2_ptr = w2; + self.aux_trunk_b2_ptr = b2; + self.aux_trunk_w3_ptr = w3; + self.aux_trunk_b3_ptr = b3; + info!( + aux_trunk_w1_ptr = w1, aux_trunk_b1_ptr = b1, + aux_trunk_w2_ptr = w2, aux_trunk_b2_ptr = b2, + aux_trunk_w3_ptr = w3, aux_trunk_b3_ptr = b3, + "Experience collector: zero-copy aux trunk param pointers set (SP14-C.5a-fixup, DEAD until C.5b)" + ); + } + /// Fill per_sample_support_buf [alloc_episodes, 4, 3] with uniform v_min/v_max/delta_z /// broadcast across all 4 branches. /// diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index 1c0dd2957..d4f8d5ccc 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -1701,6 +1701,18 @@ impl DQNTrainer { // Wire params + epsilon pointers from fused trainer (zero-copy weight access) if let Some(ref fused_ctx) = self.fused_ctx { collector.set_trainer_params_ptr(fused_ctx.params_flat_ptr()); + // SP14 Layer C Phase C.5a-fixup (2026-05-08): zero-copy + // wire of the trainer's 6 aux trunk parameter pointers + // alongside the existing main-params ptr. Cold-path, + // mirrors `set_trainer_params_ptr`'s contract — the + // collector reads the trainer's authoritative aux + // trunk weights without DtoD copy. DEAD until C.5b + // inserts the collector-side `aux_trunk_forward.launch`. + let (a_w1, a_b1, a_w2, a_b2, a_w3, a_b3) = + fused_ctx.trainer().aux_trunk_param_ptrs(); + collector.set_trainer_aux_trunk_param_ptrs( + a_w1, a_b1, a_w2, a_b2, a_w3, a_b3, + ); collector.set_per_sample_epsilon_ptr(fused_ctx.per_sample_epsilon_ptr()); // Tile per_sample_support with initial v_range (will be re-tiled each epoch) let vr = fused_ctx.eval_v_range(); @@ -2315,6 +2327,15 @@ impl DQNTrainer { // Zero-copy: wire collector to read weights directly from trainer's params_buf. if let Some(ref mut collector) = self.gpu_experience_collector { collector.set_trainer_params_ptr(ctx.params_flat_ptr()); + // SP14 Layer C Phase C.5a-fixup (2026-05-08): re-wire + // the aux trunk pointers on fused-ctx re-init (e.g. + // fold boundary). Mirrors the main-params setter's + // re-wire above. DEAD until C.5b. + let (a_w1, a_b1, a_w2, a_b2, a_w3, a_b3) = + ctx.trainer().aux_trunk_param_ptrs(); + collector.set_trainer_aux_trunk_param_ptrs( + a_w1, a_b1, a_w2, a_b2, a_w3, a_b3, + ); // Tile per_sample_support with current v_range let vr = ctx.eval_v_range(); if let Err(e) = collector.update_per_sample_support(vr[0], vr[1]) { diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index 9d1c3390c..7356ca5e4 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -7410,3 +7410,97 @@ C.5b is the genuine atomic contract migration: 5. Insert `aux_trunk_backward.launch(...)` + `launch_aux_trunk_adam_update(...)` calls. 6. Initialize `aux_trunk_lr_pinned` + `aux_trunk_grad_clip_pinned` from ISV per-step (host writes before each launch). 7. Increment `aux_trunk_adam_step` per Adam call. + +## SP14 Layer C Phase C.5a-fixup — missing scratch buffers + collector ptr scaffolding (2026-05-08) + +> Authorized 2026-05-08 after C.5b implementer flagged five gaps in C.5a: +> (1) two scratch buffers required by `aux_trunk_backward.launch` were never allocated; +> (2) collector lacked aux trunk param ptrs needed for the zero-copy contract; +> (3) graph-capture status of the C.5b new launch sites was unverified; +> (4) `aux_heads_backward` `&self` → `&mut self` ripple from host-side step counter (deferred to C.5b); +> (5) C.5b LOC ≈600-900 vs the 300 originally estimated (mitigated by removing scaffolding from C.5b's scope). +> C.5a-fixup addresses gaps 1, 2, 3 — additive-only DEAD CODE, no contract change. Gaps 4 + 5 are C.5b's concerns. + +### Gap 1 — missing dh_pre scratch buffers (`gpu_dqn_trainer.rs`) + +`aux_trunk_backward.launch` ([`gpu_aux_trunk.rs:266-267`](../crates/ml/src/cuda_pipeline/gpu_aux_trunk.rs)) requires two caller-allocated scratch buffers for the per-layer pre-activation gradient tile: +- `dh_aux1_pre_scratch` `[B_max, AUX_TRUNK_H1=256]` — Layer-1 dh_pre output of `aux_trunk_bwd_dh_pre`, consumed by `aux_trunk_bwd_dW_reduce` for `dW1 + db1`. +- `dh_aux2_pre_scratch` `[B_max, AUX_TRUNK_H2=128]` — Layer-2 dh_pre output, consumed for `dW2 + db2`. + +C.5a missed these. Allocated alongside the existing C.5a saved-fwd / accumulator / grad buffers via `alloc_f32(&stream, b * h_layer, "...")`. Both fields tagged `#[allow(dead_code)]` until C.5b atomically wires the backward chain. + +### Gap 2 — collector aux trunk ptr scaffolding (`gpu_experience_collector.rs`) + +C.5b inserts `aux_trunk_forward.launch(...)` in the collector's per-rollout-step body so the collector-native EGF chain feeds from the aux representation pipeline instead of Q's `h_s2`. That launch needs the trainer's aux trunk parameter pointers (zero-copy — same architectural contract as `trainer_params_ptr`). Added: + +- 6× `u64` fields on `GpuExperienceCollector`: `aux_trunk_{w1,b1,w2,b2,w3,b3}_ptr` (init to 0). +- `exp_aux_trunk_forward_ops: AuxTrunkForwardOps` field, populated in the collector constructor via `AuxTrunkForwardOps::new(&stream)` on the **collector's stream** (not the trainer's) — same-stream contract per `pearl_no_host_branches_in_captured_graph`. +- Public setter `set_trainer_aux_trunk_param_ptrs(w1, b1, w2, b2, w3, b3)` — mirrors `set_trainer_params_ptr`'s shape. +- Trainer-side accessor `GpuDqnTrainer::aux_trunk_param_ptrs() -> (u64×6)` returning `self.aux_trunk_{w,b}{1,2,3}.raw_ptr()`. Pattern matches the existing `params_buf_ptr() / target_params_buf_ptr()` family. +- Setter wired at both `training_loop.rs` `set_trainer_params_ptr` call sites (pre-existing fused-ctx init + fused-ctx fold-boundary re-init); the new setter call passes the unpacked tuple from the trainer's accessor. + +DEAD: ptrs are populated post-fused-ctx-init but no collector-side launch reads them in C.5a-fixup. + +### Gap 3 — Graph-capture audit: aux_heads_backward IS INSIDE captured graph + +**Result: INSIDE captured graph.** Verified by tracing the call chain: + +- `aux_heads_backward` ([`gpu_dqn_trainer.rs:17663`](../crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs)) is invoked from `launch_cublas_backward_to` ([`gpu_dqn_trainer.rs:29612`](../crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs)), which is called by `launch_cublas_backward` ([line 29354](../crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs)). +- `launch_cublas_backward` is called from `submit_forward_ops_main` ([line 27923](../crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs)). +- `submit_forward_ops_main` runs inside the `forward` child graph captured by `capture_child_graph("forward", ...)` in [`fused_training.rs:3033-3037`](../crates/ml/src/trainers/dqn/fused_training.rs). +- The single CUDA stream `begin_capture` boundary in the trainer family is `fused_training.rs:2964` (`capture_child_graph` body); confirmed by `grep -nE "begin_capture" crates/ml/src/`. + +**Implication for the existing aux_heads_backward.** The current implementation is `pub(crate) fn aux_heads_backward(&self, ...)` — entirely device-side: every operation is a kernel launch (`backward_next_bar`, `backward_regime`, 8× `param_grad_reduce` + SAXPY, 2× dh_s2 SAXPY). Zero host writes inside the function body — capture-safe by construction. + +**Implication for C.5b's new aux trunk Adam launch.** The Adam launcher (`launch_aux_trunk_adam_update`, [`gpu_dqn_trainer.rs:9866`](../crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs)) is also fully device-side (6× grad_norm partials → 1× finalize → 6× Adam update kernels). Capture-safe. + +**Critical C.5b concern: the per-step host writes.** The C.5a allocations include three pinned host pointers needing per-step host writes before each Adam launch: +- `aux_trunk_grad_clip_pinned` — host writes ISV[`AUX_TRUNK_GRAD_CLIP_INDEX`]. +- `aux_trunk_lr_pinned` — host writes ISV[`AUX_TRUNK_LR_INDEX`]. +- `aux_trunk_adam_step` (host i32 mirror) → `aux_trunk_t_pinned`. + +**These host writes CANNOT happen inside the captured `forward` child graph.** Two graph-capture-safe strategies are already established in the codebase: + +1. **Pre-capture host write pattern.** `step_selectivity_adam` ([`gpu_dqn_trainer.rs:32169`](../crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs)) demonstrates the pattern: increment host counter (`self.sel_adam_step += 1`), `unsafe { *self.sel_t_pinned = step_val }`, then issue kernel launches. **Note: `step_selectivity_adam` is not currently invoked from any captured-graph path** — it's an orphan, so it doesn't disprove the constraint. C.5b would need to issue the host writes outside the captured region (e.g., a pre-capture helper or in `process_epoch_boundary`) so the captured graph just records the device-side launches reading the pre-set pinned values. + +2. **GPU-side step counter increment kernel** (PREFERRED for graph capture). The trainer already loads an `increment_step_counters` kernel ([line 22793](../crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs)) plus a `submit_counters_ops` capture point ([line 22799](../crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs)) explicitly designed to bump all 8 main Adam step counters from inside the graph. **C.5b should add the aux trunk's `t_pinned` to the same `submit_counters_ops` chain so the step counter is incremented graph-side.** ISV-driven LR + clip values are already mapped-pinned with host-readable dev_ptrs; the host writes the LR/clip _before_ entering the captured region (cold-path per-epoch / per-fold), and the kernel reads `*lr_ptr` / `*clip_ptr` at execution time. This is the pattern the existing `lb_budget_*_dev_ptr` chain uses ([fused_training.rs:1940-1946](../crates/ml/src/trainers/dqn/fused_training.rs)) — capture records pointer reads, not values. + +**C.5b wiring strategy (binding for C.5b's implementer):** +- Insert the new Adam launch inside the captured `forward` child graph, immediately after the existing `aux_heads_backward` SAXPY into `bw_d_h_s2`. Both are device-side kernel chains, both run on the same captured stream. +- Add the aux trunk's `t_pinned` to the `submit_counters_ops` GPU-side increment chain (or a sibling captured child) — DO NOT use host-side `aux_trunk_adam_step += 1` patterns inside the captured region. +- ISV[`AUX_TRUNK_LR_INDEX`] / ISV[`AUX_TRUNK_GRAD_CLIP_INDEX`] writes happen pre-capture (cold-path host writes to `aux_trunk_lr_pinned` / `aux_trunk_grad_clip_pinned`); the captured graph reads via the pre-cached `aux_trunk_lr_dev_ptr` / `aux_trunk_grad_clip_dev_ptr`. +- The `&self` → `&mut self` ripple on `aux_heads_backward` (gap 4) is therefore avoided IF C.5b uses the GPU-side counter increment path. If C.5b instead adopts the pre-capture host-write strategy, it must restructure the call site so the host writes occur before `capture_child_graph("forward", ...)` begins — which is non-trivial because the `forward` child is composed once at capture-time (line 3033) and its body is the same closure on every replay. + +### Files modified (this commit) + +- Modify: [`crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs`](../crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs) — +2 fields (`dh_aux1_pre_scratch`, `dh_aux2_pre_scratch`), +2 allocations adjacent to C.5a's site, +`aux_trunk_param_ptrs()` accessor. +- Modify: [`crates/ml/src/cuda_pipeline/gpu_experience_collector.rs`](../crates/ml/src/cuda_pipeline/gpu_experience_collector.rs) — +6× `u64` ptr fields + 1× `AuxTrunkForwardOps` field on `GpuExperienceCollector`, +1× orchestrator construction in collector ctor, +`set_trainer_aux_trunk_param_ptrs` setter. +- Modify: [`crates/ml/src/trainers/dqn/trainer/training_loop.rs`](../crates/ml/src/trainers/dqn/trainer/training_loop.rs) — +setter calls at the two existing `set_trainer_params_ptr` sites (~line 1703 / ~line 2317). +- Modify: this audit doc. + +### Verification + +- `SQLX_OFFLINE=true CUDA_COMPUTE_CAP=86 cargo check -p ml --tests --all-targets` — clean (only pre-existing warnings; no new errors). +- `SQLX_OFFLINE=true CUDA_COMPUTE_CAP=86 cargo test -p ml --test aux_trunk_oracle_tests --test sp14_oracle_tests --release -- --ignored --nocapture` — 12/12 pass (8 aux_trunk + 4 sp14, bit-identical to C.5a baseline; no regression — pure scaffolding). +- `grep -rn "set_trainer_aux_trunk_param_ptrs\|aux_trunk_param_ptrs\|dh_aux[12]_pre_scratch" crates/ml/src/` shows the new symbols defined and called at the exact 2 setter sites + 2 trainer-side allocations + 2 trainer-side struct-literal entries. + +### Wire status — DEAD scaffolding (post-C.5a-fixup) + +| Component | Status | Lit by C.5b? | +|-----------|--------|--------------| +| `dh_aux1_pre_scratch` / `dh_aux2_pre_scratch` (trainer) | allocated, no consumers | ✓ via `aux_trunk_backward.launch` | +| `aux_trunk_param_ptrs()` accessor (trainer) | defined, called from training_loop | ✓ result already populates the setter | +| `set_trainer_aux_trunk_param_ptrs` (collector) | called, fields populated, no readers | ✓ via collector-side `aux_trunk_forward.launch` | +| `aux_trunk_w{1,2,3}_ptr` / `b{1,2,3}_ptr` (collector) | populated, no readers | ✓ | +| `exp_aux_trunk_forward_ops: AuxTrunkForwardOps` (collector) | constructed, no `.launch()` calls | ✓ | + +### Follow-up — C.5b is now scaffolding-free + +With C.5a + C.5a-fixup providing all scaffolding (buffers, scratch, ptrs, accessor, setter, AuxTrunkForwardOps in collector), C.5b becomes a true atomic contract flip: +1. Param-rename `h_s2 → h_s2_aux` in the 4 aux_heads kernels. +2. Revert zero-fills at `aux_next_bar_backward:599-638` + `aux_regime_backward:757-790`. +3. Insert `aux_trunk_forward.launch(...)` post-`forward_online_f32` (collector) + post-encoder forward (trainer). +4. Insert `aux_trunk_backward.launch(...)` + `launch_aux_trunk_adam_update(...)` in the trainer's backward chain (graph-capture-safe per the audit above). +5. Add aux trunk `t_pinned` to `submit_counters_ops` chain (GPU-side step counter increment). +6. Pre-capture host write of ISV-driven `aux_trunk_lr_pinned` / `aux_trunk_grad_clip_pinned` (cold-path per-epoch). +7. Redirect `aux_heads_fwd` input pointer from Q's `h_s2` to `h_s2_aux`.