perf(dqn): online softmax in compute_expected_q (kernel #3)
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) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -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);
|
||||
|
||||
@@ -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<CudaSlice<f32>>` 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.
|
||||
|
||||
Reference in New Issue
Block a user