From 5f26c57514edf72a8d5f05de7a0843342d3d4c7d Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Tue, 28 Apr 2026 23:33:14 +0200 Subject: [PATCH] perf(dqn): online softmax in compute_expected_q (kernel #3) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit nsys profile of multi_fold_convergence on L40S identified compute_expected_q as the #3 GPU consumer at 12.9% (207ms / 1382 calls — ~150 µs per call, compute-bound at typical batch=8192 × num_atoms=51). The kernel computed each per-action softmax over atoms in 3 passes (max → sum_exp → normalize+expected +var+entropy+util), re-reading v_row[z]+adv_a[z] 3 times per atom for each of the 4 × ~3.25 = 13 actions per sample. Replaced with online softmax (running-max + running-sum) so expected_q, sum_z_sq, and entropy all accumulate in a single forward pass over atoms. Atom utilisation still needs a 2nd pass because it depends on the final S = sum exp(logit-M) being known to compute prob_z = exp(logit_z-M)/S for the threshold test. Inner loop reduction: 3 → 2 passes, ~33% fewer global loads of branch advantage logits (which dominate because v_row[z] is reused across 13 (action, branch) pairs and may stay in L1, while adv_a[z] flips per action and is pure global). Online accumulation pattern (Page-Olshen): M, S, TZ, TZ², TLM = -inf, 0, 0, 0, 0 for z in 0..num_atoms: logit = v_row[z] + adv_a[z] if logit > M: scale = exp(M - logit) # ≤ 1, no overflow S, TZ, TZ², TLM *= scale M = logit w = exp(logit - M) S += w TZ += w * z_val TZ² += w * z_val² TLM += w * (logit - M) expected_q = TZ / S sum_z_sq = TZ² / S entropy = log(S) - TLM / S # analytical: -sum(p log p) Correctness: bit-stable when atoms processed in fixed order (z = 0..num_atoms-1). The analytical entropy form is mathematically identical to -sum(p log p): -sum(p log p) = log(S) - (1/S) * sum(exp(logit-M) * (logit-M)) = log(S) - TLM / S The previous code's `if (prob > 1e-10f)` underflow guard is no longer needed: exp(logit - M) underflows cleanly to 0 when logit << M, and the analytical form does not multiply tiny probs by very negative logs. Bit-stable per feedback_stop_on_anomaly.md — no fuzzy tolerance change. Atom-stat block-sum reduction unchanged (warp shfl + shared-mem bank, deterministic). Expected wall-clock saving: ~33% of 207ms = ~70ms across the smoke run; on production batch=8192 × 1382 calls per fold, roughly proportional. Lower-bound estimate — at low batch (smoke) the kernel may be launch-bound rather than load-bound. nsys re-profile after deployment will quantify. ABI unchanged; no Rust caller changes required. Per Invariant 7 the audit doc dqn-gpu-hot-path-audit.md is updated with Fix 19 entry. Build: SQLX_OFFLINE=true cargo check -p ml --lib clean (12 warnings, baseline). Tests: cargo test -p ml --lib --no-run clean. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../src/cuda_pipeline/experience_kernels.cu | 97 +++++++++++++------ docs/dqn-gpu-hot-path-audit.md | 38 ++++++++ 2 files changed, 106 insertions(+), 29 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/experience_kernels.cu b/crates/ml/src/cuda_pipeline/experience_kernels.cu index 13191ec7c..670cb4b59 100644 --- a/crates/ml/src/cuda_pipeline/experience_kernels.cu +++ b/crates/ml/src/cuda_pipeline/experience_kernels.cu @@ -3258,43 +3258,82 @@ extern "C" __global__ void compute_expected_q( int support_base = i * 12 + d * 3; float v_min = per_sample_support[support_base + 0]; float dz = per_sample_support[support_base + 2]; + const float* atom_pos_d = (atom_positions != NULL) + ? atom_positions + (long long)d * num_atoms + : NULL; for (int a = 0; a < n_d; a++) { const float* adv_a = branch_base + (long long)a * num_atoms; - /* Combine value + advantage. Softmax + expectation over atoms. */ - float max_logit = -1e30f; + /* ── Online softmax: 2-pass instead of 3-pass ──────────────── + * Pass 1: streaming (running_max, running_sum_exp_shifted) plus + * running accumulators for E[Z], E[Z²], E[Z*(logit - max)] — + * the last gives entropy via: + * -sum(p log p) = log(S) - (1/S) * sum(exp(logit-M)*(logit-M)) + * where S = sum_exp_shifted (reduces 3rd pass to a final + * normalisation step). + * + * Online rescale: when running_max updates by delta = old - new, + * each accumulator multiplied by exp(delta) (delta <= 0 → factor + * in [0,1], no overflow). Bit-stable when atoms are processed + * in fixed order (z = 0..num_atoms-1). + * + * Pass 2: atom utilisation (need final S = sum_exp_shifted to + * compute prob_z = exp(logit-M)/S and threshold-test). One + * extra pass instead of folding into Pass 1 because util + * depends on S which is only known after Pass 1 completes. + * + * Saving: 3*num_atoms reads of v_row[z]+adv_a[z] → 2*num_atoms. + * For num_atoms=51 the inner action loop drops by ~33% global + * loads against the inner advantage logits (which dominate + * because v_row may stay in L1 across the 4×n_d action loop; + * adv_a flips per action — pure global). */ + float M = -1e30f; + float S = 0.0f; /* sum exp(logit - M) */ + float TZ = 0.0f; /* sum exp(logit - M) * z_val */ + float TZ2 = 0.0f; /* sum exp(logit - M) * z_val^2 */ + float TLM = 0.0f; /* sum exp(logit - M) * (logit - M) */ for (int z = 0; z < num_atoms; z++) { - float logit = (v_row[z]) + (adv_a[z]); - if (logit > max_logit) max_logit = logit; - } - - float sum_exp = 0.0f; - for (int z = 0; z < num_atoms; z++) { - float logit = (v_row[z]) + (adv_a[z]); - sum_exp += expf(logit - max_logit); - } - - float expected_q = 0.0f; - float sum_z_sq = 0.0f; - for (int z = 0; z < num_atoms; z++) { - float logit = (v_row[z]) + (adv_a[z]); - float prob = expf(logit - max_logit) / sum_exp; - float z_val; - if (atom_positions != NULL) { - z_val = atom_positions[(long long)d * num_atoms + z]; - } else { - z_val = v_min + (float)z * dz; + float logit = v_row[z] + adv_a[z]; + float z_val = (atom_pos_d != NULL) + ? atom_pos_d[z] + : v_min + (float)z * dz; + if (logit > M) { + /* Rescale running accumulators down to the new max. */ + float scale = expf(M - logit); + S *= scale; + TZ *= scale; + TZ2 *= scale; + TLM *= scale; + M = logit; } - expected_q += prob * z_val; - sum_z_sq += prob * z_val * z_val; - /* Atom utilization: count atoms carrying meaningful mass */ - if (prob > util_threshold) sample_utilized++; - /* Atom entropy: -sum(p * log p) */ - if (prob > 1e-10f) sample_entropy -= prob * logf(prob); + float w = expf(logit - M); /* exp(logit_z - M_current) */ + S += w; + TZ += w * z_val; + TZ2 += w * z_val * z_val; + TLM += w * (logit - M); } - q_values[(long long)i * total_actions + action_idx] = (expected_q); + float inv_S = 1.0f / S; + float expected_q = TZ * inv_S; + float sum_z_sq = TZ2 * inv_S; + /* Entropy = -sum(p log p) = -sum(p (logit - M - log S)) + * = log(S) - sum(p (logit - M)) + * = log(S) - TLM / S */ + float branch_entropy = logf(S) - TLM * inv_S; + sample_entropy += branch_entropy; + + /* Pass 2: atom utilisation — count atoms with prob > threshold. + * prob_z = exp(logit_z - M)/S > util_threshold + * ⇔ exp(logit_z - M) > util_threshold * S + * The branch is identical to the prior 3-pass version. */ + float util_cutoff = util_threshold * S; + for (int z = 0; z < num_atoms; z++) { + float logit = v_row[z] + adv_a[z]; + if (expf(logit - M) > util_cutoff) sample_utilized++; + } + + q_values[(long long)i * total_actions + action_idx] = expected_q; /* Var[Q] = E[Z²] - E[Z]² — distributional uncertainty per action */ if (q_variance != NULL) { float var_q = fmaxf(sum_z_sq - expected_q * expected_q, 0.0f); diff --git a/docs/dqn-gpu-hot-path-audit.md b/docs/dqn-gpu-hot-path-audit.md index 675591a0c..558855b2b 100644 --- a/docs/dqn-gpu-hot-path-audit.md +++ b/docs/dqn-gpu-hot-path-audit.md @@ -209,3 +209,41 @@ Per-chunk launch reduction: from `2 * chunk_len + chunk_len * 7` to `1 + 7` for Borrow restructure: the prior single-shot `let ch_states = self.chunked_states_buf.as_ref()?` at top of `submit_dqn_step_loop_cublas` was removed because TLOB now needs `&mut self.chunked_states_buf`. Replaced with per-chunk `ch_states_base` (raw u64 device pointer extracted in a tight scope), reused for both `launch_gather_chunk` and the Phase 2+3 `compute_q_values_to` call. The pointer is stable across the chunk because the `Option>` does not reallocate. No new HtoD/DtoH paths introduced. Strictly removes 512 per-chunk DtoD copies. Per `feedback_no_partial_refactor.md`: kernel ABI is new (separate function name), prior `backtest_state_gather` retained for `evaluate()` / `evaluate_ppo()` / `evaluate_supervised()` paths that still need a single-step writer. + +### Fix 19 (perf) — Online softmax in compute_expected_q: 3-pass → 2-pass (2026-04-28) +`experience_kernels.cu`: nsys profile of `multi_fold_convergence` on L40S identified `compute_expected_q` as the #3 GPU consumer at 12.9% (207 ms / 1382 calls — ~150 µs per call, compute-bound at typical batch=8192 × num_atoms=51). The kernel computed each per-action softmax over atoms in 3 passes (max → sum_exp → normalize+expected/var/entropy/util), re-reading `v_row[z] + adv_a[z]` 3 times per atom for each of the 4×~3.25 = 13 actions per sample. + +Replaced with **online softmax** (Page-Olshen running-max + running-sum) so expected_q, sum_z_sq, and entropy all accumulate in a single forward pass over atoms. Atom utilisation still needs a 2nd pass because it depends on the final `S = sum exp(logit-M)` being known to compute `prob_z = exp(logit_z - M)/S` for the threshold test. Total: 3 → 2 passes, ~33% reduction in inner-loop global loads of branch advantage logits (which dominate because `v_row[z]` is reused across 13 (action, branch) pairs and may stay in L1, while `adv_a[z]` flips per action and is pure global). + +Online accumulation: +``` +M, S, TZ, TZ², TLM = -inf, 0, 0, 0, 0 +for z in 0..num_atoms: + logit = v_row[z] + adv_a[z] + if logit > M: + scale = exp(M - logit) # ≤ 1, no overflow + S, TZ, TZ², TLM *= scale + M = logit + w = exp(logit - M) # ≤ 1 + S += w + TZ += w * z_val + TZ² += w * z_val² + TLM += w * (logit - M) +expected_q = TZ / S +sum_z_sq = TZ² / S +entropy = log(S) - TLM / S # analytical: -sum(p log p) where p = w/S +``` + +Correctness: bit-stable when atoms processed in fixed order (z = 0..num_atoms-1). The analytical entropy form is mathematically equivalent to `-sum(p log p)`: +``` +-sum(p log p) = -sum(p (logit - M - log S)) = log(S) * sum(p) - sum(p (logit-M)) + = log(S) - (1/S) * sum(exp(logit-M) * (logit-M)) + = log(S) - TLM / S +``` +The previous code's `if (prob > 1e-10f)` underflow guard is no longer needed: `exp(logit - M)` underflows cleanly to 0 when logit << M, and the analytical form does not multiply tiny probs by very negative logs. Result is bit-stable per `feedback_stop_on_anomaly.md` — no fuzzy tolerance change. + +Atom-stat block-sum reduction unchanged (warp shfl + shared-mem bank, deterministic). + +Expected wall-clock saving on this kernel: ~33% of 207ms = ~70ms across the smoke run; on production batch=8192 × 1382 calls per fold, roughly proportional. Lower-bound estimate; actual saving depends on memory bandwidth vs compute mix — at low N (smoke) the kernel may be launch-bound rather than load-bound. nsys re-profile after deployment will quantify. + +ABI unchanged; no Rust caller changes required.