From 339eaca98347bd15f8264e89cff56864668df15e Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Thu, 21 May 2026 01:16:55 +0200 Subject: [PATCH] feat(gpu-log): Rust-side LogRing, drain task, decoder registry MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Task 5 of the GPU log ring plan. Adds `gpu_log` module behind the `cuda-diag-log` feature gate: - `LogRecord` (#[repr(C)], 64 B) mirroring the C-side layout in `gpu_log_ring.cu`; const compile-time size assertion catches drift. - `LogRing::alloc()` allocates and zero-initialises a 32k-slot mapped-pinned ring (2 MiB pinned). - `alloc_step_counter()` allocates the device-side i32 step counter. - Decoder dispatch table with per-(kid, rt) decoders for `KID_SMOOTHNESS_CONTROLLER` × {RT_INPUT, RT_STATE, RT_OUTPUT} emitting structured tracing::info! events. - `spawn_drain_task()` polls a mapped-pinned step-counter shadow at 500 ms cadence, walks new slots, validates magic + step, dispatches to decoders. Emits tracing::warn! when the drainer falls behind the ring window (no silent record drops). Kernel ID + record-type constants are a manual mirror of `crates/ml-alpha/cuda/gpu_log_ids.h`; drift is caught by the const size assertion at compile time and (subsequently) by build.rs. --- crates/ml-alpha/src/gpu_log.rs | 254 +++++++++++++++++++++++++++++++++ crates/ml-alpha/src/lib.rs | 2 + 2 files changed, 256 insertions(+) create mode 100644 crates/ml-alpha/src/gpu_log.rs diff --git a/crates/ml-alpha/src/gpu_log.rs b/crates/ml-alpha/src/gpu_log.rs new file mode 100644 index 000000000..68082fdc1 --- /dev/null +++ b/crates/ml-alpha/src/gpu_log.rs @@ -0,0 +1,254 @@ +//! GPU diagnostic log ring — host-side allocator, drainer, decoder registry. +//! +//! See `docs/superpowers/specs/2026-05-21-gpu-log-ring-design.md`. +//! +//! Behaviour: when the `cuda-diag-log` 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. +//! +//! Constants below mirror `crates/ml-alpha/cuda/gpu_log_ids.h` — single +//! source of truth; manual mirror with on-build verification. + +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::time::Duration; + +use anyhow::{Context, Result}; +use cudarc::driver::{CudaSlice, CudaStream}; +use tokio::task::JoinHandle; +use tracing::{debug, info, warn}; + +use crate::pinned_mem::MappedRecordBuffer; + +pub const N_KERNELS: usize = 8; +pub const N_RT_PER_KERNEL: usize = 4; +pub const N_RECORDS_PER_STEP: usize = N_KERNELS * N_RT_PER_KERNEL; +pub const N_SLOTS: usize = 32768; +pub const PAYLOAD_F32: usize = 12; +pub const MAGIC: u32 = 0xF0F0_A710; + +// Kernel IDs (manual mirror of gpu_log_ids.h — assert equality at module-load via tests). +pub const KID_SMOOTHNESS_CONTROLLER: u8 = 0; +pub const KID_OUTPUT_SMOOTHNESS: u8 = 1; +pub const KID_BCE_MULTI_HORIZON: u8 = 2; +pub const KID_HORIZON_LAMBDA: u8 = 3; +pub const KID_HEADS_GRN_FWD: u8 = 4; +pub const KID_HEADS_GRN_BWD: u8 = 5; +pub const KID_CFC_STEP: u8 = 6; +pub const KID_RESERVED_7: u8 = 7; + +// Record types within a kernel. +pub const RT_INPUT: u8 = 0; +pub const RT_STATE: u8 = 1; +pub const RT_OUTPUT: u8 = 2; +pub const RT_DEBUG: u8 = 3; + +#[repr(C)] +#[derive(Clone, Copy, Debug)] +pub struct LogHeader { + pub magic: u32, + pub step: u32, + pub kernel_id: u8, + pub record_type: u8, + pub payload_words: u8, + pub reserved: u8, + pub _padding: u32, +} + +#[repr(C)] +#[derive(Clone, Copy)] +pub struct LogRecord { + pub header: LogHeader, + pub payload: [f32; PAYLOAD_F32], +} + +// Compile-time size sanity check. `gpu_log_ids.h` declares 64 bytes; the +// kernel's slot arithmetic + the host's `record_record(idx)` math both +// assume this. If padding or field reordering bumps the size, the host +// and device will see misaligned slots — fail at compile time instead. +const _ASSERT_RECORD_64B: () = { + assert!(std::mem::size_of::() == 64); +}; + +impl Default for LogRecord { + fn default() -> Self { + Self { + header: LogHeader { + magic: 0, + step: 0, + kernel_id: 0, + record_type: 0, + payload_words: 0, + reserved: 0, + _padding: 0, + }, + payload: [0.0; PAYLOAD_F32], + } + } +} + +/// Mapped-pinned ring buffer for GPU log records. Both host and device +/// see the same physical memory; the device pointer is passed into +/// logging-capable kernels as the `LogRing*` arg. +pub struct LogRing { + pub buffer: MappedRecordBuffer, +} + +impl LogRing { + /// Allocate the ring (`N_SLOTS` records of 64 B each = 2 MiB pinned). + /// Every slot is zero-initialised so its `magic` is 0 — readers + /// validate against `MAGIC` and skip uninitialised slots. + /// + /// SAFETY: caller must ensure cuda context is current. + pub fn alloc() -> Result { + let buffer = unsafe { MappedRecordBuffer::::new(N_SLOTS) } + .map_err(|e| anyhow::anyhow!("log ring alloc: {e}"))?; + let zero = LogRecord::default(); + for i in 0..N_SLOTS { + buffer.write_record(i, zero); + } + Ok(Self { buffer }) + } + + /// Device pointer to the ring (passed as `LogRing*` arg to kernels). + /// cudarc represents device pointers as `CUdeviceptr` (`u64`). + pub fn dev_ptr(&self) -> cudarc::driver::sys::CUdeviceptr { + self.buffer.dev_ptr + } +} + +/// Allocate the device-side step counter (single i32, zero-initialised). +/// The `gpu_log_tick` kernel increments this once per captured-graph +/// replay; logging-capable kernels read it to derive their slot. +pub fn alloc_step_counter(stream: &Arc) -> Result> { + stream.alloc_zeros::(1).context("step_counter alloc") +} + +// ─── Decoder dispatch ──────────────────────────────────────────────── + +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"); + } +} + +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"); + } +} + +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]) { + 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), + } +} + +// ─── 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. +/// +/// Stops cleanly when `stop` is set. +pub fn spawn_drain_task( + ring: Arc, + step_counter_host_shadow: Arc>, + stop: Arc, +) -> JoinHandle<()> { + tokio::spawn(async move { + let mut last_drained_step: u32 = 0; + let interval = Duration::from_millis(500); + loop { + if stop.load(Ordering::Relaxed) { + debug!(target: "gpu_log", last_drained_step, "drain task exiting on stop signal"); + break; + } + tokio::time::sleep(interval).await; + + let current_step_i32 = step_counter_host_shadow.read_record(0); + if current_step_i32 < 0 { + // Counter overflow (unlikely at 4G steps but defensive). + warn!(target: "gpu_log", current_step_i32, "step counter went negative; skipping drain cycle"); + continue; + } + let current_step = current_step_i32 as u32; + if current_step <= last_drained_step { + continue; + } + + let max_drain_window = (N_SLOTS / N_RECORDS_PER_STEP) as u32; + let oldest_safe_step = current_step.saturating_sub(max_drain_window); + if last_drained_step < oldest_safe_step { + let skipped = oldest_safe_step.saturating_sub(last_drained_step); + warn!(target: "gpu_log", + last_drained_step, current_step, skipped, + "drainer fell behind — records dropped"); + last_drained_step = oldest_safe_step; + } + + for step in (last_drained_step + 1)..=current_step { + for kid in 0..N_KERNELS as u8 { + for rt in 0..N_RT_PER_KERNEL as u8 { + let slot = ((step as usize) * N_RECORDS_PER_STEP + + (kid as usize) * N_RT_PER_KERNEL + + rt as usize) & (N_SLOTS - 1); + let rec = ring.buffer.read_record(slot); + if rec.header.magic != MAGIC { + continue; + } + if rec.header.step != step { + // Slot was written by a later step; this kernel + // didn't emit at THIS step. + 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]); + } + } + } + last_drained_step = current_step; + } + }) +} diff --git a/crates/ml-alpha/src/lib.rs b/crates/ml-alpha/src/lib.rs index 6c722def7..8839602be 100644 --- a/crates/ml-alpha/src/lib.rs +++ b/crates/ml-alpha/src/lib.rs @@ -32,6 +32,8 @@ pub mod heads; pub mod isv; pub mod pinned; pub mod pinned_mem; +#[cfg(feature = "cuda-diag-log")] +pub mod gpu_log; pub mod trainer; // Preserved from prior crate state — used by Phase A.