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

fix(sp7): GPU dispatch for cql/c51 budget — eliminate capture-time host-branch freeze

Browse Source 
The SP7 controller wrote real budgets to ISV[BUDGET_CQL_BASE/C51_BASE] and
the activation flag correctly transitioned 0→1, but downstream SAXPY ops
still saw exactly 0.02/0.05 every step. Root cause: compute_adaptive_budgets
ran host-side Rust with `if cql_active >= 0.5 { real } else { bootstrap }`,
captured into the aux_child CUDA Graph at step 0 of each fold (when
cql_active is FoldReset sentinel 0.0). The bootstrap branch resolves to
literal 0.02/0.05 and freezes into kernel arg buffers; ~1000 graph replays
per fold use frozen scalars regardless of runtime ISV state.

Same disease as SP4 host-side EMA elimination 2026-05-01. Fix:

- New consume_lb_budget_kernel.cu (8 threads, 1 block): reads ISV slots
  for activation flag + controller budget per (head, branch); applies
  if/else dispatch on-device; writes trunk_mean + per-branch correction
  to mapped-pinned scratch buf [10] = (cql_trunk + cql_corr×4 + c51_trunk
  + c51_corr×4). Same cubin houses dqn_{scale,saxpy}_f32_dev_ptr_kernel
  variants reading alpha from a device pointer (no kernel-arg freeze).
- Buffer + launcher + cubin loading + 4 dev-ptr accessor methods
  in gpu_dqn_trainer.rs.
- apply_c51_budget_scale / apply_cql_saxpy + per-branch variants take
  alpha_dev_ptr: u64 instead of scalar f32. The "skip when ≈ 1.0"
  optimization is dropped — the value lives on-device. Existing scalar
  dqn_{saxpy,scale}_f32_kernel are NOT modified — they remain reused
  by distill, VSN dW dilution, etc.
- compute_adaptive_budgets refactored to launch the dispatch kernel
  inline (captured along with consumers; replay-time fresh values
  every step) and stop returning host-resolved CQL/C51 scalars.
  IQN/ENS keep their host-side resolution (out-of-scope per change
  scope — no activation-flag dispatch).
- HEALTH_DIAG reads from lb_budget_effective_buf via new accessor
  methods on FusedTrainingCtx; dead caches on GpuDqnTrainer
  (last_{cql,c51}_budget_eff / _per_branch) deleted.
- Audit doc Fix 33. Memory pearl out-of-tree.

Bootstrap byte-equivalence preserved: kernel-arg constants
cql_bootstrap=0.02 / c51_bootstrap=0.05 remain byte-identical to the
prior host-side CQL_BOOTSTRAP_BUDGET / C51_BOOTSTRAP_BUDGET literals
and to loss_balance_controller_kernel.cu's COLD_START_FLOOR_* anchors.

Files: consume_lb_budget_kernel.cu (new), build.rs, gpu_dqn_trainer.rs,
fused_training.rs, training_loop.rs, dqn-wire-up-audit.md.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-03 16:54:36 +02:00
parent 14af4854fb
commit 07f5ed5d74
6 changed files with 699 additions and 146 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
crates/ml/build.rs
View File

@@ -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

180
crates/ml/src/cuda_pipeline/consume_lb_budget_kernel.cu Normal file
View File

@@ -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 <cuda_runtime.h>
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];
}
}

317
crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
View File

@@ -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::<f32>();
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::<f32>();
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::<f32>() 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::<f32>() 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::<f32>() 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,

204
crates/ml/src/trainers/dqn/fused_training.rs
View File

@@ -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::<f32>() / 4.0_f32;
let iqn_trunk = iqn.iter().sum::<f32>() / 4.0_f32;
let cql_trunk = cql.iter().sum::<f32>() / 4.0_f32;
let ens_trunk = ens.iter().sum::<f32>() / 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::<f32>() / 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::<f32>() / 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).

13
crates/ml/src/trainers/dqn/trainer/training_loop.rs
View File

@@ -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.

119
docs/dqn-wire-up-audit.md
View File

@@ -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<f32>` / `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.