diff --git a/crates/ml/build.rs b/crates/ml/build.rs index 1ed09c9b1..3870f7bb2 100644 --- a/crates/ml/build.rs +++ b/crates/ml/build.rs @@ -497,6 +497,18 @@ fn main() { // ISV[LB_*_VAR_*_BASE]. Replaces Pearl 2's hardcoded-zero CQL // budget output and floor-pinned C51 budget. "loss_balance_controller_kernel.cu", + // SP7 (2026-05-03): GPU dispatch for cql/c51 budget consumer. + // Eliminates capture-time host-branch freeze in + // `compute_adaptive_budgets`. The dispatch kernel resolves + // `(active >= 0.5) ? controller_budget : bootstrap` on-device every + // step (8 threads — 2 heads × 4 branches), writing trunk_mean + + // per-branch correction to a 10 f32 mapped-pinned buffer. The same + // cubin houses two new SAXPY/scale variants that read `alpha` from + // a device pointer (instead of scalar arg), so the captured + // `apply_c51_budget_scale` / `apply_cql_saxpy` chain reads runtime + // ISV-derived values instead of step-0 frozen scalars. See + // `feedback_no_cpu_compute_strict.md` and audit Fix 33. + "consume_lb_budget_kernel.cu", // SP5 Layer D Task D1 (2026-05-02): PnL aggregation producer. // First of 3 Layer D producers replacing the host-side trade-PnL // aggregation in `compute_epoch_financials` per diff --git a/crates/ml/src/cuda_pipeline/consume_lb_budget_kernel.cu b/crates/ml/src/cuda_pipeline/consume_lb_budget_kernel.cu new file mode 100644 index 000000000..35ec57d0b --- /dev/null +++ b/crates/ml/src/cuda_pipeline/consume_lb_budget_kernel.cu @@ -0,0 +1,180 @@ +// crates/ml/src/cuda_pipeline/consume_lb_budget_kernel.cu +// +// SP7 (2026-05-03) — GPU dispatch for cql/c51 budget consumer. +// +// Background: the SP7 loss-balance controller writes the controller-derived +// budget per (head, branch) into ISV[BUDGET_{CQL,C51}_BASE+b] and the +// activation flag into ISV[LB_{CQL,C51}_ACTIVE_BASE+b]. The downstream +// consumer used to resolve `effective = (active >= 0.5) ? controller : bootstrap` +// host-side (Rust) inside `compute_adaptive_budgets`, then pass the resulting +// f32 scalars to `apply_c51_budget_scale` / `apply_cql_saxpy`. Because that +// host code was reached through the `aux_child` CUDA-Graph capture, the if/else +// resolved ONCE at capture time (step 0 of each fold, when the activation flag +// is the FoldReset sentinel `0.0`); the bootstrap branch baked the literal +// 0.02/0.05 into the captured kernel-arg buffers, freezing the budget for +// every replay regardless of what the controller wrote at runtime. Same +// disease as SP4 host-side EMA elimination 2026-05-01 (see +// `feedback_no_cpu_compute_strict.md`). +// +// Fix architecture: this kernel resolves the dispatch on-GPU each step. It +// runs inside the `aux_child` graph (kernel launch is graph-safe), reading +// ISV at replay time (current values), and writes the effective per-branch +// budget plus the trunk-mean and per-branch correction to a mapped-pinned +// scratch buffer. Downstream `dqn_scale_f32_dev_ptr_kernel` / +// `dqn_saxpy_f32_dev_ptr_kernel` (also in this translation unit) read the +// scalar alpha from device memory at execution time, so every replay sees +// fresh values — capture-safe end-to-end. +// +// Buffer layout (`lb_budget_effective_buf`, 10 f32): +// [0] cql_trunk_mean — mean(cql_effective[0..4]) +// [1] cql_correction[0] — cql_effective[0] / cql_trunk_mean +// [2] cql_correction[1] — cql_effective[1] / cql_trunk_mean +// [3] cql_correction[2] — cql_effective[2] / cql_trunk_mean +// [4] cql_correction[3] — cql_effective[3] / cql_trunk_mean +// [5] c51_trunk_mean — mean(c51_effective[0..4]) +// [6] c51_correction[0] — c51_effective[0] / c51_trunk_mean +// [7] c51_correction[1] — c51_effective[1] / c51_trunk_mean +// [8] c51_correction[2] — c51_effective[2] / c51_trunk_mean +// [9] c51_correction[3] — c51_effective[3] / c51_trunk_mean +// +// SAXPY callers (`apply_c51_budget_scale` / `apply_cql_saxpy`) issue a +// trunk-wide kernel reading `lb_buf[trunk_idx]`, then per-branch correction +// kernels reading `lb_buf[trunk_idx + 1 + b]`. Net effect: branch slice ends +// up scaled by `effective[b] = trunk × correction[b]`, identical to the prior +// host-side semantics, but every value derives from runtime ISV reads. +// +// Cold-start contract: bootstrap constants (CQL=0.02, C51=0.05) are passed as +// kernel args (not hardcoded literals) per `feedback_isv_for_adaptive_bounds`, +// keeping the values flowing through ISV/kernel-args. They MUST stay +// byte-identical to the host-side `CQL_BOOTSTRAP_BUDGET=0.02` / +// `C51_BOOTSTRAP_BUDGET=0.05` literals that lived in `compute_adaptive_budgets` +// (which themselves match `loss_balance_controller_kernel.cu`'s +// `COLD_START_FLOOR_CQL=0.02` / `COLD_START_FLOOR_C51=0.05`). +// +// Layout: single block, 8 threads (2 heads × 4 branches) for the dispatch +// kernel. No atomicAdd (per `feedback_no_atomicadd`). __threadfence_system() +// after writes so the host-visible mapped-pinned slice sees the updates after +// stream sync (HEALTH_DIAG path). + +#include + +extern "C" __global__ void lb_budget_dispatch( + // ISV signal bus (read-only). + const float* __restrict__ isv_signals, + int active_cql_isv_base, // LB_CQL_ACTIVE_BASE = 313 + int active_c51_isv_base, // LB_C51_ACTIVE_BASE = 317 + int budget_cql_isv_base, // BUDGET_CQL_BASE = 198 + int budget_c51_isv_base, // BUDGET_C51_BASE = 190 + // Cold-start bootstrap constants (passed as args, not hardcoded; matches + // host-side CQL_BOOTSTRAP_BUDGET / C51_BOOTSTRAP_BUDGET in + // compute_adaptive_budgets per the new dev-ptr architecture). + float cql_bootstrap, + float c51_bootstrap, + // Output scratch buffer (10 f32). See top-of-file layout block. + float* __restrict__ out +) { + const float ACTIVE_THRESHOLD = 0.5f; + const float SP7_EPS_DIV = 1e-8f; + const float TRUNK_EPS = 1e-6f; // matches host correction-skip threshold + + int tid = threadIdx.x; + if (blockIdx.x != 0 || tid >= 8) return; // 8 = 2 heads × 4 branches + + int head = tid / 4; // 0 = CQL, 1 = C51 + int branch = tid % 4; // 0..3 (dir, mag, ord, urg) + + // ── Read activation flag + controller budget per (head, branch) ── + int active_base = (head == 0) ? active_cql_isv_base : active_c51_isv_base; + int budget_base = (head == 0) ? budget_cql_isv_base : budget_c51_isv_base; + float bootstrap = (head == 0) ? cql_bootstrap : c51_bootstrap; + + float active = isv_signals[active_base + branch]; + float ctrl = isv_signals[budget_base + branch]; + + // ── Dispatch on activation flag ────────────────────────────────── + float effective = (active >= ACTIVE_THRESHOLD) ? fmaxf(ctrl, SP7_EPS_DIV) + : bootstrap; + + // ── Stage in shared mem so thread 0 (per head) can compute trunk mean ─ + __shared__ float effective_smem[8]; + effective_smem[tid] = effective; + __syncthreads(); + + // Two threads (one per head) compute the trunk mean over their 4 branches + // and write the trunk slot. The other six threads write only their + // correction slot. This is a fixed 8-thread pattern — no atomics needed. + int head_base_out = (head == 0) ? 0 : 5; // [0..5)=cql, [5..10)=c51 + int trunk_idx_out = head_base_out + 0; // slot 0 of each head block + int corr_idx_out = head_base_out + 1 + branch; + + if (branch == 0) { + // Thread (head, 0) computes trunk mean. + int sm_base = head * 4; + float sum = effective_smem[sm_base + 0] + + effective_smem[sm_base + 1] + + effective_smem[sm_base + 2] + + effective_smem[sm_base + 3]; + float trunk_mean = sum * 0.25f; + out[trunk_idx_out] = trunk_mean; + } + __syncthreads(); + + // ── Compute and write per-branch correction ────────────────────── + float trunk_mean = out[trunk_idx_out]; + // Mirror the host-side guard `if c51_trunk > 1e-6 { c51_branch[b] / c51_trunk } else { 1.0 }`. + float correction = (trunk_mean > TRUNK_EPS) ? (effective / trunk_mean) : 1.0f; + out[corr_idx_out] = correction; + + __threadfence_system(); +} + +/* ══════════════════════════════════════════════════════════════════════ + * F32 IN-PLACE SCALE KERNEL (alpha read from device memory) + * + * y[i] *= *alpha_ptr for i = 0..n-1 + * + * Mirror of `dqn_scale_f32_kernel` in `dqn_utility_kernels.cu`, but `alpha` + * is a device pointer (single f32) read at execution time rather than a + * scalar baked into the kernel-arg buffer at capture time. Required for + * SP7 capture-safe per-branch budget consumption. + * + * Launch config: grid=(ceil(n/256), 1, 1), block=(256, 1, 1). + * ══════════════════════════════════════════════════════════════════════ */ +extern "C" __global__ void dqn_scale_f32_dev_ptr_kernel( + float* y, + const float* __restrict__ alpha_ptr, + int n +) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) { + float alpha = *alpha_ptr; + // NaN-safe: matches dqn_scale_f32_kernel — IEEE 754 says 0*NaN=NaN; + // when alpha=0 (cold-start / MSE warmup) the y buffer can contain + // NaN, scaling by 0 must yield 0. + y[i] = (alpha == 0.0f) ? 0.0f : y[i] * alpha; + } +} + +/* ══════════════════════════════════════════════════════════════════════ + * F32 SAXPY KERNEL (alpha read from device memory) + * + * y[i] += (*alpha_ptr) * x[i] for i = 0..n-1 + * + * Mirror of `dqn_saxpy_f32_kernel` in `dqn_utility_kernels.cu`, but `alpha` + * is a device pointer (single f32) read at execution time. Used by + * `apply_cql_saxpy` and per-branch variant in the SP7 capture-safe path. + * + * Launch config: grid=(ceil(n/256), 1, 1), block=(256, 1, 1). + * ══════════════════════════════════════════════════════════════════════ */ +extern "C" __global__ void dqn_saxpy_f32_dev_ptr_kernel( + float* __restrict__ y, + const float* __restrict__ x, + const float* __restrict__ alpha_ptr, + int n +) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) { + float alpha = *alpha_ptr; + y[i] = y[i] + alpha * x[i]; + } +} diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 8c23408ac..a9e8cbaee 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -468,6 +468,21 @@ static SP5_TRAINING_METRICS_EMA_CUBIN: &[u8] = static LOSS_BALANCE_CONTROLLER_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/loss_balance_controller_kernel.cubin")); +/// SP7 (2026-05-03): GPU dispatch for cql/c51 budget consumer + dev-ptr +/// SAXPY/scale variants. Consolidates three kernels into one cubin so they +/// share a CUmodule with the SP7 dispatch buffer. See +/// `consume_lb_budget_kernel.cu` for the per-kernel contract. +/// +/// Architecture: the existing SAXPY/scale path took `alpha: f32` by value at +/// the call site, which baked the value into the kernel-arg buffer at +/// CUDA-Graph capture time. After SP7 introduced a controller writing to ISV +/// each step, the captured `aux_child` graph saw step-0 (FoldReset sentinel) +/// values frozen for every replay, defeating the controller. New kernels +/// read `alpha` from a device pointer (mapped-pinned) populated by +/// `lb_budget_dispatch` each step → capture-safe end-to-end. +static CONSUME_LB_BUDGET_CUBIN: &[u8] = + include_bytes!(concat!(env!("OUT_DIR"), "/consume_lb_budget_kernel.cubin")); + /// Plan 4 Task 3 (E.3): IQN multi-quantile diagnostic EMAs into ISV[99..103). /// 4-block (256 threads/block, shmem-reduce, no atomicAdd) kernel launched /// alongside `h_s2_rms_ema_update` from `training_loop.rs`. Reads the IQN @@ -4449,6 +4464,39 @@ pub struct GpuDqnTrainer { /// `feedback_no_partial_refactor`. Loaded from /// `loss_balance_controller_kernel.cubin`. loss_balance_controller_kernel: CudaFunction, + /// SP7 (2026-05-03): GPU dispatch kernel for the cql/c51 budget consumer. + /// 8 threads (2 heads × 4 branches), single block. Reads ISV activation + /// flag + controller budget per (head, branch); resolves + /// `(active >= 0.5) ? max(ctrl, EPS_DIV) : bootstrap` on-device; writes + /// trunk_mean + per-branch correction to `lb_budget_effective_buf` (10 + /// f32 mapped-pinned). Replaces the host-side if/else inside + /// `compute_adaptive_budgets` that froze at CUDA-Graph capture time. + /// Loaded from `consume_lb_budget_kernel.cubin`. + lb_budget_dispatch_kernel: CudaFunction, + /// SP7 (2026-05-03): SAXPY variant reading `alpha` from a device pointer. + /// `y[i] += (*alpha_ptr) * x[i]`. Used by `apply_cql_saxpy` and per-branch + /// variant when consuming `lb_budget_effective_buf` slots. Loaded from + /// `consume_lb_budget_kernel.cubin` (separate CUmodule from + /// `dqn_saxpy_f32_kernel` so existing scalar callers remain unchanged). + saxpy_f32_dev_ptr_aux: CudaFunction, + /// SP7 (2026-05-03): in-place scale variant reading `alpha` from a device + /// pointer. `y[i] *= *alpha_ptr` (NaN-safe at alpha=0 — matches + /// `dqn_scale_f32_kernel`). Used by `apply_c51_budget_scale` and + /// per-branch variant when consuming `lb_budget_effective_buf` slots. + /// Loaded from `consume_lb_budget_kernel.cubin`. + scale_f32_dev_ptr_aux: CudaFunction, + /// SP7 (2026-05-03): mapped-pinned scratch holding the GPU-resolved + /// effective per-branch budgets for cql/c51. Layout (10 f32): + /// [0] cql_trunk_mean (consumed by `apply_c51_budget_scale`'s + /// sibling apply_cql_saxpy(trunk)) + /// [1..5] cql_correction[branch] (consumed by `apply_cql_saxpy_branch`) + /// [5] c51_trunk_mean (consumed by `apply_c51_budget_scale`) + /// [6..10] c51_correction[branch] (consumed by `apply_c51_budget_scale_branch`) + /// Populated each step by `lb_budget_dispatch_kernel`; consumed by the + /// dev-ptr SAXPY/scale kernels in the same captured `aux_child` graph. + /// Mapped-pinned so HEALTH_DIAG can read host-side after stream sync — + /// no DtoH copy. + pub(crate) lb_budget_effective_buf: super::mapped_pinned::MappedF32Buffer, // ── SP5 Task A4: Pearl 4 per-group Adam β1/β2/ε kernels ───────────── /// SP5 Task A4 (2026-05-01): auxiliary gradient cosine similarity kernel. /// Single-block 8-thread kernel (one thread per SP4 param group). @@ -5084,20 +5132,17 @@ pub struct GpuDqnTrainer { /// B4/G5: Last adaptive gradient budget for IQN (0.10 + 0.30×health). /// Populated by `FusedTrainingCtx::compute_adaptive_budgets`. Default: health=1 → 0.40. pub(crate) last_iqn_budget_eff: f32, - /// B4/G5: Last adaptive gradient budget for CQL (0.10×(1−regime)×health). - /// Populated by `FusedTrainingCtx::compute_adaptive_budgets`. Default: 0.00 (regime_stability=1). - pub(crate) last_cql_budget_eff: f32, - /// B4/G5: Last adaptive gradient budget for C51 (1−iqn−cql−ens). - /// Populated by `FusedTrainingCtx::compute_adaptive_budgets`. Default: health=1 → 0.55. - pub(crate) last_c51_budget_eff: f32, /// B4/G5: Adaptive gradient budget for ensemble (constant 0.05). /// Populated by `FusedTrainingCtx::compute_adaptive_budgets`. pub(crate) last_ens_budget_eff: f32, /// SP6 Pearl 2: per-branch budget arrays — set by compute_adaptive_budgets. /// Index order: dir=0, mag=1, ord=2, urg=3. - pub(crate) last_c51_budget_per_branch: [f32; 4], + /// + /// SP7 (2026-05-03): the CQL/C51 caches were removed because the + /// host-side dispatch they relied on was the source of the capture-time + /// freeze (audit Fix 33). HEALTH_DIAG now reads the GPU-resolved values + /// from `lb_budget_effective_buf` via `FusedTrainingCtx::last_{cql,c51}_budget_per_branch`. pub(crate) last_iqn_budget_per_branch: [f32; 4], - pub(crate) last_cql_budget_per_branch: [f32; 4], pub(crate) last_ens_budget_per_branch: [f32; 4], /// D1/N1: Ring buffer of best-health weight snapshots for temporal self-distillation. @@ -9743,7 +9788,8 @@ impl GpuDqnTrainer { base > 0.0 } - /// F8/G5: Scale the C51-contributed portion of `grad_buf` by the adaptive c51_budget. + /// F8/G5 + SP7 (2026-05-03): Scale the C51-contributed portion of `grad_buf` + /// by the adaptive c51_budget read from `lb_budget_effective_buf` at runtime. /// /// Must be called AFTER `submit_forward_ops_main` (which runs C51 backward and writes /// C51 gradients into `grad_buf`), and BEFORE any CQL/IQN SAXPY accumulations. @@ -9751,21 +9797,23 @@ impl GpuDqnTrainer { /// The composite effect is: /// grad = c51_budget × C51_grad + cql_budget × CQL_grad + iqn_budget × IQN_grad /// - /// Uses `scale_f32_ungraphed` (loaded from a separate CUmodule) so it is safe to - /// call outside CUDA graph capture without corrupting the forward_child graph state. - pub fn apply_c51_budget_scale(&mut self, c51_budget: f32) -> Result<(), MLError> { - // No-op when budget ≈ 1.0 (common at health=0 — avoids unnecessary kernel launch). - if (c51_budget - 1.0).abs() < 1e-6 { - return Ok(()); - } + /// SP7 (2026-05-03): `c51_budget` is read on-device from `alpha_dev_ptr` (a u64 + /// device pointer to a single f32 — typically `&lb_budget_effective_buf[5]`, + /// the c51_trunk_mean slot). Eliminates the capture-time freeze of the prior + /// scalar-arg variant; see `consume_lb_budget_kernel.cu` and audit Fix 33. + pub fn apply_c51_budget_scale(&mut self, alpha_dev_ptr: u64) -> Result<(), MLError> { + // SP7 (2026-05-03): no host-side `if budget ≈ 1.0 { return; }` skip — the + // value lives on-device and we can't read it without DtoH (forbidden mid- + // capture). The kernel itself short-circuits the alpha=0.0 NaN-safety + // clause; an alpha=1.0 multiply is a numerically-cheap no-op. let grad_ptr = self.ptrs.grad_buf; let n = self.total_params as i32; let blocks = ((self.total_params as u32 + 255) / 256) as u32; unsafe { self.stream - .launch_builder(&self.scale_f32_ungraphed) + .launch_builder(&self.scale_f32_dev_ptr_aux) .arg(&grad_ptr) - .arg(&c51_budget) + .arg(&alpha_dev_ptr) .arg(&n) .launch(LaunchConfig { grid_dim: (blocks, 1, 1), @@ -9777,18 +9825,20 @@ impl GpuDqnTrainer { Ok(()) } - /// SP6 Pearl 2: scale only branch `branch_idx` parameter slice of grad_buf by `budget`. + /// SP6 Pearl 2 + SP7 (2026-05-03): scale only branch `branch_idx` parameter + /// slice of grad_buf by the correction factor read from + /// `lb_budget_effective_buf` at runtime. /// - /// Called after `apply_c51_budget_scale(trunk_mean)` has scaled the entire grad_buf. - /// The correction factor is `branch_budget / trunk_mean`, so the branch slice ends up - /// scaled by `branch_budget` rather than `trunk_mean`. + /// Called after `apply_c51_budget_scale(trunk_dev_ptr)` has scaled the entire + /// grad_buf by c51_trunk_mean. The correction factor (loaded from + /// `correction_dev_ptr`) is `branch_budget / trunk_mean`, so the branch slice + /// ends up scaled by `branch_budget` rather than `trunk_mean`. /// /// Branch b owns tensors `(8+4b)..(12+4b)` in the padded flat layout. /// Padding bytes in `grad_buf` are always zero so scaling them is a no-op. - pub fn apply_c51_budget_scale_branch(&mut self, branch_idx: usize, budget: f32) -> Result<(), MLError> { - if (budget - 1.0).abs() < 1e-6 { - return Ok(()); - } + pub fn apply_c51_budget_scale_branch(&mut self, branch_idx: usize, correction_dev_ptr: u64) -> Result<(), MLError> { + // SP7 (2026-05-03): no host-side `if correction ≈ 1.0 { return; }` skip + // for the same reason as the trunk variant — value lives on-device. let f32_size = std::mem::size_of::(); let param_sizes = compute_param_sizes(&self.config); let first_tensor = 8 + branch_idx * 4; @@ -9801,9 +9851,9 @@ impl GpuDqnTrainer { let blocks = ((len_elems as u32 + 255) / 256) as u32; unsafe { self.stream - .launch_builder(&self.scale_f32_ungraphed) + .launch_builder(&self.scale_f32_dev_ptr_aux) .arg(&grad_slice_ptr) - .arg(&budget) + .arg(&correction_dev_ptr) .arg(&len_elems) .launch(LaunchConfig { grid_dim: (blocks, 1, 1), @@ -9821,22 +9871,26 @@ impl GpuDqnTrainer { /// /// Called after `apply_cql_gradient` populated `cql_grad_scratch`. /// No per-component clip — single global clip in Adam handles safety. - pub fn apply_cql_saxpy(&mut self, cql_budget: f32) -> Result<(), MLError> { + /// + /// SP7 (2026-05-03): `cql_budget` is read on-device from `alpha_dev_ptr` — + /// typically `&lb_budget_effective_buf[0]`, the cql_trunk_mean slot. + /// Eliminates the capture-time freeze of the prior scalar-arg variant. + pub fn apply_cql_saxpy(&mut self, alpha_dev_ptr: u64) -> Result<(), MLError> { let total = self.total_params as i32; let blocks = self.grad_norm_blocks as u32; let scratch_ptr = self.ptrs.cql_grad_scratch; let grad_ptr = self.ptrs.grad_buf; - // B4/G5: scale CQL gradient contribution by adaptive budget (0.10×(1−regime)×health). - let alpha = cql_budget; - // Plain SAXPY: grad_buf += cql_budget * cql_scratch - // Uses aux_child-specific handle — saxpy_f32_kernel is captured in forward_child. + // Plain SAXPY: grad_buf += (*alpha_dev_ptr) * cql_scratch + // Uses aux_child-specific dev-ptr handle — saxpy_f32_kernel is captured + // in forward_child and the scalar saxpy_f32_aux is reused by other + // callers (distill, etc.). The dev-ptr variant is unique to SP7. unsafe { self.stream - .launch_builder(&self.saxpy_f32_aux) + .launch_builder(&self.saxpy_f32_dev_ptr_aux) .arg(&grad_ptr) .arg(&scratch_ptr) - .arg(&alpha) + .arg(&alpha_dev_ptr) .arg(&total) .launch(LaunchConfig { grid_dim: (blocks, 1, 1), @@ -9849,15 +9903,17 @@ impl GpuDqnTrainer { Ok(()) } - /// SP6 Pearl 2: SAXPY only branch `branch_idx` parameter slice: - /// grad_buf[branch_slice] += budget * cql_grad_scratch[branch_slice] + /// SP6 Pearl 2 + SP7 (2026-05-03): SAXPY only branch `branch_idx` parameter + /// slice: + /// grad_buf[branch_slice] += (*correction_dev_ptr) * cql_grad_scratch[branch_slice] /// - /// Called after `apply_cql_saxpy(trunk_mean)` has performed the full-buffer SAXPY. - /// The correction factor is `branch_budget / trunk_mean`, so the branch slice ends up - /// accumulating `branch_budget * cql_scratch` rather than `trunk_mean * cql_scratch`. + /// Called after `apply_cql_saxpy(trunk_dev_ptr)` has performed the full-buffer + /// SAXPY. The correction factor (loaded from `correction_dev_ptr`) is + /// `branch_budget / trunk_mean`, so the branch slice ends up accumulating + /// `branch_budget * cql_scratch` rather than `trunk_mean * cql_scratch`. /// /// Branch b owns tensors `(8+4b)..(12+4b)` in the padded flat layout. - pub fn apply_cql_saxpy_branch(&mut self, branch_idx: usize, budget: f32) -> Result<(), MLError> { + pub fn apply_cql_saxpy_branch(&mut self, branch_idx: usize, correction_dev_ptr: u64) -> Result<(), MLError> { let f32_size = std::mem::size_of::(); let param_sizes = compute_param_sizes(&self.config); let first_tensor = 8 + branch_idx * 4; @@ -9871,10 +9927,10 @@ impl GpuDqnTrainer { let blocks = ((len_elems as u32 + 255) / 256) as u32; unsafe { self.stream - .launch_builder(&self.saxpy_f32_aux) + .launch_builder(&self.saxpy_f32_dev_ptr_aux) .arg(&grad_ptr) .arg(&scratch_ptr) - .arg(&budget) + .arg(&correction_dev_ptr) .arg(&len_elems) .launch(LaunchConfig { grid_dim: (blocks, 1, 1), @@ -11516,6 +11572,145 @@ impl GpuDqnTrainer { Ok(()) } + /// SP7 (2026-05-03): launch the GPU dispatch kernel that resolves + /// `(active >= 0.5) ? max(controller, EPS_DIV) : bootstrap` for every + /// (head ∈ {CQL, C51}, branch ∈ [0..4)) on-device. + /// + /// Reads `ISV[LB_{CQL,C51}_ACTIVE_BASE+b]` (activation flag) and + /// `ISV[BUDGET_{CQL,C51}_BASE+b]` (controller-derived budget); writes the + /// trunk_mean + per-branch correction tuples to `lb_budget_effective_buf` + /// (10 f32 mapped-pinned). Bootstrap constants are passed as kernel args + /// (CQL=0.02, C51=0.05 — byte-identical to the prior host-side + /// `CQL_BOOTSTRAP_BUDGET` / `C51_BOOTSTRAP_BUDGET` literals in + /// `compute_adaptive_budgets` and to `loss_balance_controller_kernel.cu`'s + /// `COLD_START_FLOOR_*` anchors per `feedback_isv_for_adaptive_bounds`). + /// + /// Must run BEFORE `apply_c51_budget_scale` / `apply_cql_saxpy` (and their + /// per-branch variants) inside the same captured `aux_child` graph. The + /// downstream SAXPY/scale kernels read `alpha` from a device pointer + /// pointing into `lb_budget_effective_buf`, so every graph replay sees + /// fresh values written by this kernel against the runtime ISV state — + /// no capture-time freeze of the host-side if/else dispatch. + pub(crate) fn launch_lb_budget_dispatch(&self) -> Result<(), MLError> { + use crate::cuda_pipeline::sp5_isv_slots::{ + BUDGET_CQL_BASE, BUDGET_C51_BASE, + LB_CQL_ACTIVE_BASE, LB_C51_ACTIVE_BASE, + }; + + debug_assert!(self.isv_signals_dev_ptr != 0, + "launch_lb_budget_dispatch: isv_signals_dev_ptr must be allocated"); + debug_assert_eq!(self.lb_budget_effective_buf.len, 10, + "launch_lb_budget_dispatch: lb_budget_effective_buf must be 10 f32"); + + let isv_dev = self.isv_signals_dev_ptr; + let out_dev = self.lb_budget_effective_buf.dev_ptr; + + let active_cql_isv_base_i32: i32 = LB_CQL_ACTIVE_BASE as i32; + let active_c51_isv_base_i32: i32 = LB_C51_ACTIVE_BASE as i32; + let budget_cql_isv_base_i32: i32 = BUDGET_CQL_BASE as i32; + let budget_c51_isv_base_i32: i32 = BUDGET_C51_BASE as i32; + // SP7 cold-start basis — must remain byte-identical to the prior host- + // side literals in `compute_adaptive_budgets` (CQL_BOOTSTRAP_BUDGET=0.02 + // / C51_BOOTSTRAP_BUDGET=0.05). Passed as kernel args (not hardcoded + // in kernel body) per `feedback_isv_for_adaptive_bounds`. + let cql_bootstrap: f32 = 0.02_f32; + let c51_bootstrap: f32 = 0.05_f32; + + unsafe { + self.stream + .launch_builder(&self.lb_budget_dispatch_kernel) + .arg(&isv_dev) + .arg(&active_cql_isv_base_i32) + .arg(&active_c51_isv_base_i32) + .arg(&budget_cql_isv_base_i32) + .arg(&budget_c51_isv_base_i32) + .arg(&cql_bootstrap) + .arg(&c51_bootstrap) + .arg(&out_dev) + .launch(LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (8, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("lb_budget_dispatch: {e}")))?; + } + + Ok(()) + } + + /// SP7 (2026-05-03): device pointer to the `cql_trunk_mean` slot in + /// `lb_budget_effective_buf`. Pass into `apply_cql_saxpy` so the kernel + /// reads the GPU-resolved budget at execution time (not capture time). + #[inline] + pub(crate) fn lb_budget_cql_trunk_dev_ptr(&self) -> u64 { + // Slot 0 = cql_trunk_mean. + self.lb_budget_effective_buf.dev_ptr + } + + /// SP7 (2026-05-03): device pointer to the `cql_correction[branch]` slot. + /// Pass into `apply_cql_saxpy_branch` so the kernel reads the per-branch + /// correction at execution time. + #[inline] + pub(crate) fn lb_budget_cql_correction_dev_ptr(&self, branch: usize) -> u64 { + debug_assert!(branch < 4, "branch must be in [0,4)"); + // Slots 1..5 = cql_correction[0..4]. + let f32_size = std::mem::size_of::() as u64; + self.lb_budget_effective_buf.dev_ptr + (1 + branch as u64) * f32_size + } + + /// SP7 (2026-05-03): device pointer to the `c51_trunk_mean` slot. + #[inline] + pub(crate) fn lb_budget_c51_trunk_dev_ptr(&self) -> u64 { + // Slot 5 = c51_trunk_mean. + let f32_size = std::mem::size_of::() as u64; + self.lb_budget_effective_buf.dev_ptr + 5 * f32_size + } + + /// SP7 (2026-05-03): device pointer to the `c51_correction[branch]` slot. + #[inline] + pub(crate) fn lb_budget_c51_correction_dev_ptr(&self, branch: usize) -> u64 { + debug_assert!(branch < 4, "branch must be in [0,4)"); + // Slots 6..10 = c51_correction[0..4]. + let f32_size = std::mem::size_of::() as u64; + self.lb_budget_effective_buf.dev_ptr + (6 + branch as u64) * f32_size + } + + /// SP7 (2026-05-03): host-side read of the GPU-resolved per-branch + /// effective budgets and trunk means. Returns `(cql_per_branch, c51_per_branch, + /// cql_trunk, c51_trunk)`. Reconstructs per-branch effectives from the + /// stored `(trunk, correction)` decomposition: `effective[b] = trunk × correction[b]`. + /// + /// Used by HEALTH_DIAG to surface the dispatched values that actually fed + /// into SAXPY/scale this step. Caller MUST have `sync_all_streams()`'d the + /// producing stream first — mapped-pinned coherence requires the kernel to + /// have completed before host reads (otherwise reads return stale or + /// in-flight values). + pub(crate) fn read_lb_budget_effective(&self) -> ([f32; 4], [f32; 4], f32, f32) { + // Safety: dev pointer and host pointer back the same physical memory + // (cuMemHostAlloc DEVICEMAP); read_volatile defeats CPU cache. + let host = self.lb_budget_effective_buf.host_ptr; + let read = |offset: usize| -> f32 { + unsafe { std::ptr::read_volatile(host.add(offset)) } + }; + let cql_trunk = read(0); + let cql_corr = [read(1), read(2), read(3), read(4)]; + let c51_trunk = read(5); + let c51_corr = [read(6), read(7), read(8), read(9)]; + let cql_per_branch = [ + cql_trunk * cql_corr[0], + cql_trunk * cql_corr[1], + cql_trunk * cql_corr[2], + cql_trunk * cql_corr[3], + ]; + let c51_per_branch = [ + c51_trunk * c51_corr[0], + c51_trunk * c51_corr[1], + c51_trunk * c51_corr[2], + c51_trunk * c51_corr[3], + ]; + (cql_per_branch, c51_per_branch, cql_trunk, c51_trunk) + } + /// SP5 Task A4 (2026-05-01): launch the two-kernel Pearl 4 chain + /// 24 `apply_pearls_ad_kernel` launches to smooth β1/β2/ε into ISV. /// @@ -14605,6 +14800,34 @@ impl GpuDqnTrainer { .map_err(|e| MLError::ModelError(format!("loss_balance_controller_update load: {e}")))? }; + // SP7 (2026-05-03): load consume_lb_budget_kernel.cubin — three kernels + // (lb_budget_dispatch + dqn_scale_f32_dev_ptr_kernel + + // dqn_saxpy_f32_dev_ptr_kernel) sharing one CUmodule. Separate from + // `aux_module` (which holds the scalar saxpy_f32/scale_f32 reused by + // distill, VSN dW dilution, etc.) so existing scalar callers remain + // unchanged. The dispatch kernel runs each step inside the captured + // `aux_child` graph; the dev-ptr SAXPY/scale kernels are launched by + // `apply_c51_budget_scale` / `apply_cql_saxpy` (and per-branch variants) + // immediately after, reading `alpha` from `lb_budget_effective_buf`. + let consume_lb_budget_module = stream.context().load_cubin(CONSUME_LB_BUDGET_CUBIN.to_vec()) + .map_err(|e| MLError::ModelError(format!("consume_lb_budget cubin load: {e}")))?; + let lb_budget_dispatch_kernel = consume_lb_budget_module + .load_function("lb_budget_dispatch") + .map_err(|e| MLError::ModelError(format!("lb_budget_dispatch load: {e}")))?; + let scale_f32_dev_ptr_aux = consume_lb_budget_module + .load_function("dqn_scale_f32_dev_ptr_kernel") + .map_err(|e| MLError::ModelError(format!("dqn_scale_f32_dev_ptr_kernel load: {e}")))?; + let saxpy_f32_dev_ptr_aux = consume_lb_budget_module + .load_function("dqn_saxpy_f32_dev_ptr_kernel") + .map_err(|e| MLError::ModelError(format!("dqn_saxpy_f32_dev_ptr_kernel load: {e}")))?; + // SP7 (2026-05-03): allocate the 10 f32 mapped-pinned scratch holding + // GPU-resolved [cql_trunk, cql_corr×4, c51_trunk, c51_corr×4]. Zero-init + // is fine — first dispatch-kernel launch overwrites the slots before + // any consumer reads them. See `consume_lb_budget_kernel.cu` top-of-file + // block for layout. + let lb_budget_effective_buf = unsafe { MappedF32Buffer::new(10) } + .map_err(|e| MLError::ModelError(format!("SP7 lb_budget_effective_buf alloc: {e}")))?; + // SP5 Task A4 (2026-05-01): allocate grad_prev_buf_per_group [total_params f32]. // Zeroed at construction; fold-boundary reset zeroes it again (Pearl A sentinel: // zero grad_prev → cosine_sim=0 on first step → β1/β2 start at envelope midpoints). @@ -17787,6 +18010,10 @@ impl GpuDqnTrainer { pearl_3_sigma_kernel, pearl_2_budget_kernel, loss_balance_controller_kernel, + lb_budget_dispatch_kernel, + saxpy_f32_dev_ptr_aux, + scale_f32_dev_ptr_aux, + lb_budget_effective_buf, grad_cosine_sim_kernel, pearl_4_adam_hparams_kernel, q_skew_kurtosis_kernel, @@ -17988,12 +18215,8 @@ impl GpuDqnTrainer { last_cql_alpha_eff: 0.0, last_sarsa_tau_factor: 0.0, last_iqn_budget_eff: 0.40, - last_cql_budget_eff: 0.00, - last_c51_budget_eff: 0.55, last_ens_budget_eff: 0.05, - last_c51_budget_per_branch: [0.55; 4], last_iqn_budget_per_branch: [0.40; 4], - last_cql_budget_per_branch: [0.00; 4], last_ens_budget_per_branch: [0.05; 4], q_snapshots: crate::cuda_pipeline::q_snapshot::SnapshotRing::new( stream_for_snapshots, diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index c230b2122..c69b1209f 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -40,8 +40,7 @@ use crate::cuda_pipeline::gpu_dqn_trainer::{ GpuDqnTrainConfig, GpuDqnTrainer, TAU_EFF_INDEX, }; use crate::cuda_pipeline::sp5_isv_slots::{ - BUDGET_C51_BASE, BUDGET_IQN_BASE, BUDGET_CQL_BASE, BUDGET_ENS_BASE, - LB_CQL_ACTIVE_BASE, LB_C51_ACTIVE_BASE, + BUDGET_IQN_BASE, BUDGET_ENS_BASE, ADAM_BETA1_BASE, ADAM_BETA2_BASE, ADAM_EPS_BASE, IQN_TAU_BASE, }; @@ -1877,9 +1876,12 @@ impl FusedTrainingCtx { // not as norm clips. This preserves gradient magnitude variation while steering // component balance. Safety is still handled by the single global clip in Adam. - // SP6 Pearl 2: Compute adaptive budgets — returns per-branch arrays + trunk means. - let (c51_branch, iqn_branch, cql_branch, _ens_branch, - c51_trunk, _iqn_trunk, cql_trunk, _ens_trunk) = self.compute_adaptive_budgets(); + // SP6 Pearl 2 + SP7 (2026-05-03): launches the on-device dispatch + // kernel populating `lb_budget_effective_buf` with the resolved CQL/C51 + // trunk + per-branch corrections; returns IQN and ENS scalars as + // before (prompt-scoped — host-branch dispatch only applied to CQL/C51). + let (iqn_branch, _ens_branch, _iqn_trunk, _ens_trunk) = + self.compute_adaptive_budgets(); // F8/G5: Scale C51 contribution in grad_buf by c51_budget. // C51 backward already wrote into grad_buf during submit_forward_ops_main (Phase 2). @@ -1891,17 +1893,24 @@ impl FusedTrainingCtx { // When `c51_budget ≈ 1.0` the scale is a no-op and delta = 0.0. self.trainer.grad_decomp_snapshot_c51_bs() .map_err(|e| anyhow::anyhow!("Task 2.0 grad_decomp_snapshot_c51_bs: {e}"))?; - // SP6 Pearl 2: trunk/value first (mean budget), then per-branch correction sub-launches. - // Branch slice has already been scaled by c51_trunk; correction = branch[b]/trunk_mean - // so branch slice ends up scaled by c51_branch[b]. - self.trainer.apply_c51_budget_scale(c51_trunk) + // SP6 Pearl 2 + SP7 (2026-05-03): trunk/value first (mean budget), then + // per-branch correction sub-launches. Both kernels now read `alpha` + // from `lb_budget_effective_buf` at execution time — capture-safe. + // Net effect: branch slice ends up scaled by trunk × correction[b] + // = effective[b], identical to the prior host-resolved semantics but + // with values written by the dispatch kernel against the runtime ISV + // state (not frozen at capture-time step 0). The "skip when + // correction ≈ 1.0" optimization is dropped — the value lives on- + // device and can't be branched on host-side without reintroducing + // the same capture-freeze bug. Cost: 4 cheap kernel launches with + // alpha=1.0 in the steady state (multiply by 1.0 is a numeric no-op). + let c51_trunk_dev_ptr = self.trainer.lb_budget_c51_trunk_dev_ptr(); + self.trainer.apply_c51_budget_scale(c51_trunk_dev_ptr) .map_err(|e| anyhow::anyhow!("c51_budget_scale_trunk: {e}"))?; for b in 0..4_usize { - let correction = if c51_trunk > 1e-6 { c51_branch[b] / c51_trunk } else { 1.0_f32 }; - if (correction - 1.0).abs() > 1e-6 { - self.trainer.apply_c51_budget_scale_branch(b, correction) - .map_err(|e| anyhow::anyhow!("c51_budget_scale_branch[{b}]: {e}"))?; - } + let correction_dev_ptr = self.trainer.lb_budget_c51_correction_dev_ptr(b); + self.trainer.apply_c51_budget_scale_branch(b, correction_dev_ptr) + .map_err(|e| anyhow::anyhow!("c51_budget_scale_branch[{b}]: {e}"))?; } self.trainer.grad_decomp_launch_c51_bs() .map_err(|e| anyhow::anyhow!("Task 2.0 grad_decomp c51_bs: {e}"))?; @@ -2422,15 +2431,18 @@ impl FusedTrainingCtx { // `grad_buf`. self.trainer.grad_decomp_snapshot_cql_sx() .map_err(|e| anyhow::anyhow!("Task 2.0 grad_decomp_snapshot_cql_sx: {e}"))?; - // SP6 Pearl 2: trunk/value SAXPY first (mean budget), then per-branch corrections. - self.trainer.apply_cql_saxpy(cql_trunk) + // SP6 Pearl 2 + SP7 (2026-05-03): trunk/value SAXPY first + // (mean budget), then per-branch corrections. Both kernels + // read `alpha` from `lb_budget_effective_buf` (populated by + // the dispatch kernel earlier in `compute_adaptive_budgets`) + // — capture-safe, see audit Fix 33. + let cql_trunk_dev_ptr = self.trainer.lb_budget_cql_trunk_dev_ptr(); + self.trainer.apply_cql_saxpy(cql_trunk_dev_ptr) .map_err(|e| anyhow::anyhow!("CQL SAXPY trunk: {e}"))?; for b in 0..4_usize { - let correction = if cql_trunk > 1e-6 { cql_branch[b] / cql_trunk } else { 1.0_f32 }; - if (correction - 1.0).abs() > 1e-6 { - self.trainer.apply_cql_saxpy_branch(b, correction) - .map_err(|e| anyhow::anyhow!("CQL SAXPY branch[{b}]: {e}"))?; - } + let correction_dev_ptr = self.trainer.lb_budget_cql_correction_dev_ptr(b); + self.trainer.apply_cql_saxpy_branch(b, correction_dev_ptr) + .map_err(|e| anyhow::anyhow!("CQL SAXPY branch[{b}]: {e}"))?; } self.trainer.grad_decomp_launch_cql_sx() .map_err(|e| anyhow::anyhow!("Task 2.0 grad_decomp cql_sx: {e}"))?; @@ -3405,90 +3417,75 @@ impl FusedTrainingCtx { self.trainer.last_sarsa_tau_factor } - /// SP6 Pearl 2: Compute adaptive gradient budgets from ISV[190..210). + /// SP6 Pearl 2 + SP7 (2026-05-03): Compute adaptive gradient budgets. /// - /// Returns `(c51_branch, iqn_branch, cql_branch, ens_branch, c51_trunk, iqn_trunk, - /// cql_trunk, ens_trunk)`. The per-branch arrays feed per-branch SAXPY correction - /// sub-launches on branch HEAD parameter slices. The trunk scalars (mean of 4 branch - /// budgets) feed the full-buf trunk/value scaling call — preserving SP5 Layer B - /// behaviour for shared parameters (D3 decision in SP6 spec). + /// Returns `(iqn_branch, ens_branch, iqn_trunk, ens_trunk)`. The CQL and + /// C51 budgets are no longer returned as scalars — they live on-device in + /// `lb_budget_effective_buf`, populated by `launch_lb_budget_dispatch` + /// (called inline below). Callers consume them via the trainer's + /// `lb_budget_{cql,c51}_{trunk,correction}_dev_ptr` accessors and pass + /// the device pointers to the SAXPY/scale kernels, which read `alpha` + /// from device memory at runtime — capture-safe. /// - /// Cold-start bootstrap: per-(head, branch) activation flag at - /// `LB_{CQL,C51}_ACTIVE_BASE` discriminates "controller has fired" from - /// "controller hasn't fired yet". IQN keeps a structural BASE_IQN=0.11 - /// floor (reference budget, Invariant 1 carve-out). CQL/C51 dispatch on - /// the activation flag: bootstrap (0.02/0.05) when active < 0.5, - /// controller verbatim (with EPS_DIV minimum) when active >= 0.5. ENS - /// keeps the prior numeric-equality bootstrap because no controller - /// drives it (SP5 Pearl 2 stopped writing BUDGET_ENS_BASE in SP7 T6). + /// Bug fixed (commit landing this change): `cql_active` / `c51_active` + /// were resolved via host-side `if/else` here, then passed by value to + /// `apply_c51_budget_scale` / `apply_cql_saxpy`. Inside the captured + /// `aux_child` graph, the `if/else` ran ONCE at capture time (step 0 of + /// each fold, when the activation flag is the FoldReset sentinel `0.0`) + /// → bootstrap branch resolved → literal `0.02` / `0.05` baked into the + /// captured kernel-arg buffers, freezing the budget for ~1000 graph + /// replays per fold regardless of what the controller wrote at runtime. + /// Same disease as SP4 host-side EMA elimination 2026-05-01; see + /// `feedback_no_cpu_compute_strict.md` and audit Fix 33. /// - /// SP7 activation-flag fix (2026-05-03): replaces the prior - /// numeric-equality bootstrap (`raw < 1e-8 → bootstrap`) for CQL/C51, - /// which couldn't distinguish "controller said bootstrap" from - /// "controller hasn't fired yet" because `CQL_BOOTSTRAP_BUDGET=0.02` - /// numerically equaled the kernel's `COLD_START_FLOOR_CQL=0.02`. The - /// activation flag is a binary discriminator written by the controller - /// itself: 1.0 when the active-path computes a real budget update, - /// 0.0 (no-op via apply_pearls_ad short-circuit) on genuine cold-start. + /// IQN keeps a structural BASE_IQN=0.11 floor (reference budget, + /// Invariant 1 carve-out — the controller's ratio compute would + /// degenerate if IQN reached 0). ENS keeps the prior numeric-equality + /// bootstrap because no controller drives BUDGET_ENS_BASE in SP7 + /// (Pearl 2 stopped writing it in T6). Both IQN and ENS still resolve + /// host-side here; the capture-time freeze applies but is + /// out-of-scope for this fix per prompt scope. pub(crate) fn compute_adaptive_budgets( &mut self, - ) -> ([f32; 4], [f32; 4], [f32; 4], [f32; 4], f32, f32, f32, f32) { - let mut c51 = [0.0_f32; 4]; + ) -> ([f32; 4], [f32; 4], f32, f32) { + // SP7 (2026-05-03): launch the on-device dispatch kernel that resolves + // `(active >= 0.5) ? max(controller, EPS_DIV) : bootstrap` for CQL and + // C51 across all 4 branches, computes the trunk mean and per-branch + // correction, and writes them to `lb_budget_effective_buf`. Captured + // into the same `aux_child` graph as the SAXPY/scale callers below; + // every replay sees fresh runtime ISV state. + if let Err(e) = self.trainer.launch_lb_budget_dispatch() { + tracing::error!("SP7 lb_budget_dispatch launch failed: {e}"); + } + let mut iqn = [0.0_f32; 4]; - let mut cql = [0.0_f32; 4]; let mut ens = [0.0_f32; 4]; - // SP7 activation-flag fix (2026-05-03): per-(head, branch) activation - // dispatch. ACTIVE_THRESHOLD=0.5 is a binary discriminator (the slot - // takes 0.0 or 1.0); not a tuning knob. Bootstrap values must match - // the cold-start basis in loss_balance_controller_kernel.cu - // (COLD_START_FLOOR_CQL=0.02, COLD_START_FLOOR_C51=0.05). IQN keeps - // the structural BASE_IQN floor (reference budget — can't be 0 - // without breaking the ratio compute in the controller). ENS still - // uses the prior numeric-equality bootstrap because no controller - // drives BUDGET_ENS_BASE; the indistinguishable-state pathology - // doesn't apply there. + // IQN/ENS bootstrap constants. ENS_BOOTSTRAP_BUDGET kept at 0.02 to + // preserve prior numeric-equality semantics (no controller drives it + // in SP7 T6+). const SP7_EPS_DIV: f32 = 1e-8; - const ACTIVE_THRESHOLD: f32 = 0.5; - const CQL_BOOTSTRAP_BUDGET: f32 = 0.02; - const C51_BOOTSTRAP_BUDGET: f32 = 0.05; const ENS_BOOTSTRAP_BUDGET: f32 = 0.02; const BASE_IQN: f32 = 0.11; let budget_or_bootstrap = |raw: f32, bootstrap: f32| -> f32 { if raw < SP7_EPS_DIV { bootstrap } else { raw } }; for b in 0..4_usize { - let cql_active = self.read_isv_signal_at(LB_CQL_ACTIVE_BASE + b); - let c51_active = self.read_isv_signal_at(LB_C51_ACTIVE_BASE + b); - cql[b] = if cql_active >= ACTIVE_THRESHOLD { - self.read_isv_signal_at(BUDGET_CQL_BASE + b).max(SP7_EPS_DIV) - } else { - CQL_BOOTSTRAP_BUDGET - }; - c51[b] = if c51_active >= ACTIVE_THRESHOLD { - self.read_isv_signal_at(BUDGET_C51_BASE + b).max(SP7_EPS_DIV) - } else { - C51_BOOTSTRAP_BUDGET - }; iqn[b] = self.read_isv_signal_at(BUDGET_IQN_BASE + b).max(BASE_IQN); ens[b] = budget_or_bootstrap(self.read_isv_signal_at(BUDGET_ENS_BASE + b), ENS_BOOTSTRAP_BUDGET); } // Trunk/value use mean of 4 branch budgets (D3 decision in SP6 spec). - let c51_trunk = c51.iter().sum::() / 4.0_f32; let iqn_trunk = iqn.iter().sum::() / 4.0_f32; - let cql_trunk = cql.iter().sum::() / 4.0_f32; let ens_trunk = ens.iter().sum::() / 4.0_f32; - // Cache per-branch arrays for HEALTH_DIAG logging. - self.trainer.last_c51_budget_per_branch = c51; + // Cache IQN/ENS host-side scalars for HEALTH_DIAG. CQL/C51 are read + // from `lb_budget_effective_buf` at HEALTH_DIAG emit time (see + // `training_loop.rs` HEALTH_DIAG block) — same memory the kernel + // populated, no DtoH copy. self.trainer.last_iqn_budget_per_branch = iqn; - self.trainer.last_cql_budget_per_branch = cql; self.trainer.last_ens_budget_per_branch = ens; - // Cache trunk scalars for backward-compat accessors. - self.trainer.last_c51_budget_eff = c51_trunk; self.trainer.last_iqn_budget_eff = iqn_trunk; - self.trainer.last_cql_budget_eff = cql_trunk; self.trainer.last_ens_budget_eff = ens_trunk; - (c51, iqn, cql, ens, c51_trunk, iqn_trunk, cql_trunk, ens_trunk) + (iqn, ens, iqn_trunk, ens_trunk) } /// SP5 Pearl 2 (flatness-modulated): last per-step IQN budget (mean of 4 @@ -3496,22 +3493,37 @@ impl FusedTrainingCtx { /// `compute_adaptive_budgets` — IQN is the reference budget used by the /// SP7 loss-balance controller's ratio compute and cannot be driven to 0. pub(crate) fn last_iqn_budget_eff(&self) -> f32 { self.trainer.last_iqn_budget_eff } - /// SP7 (2026-05-03): last per-step CQL budget (mean of 4 per-branch - /// values driven by `loss_balance_controller_kernel`). Each per-branch - /// value dispatches on the activation flag at `LB_CQL_ACTIVE_BASE+b`: - /// `CQL_BOOTSTRAP_BUDGET=0.02` while controller has yet to fire on this - /// branch (active < 0.5), controller verbatim (clamped at SP7_EPS_DIV) - /// once activated (active >= 0.5). FoldReset zeroes the activation flag - /// so each fold re-bootstraps cleanly through Pearl A. - pub(crate) fn last_cql_budget_eff(&self) -> f32 { self.trainer.last_cql_budget_eff } - /// SP7 (2026-05-03): last per-step C51 budget (mean of 4 per-branch - /// values driven by `loss_balance_controller_kernel`). Each per-branch - /// value dispatches on the activation flag at `LB_C51_ACTIVE_BASE+b`: - /// `C51_BOOTSTRAP_BUDGET=0.05` while controller has yet to fire on this - /// branch (active < 0.5), controller verbatim (clamped at SP7_EPS_DIV) - /// once activated (active >= 0.5). FoldReset zeroes the activation flag - /// so each fold re-bootstraps cleanly through Pearl A. - pub(crate) fn last_c51_budget_eff(&self) -> f32 { self.trainer.last_c51_budget_eff } + /// SP7 (2026-05-03): last per-step CQL budget — mean of the GPU-resolved + /// per-branch effective budgets in `lb_budget_effective_buf`. Each + /// per-branch value dispatches on `ISV[LB_CQL_ACTIVE_BASE+b]` on-device: + /// `CQL_BOOTSTRAP_BUDGET=0.02` while the controller has yet to fire, + /// controller verbatim (clamped at SP7_EPS_DIV) once active. Caller MUST + /// have synced the producing stream (HEALTH_DIAG path runs at end of + /// epoch — already synced). + pub(crate) fn last_cql_budget_eff(&self) -> f32 { + let (cql, _c51, _t_cql, _t_c51) = self.trainer.read_lb_budget_effective(); + cql.iter().sum::() / 4.0_f32 + } + /// SP7 (2026-05-03): last per-step C51 budget — mean of the GPU-resolved + /// per-branch effective budgets in `lb_budget_effective_buf`. Same dispatch + /// semantics as `last_cql_budget_eff` against + /// `ISV[LB_C51_ACTIVE_BASE+b]` with `C51_BOOTSTRAP_BUDGET=0.05`. + pub(crate) fn last_c51_budget_eff(&self) -> f32 { + let (_cql, c51, _t_cql, _t_c51) = self.trainer.read_lb_budget_effective(); + c51.iter().sum::() / 4.0_f32 + } + /// SP7 (2026-05-03): per-branch CQL effective budget read from + /// `lb_budget_effective_buf`. Used by HEALTH_DIAG. + pub(crate) fn last_cql_budget_per_branch(&self) -> [f32; 4] { + let (cql, _c51, _t_cql, _t_c51) = self.trainer.read_lb_budget_effective(); + cql + } + /// SP7 (2026-05-03): per-branch C51 effective budget read from + /// `lb_budget_effective_buf`. Used by HEALTH_DIAG. + pub(crate) fn last_c51_budget_per_branch(&self) -> [f32; 4] { + let (_cql, c51, _t_cql, _t_c51) = self.trainer.read_lb_budget_effective(); + c51 + } /// SP7 (2026-05-03): last per-step ensemble budget (mean of 4 per-branch /// values). ENS has no active controller driver in SP7 (Pearl 2 no longer /// writes BUDGET_ENS_BASE; loss-balance controller covers CQL/C51 only). diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index 8fd4f49ad..8c6ac77c5 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -2809,15 +2809,22 @@ impl DQNTrainer { self.last_sarsa_tau_factor = Some(fused.last_sarsa_tau_factor()); } - // B4/G5: propagate adaptive gradient budgets for logging. + // B4/G5 + SP7 (2026-05-03): propagate adaptive gradient budgets + // for logging. CQL/C51 reads pull from `lb_budget_effective_buf` + // (the GPU-dispatched values that actually fed the SAXPY/scale + // ops this step) — not the host-cached fields, which were + // populated by `compute_adaptive_budgets` at CUDA-Graph capture + // time and would freeze at the cold-start bootstrap. IQN/ENS + // still come from host caches (out-of-scope for SP7 capture- + // freeze fix per audit Fix 33). if let Some(ref fused) = self.fused_ctx { self.last_iqn_budget_eff = Some(fused.last_iqn_budget_eff()); self.last_cql_budget_eff = Some(fused.last_cql_budget_eff()); self.last_c51_budget_eff = Some(fused.last_c51_budget_eff()); // SP6 Pearl 2: cache per-branch budget arrays. - self.last_c51_budget_per_branch = Some(fused.trainer().last_c51_budget_per_branch); + self.last_c51_budget_per_branch = Some(fused.last_c51_budget_per_branch()); self.last_iqn_budget_per_branch = Some(fused.trainer().last_iqn_budget_per_branch); - self.last_cql_budget_per_branch = Some(fused.trainer().last_cql_budget_per_branch); + self.last_cql_budget_per_branch = Some(fused.last_cql_budget_per_branch()); } // Plan 2 D.2: propagate direction-branch gamma_eff from ISV for logging. diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index 750730173..18c3a1420 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -4131,3 +4131,122 @@ After merging `sp5-magnitude-differentiation` into `main` (bringing SP4+SP5+SP6+ 2. `crates/ml/src/cuda_pipeline/gpu_tlob.rs:2215` — Fix 20's regression test `tlob_dw_layout_alignment_regression_full_chain` called `adam_step(lr, max_grad_norm, weight_clamp, weight_decay)` (4 args), but sp5's evolved Adam signature added 3 more (β1, β2, ε from ISV per SP5 Pearl 4). Added standard-default values; test purpose is layout assertion, not Adam dynamics. `cargo check -p ml`: clean post-fixup. 19 pre-existing warnings, 0 new. + +### Fix 33 — SP7 GPU dispatch for cql/c51 budget consumer (2026-05-03) + +Capture-time host-branch freeze in the SP7 loss-balance controller's downstream consumer. + +**Symptom:** the SP7 controller wrote real per-branch budgets to +`ISV[BUDGET_CQL_BASE+b]` and `ISV[BUDGET_C51_BASE+b]`, and the activation +flag at `ISV[LB_{CQL,C51}_ACTIVE_BASE+b]` correctly transitioned 0→1 +within a fold. But the budgets reaching the SAXPY/scale operations +arrived as exactly the bootstrap values (`0.02` for CQL, `0.05` for C51) +every step, regardless of what the controller wrote. + +**Root cause:** `compute_adaptive_budgets` in `fused_training.rs` ran +host-side Rust: + +```rust +let cql_active = self.read_isv_signal_at(LB_CQL_ACTIVE_BASE + b); +cql[b] = if cql_active >= ACTIVE_THRESHOLD { + self.read_isv_signal_at(BUDGET_CQL_BASE + b).max(SP7_EPS_DIV) +} else { + CQL_BOOTSTRAP_BUDGET // literal 0.02 +}; +``` + +This function was called from inside `submit_aux_ops`, which is captured +into the `aux_child` CUDA Graph. CUDA stream-capture records ONLY kernel +launches, not host-side reads/branches. The host-side `if/else` +resolved ONCE at capture time (step 0 of each fold, when `cql_active` +was the FoldReset sentinel `0.0`), the bootstrap branch was taken, and +the literal `0.02` got baked into the captured kernel-arg buffer for +`apply_c51_budget_scale` / `apply_cql_saxpy`. ~1000 graph replays per +fold then used the frozen `0.02` regardless of what the SP7 controller +wrote to ISV at runtime. Same for `c51` (`0.05`). + +Same disease as the SP4 host-side EMA elimination 2026-05-01 (see +`crates/ml/src/cuda_pipeline/sp4_wiener_ema.rs:14-16` historical +reference + `feedback_no_cpu_compute_strict.md`). The `if/else` IS host +compute inside a captured-graph path, even though the values came from +mapped-pinned ISV memory. + +**Fix architecture — move the dispatch onto GPU, route SAXPY through a +device pointer:** + +1. New CUDA translation unit `consume_lb_budget_kernel.cu` (3 kernels): + - `lb_budget_dispatch` (single block, 8 threads = 2 heads × 4 + branches): reads `ISV[LB_{CQL,C51}_ACTIVE_BASE+b]` + `ISV[BUDGET_{CQL,C51}_BASE+b]`, + applies `(active ≥ 0.5) ? max(ctrl, EPS_DIV) : bootstrap` on-device. + Bootstrap constants pass as kernel args (not hardcoded in body) per + `feedback_isv_for_adaptive_bounds`. Stages effectives in shared + memory; thread (head, 0) reduces 4-element mean to compute trunk + mean; all threads write per-branch correction = `effective[b] / + trunk_mean` (with `>1e-6` guard mirroring the host-side correction- + skip threshold). Output is a 10 f32 mapped-pinned buffer: + `[cql_trunk, cql_corr×4, c51_trunk, c51_corr×4]`. + - `dqn_scale_f32_dev_ptr_kernel`: in-place scale where `alpha` is + read from a device pointer (`*alpha_ptr * y[i]`). NaN-safe at + `alpha=0` matching `dqn_scale_f32_kernel`. + - `dqn_saxpy_f32_dev_ptr_kernel`: `y[i] += (*alpha_ptr) * x[i]`. + All three live in one cubin sharing a CUmodule. Existing scalar + `dqn_scale_f32_kernel` / `dqn_saxpy_f32_kernel` in + `dqn_utility_kernels.cu` are NOT modified — they remain reused by + distill, VSN dW dilution, attn, and other unrelated callers. + +2. `apply_c51_budget_scale` / `apply_c51_budget_scale_branch` / + `apply_cql_saxpy` / `apply_cql_saxpy_branch` modified to take + `alpha_dev_ptr: u64` (CUdeviceptr) instead of `alpha: f32`. Use the + new dev-ptr kernels. Host-side "skip when budget ≈ 1.0" + optimization dropped — the value lives on-device and can't be + branched on host without reintroducing the freeze. Cost: 4 cheap + kernel launches per step with alpha=1.0 in the steady state + (multiply-by-1.0 is a numeric no-op). + +3. `compute_adaptive_budgets` refactored: launches `lb_budget_dispatch` + inline (captured into `aux_child` graph; replay-time fresh values + every step) and returns only IQN/ENS scalars. Callers in + `submit_aux_ops` now pass `lb_budget_{cql,c51}_{trunk,correction}_dev_ptr(b)` + accessor results to the SAXPY/scale ops. IQN/ENS keep their host- + side resolution (out-of-scope per change scope — they don't have an + activation-flag dispatch, just `.max(BASE_IQN)` and a sentinel + bootstrap respectively). + +4. HEALTH_DIAG read at `training_loop.rs` migrated from the host caches + `last_{c51,cql}_budget_per_branch` (formerly populated by the now- + removed host-side dispatch in `compute_adaptive_budgets`) to + `FusedTrainingCtx::last_{c51,cql}_budget_per_branch()`, which reads + `lb_budget_effective_buf` via `read_lb_budget_effective` → + `read_volatile` on mapped-pinned host pointer (no DtoH copy). Caller + must have synced the producing stream first; HEALTH_DIAG runs at + end of epoch where the stream is already synced. The dead caches + on `GpuDqnTrainer` (`last_{c51,cql}_budget_eff` / + `last_{c51,cql}_budget_per_branch`) are deleted; the + `DqnTrainer.last_*` `Option` / `Option<[f32;4]>` slots that + feed the actual logging line stay, but are populated from the + buffer-reading accessors instead. + +**Bootstrap byte-equivalence:** the kernel arg constants `cql_bootstrap=0.02` +and `c51_bootstrap=0.05` MUST stay byte-identical to the prior host-side +`CQL_BOOTSTRAP_BUDGET=0.02` / `C51_BOOTSTRAP_BUDGET=0.05` literals (also +to `loss_balance_controller_kernel.cu`'s `COLD_START_FLOOR_*` anchors) +to preserve cold-start semantics across the producer / consumer +contract boundary. + +**Atomic commit per `feedback_no_partial_refactor`:** kernel + buffer + +launcher + dev-ptr SAXPY/scale variants + consumer rewrite + HEALTH_DIAG +fix + dead-cache deletion + audit doc all together. + +Files: `cuda_pipeline/consume_lb_budget_kernel.cu` (new, ~150 LOC), +`build.rs` (+10 LOC), `cuda_pipeline/gpu_dqn_trainer.rs` (+225 LOC, +−15 LOC: cubin static + 4 struct fields + buffer/cubin/kernel-handle +loading + launch_lb_budget_dispatch + 4 dev-ptr accessor methods + +read_lb_budget_effective + 4 modified apply_*; deleted 4 dead cache +fields + their initializers), `trainers/dqn/fused_training.rs` (−40 +LOC, +30 LOC: compute_adaptive_budgets refactor + 2 new accessor +methods on FusedTrainingCtx + SAXPY caller updates), `trainers/dqn/trainer/training_loop.rs` +(+5 LOC: HEALTH_DIAG read migration), this audit entry. No +`StateResetRegistry` changes — all SP7 ISV slots already had FoldReset +entries (the bug was downstream of the controller's outputs, not in +the slot lifecycle). Memory pearl `pearl_no_host_branches_in_captured_graph.md` +out-of-tree.