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).
30 KiB
HEALTH_DIAG Inventory — Phase 0 (read-only)
Purpose. Catalogue every numeric value emitted in the HEALTH_DIAG family of log lines, classified by where the compute lives today, ahead of the greenfield port to a single GPU-resident HealthDiagSnapshot.
Scope. All tracing::info!("HEALTH_DIAG[...]: ...", …) emissions in crates/ml/src/. Target: every numeric placeholder in those format strings.
Why. L40S observed ~70 s/epoch attributed to HEALTH_DIAG even after f815f7239 (launch_validation_loss / consume_validation_loss split). Async eval moved the host wait off the critical path; it did NOT move the compute — the per-epoch HEALTH_DIAG emit still triggers a wave of memcpy_dtoh reads of multi-million-element per-sample buffers followed by Rust-side reductions (counting, summation, variance, histograms).
Architectural rules (re-stated for cross-reference; see prompt for full list):
feedback_no_htod_htoh_only_mapped_pinned.md— mapped-pinned is the ONLY CPU↔GPU path.feedback_no_atomicadd.md— noatomicAdd; shmem tree reductions only.feedback_adaptive_not_tuned.md— block/grid sizes derive from input dims.feedback_no_partial_refactor.md— every consumer of the new struct migrates in the same commit.feedback_wire_everything_up.md— every metric in the snapshot is emitted; every metric currently emitted has a slot.
Classification key
- GPU-already — value is produced by an existing CUDA kernel that writes via mapped-pinned (
cuMemHostAlloc(DEVICEMAP|PORTABLE)) or to an ISV slot. CPU just reads. Migration: rename the source slot into the newHealthDiagSnapshotstruct field; same kernel writes; same data. - CPU-bound — value is computed on the CPU after one or more
memcpy_dtohreads. Often involvesVec::with_capacity(N)followed by per-element loops on N=1M+ buffers. Migration: write a new CUDA reduction kernel; delete the CPU path entirely. - Mixed — kernel produces a per-batch buffer (no reduction) and CPU does a final reduction in Rust. Migration: extend the kernel with a fused reduction tail OR add a small dedicated reducer kernel; delete the CPU loop.
- Host-state — value is a Rust-side state machine output (e.g.
last_anti_multafter the anti-LR controller decided something). No GPU compute involved; remains CPU. Migration: read into a host-side scratchpad and write into the snapshot viahost_ptr(allocator-free since the snapshot is a fixed-size struct). - ISV-already — already lives in the ISV signal bus (
isv_signals_dev_ptr); CPU reads viaread_isv_signal_at. Migration: copy ISV slot into snapshot in a small ISV-mirror kernel, OR keep the existing host read but write into the snapshot's host_ptr from CPU after the read; no new GPU compute needed.
Format-string anchor lines
All sites surveyed:
training_loop.rs:3229— main giantHEALTH_DIAG[{epoch}]: …(the elephant; ~50 metric placeholders).training_loop.rs:3394—HEALTH_DIAG[{epoch}]: reward_split […](6 floats).training_loop.rs:3430—HEALTH_DIAG[{epoch}]: aux […](4 floats).training_loop.rs:3456—HEALTH_DIAG[{epoch}]: aux_moe […](8 utils + ent + λ_eff = 10 floats).metrics.rs:837—HEALTH_DIAG[{epoch}]: val […](16 floats).metrics.rs:890—HEALTH_DIAG[{epoch}]: val_dir_dist […](4 floats).metrics.rs:908—HEALTH_DIAG[{epoch}]: val_picked_dir_dist […](4 floats).training_loop.rs:900—HEALTH_DIAG[{epoch}]: TERMINATED_BY_REGRESSION …(one-shot fault-path emission; not part of the per-epoch hot path).
TERMINATED_BY_REGRESSION (#8) fires once at end-of-run on regression and is excluded from the snapshot port — it's a rare fault-path string carrying already-resolved values.
Inventory table — line training_loop.rs:3229 (main HEALTH_DIAG)
| Group | Metric | Source path | Class | Lines | Cost-per-epoch | Notes |
|---|---|---|---|---|---|---|
| top | health (health_value) |
self.learning_health.value() (CPU computation over EMA components) |
Host-state | training_loop.rs (computed earlier in process_epoch_boundary) |
trivial | Pure scalar arithmetic; aggregator of the 7 component fields below. |
| components | q_gap |
self.learning_health.components.q_gap_norm |
Host-state | mod.rs / learning_health.rs | trivial | Cached EMA component, scalar. |
| components | q_var |
self.learning_health.components.q_var_norm |
Host-state | mod.rs / learning_health.rs | trivial | Cached EMA component. |
| components | atoms |
self.learning_health.components.atom_util_norm |
Host-state | mod.rs / learning_health.rs | trivial | Cached EMA component. |
| components | grad_stable |
self.learning_health.components.grad_stable |
Host-state | mod.rs / learning_health.rs | trivial | Cached EMA component. |
| components | ens_agree |
self.learning_health.components.ens_agree |
Host-state | mod.rs / learning_health.rs | trivial | Cached EMA component. |
| components | grad_cos |
self.learning_health.components.grad_consistency_norm |
Host-state | mod.rs / learning_health.rs | trivial | Cached EMA component. |
| components | spectral |
self.learning_health.components.spectral_gap_norm |
Host-state | mod.rs / learning_health.rs | trivial | Cached EMA component. |
| effective | cql_alpha |
self.last_cql_alpha_eff |
Host-state | training_loop.rs | trivial | Set by adaptive controller earlier in epoch. |
| effective | iqn_budget |
self.last_iqn_budget_eff |
Host-state | training_loop.rs | trivial | Same. |
| effective | cql_budget |
self.last_cql_budget_eff |
Host-state | training_loop.rs | trivial | Same. |
| effective | c51_budget |
self.last_c51_budget_eff |
Host-state | training_loop.rs | trivial | Same. |
| effective | tau |
self.last_tau_eff |
Host-state | training_loop.rs | trivial | Same. |
| effective | sarsa_tau |
self.last_sarsa_tau_factor |
Host-state | training_loop.rs | trivial | Same. |
| effective | gamma |
self.last_gamma_eff |
Host-state | training_loop.rs | trivial | Same. |
| effective | cf_ratio |
self.last_cf_ratio_eff |
Host-state | training_loop.rs | trivial | Set by collector.set_learning_health(). |
| novels | distill (on/off bool) |
self.last_distill_active |
Host-state | training_loop.rs | trivial | Set by D-mechanism logic earlier in epoch. |
| novels | barrier |
self.last_barrier_loss |
Host-state | training_loop.rs | trivial | Same. |
| novels | plasticity (ready/cooldown bool) |
self.last_plasticity_ready |
Host-state | training_loop.rs | trivial | Same. |
| novels | ib |
self.last_ib_penalty |
Host-state | training_loop.rs ~2410 | trivial | Pure scalar arithmetic over EMA. |
| novels | ensemble_collapse |
self.last_ensemble_collapse_score |
Host-state | training_loop.rs ~2424 | trivial | Smoothstep over last_ens_disagreement. |
| novels | contrarian (on/off bool) |
self.last_contrarian_active |
Host-state | training_loop.rs ~2480 | trivial | State-machine output. |
| novels | meta_q_pred |
self.last_meta_q_pred |
Host-state | training_loop.rs ~2531 | trivial | self.meta_q.predict(input) — small CPU MLP. Stays on CPU; meta-Q is a tiny network. |
| diag | sharpe_ema |
self.training_sharpe_ema |
Host-state | training_loop.rs | trivial | Updated from train pnl. |
| diag | action_entropy |
self.last_action_entropy |
Mixed | training_loop.rs ~2790 | small CPU loop over 12-bin counts | Currently CPU computes Shannon entropy normalised over summary.action_counts[12]. Move to GPU: extend the existing monitoring_kernel's reducer to also emit normalised entropy. |
| gems | g12_predictive |
fused.read_gem_g12_loss() |
GPU-already | gpu_dqn_trainer.rs ~13169 | trivial | Already pinned. |
| mag | q_full / q_half / q_quarter |
fused.q_magnitude_bucket_means() after update_q_mag_means_cached |
CPU-bound (HEAVY) | gpu_dqn_trainer.rs:4437 | 786 KB DtoH/epoch + B×total_actions CPU loop (B=16384, total_actions=12 → 196,608 floats summed). | Hot reduction. Migrate: small kernel reducing q_out_buf cols [b0..b0+b1] across batch into 3 floats. |
| mag | var_scale (var_scale_mean) |
collector.var_scale_epoch_mean() |
CPU-bound (HEAVY) | gpu_experience_collector.rs:2784 | DtoH ~N×4 bytes (N = alloc_episodes × alloc_timesteps, ~512K-2M); CPU loop with conditional avg. | Migrate: GPU reduction kernel summing only v > 0 slots + count, returns mean. |
| mag | kelly_f (kelly_f_mean) |
self.trade_stats_history.last() arithmetic |
Host-state | training_loop.rs ~2753 | trivial | Pure CPU formula on 6-float TradeStats. |
| mag | avg_win_ratio |
same as above | Host-state | training_loop.rs ~2753 | trivial | Pure CPU formula. |
| mag | grad_ratio_mag_dir |
grad_ratio_mag_dir_cached |
GPU-already | populated by reduction kernel into pinned struct | trivial | Cached at top of process_epoch_boundary. |
| mag | dist_q / dist_h / dist_f |
monitor.action_counts[..] summed in Rust over 12 entries |
Mixed | training_loop.rs ~2839 | trivial CPU sum of 12 ints + division | Counts come from GPU via monitoring_summary (already pinned). The CPU sum is 12 ints — cheap, but part of the cascade we're collapsing. Migrate: GPU monitoring kernel emits dist_q/h/f directly. |
| grad_split_bwd | iqn cql c51 ens |
fused.grad_mag_*_ratio() after refresh_grad_component_norms() |
GPU-already | gpu_dqn_trainer.rs (in-graph reduction kernel writes pinned slots) | trivial | Read pinned. |
| grad_split_aux | distill rec pred cql_sx c51_bs |
same family of grad_mag_*_ratio() |
GPU-already | same | trivial | Read pinned. |
| grad_trunk | 9 floats | fused.grad_trunk_norms_by_component() ([f32; 9]) |
GPU-already | same | trivial | Read pinned. |
| grad_abs | dir mag |
grad_dir_abs_cached, grad_mag_abs_cached |
GPU-already | populated by same reduction kernel | trivial | Read pinned. |
| trail | fire_q/h/f hold_q/h/f |
collector.trail_fire_and_hold_per_mag() |
CPU-bound (HEAVY) | gpu_experience_collector.rs:2063 | 4× DtoH of N-element buffers (N = alloc_episodes × alloc_timesteps); CPU loop with branch + accumulate per-mag. | Migrate: 3-block kernel (one per mag) with shmem reductions over the 4 input buffers; write 6 floats (3 rates + 3 hold means). |
| mag_stats | wr_q/h/f var_q/h/f |
collector.per_magnitude_winrate_and_variance() |
CPU-bound (HEAVY) | gpu_experience_collector.rs:2682 | 4× DtoH of N-element buffers; CPU loop computing Welford-style mean+sumsq per-mag. | Migrate: 3-block kernel; shmem reduction; emit 6 floats. |
| noisy | vsn_mag vsn_dir |
ISV slots VSN_MAG_EMA_INDEX, VSN_DIR_EMA_INDEX |
ISV-already | gpu_dqn_trainer.rs (constants); attention_focus_ema kernel writes |
trivial | Read via read_isv_signal_at. |
| noisy | sigma_mag sigma_dir |
agent.branch_noisy_sigma_mean(0/1) |
CPU-bound | branched DQN model; reads NoisyNet sigma params | medium (param DtoH per branch) | NoisyNets σ mean across fc+out layers. Migrate: dedicated kernel summing |
| noisy | drift_mag drift_dir |
ISV slots TARGET_DRIFT_MAG_EMA_INDEX, TARGET_DRIFT_DIR_EMA_INDEX |
ISV-already | target_drift_ema kernel writes |
trivial | Read via read_isv_signal_at. |
| eval_dist | eq eh ef |
self.last_eval_magnitude_dist[0..3] (computed at val time) |
GPU-already (pinned read) → CPU loop | metrics.rs ~862 calls evaluator's read_eval_action_distribution_per_magnitude |
medium (CPU loop over actions_history_pinned, ~64K-256K i32s) | Currently CPU iterates the mapped-pinned actions buffer counting (a/9) % 3 per non-negative entry. Migrate: small GPU histogram kernel reading the device pointer of the same buffer; write 3 floats. |
| reward_contrib | popart cf trail_r micro |
self.last_reward_contrib (4 floats); slot 0 = prev_reward_contrib_popart_var derived; slots 1..4 from collector.reward_contrib_fractions() |
CPU-bound (HEAVY for slots 1..4) | gpu_experience_collector.rs:2163 | 6× DtoH (cf_flip, slot_live, trail, traded, micro, total) of large per-sample buffers + multi-pass CPU summation. | Migrate: single fused reduction kernel over the 6 buffers; emit 4 floats. PopArt slot stays host-state (it's already a small scalar derived from read_popart_variance). |
| controller | anti_lr tau gamma clip cql cost (6 bools), fire_frac (f32) |
host-side detection: fire_lr etc. |
Host-state | training_loop.rs ~2902-2955 | trivial | Pure CPU comparisons against prev_controller_values. |
| explore | ent_mag ent_dir |
shannon_entropy_normalized(&[dist_q, dist_h, dist_f]) and same over [dir0..dir3] |
Mixed | training_loop.rs ~2858 | trivial CPU loop over 3-4 elements | Source values come from monitor.action_counts[12] (GPU). Migrate: have the new GPU monitoring extension emit ent_mag and ent_dir directly. |
| explore | sigma_mean |
(sigma_mag + sigma_dir) * 0.5 |
Host-state | training_loop.rs ~2640 | trivial | Trivial after sigma_mag/dir migrate. |
Hidden-cost lines that materially affect total HEALTH_DIAG cost
These are computed for HEALTH_DIAG but happen in upstream methods invoked just before the emit:
| Site | Cost | Migration |
|---|---|---|
update_q_mag_means_cached (gpu_dqn_trainer.rs:4437) |
786 KB DtoH + B×total_actions Rust loop | DELETE the CPU path; replace with kernel. |
var_scale_epoch_mean (gpu_experience_collector.rs:2784) |
N×4-byte DtoH; conditional CPU avg | Replace with kernel. |
trail_fire_and_hold_per_mag (gpu_experience_collector.rs:2063) |
4× N-element DtoH; per-mag CPU loop | Replace with kernel. |
per_magnitude_winrate_and_variance (gpu_experience_collector.rs:2682) |
4× N-element DtoH; per-mag CPU sum + sum_sq | Replace with kernel. |
reward_contrib_fractions (gpu_experience_collector.rs:2163) |
6× large-buffer DtoH; multi-pass CPU sums | Replace with kernel. |
read_eval_action_distribution_per_magnitude (gpu_backtest_evaluator.rs:1255) |
actions_history_pinned read (mapped, but CPU iterates ~64-256K entries) | Replace with kernel reading the dev pointer of the same mapped buffer; CPU just reads the resulting 3 floats from the snapshot. |
read_eval_action_distribution_per_direction (gpu_backtest_evaluator.rs:1339) |
same buffer, different histogram | Same kernel, additional 4-float output. |
read_eval_intent_magnitude_distribution (gpu_backtest_evaluator.rs:1295) |
intent_mag_pinned iter | Replace with kernel; 3-float output. |
read_chunked_actions_direction_distribution (gpu_backtest_evaluator.rs:1384) |
picked_action_history_pinned iter | Replace with kernel; 4-float output. |
Inventory table — line training_loop.rs:3394 (reward_split)
| Metric | Source | Class | Notes |
|---|---|---|---|
popart |
ISV REWARD_POPART_EMA_INDEX |
ISV-already | read_isv_signal_at. |
cf |
ISV REWARD_CF_EMA_INDEX |
ISV-already | Same. |
trail |
ISV REWARD_TRAIL_EMA_INDEX |
ISV-already | Same. |
micro |
ISV REWARD_MICRO_EMA_INDEX |
ISV-already | Same. |
opp_cost |
ISV REWARD_OPP_COST_EMA_INDEX |
ISV-already | Same. |
bonus |
ISV REWARD_BONUS_EMA_INDEX |
ISV-already | Same. |
All 6 already on GPU; ISV-already. Snapshot fields = mirror copies.
Inventory table — line training_loop.rs:3430 (aux)
| Metric | Source | Class | Notes |
|---|---|---|---|
next_bar_mse |
ISV AUX_NEXT_BAR_MSE_EMA_INDEX (113) |
ISV-already | read_isv_signal_at. |
regime_ce |
ISV AUX_REGIME_CE_EMA_INDEX (114) |
ISV-already | Same. |
w (aux_weight) |
trainer.aux_weight() |
Host-state | Plain scalar field on the trainer. |
label_scale |
ISV AUX_LABEL_SCALE_EMA_INDEX (117) |
ISV-already | Same. |
Inventory table — line training_loop.rs:3456 (aux_moe)
| Metric | Source | Class | Notes |
|---|---|---|---|
util[0..8] |
ISV MOE_EXPERT_UTIL_EMA_BASE..+8 |
ISV-already | 8 read_isv_signal_at. |
ent |
ISV MOE_GATE_ENTROPY_EMA_INDEX |
ISV-already | Same. |
λ_eff |
ISV MOE_LAMBDA_EFF_INDEX |
ISV-already | Same. |
Inventory table — line metrics.rs:837 (val)
All 16 fields originate from evaluator.consume_metrics_after_event(), which already reads from metrics_buf (mapped-pinned, populated by backtest_metrics_kernel). The CPU does a small post-process to derive trades_per_bar, active_frac, dir_entropy — pure arithmetic over scalar floats already in the metrics struct.
| Metric | Source | Class | Notes |
|---|---|---|---|
sharpe |
m.sharpe |
GPU-already | metrics_buf scalar. |
sortino |
m.sortino |
GPU-already | Same. |
win_rate |
m.win_rate |
GPU-already | Same. |
max_drawdown |
m.max_drawdown |
GPU-already | Same. |
trade_count |
m.total_trades |
GPU-already | Same. |
calmar |
m.calmar |
GPU-already | Same. |
omega_ratio |
m.omega_ratio |
GPU-already | Same. |
total_pnl |
m.total_pnl |
GPU-already | Same. |
var_95 |
m.var_95 |
GPU-already | Same. |
cvar_95 |
m.cvar_95 |
GPU-already | Same. |
trades_per_bar |
derived (3 buckets ÷ window) | Host-state (trivial) | Three scalar f32 divides. |
active_frac |
derived | Host-state (trivial) | Same. |
dir_entropy |
derived (3-prob Shannon entropy) | Host-state (trivial) | Per-3-bucket entropy. Could move to kernel to land directly in metrics_buf; trivial enough to either keep or migrate. |
sharpe_annualised |
m.sharpe (alias) |
GPU-already | Aggregator alias. |
profit_factor |
m.omega_ratio (alias) |
GPU-already | Aggregator alias. |
window_bars |
buy + sell + hold |
Host-state (trivial) | One add. |
Migration note for val [...]: the metrics_buf is already mapped-pinned, so the val emit's CPU work today is 3 scalar adds + 3 entropy multiplies. The big perf win for HEALTH_DIAG is NOT in val; it's in the main line's per-sample buffer reductions. We should still fold val's fields into HealthDiagSnapshot for layout uniformity (one struct, allocator-free formatting), but the work to extend the backtest_metrics_kernel to directly emit dir_entropy etc. is optional — it can copy from metrics_buf into the snapshot at emit time.
Inventory table — line metrics.rs:890 (val_dir_dist)
| Metric | Source | Class | Notes |
|---|---|---|---|
short hold long flat (4 floats) |
read_eval_action_distribution_per_direction (gpu_backtest_evaluator.rs:1339) |
Mixed | CPU iterates actions_history_pinned counting per direction. Migrate: kernel histogram. |
Inventory table — line metrics.rs:908 (val_picked_dir_dist)
| Metric | Source | Class | Notes |
|---|---|---|---|
short hold long flat (4 floats) |
read_chunked_actions_direction_distribution (gpu_backtest_evaluator.rs:1384) |
Mixed | CPU iterates picked_action_history_pinned. Same pattern as above. |
Aggregate cost summary (per-epoch, batch 16384)
Counting only the CPU-bound + Mixed sites whose work is dominated by Vec alloc + DtoH + per-element loops:
| Site | DtoH bytes | CPU loop work | Notes |
|---|---|---|---|
update_q_mag_means_cached |
~786 KB (B × total_actions × 4) | 196 608 float-adds + 3 div | One epoch; ✗ should be a kernel. |
var_scale_epoch_mean |
~N × 4 bytes (≈ 2-8 MB at full scale) | N ifs + adds | ✗. |
trail_fire_and_hold_per_mag |
4× N-element DtoH (~8-32 MB) | 4N ifs + accumulate | ✗. |
per_magnitude_winrate_and_variance |
4× N-element DtoH (same scale) | 4N ifs + sumsq | ✗. |
reward_contrib_fractions |
6× large DtoH (≈ 24-100 MB) | 6N passes | ✗ — biggest site. |
read_eval_action_distribution_* |
reads mapped-pinned (no DtoH but CPU iterates) | up to 256K i32 iter × 4 (3 mag + 4 dir + 3 intent + 4 picked) | ✗ — CPU iter eats into the 70 s budget. |
branch_noisy_sigma_mean(0) and (1) |
NoisyNet param DtoH | small but per-branch | ✗ at the margins. |
Total estimated DtoH per HEALTH_DIAG emit: ~50-150 MB depending on N (alloc_episodes × alloc_timesteps). At PCIe Gen4 x16 (~20 GB/s sustained) that's ~3-7 ms of pure DtoH plus the host loop overhead at 1-2 ns per element on N up to 4M, which gets us into the 30-100 ms range per emit just for these readback paths. Multiplied by other epoch-boundary work (eval consumption, controller fires, optimizer maintenance) and synchronous host syncs, the 70 s/epoch number is consistent with the per-sample-buffer reductions being the dominant cost. They are also the easiest to eliminate because every per-sample buffer is already device-resident — we just stop reducing on the host.
What stays CPU after the port
Per the prompt's explicit allowance ("CPU is now strictly a string formatter + log emitter"), these remain CPU:
- Format-string assembly —
String::with_capacity(2048)once per emit, fixed cap. NoVecon hot path. - Host-state mirrors — fields tagged
Host-stateabove (controllers, mechanism flags, EMA scalars on the LearningHealth aggregator). These are already small fields on the trainer struct; we read them with no reduction. Their inclusion inHealthDiagSnapshotis "host writes the field viahost_ptrimmediately before format" — allocator-free, no GPU compute, trivially fast. - Meta-Q predictor (
self.meta_q.predict) — tiny CPU MLP; staying CPU is a deliberate architectural choice, not a HEALTH_DIAG concern.
Everything in Mixed and CPU-bound (HEAVY) rows is in scope for the kernel port and the corresponding CPU code is deleted in the same commit chain.
Snapshot field list (proposed; finalised in Phase 1 commit)
Field order matches the format-string emit order so log determinism is preserved across the port. Counts: 7 components + 8 effective + 7 novels + 2 diag + 1 gems + 10 mag + 4 grad_split_bwd + 5 grad_split_aux + 9 grad_trunk + 2 grad_abs + 6 trail + 6 mag_stats + 6 noisy + 3 eval_dist + 4 reward_contrib + 6 fire_bools + 1 fire_frac + 3 explore + 16 val + 4 val_dir_dist + 4 val_picked_dir_dist + 6 reward_split + 4 aux + 10 aux_moe + 12 monitor.action_counts (for direction-bin computation if we keep that path adjacent) ≈ ~150 f32 fields.
The booleans become u8 slots (0/1) so the struct stays POD with no padding surprises across CPU/GPU. The struct sits in a single mapped-pinned allocation of ~600-800 bytes — well under one cache line × 16, easy to fence + read.
Phased commit plan (revised 2026-04-28 — kernel-by-kernel + parallel-shadow validation)
The original Phase-1-design recommendation of a 5-kernel family with sequential
landings has been adopted. Per feedback_no_quickfixes.md, landing all 5
kernels in one commit risks silent numerical mismatches that downstream
parsers (aggregate-multi-seed-metrics.py, smoke-test summarisers) would
mis-parse. Each kernel lands in its own commit with parallel-shadow
validation; only after all 5 are bit-identical (mod float-reordering tolerance
1e-5 relative) does Phase 4 delete the CPU paths in one shot per
feedback_no_partial_refactor.md.
| Phase | Commit | Status | Files touched |
|---|---|---|---|
| 0 | This inventory doc | LANDED | docs/health_diag_inventory.md |
| 1 | HealthDiagSnapshot + MappedHealthDiagSnapshot |
LANDED | crates/ml/src/cuda_pipeline/health_diag.rs |
| 2A | health_diag_isv_mirror kernel |
LANDED | health_diag_kernel.cu + gpu_health_diag.rs (NEW) + build.rs + trainer wiring |
| 2B | health_diag_eval_histogram kernel |
Pending | health_diag_kernel.cu (append) + launcher |
| 2C | health_diag_q_mag_reduce kernel |
Pending | health_diag_kernel.cu (append) + launcher (most complex; masking-heavy) |
| 2D | health_diag_per_sample_reduce kernel |
Pending | health_diag_kernel.cu (append) + launcher |
| 2E | health_diag_finalise kernel |
Pending | health_diag_kernel.cu (append) + launcher (single __threadfence_system()) |
| 3 | Wire kernel chain into captured graph | Pending | training_loop.rs (call sites alongside launch_h_s2_rms_ema etc.) |
| 4 | Rewrite emit path + delete CPU code | Pending | training_loop.rs, metrics.rs, gpu_experience_collector.rs, gpu_backtest_evaluator.rs, gpu_dqn_trainer.rs, agent — see "Functions to delete" below |
| 5 | Final polish + audit doc updates | Pending | docs/dqn-gpu-hot-path-audit.md, docs/dqn-wire-up-audit.md |
Phase 2A — what landed
- New file
crates/ml/src/cuda_pipeline/health_diag_kernel.cu. Singlehealth_diag_isv_mirrorkernel, single-block single-thread, ~22 ISV reads- ~22 snapshot writes. WORD-index table inline at top of
.cumatches the RustHealthDiagSnapshotfield order; static_assert pinsWORD_TOTAL = 147in lockstep withsnapshot_size_is_stabletest.
- ~22 snapshot writes. WORD-index table inline at top of
- New file
crates/ml/src/cuda_pipeline/gpu_health_diag.rs.GpuHealthDiagstruct holds the cubin module, theMappedHealthDiagSnapshot, and thehealth_diag_isv_mirrorCudaFunctionhandle. Phases 2B-2E append further handles to the same struct. Public methodlaunch_isv_mirrortakes the ISV indices by-value so the kernel never embeds them. build.rsregistershealth_diag_kernel.cuinkernels_with_common.mod.rsdeclarespub mod gpu_health_diag;.gpu_dqn_trainer.rs:- imports
super::gpu_health_diag::GpuHealthDiag; - adds
health_diag: Option<GpuHealthDiag>field; - constructor calls
GpuHealthDiag::new(&stream)?after the existing cubin loads; - adds
pub fn launch_health_diag_isv_mirror(&self) -> Result<(), MLError>that pulls every ISV index from the trainer's compile-time constants (single source of truth); - adds
pub fn health_diag_snapshot(&self) -> Option<&MappedHealthDiagSnapshot>accessor for downstream consumers.
- imports
Snapshot fields populated by Phase 2A: noisy_vsn_mag, noisy_vsn_dir,
noisy_drift_mag, noisy_drift_dir, aux_next_bar_mse, aux_regime_ce,
aux_label_scale, aux_moe_util[0..8], aux_moe_ent, aux_moe_lambda_eff,
reward_split_popart, reward_split_cf, reward_split_trail,
reward_split_micro, reward_split_opp_cost, reward_split_bonus. 22 of 147
fields. The remaining 125 are populated by Phases 2B-2D (histograms +
per-sample reductions + grad-trunk mirror).
No production callers of launch_health_diag_isv_mirror yet. Wire-up into
training_loop.rs's producer family lands in Phase 3 alongside the parallel-
shadow assertion. Producer-only by design — the orchestrator is built as
additive scaffolding the same way gpu_aux_heads.rs's AuxHeadsForwardOps
landed in Plan 4 Task 6 Commit A (handles loaded, struct in place, no
production caller yet; Commit B wires it in).
Functions to delete in Phase 4 (per Phase 0 inventory)
Tracked here so the deletion commit can be checked against the source of truth (this list) rather than against a freshly-greppable signal:
gpu_experience_collector::var_scale_epoch_meangpu_experience_collector::trail_fire_and_hold_per_maggpu_experience_collector::per_magnitude_winrate_and_variancegpu_experience_collector::reward_contrib_fractionsgpu_experience_collector::launch_reward_component_ema_inplace— KEEP (still produces ISV slots that other consumers read; only the HEALTH_DIAG-side CPU readback is in scope).- (additional
gpu_experience_collectoraux readers identified per the inventory's "6 collector methods" line — final list pinned in Phase 4 PR description after kernel 2D's masking semantics resolve which legacy methods become dead code.) gpu_backtest_evaluator::read_eval_action_distribution_per_magnitudegpu_backtest_evaluator::read_eval_action_distribution_per_directiongpu_backtest_evaluator::read_eval_intent_magnitude_distributiongpu_backtest_evaluator::read_chunked_actions_direction_distributiongpu_dqn_trainer::update_q_mag_means_cached(the 786 KB DtoH).branch_noisy_sigma_mean(agent-side method).
Open questions for Phase 1 design review
-
Single kernel vs small fused family? A single mega-kernel reading every input buffer and emitting all ~150 fields is monolithic but minimises launch overhead (one kernel, one fence). A 5-7 kernel family lets us keep block configs adapted to each input's shape (per-sample buffers want N-thread reductions; ISV mirrors are scalar copies that want 1 thread). Recommended: family of 4-5 kernels run sequentially in the captured graph —
health_diag_per_sample_reduce,health_diag_q_mag_reduce,health_diag_eval_histogram,health_diag_isv_mirror,health_diag_finalise(the last one runs__threadfence_system()once over the whole struct). -
Where does the kernel run in the captured graph? End of the per-step training graph would update the snapshot every step (waste). End of
process_epoch_boundary's producer launches (next tolaunch_h_s2_rms_ema,launch_vsn_mask_ema,launch_iqn_quantile_ema,launch_aux_heads_loss_ema) is the right cadence — once per epoch, after all upstream producers have written. -
booleans + u8 in
#[repr(C)]— confirm padding behaviour. If padding leaks bytes between[u8; 6](the controller fires) and the nextf32, the kernel must explicitly write through padding-aware offsets; safer to liftu8tou32(4-byte aligned) at a tiny ~24-byte cost. Pinned allocs round to page anyway. -
val [...]block — extendbacktest_metrics_kernelor copy frommetrics_bufinto the snapshot? The cheaper path is the latter:metrics_bufis already mapped-pinned and the consumer reads it once already. We can have a smallhealth_diag_isv_mirrorkernel pull frommetrics_buf.dev_ptrinto the snapshot's val_* slots — avoids touching the existing eval graph. This keeps the eval contract stable. -
What about
HEALTH_DIAG[...]: TERMINATED_BY_REGRESSION ...? Fault-path; values are already known scalars. Keep CPU. Out of scope.