From 2186f96e830b761445aa9b2621e092ca06fd7800 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Mon, 27 Apr 2026 08:16:01 +0200 Subject: [PATCH] test(dqn): standalone Thompson direction test kernel (Phase 0) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Implements inverse-CDF over C51 atoms, uniform-τ interpolation over IQN quantiles, and argmax of E[Q] for eval mode. Plus diagnostic σ kernels (compute_sigma_c51_test, compute_sigma_iqn_test) used by Test 0.F. Exercised only by Phase 0 tests in distributional_q_tests.rs; Phase 2 production kernel will integrate the same __device__ __forceinline__ math into experience_action_select. No production callers in this commit. Plan A Task 1 of docs/superpowers/plans/2026-04-27-distributional-rl-thompson-plan-A-phase-0.md. --- crates/ml/build.rs | 7 + .../src/cuda_pipeline/thompson_test_kernel.cu | 175 ++++++++++++++++++ docs/dqn-wire-up-audit.md | 11 ++ 3 files changed, 193 insertions(+) create mode 100644 crates/ml/src/cuda_pipeline/thompson_test_kernel.cu diff --git a/crates/ml/build.rs b/crates/ml/build.rs index 77eef469a..000ea0538 100644 --- a/crates/ml/build.rs +++ b/crates/ml/build.rs @@ -180,6 +180,13 @@ fn main() { // Producer-only; ISV[AUX_NEXT_BAR_MSE_EMA_INDEX] + // ISV[AUX_REGIME_CE_EMA_INDEX] consumer wires in Commit B. "aux_heads_loss_ema_kernel.cu", + // Plan A Phase 0 (Thompson sampling spec 2026-04-26): standalone test + // kernel exercised only by `distributional_q_tests.rs`. Implements the + // Thompson direction sampling math (inverse-CDF over C51 atoms, + // uniform-τ over IQN quantiles, argmax of E[Q]) that Phase 2 will + // integrate into `experience_action_select`. Test-only; no production + // callers in this commit. + "thompson_test_kernel.cu", ]; // ALL kernels get common header (BF16 types + wrappers) diff --git a/crates/ml/src/cuda_pipeline/thompson_test_kernel.cu b/crates/ml/src/cuda_pipeline/thompson_test_kernel.cu new file mode 100644 index 000000000..0ef3038c6 --- /dev/null +++ b/crates/ml/src/cuda_pipeline/thompson_test_kernel.cu @@ -0,0 +1,175 @@ +// Standalone test-only kernel for Phase 0 verification. +// Implements the SAME math the Phase 2 production kernel will use, +// so Phase 0 tests verify the math AND Phase 2 Test 2.B verifies the +// production kernel matches this wrapper bit-for-bit (catches drift). + +#include +#include + +#ifndef N_IQN_QUANTILES +#define N_IQN_QUANTILES 5 +#endif + +// Inverse-CDF sample from C51 atom distribution. +// p[0..n_atoms] sums to ~1.0; atom_values[0..n_atoms] are the atom positions. +__device__ __forceinline__ float sample_c51_inverse_cdf( + const float* __restrict__ p, + const float* __restrict__ atom_values, + int n_atoms, + float u +) { + float cum = 0.0f; + for (int a = 0; a < n_atoms; a++) { + cum += p[a]; + if (u < cum) return atom_values[a]; + } + return atom_values[n_atoms - 1]; // numerical fallback +} + +// Sample from IQN by uniform-τ interpolation over the 5 fixed quantiles. +// q[0..5] are the values at τ in {0.05, 0.25, 0.5, 0.75, 0.95}. +__device__ __forceinline__ float sample_iqn_quantile_interp( + const float* __restrict__ q, + float u +) { + // Quantile boundaries: 0.05, 0.25, 0.50, 0.75, 0.95 + const float taus[N_IQN_QUANTILES] = {0.05f, 0.25f, 0.50f, 0.75f, 0.95f}; + + if (u <= taus[0]) return q[0]; + if (u >= taus[N_IQN_QUANTILES - 1]) return q[N_IQN_QUANTILES - 1]; + + for (int i = 0; i < N_IQN_QUANTILES - 1; i++) { + if (u <= taus[i + 1]) { + float w = (u - taus[i]) / (taus[i + 1] - taus[i]); + return q[i] * (1.0f - w) + q[i + 1] * w; + } + } + return q[N_IQN_QUANTILES - 1]; +} + +// Compute E[Q] from C51 atom probabilities + values. +__device__ __forceinline__ float compute_e_c51( + const float* __restrict__ p, + const float* __restrict__ atom_values, + int n_atoms +) { + float e = 0.0f; + for (int a = 0; a < n_atoms; a++) { + e += p[a] * atom_values[a]; + } + return e; +} + +// Compute E[Q] from IQN quantiles (mean of fixed quantiles). +__device__ __forceinline__ float compute_e_iqn(const float* __restrict__ q) { + float s = 0.0f; + for (int i = 0; i < N_IQN_QUANTILES; i++) s += q[i]; + return s / (float)N_IQN_QUANTILES; +} + +// Linear-congruential RNG (deterministic from seed). +// Used only for tests; production uses Philox. +__device__ __forceinline__ float test_rand(unsigned int* state) { + *state = (*state) * 1664525u + 1013904223u; + return (float)(*state >> 8) / (float)(1u << 24); +} + +// Test kernel: Thompson direction sampling (training mode). +// Inputs: per-direction C51 atom probs [b0_size, n_atoms] + atom values [n_atoms] +// per-direction IQN quantiles [b0_size, N_IQN_QUANTILES] +// seed for RNG state +// Output: chosen direction index [0..b0_size) +extern "C" __global__ void thompson_direction_test( + const float* __restrict__ c51_probs, + const float* __restrict__ atom_values, + const float* __restrict__ iqn_quantiles, + int b0_size, + int n_atoms, + unsigned int seed, + int* out_dir_idx +) { + if (threadIdx.x != 0 || blockIdx.x != 0) return; + + unsigned int rng_state = seed; + float best_sample = -CUDART_INF_F; + int best_d = 0; + + for (int d = 0; d < b0_size; d++) { + float u_c51 = test_rand(&rng_state); + float s_c51 = sample_c51_inverse_cdf( + c51_probs + d * n_atoms, atom_values, n_atoms, u_c51 + ); + float u_iqn = test_rand(&rng_state); + float s_iqn = sample_iqn_quantile_interp( + iqn_quantiles + d * N_IQN_QUANTILES, u_iqn + ); + // /2 dropped: argmax(a+b) == argmax((a+b)/2) + float q_sample = s_c51 + s_iqn; + if (q_sample > best_sample) { + best_sample = q_sample; + best_d = d; + } + } + + *out_dir_idx = best_d; +} + +// Test kernel: argmax of (E_C51 + E_IQN) — eval mode. +extern "C" __global__ void argmax_eq_test( + const float* __restrict__ c51_probs, + const float* __restrict__ atom_values, + const float* __restrict__ iqn_quantiles, + int b0_size, + int n_atoms, + int* out_dir_idx +) { + if (threadIdx.x != 0 || blockIdx.x != 0) return; + + float best_q = -CUDART_INF_F; + int best_d = 0; + + for (int d = 0; d < b0_size; d++) { + float e_c51 = compute_e_c51(c51_probs + d * n_atoms, atom_values, n_atoms); + float e_iqn = compute_e_iqn(iqn_quantiles + d * N_IQN_QUANTILES); + float e_joint = e_c51 + e_iqn; + if (e_joint > best_q) { best_q = e_joint; best_d = d; } + } + + *out_dir_idx = best_d; +} + +// Diagnostic kernel: compute σ from C51 atoms (used by Test 0.F). +extern "C" __global__ void compute_sigma_c51_test( + const float* __restrict__ c51_probs, + const float* __restrict__ atom_values, + int b0_size, + int n_atoms, + float* out_sigma // [b0_size] +) { + int d = blockIdx.x * blockDim.x + threadIdx.x; + if (d >= b0_size) return; + + const float* p = c51_probs + d * n_atoms; + float mean = compute_e_c51(p, atom_values, n_atoms); + float var = 0.0f; + for (int a = 0; a < n_atoms; a++) { + float diff = atom_values[a] - mean; + var += p[a] * diff * diff; + } + out_sigma[d] = sqrtf(fmaxf(var, 0.0f)); +} + +// Diagnostic kernel: compute σ from IQN quantiles (used by Test 0.F). +extern "C" __global__ void compute_sigma_iqn_test( + const float* __restrict__ iqn_quantiles, + int b0_size, + float* out_sigma // [b0_size] +) { + int d = blockIdx.x * blockDim.x + threadIdx.x; + if (d >= b0_size) return; + + const float* q = iqn_quantiles + d * N_IQN_QUANTILES; + // IQR / 1.349 ≈ σ for Gaussian-like distributions + float iqr = q[3] - q[1]; // Q_τ=0.75 - Q_τ=0.25 + out_sigma[d] = iqr / 1.349f; +} diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index 6ee764ff0..6075251e0 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -2,6 +2,17 @@ **Status:** Populated during Plan 1 Task 6 (A.5 orphan audit). Updated on every commit per Invariant 7. +Plan A Phase 0 Thompson test kernel (2026-04-27): `cuda_pipeline/thompson_test_kernel.cu` — +standalone test-only CUDA kernel implementing the Thompson direction sampling math +(inverse-CDF over C51 atoms, uniform-τ over IQN quantiles, argmax of E[Q]) plus +diagnostic σ kernels (`compute_sigma_c51_test`, `compute_sigma_iqn_test`). Exercised +only by `distributional_q_tests.rs` (Phase 0 GPU-direct hypothesis verification). +Phase 2 will inline the SAME `__device__ __forceinline__` math into +`experience_kernels.cu::experience_action_select` for the production direction-branch +action selector; this isolated test kernel is the reference that Phase 2 Test 2.B +verifies the production kernel matches bit-for-bit. No production callers in this +commit. Spec: `docs/superpowers/specs/2026-04-26-distributional-rl-aggregation-design.md`. + Persistent picked_action_history scatter (2026-04-26): `gpu_backtest_evaluator.rs` gains a window-major `picked_action_history_buf` (parallel layout to `actions_history_buf`) populated by an additional `scatter_intent_chunk`