gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(kernel-step-trace): runtime --kernel-step-trace CLI flag + JSONL drain

Browse Source 
The compile-time feature kernel-step-trace gates inclusion of the ring
code. NEW: --kernel-step-trace <PATH> CLI flag on alpha_train gates
RUNTIME activation:

- Path provided   -> ring allocated, drain spawns, JSONL records written
                    to <PATH> (truncated). One record per kernel emission.
- Path omitted    -> ring not allocated, drain not spawned, zero overhead
                    even with the feature compiled in.

JSONL schema: {"step": u32, "kid": u8, "kname": str, "rt": u8,
"rt_name": str, "payload": {field: f32, ...}}. The payload field names
come from the existing per-(kid, rt) decoders in gpu_log.rs (ported to
serde_json::json! in decode_to_json).

Trainer ring fields are now Option<_>; populated only when the runtime
trace path is Some. Tick kernel, step-counter shadow, smoothness
controller pointer-passing, and Drop all become path-conditional.

The tracing::info!-based drain variant is removed (file-writer is
strictly more useful; tests migrated to assert JSONL file contents).

Argo template:
- New parameter kernel-step-trace-enable (build-time feature opt-in)
- New parameter kernel-step-trace-path (runtime CLI value)
- ensure-binary cache-busts when feature toggles
- training step passes --kernel-step-trace flag conditionally

Per feedback_no_feature_flags: compile-time gate retained because the
ring carries real memory cost (~2 MiB pinned) and per-step write overhead;
the specific name kernel-step-trace narrows scope to this mechanism.
Runtime gate is Option<PathBuf>, not a boolean enable_*.
This commit is contained in:
jgrusewski
2026-05-21 01:55:57 +02:00
parent fa41b39ca4
commit 2d5a66f6e4
6 changed files with 345 additions and 142 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
crates/ml-alpha/examples/alpha_train.rs
View File

@@ -161,6 +161,17 @@ struct Cli {
/// `docs/superpowers/specs/2026-05-20-crt-train-output-smoothness-design.md`.
#[arg(long, default_value_t = 0.0)]
smoothness_base_lambda: f32,
/// Enable per-step kernel-state JSONL trace to the given file.
///
/// When set, the trainer allocates the GPU log ring + spawns a drain
/// task that writes one JSONL record per kernel emission to this path.
/// Requires the `kernel-step-trace` Cargo feature at compile time.
/// When unset (default), no ring allocated, zero overhead.
///
/// Example: `--kernel-step-trace /feature-cache/.../step_trace.jsonl`
#[arg(long)]
kernel_step_trace: Option<std::path::PathBuf>,
}
#[derive(Serialize, serde::Deserialize, Default)]
@@ -275,6 +286,7 @@ fn main() -> Result<()> {
horizon_weights,
n_batch: cli.batch_size,
smoothness_base_lambda: cli.smoothness_base_lambda,
kernel_step_trace_path: cli.kernel_step_trace.clone(),
};
let mut trainer = PerceptionTrainer::new(&dev, &trainer_cfg).context("trainer init")?;

196
crates/ml-alpha/src/gpu_log.rs
View File

