From c2210e8b95bb1b4986e38c6e01e3ecfc6ed9ef6e Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Mon, 27 Apr 2026 08:20:45 +0200 Subject: [PATCH] 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. --- .../trainers/dqn/distributional_q_tests.rs | 102 ++++++++++++++++++ crates/ml/src/trainers/dqn/mod.rs | 3 + docs/dqn-wire-up-audit.md | 9 ++ 3 files changed, 114 insertions(+) create mode 100644 crates/ml/src/trainers/dqn/distributional_q_tests.rs diff --git a/crates/ml/src/trainers/dqn/distributional_q_tests.rs b/crates/ml/src/trainers/dqn/distributional_q_tests.rs new file mode 100644 index 000000000..9a7997e1e --- /dev/null +++ b/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, + c51_probs: &CudaSlice, // [B0_SIZE * n_atoms] + atom_values: &CudaSlice, // [n_atoms] + iqn_quantiles: &CudaSlice, // [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::(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 = 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, + c51_probs: &CudaSlice, + atom_values: &CudaSlice, + iqn_quantiles: &CudaSlice, + 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::(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 = stream.clone_dtoh(&out).expect("dtoh out_dir"); + host[0] as u8 +} + +/// Helper: build CUDA stream + context. +fn make_test_stream() -> Arc { + let ctx = CudaContext::new(0).expect("CUDA context — is GPU available?"); + ctx.default_stream() +} + +/// Helper: upload a flat Vec to GPU. +fn upload(stream: &Arc, host: &[f32]) -> CudaSlice { + stream.clone_htod(host).expect("htod upload") +} + +// Tests follow in Tasks 3-8. diff --git a/crates/ml/src/trainers/dqn/mod.rs b/crates/ml/src/trainers/dqn/mod.rs index ef18bc7bd..0595098e8 100644 --- a/crates/ml/src/trainers/dqn/mod.rs +++ b/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 diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index 6075251e0..a9c205e5e 100644 --- a/docs/dqn-wire-up-audit.md +++ b/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`