refactor(sp4): A13 code-quality pass — DRY helper, named constant, doc fix, param cleanup

Addresses 5 IMPORTANT items from A13 code-quality review:

1. apply_pearls_to_slot helper extracted into sp4_wiener_ema.rs.
   Collapses ~30 lines of read_volatile / pearls_ad_update / write_volatile
   per-slot block into a single unsafe fn. 12 launchers now consume the
   helper (5 SP4 producers A5-A9, 6 A13 retrofits A13.0-A13.5, plus the
   inline label-scale block in aux_heads_forward Step 2b — including 2
   apply_pearls closures inside multi-slot launchers and the cross-boundary
   GpuExperienceCollector consumer). Pearl C rate_deficit site untouched
   (different mapped-pinned buffer + Rust ema array, doesn't share the
   helper's pointer-based contract).

2. REWARD_COMPONENT_COUNT named constant added next to
   REWARD_POPART_EMA_INDEX in gpu_dqn_trainer.rs. Replaces 3 hardcoded `6`
   literals across collector launcher (block_dim, Pearls A+D loop) and
   training_loop fold-reset range arithmetic. Mirrors MOE_NUM_EXPERTS /
   SL_NUM_FEATURE_GROUPS invariant-guard pattern with debug_assert_eq! in
   the collector launcher. Kernel literal `6` retained (allows nvcc full
   unroll); kernel comment now documents the host-side invariant.

3. SP4_PRODUCER_COUNT, SP4_WIENER_FLOATS_PER_SLOT, SP4_WIENER_TOTAL_FLOATS
   promoted from fn-local consts inside `pub fn new` to module-level
   `pub const`s in gpu_dqn_trainer.rs, re-exported via cuda_pipeline::mod.
   All 13 redeclarations in tests/sp4_producer_unit_tests.rs replaced
   with single `use ml::cuda_pipeline::SP4_PRODUCER_COUNT;` import. Future
   buffer growth requires single-file edit.

4. _ema_alpha_unused: f32 caller-compat shim removed from 6 retrofitted
   launchers (launch_h_s2_rms_ema, launch_aux_heads_loss_ema,
   launch_vsn_mask_ema, launch_moe_expert_util_ema, launch_iqn_quantile_ema,
   launch_reward_component_ema_inplace) and the FusedTrainingCtx proxy.
   All callers in training_loop.rs + fused_training.rs updated to drop
   the unused argument per feedback_no_legacy_aliases (no soft-deprecated
   wrappers; rename all call sites directly). Doc-comments updated from
   "α dropped per SP4 — argument preserved so callers compile unchanged"
   to "α derived adaptively from per-slot Pearls A+D Wiener state — see
   sp4_wiener_ema::pearls_ad_update".

5. Stale doc reference fixed at gpu_aux_heads.rs:391 — comment referenced
   non-existent launch_label_scale_ema_with_pearls function. Now correctly
   points at the inline Pearls A+D block in GpuDqnTrainer::aux_heads_forward
   Step 2b (which now consumes apply_pearls_to_slot).

Pure refactor — no spec or behaviour change. Same kernel launches, same
Pearls A+D semantics, same ISV slot writes.

cargo check -p ml --offline clean (12 pre-existing warnings, no new);
cargo test -p ml --lib --offline sp4_wiener_ema 7/7 passing
(6 originals + apply_pearls_to_slot_pearl_a_bootstrap_path);
cargo test -p ml --lib --offline state_reset_registry 3/3 passing.

Refs: A13 review (commits aada419de..c5add566d).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-01 08:42:40 +02:00
parent c5add566db
commit 4c231fa812
10 changed files with 362 additions and 353 deletions

View File

@@ -2337,4 +2337,6 @@ SP4 Layer A Task A13.4 — iqn_quantile_ema retrofit (Pearls A+D) (2026-05-01):
SP4 Layer A Task A13 follow-up — fold-reset sentinel contract restored (2026-05-01): A13.0..A13.5 zeroed the Wiener triples at fold boundary via the bulk `sp4_wiener_state` reset but left the companion ISV-slot dispatch arms in `state_reset_registry::reset_named_state` writing pre-SP4 cold-start values (1.0, 1/6, 1/8, ln(8)) — defeating Pearl A's sentinel branch (`prev_x_mean == 0 AND state.x_lag == 0`). Updated the 5 retrofit dispatch arms (`isv_h_s2_rms_ema`, `isv_vsn_mask_g{0..5}_ema`, `isv_aux_label_scale_ema`, `isv_moe_expert_util_ema`, `isv_moe_gate_entropy_ema`) to write 0.0 in lockstep with the Wiener-state half. Registry descriptions updated in the same commit per `feedback_no_partial_refactor.md` (doc + dispatch are two halves of the same contract). Stale `141`/`2048`/`47` Wiener-buffer comments updated to `207`/`2816`/`69` in `state_reset_registry.rs`, `training_loop.rs`, and `gpu_dqn_trainer.rs`. `cargo check -p ml --offline` clean; `sp4_wiener_ema` 6/6 + `state_reset_registry` 3/3 tests passing.
SP4 Layer A Task A13 code-quality refactor — apply_pearls_to_slot helper + REWARD_COMPONENT_COUNT named constant + SP4_PRODUCER_COUNT module-level promotion + _ema_alpha_unused removal + label-scale doc fix (2026-05-01): five IMPORTANT items from the A13 code-quality review addressed in a single coordinated commit per `feedback_no_partial_refactor.md`. (1) New `apply_pearls_to_slot(scratch_host, scratch_idx, isv_pinned, isv_idx, wiener_host, wiener_offset)` helper in `sp4_wiener_ema.rs` collapses the ~30-line read_volatile/`pearls_ad_update`/write_volatile per-slot block into a single `unsafe fn`; consumed by 11 launchers (SP4 producers `launch_sp4_target_q_p99`, `launch_sp4_atom_pos_p99_all_branches`, `launch_sp4_param_group_oracles_all_groups` (closure), `launch_sp4_grad_norm_p99`, `launch_sp4_h_s2_p99`; A13 retrofit `launch_h_s2_rms_ema`, `launch_aux_heads_loss_ema`, `launch_vsn_mask_ema`, `launch_moe_expert_util_ema` (closure), `launch_iqn_quantile_ema`; cross-boundary `GpuExperienceCollector::launch_reward_component_ema_inplace`) plus the inline label-scale block in `aux_heads_forward` Step 2b. Pearl C `pearl_c_rate_deficit_state_buf` site at `gpu_dqn_trainer.rs:21436` left untouched — different mapped-pinned buffer + Rust-side ema array means it doesn't share the helper's signature contract; refactoring would force a different abstraction. New unit test `apply_pearls_to_slot_pearl_a_bootstrap_path` exercises the helper with heap arrays standing in for mapped-pinned pointers (`*const f32`/`*mut f32` ABI is identical). (2) New `pub const REWARD_COMPONENT_COUNT: usize = 6` constant declared next to `REWARD_POPART_EMA_INDEX` replaces 4 hardcoded `6` literals: `block_dim: (REWARD_COMPONENT_COUNT as u32, 1, 1)` in `launch_reward_component_ema_inplace`, `for c in 0..REWARD_COMPONENT_COUNT` Pearls A+D loop, `REWARD_POPART_EMA_INDEX + REWARD_COMPONENT_COUNT` in `state_reset_registry::reset_named_state` ISV-slot range, plus a `debug_assert_eq!(REWARD_COMPONENT_COUNT, 6, "reward-component layout invariant — kernel block_dim and reward_components_per_sample stride must update together")` invariant guard mirroring `MOE_NUM_EXPERTS` / `SL_NUM_FEATURE_GROUPS` style. Kernel `reward_component_ema_kernel.cu` retains the literal `6` (allows nvcc to fully unroll the row-stride arithmetic) but its comment now documents the host-side `REWARD_COMPONENT_COUNT` invariant. (3) `SP4_PRODUCER_COUNT` / `SP4_WIENER_FLOATS_PER_SLOT` / `SP4_WIENER_TOTAL_FLOATS` promoted from fn-local consts inside `pub fn new` to module-level `pub const`s in `gpu_dqn_trainer.rs`, re-exported via `cuda_pipeline::mod.rs` so `crates/ml/tests/sp4_producer_unit_tests.rs` can `use ml::cuda_pipeline::SP4_PRODUCER_COUNT;` instead of redeclaring the value at 13 test sites. All 13 redeclarations deleted. (4) `_ema_alpha_unused: f32` parameter removed from 6 retrofit launchers (`launch_h_s2_rms_ema`, `launch_aux_heads_loss_ema`, `launch_vsn_mask_ema`, `launch_moe_expert_util_ema`, `launch_iqn_quantile_ema`, `launch_reward_component_ema_inplace`) and the `FusedTrainingCtx::launch_iqn_quantile_ema` proxy; all callers in `training_loop.rs` updated per `feedback_no_legacy_aliases.md` (no soft-deprecated wrappers — rename all call sites directly). Doc-comments on each launcher updated from "α dropped per SP4 — Pearls A+D adapt α from per-slot signal-vs-noise variance. The `_ema_alpha_unused` argument is preserved so callers compile unchanged" to "α is derived adaptively from per-slot Pearls A+D Wiener state — see `sp4_wiener_ema::pearls_ad_update`". (5) Stale doc reference at `gpu_aux_heads.rs::launch_label_scale_ema:391` referencing non-existent `launch_label_scale_ema_with_pearls` corrected to point at the actual inline Pearls A+D block in `GpuDqnTrainer::aux_heads_forward` Step 2b (which now consumes `apply_pearls_to_slot`). Pure refactor — no spec or behaviour change; same kernel launches, same Pearls A+D semantics, same ISV slot writes. `cargo check -p ml --offline` clean (12 pre-existing warnings, no new); `cargo test -p ml --lib --offline sp4_wiener_ema` 7/7 passing (6 originals + 1 new helper test); `cargo test -p ml --lib --offline state_reset_registry` 3/3 passing.
SP4 Layer A Task A13.5 — reward_component_ema retrofit (Pearls A+D) + cross-boundary wiring + orphan deletion (2026-05-01): retrofit the existing per-step reward-component EMA producer in `crates/ml/src/cuda_pipeline/reward_component_ema_kernel.cu` AND wire the host-side Pearls A+D update across the trainer/collector boundary (mirroring the A14/A15 Pearl C wiring path) AND delete the trainer's orphan `launch_reward_component_ema` per `feedback_wire_everything_up.md`. Kernel signature changed: drops `(isv_out, ema_alpha, isv_reward_base_slot)` for `(scratch_buf, scratch_first_index=63)`. Single-block 6-thread (one per component); each thread reduces mean |r_c| over `n_samples` samples and writes the step observation to `scratch_buf[scratch_first_index + c]` for c in 0..6 with `__threadfence_system()`. **6 ISV slots wired with Pearls A+D**: ISV[REWARD_POPART_EMA_INDEX..+6) = ISV[63..69), with Wiener offsets `(63+c)*3 = 189..207` for c in 0..6 — these are the LAST slots in the post-A13 207-float `wiener_state_buf`. **Cross-boundary wiring (mirrors A14/A15 Pearl C precedent)**: 4 new fields on `GpuExperienceCollector``reward_component_pearls_wiener_host_ptr: *mut f32`, `reward_component_pearls_scratch_dev_ptr: u64`, `reward_component_pearls_scratch_host_ptr: *mut f32`, `reward_component_pearls_isv_pinned_ptr: *mut f32` — all `NULL`/`0` until wired. New setter `set_reward_component_pearls_buffers(wiener_host, scratch_dev, scratch_host, isv_pinned)` on the collector. New accessors on `GpuDqnTrainer`: `wiener_state_buf_host_ptr()`, `producer_step_scratch_buf_dev_ptr()`, `producer_step_scratch_buf_host_ptr()`, `isv_signals_pinned_ptr()`. New wire helper `FusedTrainingCtx::wire_reward_component_pearls_buffers(&mut self, &mut GpuExperienceCollector)` pulls all four pointers from the trainer and pushes them into the collector. Wired once in `training_loop.rs::init_gpu_experience_collector` immediately after `set_curiosity_pearl_c_buffers` (mirrors the A14/A15 wire call site). The retrofitted `launch_reward_component_ema_inplace` on the collector now: launches the kernel with `(rewards_ptr, n_samples, scratch_dev, scratch_first_index=63)`, syncs the stream, applies Pearls A+D in a `for c in 0..6` loop (read prev_x_mean from `isv_pinned`, read Wiener triple from `wiener_host`, call `pearls_ad_update`, write back ISV + Wiener) — all via mapped-pinned host pointers (no HtoD/DtoH per `feedback_no_htod_htoh_only_mapped_pinned.md`), then memsets `reward_components_per_sample` to zero (preserving original behaviour). Degenerate-zero short-circuit per slot — covers the always-zero placeholder components (c=2 trail, c=5 bonus) plus cold-start before first `collect_experiences_gpu`. **Orphan deletion** per `feedback_wire_everything_up.md`: removed `GpuDqnTrainer::launch_reward_component_ema` (lines 8775-8796 — zero call sites pre-deletion), removed the trainer-side `reward_component_ema_kernel: CudaFunction` field (zero consumers post-launcher-deletion), removed the cubin loader at line 11591-11596, removed the field from the constructor's struct-init at line 14677. The `REWARD_COMPONENT_EMA_CUBIN` static remains because the collector still loads from it. **Unit test** `sp4_reward_component_ema_writes_step_obs_via_pearl_a_then_converges_pearl_d` (`#[ignore]`-gated for GPU): drives kernel with N=128 reward_components where `r[i*6+c] = sign(i) × (c+1)` (analytical mean|r_c| = c+1 for c in 0..6). Asserts each scratch slot ∈ ±1e-5 of (c+1), non-target slots remain 0; verifies Pearl A bootstrap on component 0 + Pearl D convergence (1000 stationary observations within 1% of 1.0). The cross-boundary wiring path is exercised in production by the integration smoke harness; this kernel-direct test isolates the kernel signature retrofit + Pearls A+D semantics. Per `feedback_no_atomicadd.md`, `feedback_no_htod_htoh_only_mapped_pinned.md`, `feedback_no_partial_refactor.md`, `feedback_wire_everything_up.md`. `cargo check -p ml --lib --tests --offline` clean (11 pre-existing warnings, no new warnings); `cargo test -p ml --lib sp4_wiener_ema --offline` 6/6 passing.