feat(gpu-log): wire smoothness_lambda_controller as first ring producer
Atomic refactor wiring the GPU log ring's first producer end-to-end: - cuda/gpu_log_helpers.cuh (new): extract LogHeader/LogRecord/LogRing struct definitions + the log_record() device __forceinline__ helper into a shared header. Single source of truth — other kernels include this rather than redefining the structs. - cuda/gpu_log_ring.cu: refactor to include gpu_log_helpers.cuh; retain only the gpu_log_tick kernel. - cuda/smoothness_lambda_controller.cu: include gpu_log_helpers.cuh, add trailing (LogRing*, const int*) args, capture pre-EMA jitter + post-EMA jitter + target + excess_ratio into shared mem, and emit three records per call (RT_INPUT, RT_STATE, RT_OUTPUT) gated on non-null ring pointer. - trainer/perception.rs: feature-gated (cuda-diag-log) ring allocation + step counter (mapped-pinned host shadow), gpu_log_tick launch FIRST inside the captured graph, two new pointer args on the smoothness controller launch (null when feature off), step-counter DtoD shadow alongside the other telemetry shadows, drain task spawn (skipped when no tokio runtime — sync tests still construct the trainer), and a feature-gated Drop impl that aborts the drain task on shutdown. - tests/smoothness_lambda_controller_invariants.rs: pass null pointers for the two new kernel args; the kernel's nullptr guard preserves pre-existing behaviour. - build.rs: rerun-if-changed on cuda/gpu_log_ids.h and cuda/gpu_log_helpers.cuh so header edits trigger cubin rebuilds. Verified: cargo build / check --all-targets clean both with and without the cuda-diag-log feature; 4/4 smoothness controller invariant tests pass; 9/9 perception_overfit integration tests pass under the feature. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -26,10 +26,15 @@ const KERNELS: &[&str] = &[
|
||||
"gpu_log_ring", // GPU diagnostic log ring — tick kernel + log_record helper
|
||||
];
|
||||
|
||||
// Cache bust v14 (2026-05-21): gpu_log_ring.cu added — unified in-kernel diagnostic logging.
|
||||
// Cache bust v15 (2026-05-21): gpu_log_helpers.cuh extracted + smoothness_lambda_controller emits records.
|
||||
|
||||
fn main() {
|
||||
println!("cargo:rerun-if-changed=build.rs");
|
||||
// Track shared headers so .cuh / .h edits trigger rebuilds of every
|
||||
// .cu that #includes them. Without these, an edit to a helper header
|
||||
// leaves a stale cubin.
|
||||
println!("cargo:rerun-if-changed=cuda/gpu_log_ids.h");
|
||||
println!("cargo:rerun-if-changed=cuda/gpu_log_helpers.cuh");
|
||||
|
||||
println!("cargo:rerun-if-env-changed=CARGO_FEATURE_CUDA");
|
||||
if std::env::var("CARGO_FEATURE_CUDA").is_err() {
|
||||
|
||||
80
crates/ml-alpha/cuda/gpu_log_helpers.cuh
Normal file
80
crates/ml-alpha/cuda/gpu_log_helpers.cuh
Normal file
@@ -0,0 +1,80 @@
|
||||
// gpu_log_helpers.cuh — header-only device helpers for the GPU log ring.
|
||||
//
|
||||
// Include this from any .cu file that needs to call `log_record()`.
|
||||
// The `gpu_log_tick` kernel itself lives in `gpu_log_ring.cu`.
|
||||
//
|
||||
// Per `feedback_single_source_of_truth_no_duplicates`: this is the ONE
|
||||
// definition of `LogHeader` / `LogRecord` / `LogRing` and `log_record`;
|
||||
// other kernels must include this header instead of duplicating the
|
||||
// struct layouts. The IDs (kernel_id / record_type / magic / slot math
|
||||
// constants) live in `gpu_log_ids.h` and are included here.
|
||||
|
||||
#ifndef GPU_LOG_HELPERS_CUH
|
||||
#define GPU_LOG_HELPERS_CUH
|
||||
|
||||
#include "gpu_log_ids.h"
|
||||
#include <cstdint>
|
||||
|
||||
struct LogHeader {
|
||||
uint32_t magic;
|
||||
uint32_t step;
|
||||
uint8_t kernel_id;
|
||||
uint8_t record_type;
|
||||
uint8_t payload_words;
|
||||
uint8_t reserved;
|
||||
uint32_t _padding;
|
||||
};
|
||||
|
||||
struct LogRecord {
|
||||
LogHeader header;
|
||||
float payload[GPU_LOG_PAYLOAD_F32];
|
||||
};
|
||||
|
||||
struct LogRing {
|
||||
LogRecord records[GPU_LOG_N_SLOTS];
|
||||
};
|
||||
|
||||
// log_record — emit a diagnostic record from a kernel.
|
||||
//
|
||||
// Single-writer per (step, kernel_id, record_type) slot. Thread (0,0) of
|
||||
// block 0 writes; all other threads no-op. Header-last commit ensures
|
||||
// torn-record reads on the host side are detectable via magic mismatch.
|
||||
//
|
||||
// Caller responsibility:
|
||||
// - g_log_ring == nullptr → no-op (logging disabled).
|
||||
// - g_step_counter is the device step counter (incremented by gpu_log_tick).
|
||||
// - payload_words ≤ GPU_LOG_PAYLOAD_F32 (12); kernel must respect this.
|
||||
__device__ __forceinline__ void log_record(
|
||||
LogRing* g_log_ring,
|
||||
const int* g_step_counter,
|
||||
uint8_t kernel_id,
|
||||
uint8_t record_type,
|
||||
const float* payload,
|
||||
int payload_words
|
||||
) {
|
||||
if (g_log_ring == nullptr) return;
|
||||
if (blockIdx.x != 0 || threadIdx.x != 0) return;
|
||||
|
||||
const int step = g_step_counter[0];
|
||||
const int slot = (step * GPU_LOG_N_RECORDS_PER_STEP
|
||||
+ (int)kernel_id * GPU_LOG_N_RT_PER_KERNEL
|
||||
+ (int)record_type) & (GPU_LOG_N_SLOTS - 1);
|
||||
LogRecord* rec = &g_log_ring->records[slot];
|
||||
|
||||
const int n = payload_words < GPU_LOG_PAYLOAD_F32 ? payload_words : GPU_LOG_PAYLOAD_F32;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < GPU_LOG_PAYLOAD_F32; ++i) {
|
||||
rec->payload[i] = (i < n) ? payload[i] : 0.0f;
|
||||
}
|
||||
rec->header.step = (uint32_t)step;
|
||||
rec->header.kernel_id = kernel_id;
|
||||
rec->header.record_type = record_type;
|
||||
rec->header.payload_words = (uint8_t)n;
|
||||
rec->header.reserved = 0;
|
||||
rec->header._padding = 0;
|
||||
|
||||
__threadfence(); // make payload + header fields visible …
|
||||
rec->header.magic = GPU_LOG_MAGIC; // … BEFORE the commit sentinel.
|
||||
}
|
||||
|
||||
#endif // GPU_LOG_HELPERS_CUH
|
||||
@@ -1,40 +1,21 @@
|
||||
// gpu_log_ring.cu — unified GPU diagnostic log ring (device-side).
|
||||
// gpu_log_ring.cu — unified GPU diagnostic log ring (tick kernel).
|
||||
//
|
||||
// Provides `log_record()` (a device __forceinline__ helper) and the
|
||||
// `gpu_log_tick` kernel that increments the step counter once per
|
||||
// captured-graph step.
|
||||
// Provides the `gpu_log_tick` kernel that increments the step counter
|
||||
// once per captured-graph step.
|
||||
//
|
||||
// Constraints:
|
||||
// Struct definitions (`LogHeader`, `LogRecord`, `LogRing`) and the
|
||||
// `log_record` device helper live in `gpu_log_helpers.cuh` so other
|
||||
// kernels can include the helpers without re-compiling them as
|
||||
// separate symbols. Single source of truth — see
|
||||
// `feedback_single_source_of_truth_no_duplicates`.
|
||||
//
|
||||
// Constraints (still apply to anything including the helper header):
|
||||
// - No atomicAdd (slot reservation by deterministic step/kid/rt math).
|
||||
// - No host branches; thread-id gating only.
|
||||
// - Header-last commit (magic field written after payload + other header).
|
||||
// - Mapped-pinned ring; host reads via volatile, no DtoH memcpy.
|
||||
//
|
||||
// Per `feedback_nvidia_grade_perf_for_kernels`: single-thread writer for
|
||||
// the slot's 64 B cacheline, one coalesced store.
|
||||
|
||||
#include "gpu_log_ids.h"
|
||||
|
||||
#include <cstdint>
|
||||
|
||||
struct LogHeader {
|
||||
uint32_t magic;
|
||||
uint32_t step;
|
||||
uint8_t kernel_id;
|
||||
uint8_t record_type;
|
||||
uint8_t payload_words;
|
||||
uint8_t reserved;
|
||||
uint32_t _padding;
|
||||
};
|
||||
|
||||
struct LogRecord {
|
||||
LogHeader header;
|
||||
float payload[GPU_LOG_PAYLOAD_F32];
|
||||
};
|
||||
|
||||
struct LogRing {
|
||||
LogRecord records[GPU_LOG_N_SLOTS];
|
||||
};
|
||||
#include "gpu_log_helpers.cuh"
|
||||
|
||||
// Tick kernel — 1 thread, 1 block. Increments step counter; runs FIRST in
|
||||
// every captured-graph replay so all subsequent kernels see the new step.
|
||||
@@ -43,46 +24,3 @@ extern "C" __global__ void gpu_log_tick(int* step_counter) {
|
||||
step_counter[0] = step_counter[0] + 1;
|
||||
}
|
||||
}
|
||||
|
||||
// log_record — emit a diagnostic record from a kernel.
|
||||
//
|
||||
// Single-writer per (step, kernel_id, record_type) slot. Thread (0,0) of
|
||||
// block 0 writes; all other threads no-op. Header-last commit ensures
|
||||
// torn-record reads on the host side are detectable via magic mismatch.
|
||||
//
|
||||
// Caller responsibility:
|
||||
// - g_log_ring == nullptr → no-op (logging disabled).
|
||||
// - g_step_counter is the device step counter (incremented by gpu_log_tick).
|
||||
// - payload_words ≤ GPU_LOG_PAYLOAD_F32 (12); kernel must respect this.
|
||||
__device__ __forceinline__ void log_record(
|
||||
LogRing* g_log_ring,
|
||||
const int* g_step_counter,
|
||||
uint8_t kernel_id,
|
||||
uint8_t record_type,
|
||||
const float* payload,
|
||||
int payload_words
|
||||
) {
|
||||
if (g_log_ring == nullptr) return;
|
||||
if (blockIdx.x != 0 || threadIdx.x != 0) return;
|
||||
|
||||
const int step = g_step_counter[0];
|
||||
const int slot = (step * GPU_LOG_N_RECORDS_PER_STEP
|
||||
+ (int)kernel_id * GPU_LOG_N_RT_PER_KERNEL
|
||||
+ (int)record_type) & (GPU_LOG_N_SLOTS - 1);
|
||||
LogRecord* rec = &g_log_ring->records[slot];
|
||||
|
||||
const int n = payload_words < GPU_LOG_PAYLOAD_F32 ? payload_words : GPU_LOG_PAYLOAD_F32;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < GPU_LOG_PAYLOAD_F32; ++i) {
|
||||
rec->payload[i] = (i < n) ? payload[i] : 0.0f;
|
||||
}
|
||||
rec->header.step = (uint32_t)step;
|
||||
rec->header.kernel_id = kernel_id;
|
||||
rec->header.record_type = record_type;
|
||||
rec->header.payload_words = (uint8_t)n;
|
||||
rec->header.reserved = 0;
|
||||
rec->header._padding = 0;
|
||||
|
||||
__threadfence(); // make payload + header fields visible …
|
||||
rec->header.magic = GPU_LOG_MAGIC; // … BEFORE the commit sentinel.
|
||||
}
|
||||
|
||||
@@ -18,6 +18,14 @@
|
||||
// Per `feedback_no_atomicadd`: 5 threads, single block, single writer
|
||||
// per (h) slot; no atomics. Per `pearl_no_host_branches_in_captured_graph`:
|
||||
// no host branching; thread-id gating only.
|
||||
//
|
||||
// GPU log ring producer: emits three records per call when a non-null
|
||||
// `g_log_ring` is passed (cuda-diag-log feature on the Rust side):
|
||||
// RT_INPUT — raw_per_h[5] + jitter_in[5] (pre-EMA state)
|
||||
// RT_STATE — jitter_ema[5] + target[5] (post-EMA state + derived target)
|
||||
// RT_OUTPUT — excess[5] + lambda_out[5] (control signal + emitted λ)
|
||||
|
||||
#include "gpu_log_helpers.cuh"
|
||||
|
||||
#define SLC_N_HORIZONS 5
|
||||
#define SLC_LAMBDA_FLOOR 1.0e-4f
|
||||
@@ -40,16 +48,25 @@ extern "C" __global__ void smoothness_lambda_controller(
|
||||
float* __restrict__ jitter_ema, // [5] in/out — EMA state
|
||||
int* __restrict__ first_obs, // [1] in/out — sentinel
|
||||
float base_lambda, // scalar — amplitude knob
|
||||
float* __restrict__ lambda_out // [5] output — λ for next step
|
||||
float* __restrict__ lambda_out, // [5] output — λ for next step
|
||||
LogRing* g_log_ring, // nullable — log ring (cuda-diag-log)
|
||||
const int* g_step_counter // nullable — device step counter
|
||||
) {
|
||||
const int h = threadIdx.x;
|
||||
if (h >= SLC_N_HORIZONS) return;
|
||||
|
||||
__shared__ float s_jitter_after[SLC_N_HORIZONS];
|
||||
__shared__ float s_jitter_in[SLC_N_HORIZONS]; // pre-EMA EMA reading (0 on bootstrap)
|
||||
__shared__ float s_jitter_after[SLC_N_HORIZONS]; // post-EMA value
|
||||
__shared__ float s_target[SLC_N_HORIZONS]; // derived target
|
||||
__shared__ float s_excess[SLC_N_HORIZONS]; // excess_ratio for log
|
||||
|
||||
// Pass 1: EMA update with sentinel bootstrap.
|
||||
// Pass 1: read raw + sentinel; capture pre-EMA state for logging
|
||||
// (zero on bootstrap step — first_obs gates).
|
||||
const float raw_h = raw_per_h[h];
|
||||
const int sentinel = first_obs[0];
|
||||
s_jitter_in[h] = (sentinel == 0) ? 0.0f : jitter_ema[h];
|
||||
|
||||
// EMA update with sentinel bootstrap.
|
||||
float jitter_h;
|
||||
if (sentinel == 0) {
|
||||
jitter_h = raw_h;
|
||||
@@ -71,6 +88,7 @@ extern "C" __global__ void smoothness_lambda_controller(
|
||||
// < jitter_ema[0] for h>0
|
||||
const float jitter_h0 = s_jitter_after[0];
|
||||
const float target_h = jitter_h0 * TARGET_K_RATIO[h];
|
||||
s_target[h] = target_h;
|
||||
|
||||
// Excess controller: ratio - 1, clamped at zero (only push UP).
|
||||
// When observed > target: excess > 0 → λ grows
|
||||
@@ -78,6 +96,41 @@ extern "C" __global__ void smoothness_lambda_controller(
|
||||
// Floor: λ ≥ LAMBDA_FLOOR.
|
||||
const float safe_target = fmaxf(target_h, SLC_TARGET_EPS);
|
||||
const float excess_ratio = fmaxf(0.0f, jitter_h / safe_target - 1.0f);
|
||||
s_excess[h] = excess_ratio;
|
||||
|
||||
const float lambda_new = base_lambda * (1.0f + excess_ratio);
|
||||
lambda_out[h] = fmaxf(SLC_LAMBDA_FLOOR, lambda_new);
|
||||
|
||||
// Wait for all 5 threads to populate shared mem AND publish their
|
||||
// lambda_out[h] store before thread 0 gathers all 5 for logging.
|
||||
__syncthreads();
|
||||
|
||||
if (h == 0) {
|
||||
// RT_INPUT: raw_per_h[0..5] + jitter_in[0..5].
|
||||
float in_payload[10] = {
|
||||
raw_per_h[0], raw_per_h[1], raw_per_h[2], raw_per_h[3], raw_per_h[4],
|
||||
s_jitter_in[0], s_jitter_in[1], s_jitter_in[2], s_jitter_in[3], s_jitter_in[4],
|
||||
};
|
||||
log_record(g_log_ring, g_step_counter,
|
||||
KID_SMOOTHNESS_CONTROLLER, RT_INPUT,
|
||||
in_payload, 10);
|
||||
|
||||
// RT_STATE: jitter_ema_out[0..5] + target[0..5].
|
||||
float state_payload[10] = {
|
||||
s_jitter_after[0], s_jitter_after[1], s_jitter_after[2], s_jitter_after[3], s_jitter_after[4],
|
||||
s_target[0], s_target[1], s_target[2], s_target[3], s_target[4],
|
||||
};
|
||||
log_record(g_log_ring, g_step_counter,
|
||||
KID_SMOOTHNESS_CONTROLLER, RT_STATE,
|
||||
state_payload, 10);
|
||||
|
||||
// RT_OUTPUT: excess_ratio[0..5] + lambda_out[0..5].
|
||||
float out_payload[10] = {
|
||||
s_excess[0], s_excess[1], s_excess[2], s_excess[3], s_excess[4],
|
||||
lambda_out[0], lambda_out[1], lambda_out[2], lambda_out[3], lambda_out[4],
|
||||
};
|
||||
log_record(g_log_ring, g_step_counter,
|
||||
KID_SMOOTHNESS_CONTROLLER, RT_OUTPUT,
|
||||
out_payload, 10);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -316,6 +316,29 @@ pub struct PerceptionTrainer {
|
||||
/// Cached handle for `smoothness_lambda_controller`.
|
||||
smoothness_controller_fn: CudaFunction,
|
||||
_smoothness_controller_module: Arc<CudaModule>,
|
||||
// ── GPU log ring (feature: cuda-diag-log) ──
|
||||
// When enabled, 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`.
|
||||
#[cfg(feature = "cuda-diag-log")]
|
||||
log_ring: std::sync::Arc<crate::gpu_log::LogRing>,
|
||||
#[cfg(feature = "cuda-diag-log")]
|
||||
log_step_counter_d: 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 = "cuda-diag-log")]
|
||||
log_step_counter_host_shadow: std::sync::Arc<crate::pinned_mem::MappedRecordBuffer<i32>>,
|
||||
#[cfg(feature = "cuda-diag-log")]
|
||||
gpu_log_tick_fn: CudaFunction,
|
||||
#[cfg(feature = "cuda-diag-log")]
|
||||
_gpu_log_module: Arc<CudaModule>,
|
||||
#[cfg(feature = "cuda-diag-log")]
|
||||
log_drain_stop: std::sync::Arc<std::sync::atomic::AtomicBool>,
|
||||
#[cfg(feature = "cuda-diag-log")]
|
||||
log_drain_handle: Option<tokio::task::JoinHandle<()>>,
|
||||
// X6: LN_b weights (formerly ln_gain_d / ln_bias_d) moved to
|
||||
// `self.trunk.ln_b_gain_d` / `ln_b_bias_d`. LN_b is the LayerNorm
|
||||
// AFTER Mamba2 stack 2 — stabilises the trunk-to-CfC distribution.
|
||||
@@ -1106,6 +1129,53 @@ impl PerceptionTrainer {
|
||||
let smoothness_jitter_first_obs_d = stream.alloc_zeros::<i32>(1)
|
||||
.context("smoothness_jitter_first_obs_d alloc")?;
|
||||
|
||||
// ── GPU log ring allocation (feature: cuda-diag-log) ──
|
||||
// Per `feedback_no_htod_htoh_only_mapped_pinned`: ring is mapped-
|
||||
// pinned; step counter shadow is mapped-pinned; the on-stream
|
||||
// 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.
|
||||
#[cfg(feature = "cuda-diag-log")]
|
||||
let log_ring = std::sync::Arc::new(crate::gpu_log::LogRing::alloc()?);
|
||||
#[cfg(feature = "cuda-diag-log")]
|
||||
let log_step_counter_d = crate::gpu_log::alloc_step_counter(&stream)?;
|
||||
#[cfg(feature = "cuda-diag-log")]
|
||||
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 = "cuda-diag-log")]
|
||||
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 = "cuda-diag-log")]
|
||||
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 = "cuda-diag-log")]
|
||||
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(),
|
||||
))
|
||||
} else {
|
||||
tracing::debug!(target: "gpu_log",
|
||||
"no tokio runtime available — skipping drain task spawn (records still emitted to ring)");
|
||||
None
|
||||
};
|
||||
|
||||
let k = cfg.seq_len;
|
||||
Ok(Self {
|
||||
cfg: cfg.clone(),
|
||||
@@ -1128,6 +1198,20 @@ impl PerceptionTrainer {
|
||||
smoothness_jitter_first_obs_d,
|
||||
smoothness_controller_fn,
|
||||
_smoothness_controller_module: smoothness_controller_module,
|
||||
#[cfg(feature = "cuda-diag-log")]
|
||||
log_ring,
|
||||
#[cfg(feature = "cuda-diag-log")]
|
||||
log_step_counter_d,
|
||||
#[cfg(feature = "cuda-diag-log")]
|
||||
log_step_counter_host_shadow,
|
||||
#[cfg(feature = "cuda-diag-log")]
|
||||
gpu_log_tick_fn,
|
||||
#[cfg(feature = "cuda-diag-log")]
|
||||
_gpu_log_module: gpu_log_module,
|
||||
#[cfg(feature = "cuda-diag-log")]
|
||||
log_drain_stop,
|
||||
#[cfg(feature = "cuda-diag-log")]
|
||||
log_drain_handle,
|
||||
// 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
|
||||
@@ -1584,6 +1668,24 @@ impl PerceptionTrainer {
|
||||
k_seq: usize,
|
||||
total_snaps: usize,
|
||||
) -> Result<()> {
|
||||
// ── GPU log ring tick (feature: cuda-diag-log) ──
|
||||
// Runs FIRST inside the captured graph so every downstream
|
||||
// logging-capable kernel sees the new step counter. Single-
|
||||
// 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`).
|
||||
#[cfg(feature = "cuda-diag-log")]
|
||||
{
|
||||
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);
|
||||
unsafe { launch.launch(tick_cfg).context("gpu_log_tick launch")?; }
|
||||
}
|
||||
|
||||
// DtoD: staging (mapped-pinned, device-visible) → device buffers.
|
||||
let n10 = total_snaps * 10;
|
||||
let n6 = total_snaps * REGIME_DIM;
|
||||
@@ -2052,6 +2154,20 @@ 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.
|
||||
#[cfg(feature = "cuda-diag-log")]
|
||||
let log_ring_ptr: cudarc::driver::sys::CUdeviceptr = self.log_ring.dev_ptr();
|
||||
#[cfg(not(feature = "cuda-diag-log"))]
|
||||
let log_ring_ptr: cudarc::driver::sys::CUdeviceptr = 0;
|
||||
#[cfg(feature = "cuda-diag-log")]
|
||||
let log_step_ptr: cudarc::driver::sys::CUdeviceptr = {
|
||||
let (p, _g) = self.log_step_counter_d.device_ptr(&self.stream);
|
||||
p
|
||||
};
|
||||
#[cfg(not(feature = "cuda-diag-log"))]
|
||||
let log_step_ptr: cudarc::driver::sys::CUdeviceptr = 0;
|
||||
{
|
||||
let mut launch = self.stream.launch_builder(&self.smoothness_controller_fn);
|
||||
launch
|
||||
@@ -2059,7 +2175,9 @@ impl PerceptionTrainer {
|
||||
.arg(&mut self.smoothness_jitter_ema_d)
|
||||
.arg(&mut self.smoothness_jitter_first_obs_d)
|
||||
.arg(&base_lambda)
|
||||
.arg(&mut self.smoothness_lambda_d);
|
||||
.arg(&mut self.smoothness_lambda_d)
|
||||
.arg(&log_ring_ptr)
|
||||
.arg(&log_step_ptr);
|
||||
unsafe { launch.launch(smooth_ctrl_cfg).context("smoothness_lambda_controller launch")?; }
|
||||
}
|
||||
|
||||
@@ -2632,6 +2750,21 @@ impl PerceptionTrainer {
|
||||
self.stream.cu_stream(),
|
||||
).context("smoothness_loss_per_horizon_host_d shadow")?;
|
||||
}
|
||||
// ── Shadow device step counter into mapped-pinned host shadow ──
|
||||
// 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.
|
||||
#[cfg(feature = "cuda-diag-log")]
|
||||
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")?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -4039,3 +4172,20 @@ fn download(stream: &Arc<CudaStream>, src: &CudaSlice<f32>) -> Result<Vec<f32>>
|
||||
stream.synchronize().context("stacked download sync")?;
|
||||
Ok(staging.read_all())
|
||||
}
|
||||
|
||||
// ── Drop: shut down the GPU log drain task (feature-gated) ──
|
||||
// Setting the stop flag lets the drainer exit cleanly on its next 500ms
|
||||
// tick; `.abort()` is the belt-and-braces guard against a dropped
|
||||
// runtime / hanging shutdown. Drop cannot be async, so we cannot await
|
||||
// the handle — abort is the canonical pattern for tokio JoinHandle in
|
||||
// non-async destructors.
|
||||
#[cfg(feature = "cuda-diag-log")]
|
||||
impl Drop for PerceptionTrainer {
|
||||
fn drop(&mut self) {
|
||||
use std::sync::atomic::Ordering;
|
||||
self.log_drain_stop.store(true, Ordering::Relaxed);
|
||||
if let Some(h) = self.log_drain_handle.take() {
|
||||
h.abort();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -51,13 +51,22 @@ fn run_controller(
|
||||
block_dim: (N_HORIZONS as u32, 1, 1),
|
||||
shared_mem_bytes: 0,
|
||||
};
|
||||
// The kernel now accepts two trailing pointer args for the GPU log
|
||||
// ring (`LogRing*` + `int* g_step_counter`). Passing null pointers
|
||||
// makes `log_record()` short-circuit at the `g_log_ring == nullptr`
|
||||
// guard — behaviorally identical to the pre-log signature. We must
|
||||
// still push the args because launch_builder validates the count.
|
||||
let null_ring: cudarc::driver::sys::CUdeviceptr = 0;
|
||||
let null_step: cudarc::driver::sys::CUdeviceptr = 0;
|
||||
let mut launch = stream.launch_builder(&func);
|
||||
launch
|
||||
.arg(&raw_d)
|
||||
.arg(&mut jitter_d)
|
||||
.arg(&mut first_obs_d)
|
||||
.arg(&base_lambda)
|
||||
.arg(&mut lambda_d);
|
||||
.arg(&mut lambda_d)
|
||||
.arg(&null_ring)
|
||||
.arg(&null_step);
|
||||
unsafe {
|
||||
launch.launch(cfg).context("controller launch")?;
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user