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:
@@ -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,
|
||||
|
||||
188
crates/ml/src/cuda_pipeline/phase_e_kernels.rs
Normal file
188
crates/ml/src/cuda_pipeline/phase_e_kernels.rs
Normal 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);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user