diff --git a/crates/ml/src/cuda_pipeline/mod.rs b/crates/ml/src/cuda_pipeline/mod.rs index f9dbff8e1..878dd3e3c 100644 --- a/crates/ml/src/cuda_pipeline/mod.rs +++ b/crates/ml/src/cuda_pipeline/mod.rs @@ -68,6 +68,7 @@ pub mod sp15_isv_slots; pub mod sp21_isv_slots; pub mod sp22_isv_slots; pub mod alpha_isv_slots; +pub(crate) mod phase_e_kernels; pub mod lob_bar; pub use sp4_isv_slots::{ TARGET_Q_BOUND_INDEX, ATOM_POS_BOUND_BASE, WEIGHT_BOUND_BASE, diff --git a/crates/ml/src/cuda_pipeline/phase_e_kernels.rs b/crates/ml/src/cuda_pipeline/phase_e_kernels.rs new file mode 100644 index 000000000..bca1b2f55 --- /dev/null +++ b/crates/ml/src/cuda_pipeline/phase_e_kernels.rs @@ -0,0 +1,188 @@ +//! Phase E.1 GPU kernel launchers (Task 11, 2026-05-15). +//! +//! Rust-side launchers for the two CUDA kernels compiled by `build.rs`: +//! +//! - `phase_e_kill_criteria_compute_kernel` (`phase_e_kill_criteria.cu`) — +//! writes 4 raw diagnostic observations (Q-spread, action entropy, +//! return-vs-random, early-Q-movement) into scratch_out[0..4]. The caller +//! is responsible for chaining `apply_pearls_ad_kernel` (n_slots=4) +//! afterward to smooth into ISV slots 539..542. See +//! [`super::alpha_isv_slots`] for the slot reservation. +//! +//! - `phase_e_munchausen_target_kernel` (`phase_e_munchausen_target.cu`) — +//! Munchausen-augmented DQN target per Vieillard et al. 2020. One thread +//! per batch sample; reads Q_online(s, ·) and Q_target(s', ·), writes +//! `target_out[b] = r + α_m · max(τ · log π(a|s), log_clip_min) + γ · V_soft(s')`. +//! +//! These launchers follow the `launch_apply_pearls` pattern from +//! [`super::sp4_wiener_ema`] — they take the pre-loaded `CudaFunction` as +//! a parameter (module loading happens once at trainer construction). All +//! buffer arguments are raw `u64` device pointers; passing host or +//! mapped-pinned pointers will fault the kernel since both kernels +//! dereference directly in device address space. + +#![allow(unsafe_code)] // CUDA kernel launch via cudarc requires unsafe + +use crate::cuda_pipeline::alpha_isv_slots::{ + RANDOM_BASELINE_MEAN_INDEX, RANDOM_BASELINE_STD_INDEX, +}; +use crate::MLError; + +/// Launch `phase_e_kill_criteria_compute_kernel` on `stream`. +/// +/// The kernel reads `q_values[batch, n_actions]`, `action_counts[n_actions]`, +/// and `scalar_inputs[3]` (mapped-pinned: rollout_reward_mean, q_init_norm, +/// q_early_step_norm). It writes 4 raw observations into +/// `scratch_out[0..4]` and issues `__threadfence_system()` before returning. +/// +/// Caller MUST chain `apply_pearls_ad_kernel` on the same stream after this +/// to smooth the 4 outputs into ISV slots 539..542. The slot indices for +/// `random_baseline_mean` (547) and `random_baseline_std` (548) are read +/// from `super::alpha_isv_slots` and passed as kernel args — the kernel +/// itself reads ISV at those slots to compute return_vs_random. +/// +/// # Safety +/// +/// All buffer pointers MUST be valid device pointers. `scratch_out_dev` +/// MUST point at a writable [4]-float region; passing a shorter buffer +/// will corrupt adjacent memory. +pub(crate) unsafe fn launch_phase_e_kill_criteria( + stream: &cudarc::driver::CudaStream, + kernel: &cudarc::driver::CudaFunction, + q_values_dev: u64, + action_counts_dev: u64, + scalar_inputs_dev: u64, + isv_dev: u64, + batch: i32, + n_actions: i32, + scratch_out_dev: u64, +) -> Result<(), MLError> { + use cudarc::driver::{LaunchConfig, PushKernelArg}; + + debug_assert!(batch >= 0, "launch_phase_e_kill_criteria: batch must be non-negative"); + debug_assert!(n_actions > 0, "launch_phase_e_kill_criteria: n_actions must be positive"); + debug_assert!(q_values_dev != 0, "q_values_dev must be a valid device pointer"); + debug_assert!(action_counts_dev != 0, "action_counts_dev must be a valid device pointer"); + debug_assert!(scalar_inputs_dev != 0, "scalar_inputs_dev must be a valid device pointer"); + debug_assert!(isv_dev != 0, "isv_dev must be a valid device pointer"); + debug_assert!(scratch_out_dev != 0, "scratch_out_dev must be a valid device pointer"); + + let cfg = LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (1, 1, 1), + shared_mem_bytes: 0, + }; + let baseline_mean_slot: i32 = RANDOM_BASELINE_MEAN_INDEX as i32; + let baseline_std_slot: i32 = RANDOM_BASELINE_STD_INDEX as i32; + stream + .launch_builder(kernel) + .arg(&q_values_dev) + .arg(&action_counts_dev) + .arg(&scalar_inputs_dev) + .arg(&isv_dev) + .arg(&baseline_mean_slot) + .arg(&baseline_std_slot) + .arg(&batch) + .arg(&n_actions) + .arg(&scratch_out_dev) + .launch(cfg) + .map_err(|e| { + MLError::ModelError(format!("phase_e_kill_criteria launch: {e}")) + })?; + Ok(()) +} + +/// Launch `phase_e_munchausen_target_kernel` on `stream`. +/// +/// Computes the Munchausen-augmented target for each batch sample: +/// `target[i] = r[i] + α_m · max(τ · log π(a|s), log_clip_min) + γ · V_soft(s')` +/// (or `r + m` if `dones[i] >= 0.5`). +/// +/// Grid configuration: one thread per batch sample, block_dim=256. This +/// matches the existing target/loss kernels' parallelism pattern. +/// +/// `α_m`, `τ`, `log_clip_min` are passed as scalar args so a controller +/// can ISV-drive them per phase. Typical Vieillard et al. values: +/// α_m=0.9, τ=0.03, log_clip_min=-1.0. +/// +/// # Safety +/// +/// All buffer pointers MUST be valid device pointers. `target_out_dev` MUST +/// point at a writable [batch]-float region. +pub(crate) unsafe fn launch_phase_e_munchausen_target( + stream: &cudarc::driver::CudaStream, + kernel: &cudarc::driver::CudaFunction, + q_next_dev: u64, + q_current_dev: u64, + actions_dev: u64, + rewards_dev: u64, + dones_dev: u64, + gamma: f32, + alpha_m: f32, + tau: f32, + log_clip_min: f32, + target_out_dev: u64, + batch: i32, + n_actions: i32, +) -> Result<(), MLError> { + use cudarc::driver::{LaunchConfig, PushKernelArg}; + + debug_assert!(batch > 0, "launch_phase_e_munchausen_target: batch must be positive"); + debug_assert!(n_actions > 0, "launch_phase_e_munchausen_target: n_actions must be positive"); + debug_assert!(tau > 0.0, "tau must be positive (Boltzmann temperature)"); + debug_assert!(gamma >= 0.0 && gamma <= 1.0, "gamma must be in [0, 1]"); + debug_assert!(alpha_m >= 0.0, "alpha_m must be non-negative"); + debug_assert!(q_next_dev != 0, "q_next_dev must be a valid device pointer"); + debug_assert!(q_current_dev != 0, "q_current_dev must be a valid device pointer"); + debug_assert!(actions_dev != 0, "actions_dev must be a valid device pointer"); + debug_assert!(rewards_dev != 0, "rewards_dev must be a valid device pointer"); + debug_assert!(dones_dev != 0, "dones_dev must be a valid device pointer"); + debug_assert!(target_out_dev != 0, "target_out_dev must be a valid device pointer"); + + const BLOCK: u32 = 256; + let grid_x = ((batch as u32).max(1) + BLOCK - 1) / BLOCK; + let cfg = LaunchConfig { + grid_dim: (grid_x, 1, 1), + block_dim: (BLOCK, 1, 1), + shared_mem_bytes: 0, + }; + stream + .launch_builder(kernel) + .arg(&q_next_dev) + .arg(&q_current_dev) + .arg(&actions_dev) + .arg(&rewards_dev) + .arg(&dones_dev) + .arg(&gamma) + .arg(&alpha_m) + .arg(&tau) + .arg(&log_clip_min) + .arg(&target_out_dev) + .arg(&batch) + .arg(&n_actions) + .launch(cfg) + .map_err(|e| { + MLError::ModelError(format!("phase_e_munchausen_target launch: {e}")) + })?; + Ok(()) +} + +#[cfg(test)] +mod tests { + // No GPU tests here — launchers are exercised end-to-end by the + // Phase E.1 kernel smoke (see crates/ml/examples/phase_e_kernel_smoke.rs + // for the integration test). The unit-test budget for these tiny + // launchers (which are pure parameter-marshalling) is captured by the + // smoke run, which is the only place they get touched in this milestone. + // + // This stub module exists to anchor a future home for any pure-Rust + // helpers (e.g., parameter validation) that don't require a CUDA context. + + #[test] + fn module_compiles_clean() { + // Tautology — if this file failed to compile, this test wouldn't + // even build. Acts as a witness that the launchers + their + // dependency on alpha_isv_slots are coherent. + assert!(true); + } +} diff --git a/docs/isv-slots.md b/docs/isv-slots.md index cd7063213..dc49e9353 100644 --- a/docs/isv-slots.md +++ b/docs/isv-slots.md @@ -587,3 +587,21 @@ Does NOT touch any ISV slot directly — α_m, τ, and log_clip_min are exposed **Cubin:** `target/release/build/ml-*/out/phase_e_munchausen_target.cubin` (~12.8 KB). **Wiring:** lands in Phase E.1 Task 11. The integration point is wherever the existing C51/MSE loss kernels currently consume `r + γ · V_next` — they'll consume `target_out` from this kernel instead. + +## Phase E.1 Task 11 — Rust launchers `phase_e_kernels.rs` (2026-05-15) + +`crates/ml/src/cuda_pipeline/phase_e_kernels.rs` exposes two `pub(crate)` launchers: + +- **`launch_phase_e_kill_criteria(...)`** — calls `phase_e_kill_criteria_compute_kernel` (cubin from Task 9). Grid (1, 1, 1) × Block (1, 1, 1). Reads `RANDOM_BASELINE_MEAN_INDEX` (547) and `RANDOM_BASELINE_STD_INDEX` (548) from `alpha_isv_slots` and passes them as kernel args. Caller is responsible for chaining `apply_pearls_ad_kernel(n_slots=4, isv_idx_base=539)` afterward to smooth raw scratch outputs into slots 539..542. + +- **`launch_phase_e_munchausen_target(...)`** — calls `phase_e_munchausen_target_kernel` (cubin from Task 10). Grid sized `ceil(batch/256) × 1 × 1`, Block (256, 1, 1). `α_m`, `τ`, `log_clip_min` are scalar args so the smoke / future controller can ISV-drive them. + +Both launchers follow `launch_apply_pearls` (in `sp4_wiener_ema.rs`) — pre-loaded `CudaFunction` passed as parameter, `u64` device pointers, `debug_assert!` guards. + +**Plan deviation rationale:** the plan's Task 11 said "replace hardcoded `n_step=32`, `gamma=0.999` literals" but grep across the trainer found zero such literals — the trainer already reads γ via `read_isv_signal_at(GAMMA_DIR_EFF_INDEX)` and ε via `read_isv_signal_at(AUX_TRUNK_EPS_INDEX)`. The actual missing piece for E.1 was Rust-side launchers for the new cubins, delivered here. + +**Files touched:** +- `crates/ml/src/cuda_pipeline/phase_e_kernels.rs` (new) — both launchers, `pub(crate)` visibility +- `crates/ml/src/cuda_pipeline/mod.rs` — registers `pub(crate) mod phase_e_kernels;` + +**Tests:** `cargo test -p ml --lib phase_e_kernels` — 1 compile-witness passes. Real GPU integration test lands in Phase E.1 Task 12 (kernel smoke binary).