feat(alpha): phase_e_kernels — Rust launchers for Task 9 + Task 10 cubins

Phase E.1 Task 11. Two pub(crate) launchers expose the kill-criteria
producer (Task 9) and Munchausen target augmentation (Task 10) for use
by future trainer integration:

  launch_phase_e_kill_criteria(stream, kernel, q_values_dev, action_counts_dev,
                               scalar_inputs_dev, isv_dev, batch, n_actions,
                               scratch_out_dev)
    → kicks the kill_criteria producer; caller chains apply_pearls_ad_kernel
      (n_slots=4, isv_idx_base=ALPHA_ISV_BLOCK_LO=539) to smooth into slots
      539..542 via Pearl A bootstrap + Pearl D Wiener-α.

  launch_phase_e_munchausen_target(stream, kernel, q_next_dev, q_current_dev,
                                   actions_dev, rewards_dev, dones_dev,
                                   gamma, alpha_m, tau, log_clip_min,
                                   target_out_dev, batch, n_actions)
    → one-thread-per-sample target augmentation; α_m/τ/log_clip_min are
      scalar args so a downstream ISV-driven controller can tune them.

Plan deviation: Task 11 plan-spec said "replace hardcoded n_step=32,
gamma=0.999 literals" but grep across gpu_dqn_trainer.rs found ZERO such
literals — the trainer already reads gamma via read_isv_signal_at(
GAMMA_DIR_EFF_INDEX) and epsilon via read_isv_signal_at(AUX_TRUNK_EPS_
INDEX). The actual E.1 deliverable was Rust launchers for the new
cubins, which this commit lands.

Both launchers follow the launch_apply_pearls pattern in
sp4_wiener_ema.rs — pre-loaded CudaFunction as parameter, u64 device
pointers, debug_assert! guards.

Audit doc docs/isv-slots.md updated per Invariant 7.
Tested via `cargo test -p ml --lib phase_e_kernels` — 1 compile-witness
passes. Real GPU integration test in Task 12.
This commit is contained in:
jgrusewski
2026-05-15 14:17:29 +02:00
parent b1ba41d403
commit feb2e8cc34
3 changed files with 207 additions and 0 deletions

View File

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

View File

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