@@ -2,23 +2,28 @@
//!
//! See `docs/superpowers/specs/2026-05-21-gpu-log-ring-design.md`.
//!
//! Behaviour: when the `kernel-step-trace` feature is enabled, the trainer
//! allocates a mapped-pinned ring buffer + a 1-element device step counter,
//! passes their pointers into logging-capable kernels, and starts a tokio
//! drain task that polls the ring at 500 ms cadence and emits structured
//! tracing events.
//! Behaviour: when the `kernel-step-trace` feature is enabled AND the
//! trainer is invoked with `--kernel-step-trace <PATH>`, the trainer
//! allocates a mapped-pinned ring buffer + a 1-element device step
//! counter, passes their pointers into logging-capable kernels, and
//! starts a tokio drain task that polls the ring at 500 ms cadence and
//! writes structured JSONL records to `<PATH>`.
//!
//! Constants below mirror `crates/ml-alpha/cuda/gpu_log_ids.h` — single
//! source of truth; manual mirror with on-build verification.
use std::fs::File;
use std::io::{BufWriter, Write};
use std::path::PathBuf;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::Duration;
use std::time::{Duration, Instant};
use anyhow::{Context, Result};
use cudarc::driver::{CudaSlice, CudaStream};
use serde_json::{json, Value};
use tokio::task::JoinHandle;
use tracing::{debug, info, warn};
use tracing::{debug, warn};
use crate::pinned_mem::MappedRecordBuffer;
@@ -126,82 +131,126 @@ pub fn alloc_step_counter(stream: &Arc<CudaStream>) -> Result<CudaSlice<i32>> {
stream.alloc_zeros::<i32>(1).context("step_counter alloc")
}
// ─── Decoder dispatch ────────────────────────────────────────────────
// ─── JSONL serializer ────────────────────────────────────────────────
//
// Each (kernel_id, record_type) pair maps to a fixed schema; the
// `payload` field names below come from the kernel-side emission
// sites (mirror of `smoothness_lambda_controller.cu`'s `log_record()`
// calls). Adding a new logging-capable kernel = adding a new arm.
fn default_decoder(step: u32, kernel_id: u8, record_type: u8, payload: &[f32]) {
debug!(target: "gpu_log", step, kernel_id, record_type, ?payload, "unknown record");
}
fn smoothness_input(step: u32, payload: &[f32]) {
if payload.len() >= 10 {
info!(target: "gpu_log",
step,
raw_h30 = payload[0], raw_h100 = payload[1], raw_h300 = payload[2],
raw_h1000 = payload[3], raw_h6000 = payload[4],
jitter_in_h30 = payload[5], jitter_in_h100 = payload[6],
jitter_in_h300 = payload[7], jitter_in_h1000 = payload[8],
jitter_in_h6000 = payload[9],
"smoothness_ctrl RT_INPUT");
/// Kernel-name mirror for the `kname` JSONL field.
fn kernel_name(kernel_id: u8) -> &'static str {
match kernel_id {
KID_SMOOTHNESS_CONTROLLER => "smoothness_controller",
KID_OUTPUT_SMOOTHNESS => "output_smoothness",
KID_BCE_MULTI_HORIZON => "bce_multi_horizon",
KID_HORIZON_LAMBDA => "horizon_lambda",
KID_HEADS_GRN_FWD => "heads_grn_fwd",
KID_HEADS_GRN_BWD => "heads_grn_bwd",
KID_CFC_STEP => "cfc_step",
KID_RESERVED_7 => "reserved_7",
_ => "unknown",
}
}
fn smoothness_state(step: u32, payload: &[f32]) {
if payload.len() >= 10 {
info!(target: "gpu_log",
step,
ema_h30 = payload[0], ema_h100 = payload[1], ema_h300 = payload[2],
ema_h1000 = payload[3], ema_h6000 = payload[4],
target_h30 = payload[5], target_h100 = payload[6], target_h300 = payload[7],
target_h1000 = payload[8], target_h6000 = payload[9],
"smoothness_ctrl RT_STATE");
/// Record-type-name mirror for the `rt_name` JSONL field.
fn rt_name(record_type: u8) -> &'static str {
match record_type {
RT_INPUT => "input",
RT_STATE => "state",
RT_OUTPUT => "output",
RT_DEBUG => "debug",
_ => "unknown",
}
}
fn smoothness_output(step: u32, payload: &[f32]) {
if payload.len() >= 10 {
info!(target: "gpu_log",
step,
excess_h30 = payload[0], excess_h100 = payload[1], excess_h300 = payload[2],
excess_h1000 = payload[3], excess_h6000 = payload[4],
lambda_h30 = payload[5], lambda_h100 = payload[6], lambda_h300 = payload[7],
lambda_h1000 = payload[8], lambda_h6000 = payload[9],
"smoothness_ctrl RT_OUTPUT");
}
}
/// Static dispatch table indexed by (kernel_id, record_type).
/// Add per-(kid, rt) decoders here as new kernels opt into the ring.
pub fn decode(kernel_id: u8, record_type: u8, step: u32, payload: &[f32]) {
/// Build the per-(kid, rt) payload map. Returns a JSON object with
/// field names matching the kernel-side emission contract. Unknown
/// (kid, rt) pairs fall through to an array-of-floats payload so no
/// records are silently dropped.
fn payload_json(kernel_id: u8, record_type: u8, payload: &[f32]) -> Value {
match (kernel_id, record_type) {
(KID_SMOOTHNESS_CONTROLLER, RT_INPUT) => smoothness_input(step, payload),
(KID_SMOOTHNESS_CONTROLLER, RT_STATE) => smoothness_state(step, payload),
(KID_SMOOTHNESS_CONTROLLER, RT_OUTPUT) => smoothness_output(step, payload),
_ => default_decoder(step, kernel_id, record_type, payload),
(KID_SMOOTHNESS_CONTROLLER, RT_INPUT) if payload.len() >= 10 => json!({
"raw_h30": payload[0], "raw_h100": payload[1], "raw_h300": payload[2],
"raw_h1000": payload[3], "raw_h6000": payload[4],
"jitter_in_h30": payload[5], "jitter_in_h100": payload[6],
"jitter_in_h300": payload[7], "jitter_in_h1000": payload[8],
"jitter_in_h6000": payload[9],
}),
(KID_SMOOTHNESS_CONTROLLER, RT_STATE) if payload.len() >= 10 => json!({
"ema_h30": payload[0], "ema_h100": payload[1], "ema_h300": payload[2],
"ema_h1000": payload[3], "ema_h6000": payload[4],
"target_h30": payload[5], "target_h100": payload[6], "target_h300": payload[7],
"target_h1000": payload[8], "target_h6000": payload[9],
}),
(KID_SMOOTHNESS_CONTROLLER, RT_OUTPUT) if payload.len() >= 10 => json!({
"excess_h30": payload[0], "excess_h100": payload[1], "excess_h300": payload[2],
"excess_h1000": payload[3], "excess_h6000": payload[4],
"lambda_h30": payload[5], "lambda_h100": payload[6], "lambda_h300": payload[7],
"lambda_h1000": payload[8], "lambda_h6000": payload[9],
}),
_ => {
// Unknown (kid, rt) — preserve the raw payload as a "values"
// array so the JSONL line is still well-formed and the
// record isn't silently lost.
let values: Vec<Value> = payload.iter().map(|v| json!(*v)).collect();
json!({ "values": values })
}
}
}
/// Build the full JSONL record for one (kid, rt, step, payload) tuple.
pub fn decode_to_json(kernel_id: u8, record_type: u8, step: u32, payload: &[f32]) -> Value {
json!({
"step": step,
"kid": kernel_id,
"kname": kernel_name(kernel_id),
"rt": record_type,
"rt_name": rt_name(record_type),
"payload": payload_json(kernel_id, record_type, payload),
})
}
// ─── Drain task ──────────────────────────────────────────────────────
/// Spawn the drain task. Polls the step counter (via a mapped-pinned host
/// shadow that some other component updates each tick — typically a
/// secondary memcpy_dtod_async in `step_batched`, or a separate background
/// reader) at 500 ms cadence; for each new step in
/// `[last_drained..current]`, walks `N_RECORDS_PER_STEP` slots, validates
/// `magic == MAGIC` and `header.step == expected_step`, dispatches to
/// `decode()`. Drops + warns if the drainer falls behind the ring window.
/// Spawn the drain task with a JSONL file sink. Opens `output_path`
/// (truncating any existing file), wraps in a `BufWriter`, and writes
/// one JSON object per line for every valid record observed. Polls
/// the step counter (via a mapped-pinned host shadow updated each
/// captured-graph replay) at 500 ms cadence; for each new step in
/// `[last_drained..current]`, walks `N_RECORDS_PER_STEP` slots,
/// validates `magic == MAGIC` and `header.step == expected_step`,
/// and serializes to the writer. Drops + warns if the drainer falls
/// behind the ring window.
///
/// The writer is flushed every 1024 records or every 500 ms,
/// whichever comes first, and flushed one final time on stop.
///
/// Stops cleanly when `stop` is set.
pub fn spawn_drain_task(
pub fn spawn_drain_task_to_file(
ring: Arc<LogRing>,
step_counter_host_shadow: Arc<MappedRecordBuffer<i32>>,
stop: Arc<AtomicBool>,
) -> JoinHandle<()> {
tokio::spawn(async move {
output_path: PathBuf,
) -> Result<JoinHandle<()>> {
let file = File::create(&output_path)
.with_context(|| format!("kernel-step-trace: create {}", output_path.display()))?;
let mut writer = BufWriter::new(file);
Ok(tokio::spawn(async move {
let mut last_drained_step: u32 = 0;
let interval = Duration::from_millis(500);
let mut records_since_flush: usize = 0;
let mut last_flush = Instant::now();
const FLUSH_RECORDS: usize = 1024;
const FLUSH_INTERVAL: Duration = Duration::from_millis(500);
loop {
if stop.load(Ordering::Relaxed) {
debug!(target: "gpu_log", last_drained_step, "drain task exiting on stop signal");
if let Err(e) = writer.flush() {
warn!(target: "gpu_log", error = %e, "final flush failed");
}
break;
}
tokio::time::sleep(interval).await;
@@ -243,12 +292,33 @@ pub fn spawn_drain_task(
continue;
}
let n = (rec.header.payload_words as usize).min(PAYLOAD_F32);
decode(rec.header.kernel_id, rec.header.record_type,
rec.header.step, &rec.payload[..n]);
let value = decode_to_json(
rec.header.kernel_id,
rec.header.record_type,
rec.header.step,
&rec.payload[..n],
);
if let Err(e) = serde_json::to_writer(&mut writer, &value) {
warn!(target: "gpu_log", error = %e, "JSONL serialize failed");
continue;
}
if let Err(e) = writeln!(&mut writer) {
warn!(target: "gpu_log", error = %e, "JSONL newline write failed");
continue;
}
records_since_flush += 1;
}
}
}
last_drained_step = current_step;
if records_since_flush >= FLUSH_RECORDS || last_flush.elapsed() >= FLUSH_INTERVAL {
if let Err(e) = writer.flush() {
warn!(target: "gpu_log", error = %e, "periodic flush failed");
}
records_since_flush = 0;
last_flush = Instant::now();
}
}
})
}))
}

191
crates/ml-alpha/src/trainer/perception.rs
View File

@@ -94,6 +94,13 @@ pub struct PerceptionTrainerConfig {
/// regularization is negligible). Set to e.g. 0.01 to enable.
/// See `docs/superpowers/specs/2026-05-20-crt-train-output-smoothness-design.md`.
pub smoothness_base_lambda: f32,
/// Optional path for the kernel-step-trace JSONL drain. None →
/// ring not allocated, drain task not spawned. Some(path) → ring
/// allocated + drain writes JSONL records to `path` (truncated).
/// Per `feedback_no_feature_flags`: gated by the compile-time
/// `kernel-step-trace` feature; specific name justifies the gate.
pub kernel_step_trace_path: Option<std::path::PathBuf>,
}
impl Default for PerceptionTrainerConfig {
@@ -107,6 +114,7 @@ impl Default for PerceptionTrainerConfig {
horizon_weights: [1.0; N_HORIZONS],
n_batch: 1,
smoothness_base_lambda: 0.0,
kernel_step_trace_path: None,
}
}
}
@@ -317,26 +325,30 @@ pub struct PerceptionTrainer {
smoothness_controller_fn: CudaFunction,
_smoothness_controller_module: Arc<CudaModule>,
// ── GPU log ring (feature: kernel-step-trace) ──
// When enabled, the trainer allocates a mapped-pinned ring + a device
// When the feature is compiled AND `cfg.kernel_step_trace_path` is
// Some(path), the trainer allocates a mapped-pinned ring + a device
// step counter, passes their pointers into logging-capable kernels
// (currently `smoothness_lambda_controller`), and spawns a tokio drain
// task that polls the ring at 500 ms cadence and emits structured
// tracing events. See `docs/superpowers/specs/2026-05-21-gpu-log-ring-design.md`.
// task that polls the ring at 500 ms cadence and writes JSONL records
// to `path`. When the path is None, every Option below stays None —
// no allocations, no tick kernel, no drain task, zero overhead even
// with the feature compiled in.
// See `docs/superpowers/specs/2026-05-21-gpu-log-ring-design.md`.
#[cfg(feature = "kernel-step-trace")]
log_ring: std::sync::Arc<crate::gpu_log::LogRing>,
log_ring: Option<std::sync::Arc<crate::gpu_log::LogRing>>,
#[cfg(feature = "kernel-step-trace")]
log_step_counter_d: CudaSlice<i32>,
log_step_counter_d: Option<CudaSlice<i32>>,
/// Mapped-pinned host shadow of `log_step_counter_d`. Updated via
/// memcpy_dtod_async inside the captured region; read by the drain
/// task each tick.
#[cfg(feature = "kernel-step-trace")]
log_step_counter_host_shadow: std::sync::Arc<crate::pinned_mem::MappedRecordBuffer<i32>>,
log_step_counter_host_shadow: Option<std::sync::Arc<crate::pinned_mem::MappedRecordBuffer<i32>>>,
#[cfg(feature = "kernel-step-trace")]
gpu_log_tick_fn: CudaFunction,
gpu_log_tick_fn: Option<CudaFunction>,
#[cfg(feature = "kernel-step-trace")]
_gpu_log_module: Arc<CudaModule>,
_gpu_log_module: Option<Arc<CudaModule>>,
#[cfg(feature = "kernel-step-trace")]
log_drain_stop: std::sync::Arc<std::sync::atomic::AtomicBool>,
log_drain_stop: Option<std::sync::Arc<std::sync::atomic::AtomicBool>>,
#[cfg(feature = "kernel-step-trace")]
log_drain_handle: Option<tokio::task::JoinHandle<()>>,
// X6: LN_b weights (formerly ln_gain_d / ln_bias_d) moved to
@@ -1135,45 +1147,67 @@ impl PerceptionTrainer {
// DtoD memcpy of the device step counter into the host shadow
// is captured in the graph, so the drainer sees the latest
// counter without an extra sync.
//
// Runtime activation: only allocate + spawn when
// `cfg.kernel_step_trace_path` is Some(path). Otherwise every
// field below stays None and the trainer carries zero overhead.
#[cfg(feature = "kernel-step-trace")]
let log_ring = std::sync::Arc::new(crate::gpu_log::LogRing::alloc()?);
#[cfg(feature = "kernel-step-trace")]
let log_step_counter_d = crate::gpu_log::alloc_step_counter(&stream)?;
#[cfg(feature = "kernel-step-trace")]
let log_step_counter_host_shadow = {
let buf = unsafe { crate::pinned_mem::MappedRecordBuffer::<i32>::new(1) }
.map_err(|e| anyhow::anyhow!("log step counter shadow: {e}"))?;
buf.write_record(0, 0);
std::sync::Arc::new(buf)
};
#[cfg(feature = "kernel-step-trace")]
let (gpu_log_tick_fn, gpu_log_module) = {
const GPU_LOG_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/gpu_log_ring.cubin"));
let module = ctx.load_cubin(GPU_LOG_CUBIN.to_vec())
.context("load gpu_log_ring cubin")?;
let f = module.load_function("gpu_log_tick").context("load gpu_log_tick")?;
(f, module)
};
#[cfg(feature = "kernel-step-trace")]
let log_drain_stop = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false));
// The drain task uses `tokio::spawn`, which requires a runtime
// handle on the current thread. In production the trainer runs
// inside a tokio runtime (binaries are `#[tokio::main]`). In
// sync `#[test]` units there is no runtime — we skip the drain
// (records still land in the ring; the test simply doesn't read
// them through the structured tracing path). This is the
// canonical pattern for opt-in async telemetry.
#[cfg(feature = "kernel-step-trace")]
let log_drain_handle = if tokio::runtime::Handle::try_current().is_ok() {
Some(crate::gpu_log::spawn_drain_task(
log_ring.clone(),
log_step_counter_host_shadow.clone(),
log_drain_stop.clone(),
))
let (
log_ring,
log_step_counter_d,
log_step_counter_host_shadow,
gpu_log_tick_fn,
gpu_log_module,
log_drain_stop,
log_drain_handle,
) = if let Some(trace_path) = cfg.kernel_step_trace_path.clone() {
let log_ring = std::sync::Arc::new(crate::gpu_log::LogRing::alloc()?);
let log_step_counter_d = crate::gpu_log::alloc_step_counter(&stream)?;
let log_step_counter_host_shadow = {
let buf = unsafe { crate::pinned_mem::MappedRecordBuffer::<i32>::new(1) }
.map_err(|e| anyhow::anyhow!("log step counter shadow: {e}"))?;
buf.write_record(0, 0);
std::sync::Arc::new(buf)
};
let (gpu_log_tick_fn, gpu_log_module) = {
const GPU_LOG_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/gpu_log_ring.cubin"));
let module = ctx.load_cubin(GPU_LOG_CUBIN.to_vec())
.context("load gpu_log_ring cubin")?;
let f = module.load_function("gpu_log_tick").context("load gpu_log_tick")?;
(f, module)
};
let log_drain_stop = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false));
// The drain task uses `tokio::spawn`, which requires a runtime
// handle on the current thread. In production the trainer runs
// inside a tokio runtime (binaries are `#[tokio::main]`). In
// sync `#[test]` units there is no runtime — we skip the drain
// (records still land in the ring; the test simply doesn't read
// them via the JSONL file path). This is the canonical pattern
// for opt-in async telemetry.
let log_drain_handle = if tokio::runtime::Handle::try_current().is_ok() {
Some(crate::gpu_log::spawn_drain_task_to_file(
log_ring.clone(),
log_step_counter_host_shadow.clone(),
log_drain_stop.clone(),
trace_path.clone(),
)?)
} else {
tracing::debug!(target: "gpu_log",
path = %trace_path.display(),
"no tokio runtime available — skipping drain task spawn (records still emitted to ring)");
None
};
(
Some(log_ring),
Some(log_step_counter_d),
Some(log_step_counter_host_shadow),
Some(gpu_log_tick_fn),
Some(gpu_log_module),
Some(log_drain_stop),
log_drain_handle,
)
} else {
tracing::debug!(target: "gpu_log",
"no tokio runtime available — skipping drain task spawn (records still emitted to ring)");
None
(None, None, None, None, None, None, None)
};
let k = cfg.seq_len;
@@ -1211,7 +1245,7 @@ impl PerceptionTrainer {
#[cfg(feature = "kernel-step-trace")]
log_drain_stop,
#[cfg(feature = "kernel-step-trace")]
log_drain_handle,
log_drain_handle, // already Option<_>; field type unchanged
// LayerNorm state. `ln_gain_d` / `ln_bias_d` constructed
// above via `upload` (gain=1, bias=0). LN_HIDDEN matches
// HIDDEN_DIM (128) by construction — checked at kernel
@@ -1674,15 +1708,23 @@ impl PerceptionTrainer {
// thread/single-block kernel: no host malloc, no host branch
// (per `pearl_cudarc_disable_event_tracking_for_graph_capture`
// and `pearl_no_host_branches_in_captured_graph`).
//
// Activated only when the runtime trace path is Some — i.e.
// when the trainer was constructed with
// `cfg.kernel_step_trace_path = Some(_)`. With path = None all
// the trace fields are None and this block is skipped, so the
// tick kernel never launches and zero overhead is incurred.
#[cfg(feature = "kernel-step-trace")]
if let (Some(tick_fn), Some(step_counter_d)) =
(self.gpu_log_tick_fn.as_ref(), self.log_step_counter_d.as_mut())
{
let tick_cfg = LaunchConfig {
grid_dim: (1, 1, 1),
block_dim: (1, 1, 1),
shared_mem_bytes: 0,
};
let mut launch = self.stream.launch_builder(&self.gpu_log_tick_fn);
launch.arg(&mut self.log_step_counter_d);
let mut launch = self.stream.launch_builder(tick_fn);
launch.arg(step_counter_d);
unsafe { launch.launch(tick_cfg).context("gpu_log_tick launch")?; }
}
@@ -2154,17 +2196,27 @@ impl PerceptionTrainer {
block_dim: (N_HORIZONS as u32, 1, 1),
shared_mem_bytes: 0,
};
// GPU log ring pointers. With feature off, both are 0 (null) and
// the kernel's `log_record()` short-circuits at the nullptr guard
// — behaviorally identical to the pre-log kernel.
// GPU log ring pointers. When the feature is OFF, both are 0
// (null). When the feature is ON but the runtime trace path is
// None (no `--kernel-step-trace`), the ring + step counter
// Options are also None — pass null and the kernel's
// `log_record()` short-circuits at the nullptr guard, behaviorally
// identical to the pre-log kernel.
#[cfg(feature = "kernel-step-trace")]
let log_ring_ptr: cudarc::driver::sys::CUdeviceptr = self.log_ring.dev_ptr();
let log_ring_ptr: cudarc::driver::sys::CUdeviceptr = self
.log_ring
.as_ref()
.map(|r| r.dev_ptr())
.unwrap_or(0);
#[cfg(not(feature = "kernel-step-trace"))]
let log_ring_ptr: cudarc::driver::sys::CUdeviceptr = 0;
#[cfg(feature = "kernel-step-trace")]
let log_step_ptr: cudarc::driver::sys::CUdeviceptr = {
let (p, _g) = self.log_step_counter_d.device_ptr(&self.stream);
p
let log_step_ptr: cudarc::driver::sys::CUdeviceptr = match self.log_step_counter_d.as_ref() {
Some(slice) => {
let (p, _g) = slice.device_ptr(&self.stream);
p
}
None => 0,
};
#[cfg(not(feature = "kernel-step-trace"))]
let log_step_ptr: cudarc::driver::sys::CUdeviceptr = 0;
@@ -2754,16 +2806,21 @@ impl PerceptionTrainer {
// The drain task reads the host shadow at 500ms cadence to know
// how many steps have completed since its last drain. Captured
// inside the same graph as the other shadow memcpys, so a single
// end-of-step sync makes all of them visible.
// end-of-step sync makes all of them visible. Skipped when the
// ring isn't allocated (runtime trace path is None).
#[cfg(feature = "kernel-step-trace")]
unsafe {
let (src_ptr, _g) = self.log_step_counter_d.device_ptr(&self.stream);
cudarc::driver::result::memcpy_dtod_async(
self.log_step_counter_host_shadow.dev_ptr,
src_ptr,
std::mem::size_of::<i32>(),
self.stream.cu_stream(),
).context("log_step_counter_host_shadow shadow")?;
if let (Some(step_counter_d), Some(shadow)) =
(self.log_step_counter_d.as_ref(), self.log_step_counter_host_shadow.as_ref())
{
unsafe {
let (src_ptr, _g) = step_counter_d.device_ptr(&self.stream);
cudarc::driver::result::memcpy_dtod_async(
shadow.dev_ptr,
src_ptr,
std::mem::size_of::<i32>(),
self.stream.cu_stream(),
).context("log_step_counter_host_shadow shadow")?;
}
}
Ok(())
}
@@ -4183,7 +4240,9 @@ fn download(stream: &Arc<CudaStream>, src: &CudaSlice<f32>) -> Result<Vec<f32>>
impl Drop for PerceptionTrainer {
fn drop(&mut self) {
use std::sync::atomic::Ordering;
self.log_drain_stop.store(true, Ordering::Relaxed);
if let Some(stop) = self.log_drain_stop.as_ref() {
stop.store(true, Ordering::Relaxed);
}
if let Some(h) = self.log_drain_handle.take() {
h.abort();
}

56
crates/ml-alpha/tests/gpu_log_ring_invariants.rs
View File

@@ -17,11 +17,12 @@ use std::sync::atomic::{AtomicBool, Ordering};
use std::time::Duration;
use ml_alpha::gpu_log::{
spawn_drain_task, LogHeader, LogRecord, LogRing, KID_SMOOTHNESS_CONTROLLER, MAGIC,
spawn_drain_task_to_file, LogHeader, LogRecord, LogRing, KID_SMOOTHNESS_CONTROLLER, MAGIC,
N_RECORDS_PER_STEP, N_RT_PER_KERNEL, N_SLOTS, PAYLOAD_F32, RT_INPUT, RT_OUTPUT, RT_STATE,
};
use ml_alpha::pinned_mem::MappedRecordBuffer;
use ml_core::device::MlDevice;
use tempfile::NamedTempFile;
// ── helpers ──────────────────────────────────────────────────────────
@@ -147,15 +148,29 @@ async fn magic_validation_skips_torn() {
counter.write_record(0, 1);
let h = spawn_drain_task(ring.clone(), counter.clone(), stop.clone());
let tmp = NamedTempFile::new().expect("tempfile");
let path = tmp.path().to_path_buf();
let h = spawn_drain_task_to_file(ring.clone(), counter.clone(), stop.clone(), path.clone())
.expect("spawn drain");
// Give drainer one tick (>500 ms cadence + a bit of slack).
tokio::time::sleep(Duration::from_millis(700)).await;
stop.store(true, Ordering::Relaxed);
// Wait for the next sleep boundary so the drain task observes the stop flag.
let _ = tokio::time::timeout(Duration::from_secs(2), h).await;
// Pass if no panic — torn slot was skipped, mismatch slot was skipped.
// (We don't intercept tracing output here; the absence of panic is the
// pass condition — drainer must defensively handle both cases.)
// The valid record should produce exactly one JSONL line; the torn
// (magic=0) and mismatched-step records must NOT appear.
let contents = std::fs::read_to_string(&path).expect("read trace file");
let lines: Vec<&str> = contents.lines().filter(|l| !l.is_empty()).collect();
assert_eq!(lines.len(), 1, "expected exactly 1 valid record, got {}", lines.len());
let v: serde_json::Value = serde_json::from_str(lines[0]).expect("parse JSONL");
assert_eq!(v["step"], 1);
assert_eq!(v["kid"], KID_SMOOTHNESS_CONTROLLER);
assert_eq!(v["rt"], RT_INPUT);
assert_eq!(v["kname"], "smoothness_controller");
assert_eq!(v["rt_name"], "input");
// Payload field comes from the smoothness_input schema.
assert!((v["payload"]["raw_h30"].as_f64().unwrap() - 100.0).abs() < 1e-6);
}
#[test]
@@ -199,7 +214,14 @@ async fn step_gap_detection_emits_warn() {
let big_jump = ((N_SLOTS / N_RECORDS_PER_STEP) + 10) as i32;
counter.write_record(0, big_jump);
let h = spawn_drain_task(ring.clone(), counter.clone(), stop.clone());
let tmp = NamedTempFile::new().expect("tempfile");
let h = spawn_drain_task_to_file(
ring.clone(),
counter.clone(),
stop.clone(),
tmp.path().to_path_buf(),
)
.expect("spawn drain");
tokio::time::sleep(Duration::from_millis(700)).await;
stop.store(true, Ordering::Relaxed);
let _ = tokio::time::timeout(Duration::from_secs(2), h).await;
@@ -226,17 +248,27 @@ async fn decoder_dispatch_unknown_falls_through() {
ring.buffer.write_record(slot, rec);
counter.write_record(0, 1);
let h = spawn_drain_task(ring.clone(), counter.clone(), stop.clone());
let tmp = NamedTempFile::new().expect("tempfile");
let h = spawn_drain_task_to_file(
ring.clone(),
counter.clone(),
stop.clone(),
tmp.path().to_path_buf(),
)
.expect("spawn drain");
tokio::time::sleep(Duration::from_millis(700)).await;
stop.store(true, Ordering::Relaxed);
let _ = tokio::time::timeout(Duration::from_secs(2), h).await;
// The drainer's iteration is `for kid in 0..N_KERNELS` — so kid=99 will
// never be visited by the drainer's slot walk. But the `decode()` function
// itself MUST handle unknown kernel_id without panic. Exercise it directly:
ml_alpha::gpu_log::decode(unknown_kid, RT_INPUT, 1, &[7.0; 10]);
ml_alpha::gpu_log::decode(0, 99, 1, &[7.0; 10]); // unknown record_type
// Pass if no panic.
// never be visited by the drainer's slot walk. But the `decode_to_json()`
// function itself MUST handle unknown kernel_id without panic. Exercise
// it directly — unknown (kid, rt) pairs fall through to a "values" array.
let v1 = ml_alpha::gpu_log::decode_to_json(unknown_kid, RT_INPUT, 1, &[7.0; 10]);
assert_eq!(v1["kname"], "unknown");
assert!(v1["payload"]["values"].is_array());
let v2 = ml_alpha::gpu_log::decode_to_json(0, 99, 1, &[7.0; 10]); // unknown record_type
assert_eq!(v2["rt_name"], "unknown");
}
// ── CUDA integration test #7 ─────────────────────────────────────────

8
crates/ml-alpha/tests/perception_overfit.rs
View File

@@ -73,6 +73,7 @@ fn stacked_trainer_loss_shrinks_on_constant_signal() {
horizon_weights: [1.0; 5],
n_batch: 1,
smoothness_base_lambda: 0.0,
kernel_step_trace_path: None,
};
let mut trainer = PerceptionTrainer::new(&dev, &cfg).expect("init");
@@ -147,6 +148,7 @@ fn stacked_trainer_loss_shrinks_with_stride_4() {
horizon_weights: [1.0; 5],
n_batch: 1,
smoothness_base_lambda: 0.0,
kernel_step_trace_path: None,
};
let mut trainer = PerceptionTrainer::new(&dev, &cfg).expect("init");
@@ -201,6 +203,7 @@ fn stacked_trainer_loss_shrinks_at_batch_32() {
horizon_weights: [1.0; 5],
n_batch: 32,
smoothness_base_lambda: 0.0,
kernel_step_trace_path: None,
};
let mut trainer = PerceptionTrainer::new(&dev, &cfg).expect("init");
@@ -299,6 +302,7 @@ fn evaluate_alone_succeeds() {
horizon_weights: [1.0; 5],
n_batch: 1,
smoothness_base_lambda: 0.0,
kernel_step_trace_path: None,
};
let mut trainer = PerceptionTrainer::new(&dev, &cfg).expect("init");
let ts = 1_000_000u64;
@@ -327,6 +331,7 @@ fn evaluate_works_after_captured_training_step() {
horizon_weights: [1.0; 5],
n_batch: 1,
smoothness_base_lambda: 0.0,
kernel_step_trace_path: None,
};
let mut trainer = PerceptionTrainer::new(&dev, &cfg).expect("init");
@@ -362,6 +367,7 @@ fn evaluate_works_after_capture_no_replay() {
horizon_weights: [1.0; 5],
n_batch: 1,
smoothness_base_lambda: 0.0,
kernel_step_trace_path: None,
};
let mut trainer = PerceptionTrainer::new(&dev, &cfg).expect("init");
let ts = 1_000_000u64;
@@ -392,6 +398,7 @@ fn horizon_ema_and_lambda_track_after_training() {
horizon_weights: [1.0; 5],
n_batch: 1,
smoothness_base_lambda: 0.0,
kernel_step_trace_path: None,
};
let mut trainer = PerceptionTrainer::new(&dev, &cfg).expect("init");
@@ -448,6 +455,7 @@ fn evaluate_works_after_warmup_only() {
horizon_weights: [1.0; 5],
n_batch: 1,
smoothness_base_lambda: 0.0,
kernel_step_trace_path: None,
};
let mut trainer = PerceptionTrainer::new(&dev, &cfg).expect("init");
let ts = 1_000_000u64;

24
infra/k8s/argo/alpha-perception-template.yaml
View File

@@ -80,6 +80,10 @@ spec:
value: "1"
- name: smoothness-base-lambda
value: "0.0"
- name: kernel-step-trace-enable
value: "false"
- name: kernel-step-trace-path
value: ""
volumes:
- name: git-ssh-key
@@ -299,8 +303,20 @@ spec:
fi
export PATH="${CARGO_HOME}/bin:${PATH}"
FEATURES_FLAG=""
if [ "{{workflow.parameters.kernel-step-trace-enable}}" = "true" ]; then
FEATURES_FLAG="--features kernel-step-trace"
echo "Building with --features kernel-step-trace"
# Bust cache: feature-on binary differs from feature-off at same SHA.
if [ -x "$BIN_DIR/alpha_train" ]; then
echo "Removing cached binary (feature change requires rebuild)"
rm -f "$BIN_DIR/alpha_train"
fi
fi
echo "Building alpha_train (ml-alpha)..."
cargo build --release -p ml-alpha --example alpha_train
cargo build --release $FEATURES_FLAG -p ml-alpha --example alpha_train
mkdir -p "$BIN_DIR"
cp "$CARGO_TARGET_DIR/release/examples/alpha_train" "$BIN_DIR/"
@@ -432,6 +448,11 @@ spec:
EXTRA_FLAGS="$EXTRA_FLAGS --auto-horizon-weights"
fi
TRACE_FLAG=""
if [ -n "{{workflow.parameters.kernel-step-trace-path}}" ]; then
TRACE_FLAG="--kernel-step-trace {{workflow.parameters.kernel-step-trace-path}}"
fi
"$BIN" \
--mbp10-data-dir "$MBP10_DIR" \
--predecoded-dir "$PREDECODED_DIR" \
@@ -452,6 +473,7 @@ spec:
--cv-train-window {{workflow.parameters.cv-train-window}} \
--decision-stride {{workflow.parameters.decision-stride}} \
--smoothness-base-lambda {{workflow.parameters.smoothness-base-lambda}} \
${TRACE_FLAG} \
$EXTRA_FLAGS
echo "=== Training complete ==="