test(dqn): Phase 0 test module skeleton + kernel wrappers

Adds distributional_q_tests.rs with launch_thompson_direction and
launch_argmax_eq wrappers around thompson_test_kernel.cu. Tests follow
in subsequent tasks.

crates/ml/src/trainers/dqn/distributional_q_tests.rs

@@ -0,0 +1,102 @@
+//! Phase 0 GPU-direct verification tests for the Thompson sampling
+//! design (spec: docs/superpowers/specs/2026-04-26-distributional-rl-aggregation-design.md).
+//!
+//! All tests in this module are `#[ignore]` because they require a CUDA
+//! GPU. Run them with:
+//!     cargo test -p ml --lib distributional_q_tests -- --ignored --nocapture
+
+use std::sync::Arc;
+use cudarc::driver::{CudaContext, CudaSlice, CudaStream, LaunchConfig, PushKernelArg};
+
+const TEST_KERNEL_CUBIN: &[u8] =
+    include_bytes!(concat!(env!("OUT_DIR"), "/thompson_test_kernel.cubin"));
+
+const N_IQN_QUANTILES: usize = 5;
+const B0_SIZE: usize = 4;  // 4-direction action space
+
+/// Launch `thompson_direction_test` with provided inputs; return chosen direction.
+fn launch_thompson_direction(
+    stream: &Arc<CudaStream>,
+    c51_probs: &CudaSlice<f32>,    // [B0_SIZE * n_atoms]
+    atom_values: &CudaSlice<f32>,  // [n_atoms]
+    iqn_quantiles: &CudaSlice<f32>, // [B0_SIZE * N_IQN_QUANTILES]
+    n_atoms: i32,
+    seed: u32,
+) -> u8 {
+    let module = stream.context()
+        .load_cubin(TEST_KERNEL_CUBIN.to_vec())
+        .expect("load thompson_test_kernel cubin");
+    let kernel = module.load_function("thompson_direction_test")
+        .expect("load thompson_direction_test");
+
+    let mut out = stream.alloc_zeros::<i32>(1).expect("alloc out_dir");
+
+    unsafe {
+        stream.launch_builder(&kernel)
+            .arg(c51_probs)
+            .arg(atom_values)
+            .arg(iqn_quantiles)
+            .arg(&(B0_SIZE as i32))
+            .arg(&n_atoms)
+            .arg(&seed)
+            .arg(&mut out)
+            .launch(LaunchConfig {
+                grid_dim: (1, 1, 1),
+                block_dim: (1, 1, 1),
+                shared_mem_bytes: 0,
+            })
+            .expect("launch thompson_direction_test");
+    }
+
+    let host: Vec<i32> = stream.clone_dtoh(&out).expect("dtoh out_dir");
+    host[0] as u8
+}
+
+/// Launch `argmax_eq_test`; return chosen direction (eval mode).
+fn launch_argmax_eq(
+    stream: &Arc<CudaStream>,
+    c51_probs: &CudaSlice<f32>,
+    atom_values: &CudaSlice<f32>,
+    iqn_quantiles: &CudaSlice<f32>,
+    n_atoms: i32,
+) -> u8 {
+    let module = stream.context()
+        .load_cubin(TEST_KERNEL_CUBIN.to_vec())
+        .expect("load thompson_test_kernel cubin");
+    let kernel = module.load_function("argmax_eq_test")
+        .expect("load argmax_eq_test");
+
+    let mut out = stream.alloc_zeros::<i32>(1).expect("alloc out_dir");
+
+    unsafe {
+        stream.launch_builder(&kernel)
+            .arg(c51_probs)
+            .arg(atom_values)
+            .arg(iqn_quantiles)
+            .arg(&(B0_SIZE as i32))
+            .arg(&n_atoms)
+            .arg(&mut out)
+            .launch(LaunchConfig {
+                grid_dim: (1, 1, 1),
+                block_dim: (1, 1, 1),
+                shared_mem_bytes: 0,
+            })
+            .expect("launch argmax_eq_test");
+    }
+
+    let host: Vec<i32> = stream.clone_dtoh(&out).expect("dtoh out_dir");
+    host[0] as u8
+}
+
+/// Helper: build CUDA stream + context.
+fn make_test_stream() -> Arc<CudaStream> {
+    let ctx = CudaContext::new(0).expect("CUDA context — is GPU available?");
+    ctx.default_stream()
+}
+
+/// Helper: upload a flat Vec<f32> to GPU.
+fn upload(stream: &Arc<CudaStream>, host: &[f32]) -> CudaSlice<f32> {
+    stream.clone_htod(host).expect("htod upload")
+}
+
+// Tests follow in Tasks 3-8.

crates/ml/src/trainers/dqn/mod.rs

@@ -35,6 +35,9 @@ mod statistics;
 mod trainer;
 mod smoke_tests;
 pub mod state_reset_registry;
+
+#[cfg(test)]
+pub mod distributional_q_tests;
 pub use state_reset_registry::{ResetCategory, RegistryEntry, StateResetRegistry};
 
 // Re-export all public items for backward compatibility

docs/dqn-wire-up-audit.md

@@ -13,6 +13,15 @@ action selector; this isolated test kernel is the reference that Phase 2 Test 2.
 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`.
 
+Plan A Phase 0 Task 2 test wrappers (2026-04-27): `trainers/dqn/distributional_q_tests.rs` —
+test-only Rust module gated `#[cfg(test)] pub mod` in `trainers/dqn/mod.rs`. Wraps the
+Task 1 cubin via `include_bytes!(concat!(env!("OUT_DIR"), "/thompson_test_kernel.cubin"))`
+with two launchers (`launch_thompson_direction`, `launch_argmax_eq`) plus
+`make_test_stream` / `upload` helpers. Skeleton only; actual `#[ignore]`-gated
+verification tests (0.A bias-reproduces, 0.B inverse-CDF, 0.C IQN symmetry, 0.D
+Thompson-reverses, 0.E synthetic edge, 0.F converged-checkpoint) land in Tasks 3-8.
+No production callers. Plan: `docs/superpowers/plans/2026-04-27-distributional-rl-thompson-plan-A-phase-0.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`