Phase 2A of the HEALTH_DIAG GPU port. Lands the simplest member of the
kernel family (single-block single-thread ISV scalar copies) plus the
orchestrator scaffolding (cubin loader, mapped-pinned snapshot owner,
launcher) that subsequent phases will reuse.
Why simplest first: 5 kernels each with different masking semantics
landing in one commit risks silent numerical mismatches that downstream
parsers (aggregate-multi-seed-metrics.py, smoke summarisers) would
mis-parse. Per feedback_no_quickfixes.md, kernel-by-kernel with
parallel-shadow validation. Phase 4 deletes CPU paths in one commit
once all kernels are bit-identical.
What this kernel does: copies 22 ISV signal-bus slots into the matching
HealthDiagSnapshot fields — reward-component EMAs (slots 63-68), VSN
attention focus EMAs (87, 88), target-drift EMAs (92, 93), aux-head
loss EMAs (113, 114, 117), MoE expert utilisation (118-126), gate
entropy (126), and λ_eff (128). 22 of 147 snapshot words populated
(reward_split[6] + noisy_vsn[2] + noisy_drift[2] + aux[3] + aux_moe[10]).
Producer-only — no consumer reads health_diag.snapshot() yet; Phase 3
wires the call site, Phase 4 makes the snapshot the sole emit source.
Architecture rules upheld:
- No HtoD / DtoH (writes through mapped-pinned device pointer).
- No atomicAdd (single-thread kernel; family-wide rule for 2B-2E too).
- No DtoD memcpy on the diag path (kernel pointer-load + store).
- Kernel WORD-index table inline; static_assert pins WORD_TOTAL=147 in
lockstep with snapshot_size_is_stable test.
- ISV indices passed as kernel args (not #define'd) so the kernel
never embeds the numbering — single source of truth lives in
gpu_dqn_trainer.rs constants.
Changes:
- new crates/ml/src/cuda_pipeline/health_diag_kernel.cu (311 lines).
- new crates/ml/src/cuda_pipeline/gpu_health_diag.rs (190 lines)
holds GpuHealthDiag orchestrator (cubin module + isv_mirror_kernel
handle + MappedHealthDiagSnapshot).
- mod.rs: pub mod gpu_health_diag.
- build.rs: registers health_diag_kernel.cu in kernels_with_common.
- gpu_dqn_trainer.rs:
- imports super::gpu_health_diag::GpuHealthDiag;
- adds health_diag: Option<GpuHealthDiag> field next to aux_heads_fwd;
- constructor calls GpuHealthDiag::new() after the cubin loads;
- adds pub fn launch_health_diag_isv_mirror() (uses compile-time
ISV index constants, single source of truth);
- adds pub fn health_diag_snapshot() accessor.
- docs/health_diag_inventory.md: Phase 2A entry + Phase 4 deletion
list pinned per Invariant 7.
- docs/dqn-wire-up-audit.md: new row for health_diag_kernel.cu under
"CUDA Kernels (.cu files)" — classified Pending wire-up (Phase 3).
Verification:
- cargo check -p ml: 0 new warnings (12 lib warnings pre-existing,
none touch health_diag/gpu_health_diag).
- cargo test -p ml --no-run: clean.
- cargo test -p ml --lib cuda_pipeline::health_diag::tests: 3/3 pass
(snapshot_size_is_stable, default_is_zeroed, alignment_is_4_bytes
— the WORD_TOTAL=147 static_assert in the kernel matches the host
size_of test bit-for-bit).
Scope discipline (per prompt's "if kernel #3 takes too long, stop"):
This commit lands the orchestrator scaffolding + simplest kernel only.
Phases 2B-2E (eval-histogram, q-mag-reduce, per-sample-reduce,
finalise) intentionally deferred — masking semantics for q_mag_reduce
in particular need careful per-CPU-getter inspection before kernel-side
implementation. Producer-only by design; no production caller of
launch_health_diag_isv_mirror in this commit (matches AuxHeadsForwardOps
landing pattern from Plan 4 Task 6 Commit A).
Catalogue every numeric field in the HEALTH_DIAG family of log lines
(7 emit sites across training_loop.rs + metrics.rs) and classify each
source as GPU-already / CPU-bound / Mixed / Host-state / ISV-already.
Identifies the dominant cost contributors driving the observed
~70 s/epoch HEALTH_DIAG overhead on L40S after the eval async-split
(commit f815f7239) — they are all CPU-side reductions over multi-million
element per-sample buffers fetched via memcpy_dtoh:
- update_q_mag_means_cached (~786 KB DtoH + B×total_actions loop)
- var_scale_epoch_mean (N-element DtoH + conditional avg)
- trail_fire_and_hold_per_mag (4× N-element DtoH + per-mag loop)
- per_magnitude_winrate_and_variance (4× DtoH + sumsq loop)
- reward_contrib_fractions (6× large-buffer DtoH + 6 passes)
- read_eval_action_distribution_* (4 sites; CPU iter on pinned)
- branch_noisy_sigma_mean (NoisyNet param DtoH per branch)
Aggregate per-emit DtoH ≈ 50-150 MB depending on alloc_episodes ×
alloc_timesteps; fully consistent with the 70 s/epoch observation.
All inputs are already device-resident, so the migration is "stop
reducing on the host" plus a small kernel family writing into a single
mapped-pinned HealthDiagSnapshot struct.
Phase 0 deliverable per the dispatching prompt; subsequent commits
(struct + mapped wrapper, kernel family, CPU emit rewrite, audit-doc
updates) follow the phased plan documented at the bottom of the file
and gate-review on this inventory before touching code.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
