test(dqn): standalone Thompson direction test kernel (Phase 0)

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

@@ -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)

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 <cuda_runtime.h>
+#include <math_constants.h>
+
+#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;
+}

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`