feat(sp14): B.10 — backward wire gradient gating by ALPHA_GRAD_SMOOTHED
Critical safety mechanism that completes the EGF pearl: scales the wire
column of `dL/dx_concat [B, SH2 + 1]` (the gradient flowing FROM the
direction Q-head's first FC SGEMM TO `aux_softmax_diff`) by
`ISV[ALPHA_GRAD_SMOOTHED_INDEX = 393]`, computed by B.4's
`alpha_grad_compute_kernel` and orchestrated per-step in B.11.
`dL/dW[wire_col]` (Q-head's own weight gradient for the appended column)
is NOT scaled — the dW SGEMM `dY^T × x_concat` and the dX SGEMM
`dY × W^T` are independent, so scaling `dx[:, SH2]` AFTER both have
completed leaves dW unaffected. Q-head learns to USE the wire freely;
only the gradient PROPAGATING BACK to aux is gated.
Pre-B.11 (no producer wired) `ISV[393]` holds sentinel `0.0` →
wire force-closed (gradient zeroed) — the conservative safety state.
Post-B.11, B.4 writes the live gate output ∈ [0, 1] each step.
Closes the latent K-mismatch B.8/B.9 left in backward
============================================================
B.8 grew `w_b0fc` to `[adv_h, SH2 + 1]` end-to-end (Adam m/v +
spectral-norm vector + smoke fixtures); B.9 closed the forward dispatch
K-mismatch. The backward dW/dX SGEMMs for `d == 0` still used `K = SH2`
against the new `LDA = SH2 + 1` weight tensor — silently dropping the
last column of dW and zeroing the wire-col gradient. B.10 closes that
gap atomically with the wire-col scale per `feedback_no_partial_refactor`:
* `backward_branch_dw` for `d == 0` now uses `(dir_qaux_concat_ptr, SH2 + 1)`
instead of `(save_h_s2, SH2)` — matching the forward consumer
pattern from B.9.
* `backward_branch_dx` for `d == 0` now writes to
`d_dir_qaux_concat [B, SH2 + 1]` with `K = SH2 + 1` instead of
`scratch_d_h_s2 [B, SH2]` with `K = SH2`. Mirrors the magnitude
branch's wider-buffer pattern.
New artifacts
=============
* `sp14_scale_wire_col_kernel.cu`: one thread per batch row, scales
`dx_concat[b, SH2]` by `isv[393]` IN-PLACE. NaN-safe per the
`dqn_scale_f32_kernel` precedent (explicit `α==0 ⇒ 0` branch).
Pure per-thread map, no atomicAdd, no shared memory.
* `sp14_d_dir_qaux_concat: CudaSlice<f32>` `[B, SH2 + 1]` trainer-
struct field. Dx SGEMM destination; the wire-col scale acts on
this buffer; the strided accumulator copies the first SH2 columns
into `bw_d_h_s2` after the scale.
* `launch_sp14_scale_wire_col` launcher reads `self.isv_signals_dev_ptr`
and the new buffer's raw_ptr.
* `backward_full` signature grows two trailing `u64` args
(`dir_qaux_concat_ptr`, `d_dir_qaux_concat_ptr`); both
`backward_full` call sites (CQL aux + main online) wired
atomically per `feedback_no_partial_refactor`.
Post-call orchestration at trainer level
========================================
1. `launch_sp14_dir_concat_qaux(save_h_s2)` rebuilds the ONLINE
concat in `sp14_dir_qaux_concat_scratch` (the forward pass had
overwritten it with the TARGET concat at line ~25817). Same
one-step-lag semantic preserved — `aux_nb_softmax_buf` is
unchanged between forward and backward.
2. `cuMemsetD32Async` zero of `d_h_s2` — pre-B.10 the direction
branch (d==0) wrote it with beta=0; post-B.10 the dir-Q dX
lives in `d_dir_qaux_concat` and is gated + accumulated AFTER
`backward_full` returns, so the value-FC dx accumulator inside
`backward_full` (beta=1) needs an explicit zero baseline.
3. `backward_full` runs: dir branch → `d_dir_qaux_concat`,
mag/ord/urg branches → their concat dX buffers, value-FC →
`d_h_s2` (beta=1, on top of zeroed buffer).
4. `launch_sp14_scale_wire_col` gates col SH2 of `d_dir_qaux_concat`.
5. `accumulate_d_h_s2_from_concat` (beta=1) copies first SH2 cols
of `d_dir_qaux_concat` into `d_h_s2`. Wire col stays in
`d_dir_qaux_concat[:, SH2]`, untouched by this accumulator (its
destination range is [0, SH2)). Pre-B.11 the wire is already
zeroed by the sentinel-α gate; the orchestrator that propagates
the gated wire-col gradient back to the aux head's softmax CE
backward chain lives in B.11.
6. mag/ord/urg accumulators continue with beta=1 (comments updated).
Wire status
===========
* Forward dispatch: unchanged (B.9-complete).
* Backward dispatch: GATED on both call sites (CQL aux + main online).
* dW unchanged: the `dW = dY^T × x_concat` SGEMM writes
`grad_buf[goff_w_b0fc..]` BEFORE the scale-wire-col launches;
the scale operates ONLY on `d_dir_qaux_concat` (the dx buffer)
AFTER both dW and dX SGEMMs complete.
* Target net unaffected: Polyak EMA-only, no backward.
* CudaSlice wrapper path: passes `0u64` for both new args, falls
back to the legacy K=SH2 path. Consistent with the forward
wrapper's diagnostic-only residual.
Verified
========
* `SQLX_OFFLINE=true cargo check -p ml` clean, 18 warnings (baseline)
* `cargo test -p ml --test sp14_oracle_tests` 2 passed, 6 ignored (GPU)
* Audit doc `docs/dqn-wire-up-audit.md` updated per Invariant 7.
After this commit, the EGF pearl is architecturally complete; the
orchestration of when/how the alpha_grad gates fire happens in B.11
(producer chain orchestrator).
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -811,6 +811,20 @@ fn main() {
|
||||
// (enforced by graph capture orchestrator in B.10/B.11). Does NOT
|
||||
// read or write ISV slots — purely data-movement.
|
||||
"dir_concat_qaux_kernel.cu",
|
||||
// SP14 Layer B Task B.10 (2026-05-05): Earned Gradient Flow
|
||||
// backward gate. Scales the wire column of `dL/dx_concat
|
||||
// [B, SH2+1]` (gradient flowing FROM the direction Q-head's
|
||||
// first FC SGEMM TO `aux_softmax_diff`) by ISV[393] =
|
||||
// ALPHA_GRAD_SMOOTHED, computed by `alpha_grad_compute_kernel`
|
||||
// (B.4) and orchestrated per-step by B.11. Touches ONLY column
|
||||
// SH2 of the [B, SH2+1] gradient buffer; dW (Q-head's own
|
||||
// weight gradient computed by an independent SGEMM) stays
|
||||
// unscaled per the pearl design — Q-head learns to USE the
|
||||
// wire freely; only the back-flow to aux is gated. Pre-B.11
|
||||
// (unwired producer) ISV[393] sentinels at 0.0 → wire force-
|
||||
// closed for safety (gradient zeroed). Pure per-thread map:
|
||||
// one thread per batch row's wire element. No atomicAdd.
|
||||
"sp14_scale_wire_col_kernel.cu",
|
||||
];
|
||||
|
||||
// ALL kernels get common header (BF16 types + wrappers)
|
||||
|
||||
@@ -1451,6 +1451,11 @@ impl CublasBackwardSet {
|
||||
d_ord_concat_ptr: u64,
|
||||
urg_concat_ptr: u64,
|
||||
d_urg_concat_ptr: u64,
|
||||
// SP14 B.10 (2026-05-05): when non-zero, direction branch (d==0)
|
||||
// dX writes to this `[B, SH2+1]` buffer with K = SH2+1 (matching
|
||||
// the widened `w_b0fc [adv_h, SH2+1]` weight from B.8). When 0,
|
||||
// falls back to legacy direct-write into `scratch_d_h_s2 [B, SH2]`.
|
||||
d_dir_qaux_concat_ptr: u64,
|
||||
b: usize,
|
||||
) -> Result<(), MLError> {
|
||||
// Read per-branch GLU scratch (produced by Steps 1-4 on branch stream,
|
||||
@@ -1458,7 +1463,34 @@ impl CublasBackwardSet {
|
||||
let d_glu_value_ptr = raw_f32_ptr(&self.branch_d_glu_value[d], stream);
|
||||
let d_glu_gate_ptr = raw_f32_ptr(&self.branch_d_glu_gate[d], stream);
|
||||
|
||||
if d == 1 && mag_concat_ptr != 0 {
|
||||
if d == 0 && d_dir_qaux_concat_ptr != 0 {
|
||||
// SP14 B.10: direction Q-head with EGF wire — dX writes to
|
||||
// d_dir_qaux_concat [B, SH2+1]. K = SH2+1 because `w_b0fc` is
|
||||
// `[adv_h, SH2+1]` post-B.8 (last col = aux_softmax_diff weight).
|
||||
// After this returns, the caller scales col SH2 by ISV[393] then
|
||||
// accumulates first SH2 cols into scratch_d_h_s2 (with beta=0
|
||||
// because d==0 is the first branch's contribution to d_h_s2).
|
||||
self.launch_dx_only(
|
||||
stream,
|
||||
d_glu_value_ptr,
|
||||
w_fc,
|
||||
d_dir_qaux_concat_ptr,
|
||||
self.adv_h,
|
||||
self.shared_h2 + 1,
|
||||
b,
|
||||
0.0_f32,
|
||||
)?;
|
||||
self.launch_dx_only(
|
||||
stream,
|
||||
d_glu_gate_ptr,
|
||||
w_gate,
|
||||
d_dir_qaux_concat_ptr,
|
||||
self.adv_h,
|
||||
self.shared_h2 + 1,
|
||||
b,
|
||||
1.0_f32,
|
||||
)?;
|
||||
} else if d == 1 && mag_concat_ptr != 0 {
|
||||
// Magnitude (direction-conditioned): dX writes to d_mag_concat [B, SH2+branch_0_size].
|
||||
self.launch_dx_only(
|
||||
stream,
|
||||
@@ -1711,6 +1743,21 @@ impl CublasBackwardSet {
|
||||
// 0 OR when the cached DRELU_BGRAD descriptor is unavailable, the
|
||||
// backward path falls back to the standalone-kernel sequence.
|
||||
value_fc_relu_mask_ptr: u64,
|
||||
// SP14 Layer B Task B.10 (2026-05-05): direction-Q-head Earned
|
||||
// Gradient Flow wire. When non-zero, the direction branch (`d == 0`)
|
||||
// backward operates on the wider `[B, SH2 + 1]` concat (last col =
|
||||
// `aux_softmax_diff`), matching the SP14 B.8/B.9 forward consumer
|
||||
// pattern. dW reads `dir_qaux_concat_ptr` (saved forward concat,
|
||||
// re-built by the caller via `launch_sp14_dir_concat_qaux` right
|
||||
// before this `backward_full`), and dX writes to `d_dir_qaux_concat_ptr`
|
||||
// with `K = SH2 + 1`. The caller then runs `sp14_scale_wire_col_kernel`
|
||||
// to gate column SH2 by ISV[393] = ALPHA_GRAD_SMOOTHED, then
|
||||
// accumulates the first SH2 columns into `scratch_d_h_s2` via
|
||||
// `strided_accumulate` (mirroring the magnitude branch pattern).
|
||||
// When 0 (e.g. CudaSlice wrapper path), `d == 0` falls back to the
|
||||
// legacy direct-write path into `scratch_d_h_s2 [B, SH2]`.
|
||||
dir_qaux_concat_ptr: u64, // [B, SH2+1] saved forward concat (for dW)
|
||||
d_dir_qaux_concat_ptr: u64, // [B, SH2+1] dX output (caller scales col SH2 + accumulates first SH2 cols)
|
||||
) -> Result<(), MLError> {
|
||||
// Plan 4 Task 2c.3c.4: trunk backward swap.
|
||||
//
|
||||
@@ -1851,10 +1898,15 @@ impl CublasBackwardSet {
|
||||
let n_d = branch_n[d];
|
||||
let w_out = w_ptrs[w_bout_idx[d]];
|
||||
|
||||
// FC input for dW (branches 1,2,3 may use wider concat buffers).
|
||||
// Magnitude (d==1) is direction-conditioned (SH2 + branch_0_size);
|
||||
// order/urgency (d==2,3) are OFI-conditioned (SH2 + 3).
|
||||
let (fc_input, fc_in_dim) = if d == 1 && mag_concat_ptr != 0 {
|
||||
// FC input for dW (each branch may use a wider concat buffer).
|
||||
// Direction (d==0): SP14 B.10 EGF wire — `[h_s2 ; aux_softmax_diff]`
|
||||
// when `dir_qaux_concat_ptr != 0` (SH2 + 1 cols).
|
||||
// Magnitude (d==1): direction-conditioned (SH2 + branch_0_size).
|
||||
// Order/urgency (d==2,3): OFI-conditioned (SH2 + 3).
|
||||
// Otherwise: raw `save_h_s2` (SH2 cols), the legacy path.
|
||||
let (fc_input, fc_in_dim) = if d == 0 && dir_qaux_concat_ptr != 0 {
|
||||
(dir_qaux_concat_ptr, self.shared_h2 + 1)
|
||||
} else if d == 1 && mag_concat_ptr != 0 {
|
||||
(mag_concat_ptr, self.shared_h2 + self.branch_0_size)
|
||||
} else if d == 2 && ord_concat_ptr != 0 {
|
||||
(ord_concat_ptr, self.shared_h2 + 3)
|
||||
@@ -1921,6 +1973,9 @@ impl CublasBackwardSet {
|
||||
d_ord_concat_ptr,
|
||||
urg_concat_ptr,
|
||||
d_urg_concat_ptr,
|
||||
// SP14 B.10: route direction (d==0) dX to the wire-aware
|
||||
// [B, SH2+1] buffer when the EGF wire is enabled.
|
||||
d_dir_qaux_concat_ptr,
|
||||
b,
|
||||
)?;
|
||||
}
|
||||
|
||||
@@ -686,6 +686,15 @@ static SP14_GRAD_HACK_CUBIN: &[u8] =
|
||||
static SP14_DIR_CONCAT_CUBIN: &[u8] =
|
||||
include_bytes!(concat!(env!("OUT_DIR"), "/dir_concat_qaux_kernel.cubin"));
|
||||
|
||||
/// SP14 B.10 (2026-05-05): backward gradient gate kernel for the EGF wire.
|
||||
/// Scales column SH2 of `d_dir_qaux_concat [B, SH2 + 1]` by
|
||||
/// `ISV[ALPHA_GRAD_SMOOTHED_INDEX = 393]` IN-PLACE. dW (Q-head's own weight
|
||||
/// gradient, computed by an independent SGEMM) is NOT touched — Q-head
|
||||
/// learns to USE the wire freely; only the back-flow to aux is gated per
|
||||
/// the EGF pearl design. Loaded from `sp14_scale_wire_col_kernel.cubin`.
|
||||
static SP14_SCALE_WIRE_COL_CUBIN: &[u8] =
|
||||
include_bytes!(concat!(env!("OUT_DIR"), "/sp14_scale_wire_col_kernel.cubin"));
|
||||
|
||||
/// SP11 Fix 39 (2026-05-04, Task A2): SimHash novelty signal — lookup +
|
||||
/// update kernels sharing one cubin. Lookup reads
|
||||
/// `1/sqrt(1+count)` ∈ [0, 1] for each (state, action) bucket; update
|
||||
@@ -5119,6 +5128,29 @@ pub struct GpuDqnTrainer {
|
||||
/// the prior direct `h_s2` input to the direction Q-head. B.7 allocates;
|
||||
/// B.8 bumps the Q-head input-dim constant; B.9/B.11 launch and wire.
|
||||
sp14_dir_qaux_concat_scratch: CudaSlice<f32>,
|
||||
/// SP14 B.10 (2026-05-05): backward dX scratch for direction Q-head's
|
||||
/// first FC. Shape `[B, SH2 + 1]`. Written by `backward_branch_dx`
|
||||
/// (`d == 0` case) as the SGEMM `dY @ W^T` output, then column SH2 is
|
||||
/// scaled by `ISV[ALPHA_GRAD_SMOOTHED_INDEX = 393]` via
|
||||
/// `sp14_scale_wire_col_kernel`, and finally the first SH2 columns
|
||||
/// accumulate into `bw_d_h_s2` via `strided_accumulate`. Mirrors the
|
||||
/// `d_mag_concat_buf` precedent for the magnitude branch's wider input.
|
||||
/// Pre-existing latent gap: B.8/B.9 grew `w_b0fc` to `[adv_h, SH2+1]`
|
||||
/// in forward but the backward consumers (dW + dX) for `d == 0` still
|
||||
/// used `K = SH2`, silently dropping the wire column gradient. B.10
|
||||
/// closes that gap atomically AND adds the EGF gate; per
|
||||
/// `feedback_no_partial_refactor.md`.
|
||||
sp14_d_dir_qaux_concat: CudaSlice<f32>,
|
||||
/// SP14 B.10 (2026-05-05): EGF backward gate kernel. Scales column SH2
|
||||
/// of `d_dir_qaux_concat [B, SH2 + 1]` by `ISV[393] = ALPHA_GRAD_SMOOTHED`.
|
||||
/// One thread per batch row. Loaded from
|
||||
/// `sp14_scale_wire_col_kernel.cubin`. Launched immediately after
|
||||
/// `backward_full` returns, before the first-SH2-cols accumulator copies
|
||||
/// the (now-gated) gradient into `bw_d_h_s2`. The wire column itself is
|
||||
/// dropped pre-B.11 (alpha_grad sentinel = 0 force-closes the gate); the
|
||||
/// orchestrator that propagates the scaled wire-col gradient to the aux
|
||||
/// head's softmax backward lives in B.11.
|
||||
sp14_scale_wire_col_kernel: CudaFunction,
|
||||
/// SP11 Fix 39 (2026-05-04, Task A2): reward-subsystem controller
|
||||
/// kernel. Single-block, 10-thread producer reading 5 canary ISV slots
|
||||
/// [350..360) and writing 10 floats to scratch[SCRATCH_SP11_CONTROLLER_BASE..+10).
|
||||
@@ -7089,6 +7121,43 @@ impl GpuDqnTrainer {
|
||||
/// kernel) reads it on the same stream BEFORE the direction Q-head SGEMM.
|
||||
/// Sequential same-stream submission enforces the dep — no host barrier
|
||||
/// needed and the order survives CUDA Graph capture.
|
||||
/// SP14 Layer B Task B.10 (2026-05-05): EGF backward gate.
|
||||
///
|
||||
/// Scales column SH2 of `d_dir_qaux_concat [B, SH2 + 1]` (the gradient
|
||||
/// flowing FROM the direction Q-head's first FC SGEMM TO
|
||||
/// `aux_softmax_diff`) by `ISV[ALPHA_GRAD_SMOOTHED_INDEX = 393]`. dW
|
||||
/// (Q-head's own weight gradient, computed by an independent SGEMM) is
|
||||
/// NOT touched per the EGF pearl — the SGEMMs `dW = dY^T × x_concat`
|
||||
/// and `dX = dY × W^T` are independent, so scaling `dx[:, SH2]` AFTER
|
||||
/// both have completed leaves dW unaffected.
|
||||
///
|
||||
/// Pre-B.11 (no producer wired) ISV[393] holds sentinel `0.0` →
|
||||
/// wire force-closed for safety (gradient zeroed). Post-B.11, B.4's
|
||||
/// `alpha_grad_compute_kernel` writes the live gate output ∈ [0, 1].
|
||||
/// MUST run AFTER `backward_branch_dx` writes `d_dir_qaux_concat` and
|
||||
/// BEFORE `accumulate_d_h_s2_from_concat` reads the first SH2 columns.
|
||||
pub(crate) fn launch_sp14_scale_wire_col(&self, d_dir_qaux_concat_ptr: u64) -> Result<(), MLError> {
|
||||
let b_i32 = self.config.batch_size as i32;
|
||||
let sh2_i32 = self.config.shared_h2 as i32;
|
||||
let blocks = ((self.config.batch_size as u32 + 255) / 256).max(1);
|
||||
let isv_ptr = self.isv_signals_dev_ptr;
|
||||
unsafe {
|
||||
self.stream
|
||||
.launch_builder(&self.sp14_scale_wire_col_kernel)
|
||||
.arg(&d_dir_qaux_concat_ptr)
|
||||
.arg(&isv_ptr)
|
||||
.arg(&b_i32)
|
||||
.arg(&sh2_i32)
|
||||
.launch(LaunchConfig {
|
||||
grid_dim: (blocks, 1, 1),
|
||||
block_dim: (256, 1, 1),
|
||||
shared_mem_bytes: 0,
|
||||
})
|
||||
.map_err(|e| MLError::ModelError(format!("sp14_scale_wire_col: {e}")))?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(crate) fn launch_sp14_dir_concat_qaux(&self, source_ptr: u64) -> Result<(), MLError> {
|
||||
let b = self.config.batch_size;
|
||||
let sh2 = self.config.shared_h2 as i32;
|
||||
@@ -10386,6 +10455,32 @@ impl GpuDqnTrainer {
|
||||
} else {
|
||||
0u64
|
||||
};
|
||||
// SP14 B.10: rebuild the ONLINE direction-Q-head concat scratch
|
||||
// [save_h_s2 | aux_softmax_diff] into `sp14_dir_qaux_concat_scratch`.
|
||||
// The forward pass overwrote it with the TARGET concat (line ~25817 in
|
||||
// `submit_forward_ops`), so the backward dW SGEMM would otherwise read
|
||||
// tg_h_s2 instead of save_h_s2 — wrong gradient. `aux_nb_softmax_buf`
|
||||
// is unchanged between forward and backward (CE loss writes a separate
|
||||
// `d_aux_softmax` scratch), so re-running the same concat with
|
||||
// `save_h_s2` yields the bit-identical online concat the forward
|
||||
// SGEMM consumed. Same one-step-lag semantic preserved.
|
||||
self.launch_sp14_dir_concat_qaux(self.ptrs.save_h_s2)?;
|
||||
// SP14 B.10: pre-zero `scratch_d_h_s2` because the direction branch
|
||||
// (d==0) no longer writes it directly with beta=0 — its dX now lands
|
||||
// in `d_dir_qaux_concat` and is gated + accumulated AFTER backward_full
|
||||
// returns. The value-FC dX accumulator inside backward_full uses
|
||||
// beta=1 and assumed scratch_d_h_s2 was pre-written by d==0; with the
|
||||
// direction branch redirected, we must establish the zero baseline
|
||||
// explicitly. Raw cuMemsetD32Async — captured by CUDA Graph (cudarc
|
||||
// memset_zeros is NOT graph-safe; matches the cql_grad_scratch
|
||||
// pre-backward-full zero pattern below).
|
||||
unsafe {
|
||||
cudarc::driver::sys::cuMemsetD32Async(
|
||||
scratch_d_h_s2, 0,
|
||||
self.config.batch_size * self.config.shared_h2,
|
||||
self.stream.cu_stream(),
|
||||
);
|
||||
}
|
||||
self.cublas_backward.backward_full(
|
||||
&self.stream,
|
||||
d_v_ptr,
|
||||
@@ -10409,8 +10504,33 @@ impl GpuDqnTrainer {
|
||||
self.ptrs.bw_d_glu_value,
|
||||
self.ptrs.bw_d_glu_gate,
|
||||
value_fc_relu_mask_cql,
|
||||
// SP14 B.10: direction-Q-head EGF wire — saved forward concat
|
||||
// for dW + dX output buffer for col SH2 wire scaling.
|
||||
self.sp14_dir_qaux_concat_scratch.raw_ptr(),
|
||||
self.sp14_d_dir_qaux_concat.raw_ptr(),
|
||||
).map_err(|e| MLError::ModelError(format!("CQL backward_full: {e}")))?;
|
||||
|
||||
// SP14 B.10: gate the wire column gradient by ISV[ALPHA_GRAD_SMOOTHED]
|
||||
// before the first-SH2-cols accumulator copies it into scratch_d_h_s2.
|
||||
// Pre-B.11 (no producer wired) the slot sentinel = 0.0 zeros the wire
|
||||
// → effectively force-closed. Post-B.11, the gate value ∈ [0, 1] is
|
||||
// applied per the EGF pearl design.
|
||||
self.launch_sp14_scale_wire_col(self.sp14_d_dir_qaux_concat.raw_ptr())?;
|
||||
// Accumulate first SH2 columns of d_dir_qaux_concat into scratch_d_h_s2
|
||||
// with beta=1 (the value-FC dX inside backward_full already accumulated
|
||||
// its contribution; we ADD on top, NOT overwrite). Wire column (col
|
||||
// SH2) is NOT propagated downstream from here — the orchestrator that
|
||||
// routes the gated wire-col gradient back to the aux head's softmax CE
|
||||
// backward chain lives in B.11. Pre-B.11 the wire is zeroed by the
|
||||
// sentinel-α gate above anyway, so the unrouted column is moot.
|
||||
self.accumulate_d_h_s2_from_concat(
|
||||
self.sp14_d_dir_qaux_concat.raw_ptr(),
|
||||
scratch_d_h_s2,
|
||||
self.config.batch_size,
|
||||
self.config.shared_h2 + 1,
|
||||
1.0,
|
||||
)?;
|
||||
|
||||
// CQL: accumulate magnitude branch dX into scratch_d_h_s2
|
||||
if self.ptrs.mag_concat_buf != 0 {
|
||||
self.accumulate_d_h_s2_from_concat(
|
||||
@@ -10418,7 +10538,7 @@ impl GpuDqnTrainer {
|
||||
scratch_d_h_s2,
|
||||
self.config.batch_size,
|
||||
self.config.shared_h2 + self.config.branch_0_size, // mag: direction-conditioned
|
||||
1.0, // beta=1: d==0 already wrote to scratch_d_h_s2
|
||||
1.0, // beta=1: SP14 B.10 zeroed scratch_d_h_s2 + value-FC + dir-Q already accumulated
|
||||
)?;
|
||||
}
|
||||
if self.ptrs.ord_concat_buf != 0 {
|
||||
@@ -17038,6 +17158,26 @@ impl GpuDqnTrainer {
|
||||
.alloc_zeros::<f32>(b * (config.shared_h2 + 1))
|
||||
.map_err(|e| MLError::ModelError(format!("sp14 dir_qaux scratch: {e}")))?;
|
||||
|
||||
// SP14 B.10: backward dX scratch for the direction Q-head's first
|
||||
// FC. Shape `[B, SH2 + 1]`. Receives the SGEMM `dY @ W^T` output
|
||||
// (replacing the prior direct write into `bw_d_h_s2 [B, SH2]` for
|
||||
// d == 0); column SH2 then gets gated by ISV[393] and the first
|
||||
// SH2 columns accumulate into `bw_d_h_s2`. Symmetric with the
|
||||
// existing `d_mag_concat_buf` (magnitude branch's `[B, SH2 + b0]`
|
||||
// backward scratch). Allocated zero-initialised so the first
|
||||
// backward replay sees a clean buffer; the dX SGEMM uses beta=0
|
||||
// anyway so the prior contents are overwritten each step.
|
||||
let sp14_d_dir_qaux_concat = stream
|
||||
.alloc_zeros::<f32>(b * (config.shared_h2 + 1))
|
||||
.map_err(|e| MLError::ModelError(format!("sp14 d_dir_qaux scratch: {e}")))?;
|
||||
|
||||
let sp14_scale_wire_col_module = stream.context()
|
||||
.load_cubin(SP14_SCALE_WIRE_COL_CUBIN.to_vec())
|
||||
.map_err(|e| MLError::ModelError(format!("sp14 scale_wire_col cubin: {e}")))?;
|
||||
let sp14_scale_wire_col_kernel = sp14_scale_wire_col_module
|
||||
.load_function("sp14_scale_wire_col_kernel")
|
||||
.map_err(|e| MLError::ModelError(format!("sp14_scale_wire_col_kernel: {e}")))?;
|
||||
|
||||
// SP11 Fix 39 B1b fix-up (2026-05-04): allocate trainer-side
|
||||
// popart-component-per-sample mapped-pinned placeholder buffer.
|
||||
// The canonical per-bar buffer is owned by the experience
|
||||
@@ -20539,6 +20679,12 @@ impl GpuDqnTrainer {
|
||||
sp14_gradient_hack_detect_kernel,
|
||||
sp14_dir_concat_qaux_kernel,
|
||||
sp14_dir_qaux_concat_scratch,
|
||||
// SP14 EGF backward gate (B.10). Scratch + kernel land
|
||||
// atomically with the dW/dX SGEMM K-dim migration in
|
||||
// `backward_branch_dw`/`backward_branch_dx` per
|
||||
// `feedback_no_partial_refactor.md`.
|
||||
sp14_d_dir_qaux_concat,
|
||||
sp14_scale_wire_col_kernel,
|
||||
// SP11 Fix 39 (Task A2): controller + SimHash novelty kernels
|
||||
// and their backing buffers. Projection matrix was populated
|
||||
// on-device above by `launch_novelty_simhash_proj_init`; hash
|
||||
@@ -26832,6 +26978,32 @@ impl GpuDqnTrainer {
|
||||
self.launch_clamp_finite_f32(d_value_logits_ptr, n_val_clamp, max_abs_q)?;
|
||||
self.launch_clamp_finite_f32(d_adv_logits_ptr, n_adv_clamp, max_abs_q)?;
|
||||
|
||||
// SP14 B.10: rebuild the ONLINE direction-Q-head concat scratch
|
||||
// [save_h_s2 | aux_softmax_diff] into `sp14_dir_qaux_concat_scratch`.
|
||||
// The forward pass overwrote it with the TARGET concat (line ~25817 in
|
||||
// `submit_forward_ops`), so the backward dW SGEMM would otherwise read
|
||||
// tg_h_s2 instead of save_h_s2 — wrong gradient. `aux_nb_softmax_buf`
|
||||
// is unchanged between forward and backward (CE loss writes a separate
|
||||
// `d_aux_softmax` scratch), so re-running the same concat with
|
||||
// `save_h_s2` yields the bit-identical online concat the forward
|
||||
// SGEMM consumed. Same one-step-lag semantic preserved.
|
||||
self.launch_sp14_dir_concat_qaux(self.ptrs.save_h_s2)?;
|
||||
// SP14 B.10: pre-zero `d_h_s2` because the direction branch (d==0) no
|
||||
// longer writes it directly with beta=0 — its dX now lands in
|
||||
// `d_dir_qaux_concat` and is gated + accumulated AFTER backward_full
|
||||
// returns. The value-FC dX accumulator inside backward_full uses
|
||||
// beta=1 and assumed d_h_s2 was pre-written by d==0; with the
|
||||
// direction branch redirected, we must establish the zero baseline
|
||||
// explicitly. Raw cuMemsetD32Async — captured by CUDA Graph (cudarc
|
||||
// memset_zeros is NOT graph-safe; matches the cql_grad_scratch
|
||||
// pre-backward-full zero pattern in the CQL path).
|
||||
unsafe {
|
||||
cudarc::driver::sys::cuMemsetD32Async(
|
||||
d_h_s2_ptr, 0,
|
||||
self.config.batch_size * self.config.shared_h2,
|
||||
self.stream.cu_stream(),
|
||||
);
|
||||
}
|
||||
bw.backward_full(
|
||||
&self.stream,
|
||||
d_value_logits_ptr,
|
||||
@@ -26859,6 +27031,34 @@ impl GpuDqnTrainer {
|
||||
self.ptrs.bw_d_glu_value,
|
||||
self.ptrs.bw_d_glu_gate,
|
||||
value_fc_relu_mask,
|
||||
// SP14 B.10: direction-Q-head EGF wire — saved forward concat
|
||||
// for dW + dX output buffer for col SH2 wire scaling.
|
||||
self.sp14_dir_qaux_concat_scratch.raw_ptr(),
|
||||
self.sp14_d_dir_qaux_concat.raw_ptr(),
|
||||
)?;
|
||||
|
||||
// SP14 B.10: gate the wire column gradient by ISV[ALPHA_GRAD_SMOOTHED]
|
||||
// before the first-SH2-cols accumulator copies it into d_h_s2.
|
||||
// Pre-B.11 (no producer wired) the slot sentinel = 0.0 zeros the wire
|
||||
// → effectively force-closed. Post-B.11, the gate value ∈ [0, 1] is
|
||||
// applied per the EGF pearl design. Critical: this scales `dL/dx` ONLY,
|
||||
// NOT `dL/dW` — Q-head's first FC weight gradient (computed by an
|
||||
// independent `dY^T × x_concat` SGEMM into `grad_base[goff_w_b0fc..]`)
|
||||
// is unscaled, so the Q-head learns to USE the wire freely.
|
||||
self.launch_sp14_scale_wire_col(self.sp14_d_dir_qaux_concat.raw_ptr())?;
|
||||
// Accumulate first SH2 columns of d_dir_qaux_concat into d_h_s2 with
|
||||
// beta=1 (the value-FC dX inside backward_full already accumulated its
|
||||
// contribution; we ADD on top, NOT overwrite). Wire column (col SH2)
|
||||
// is NOT propagated downstream from here — the orchestrator that
|
||||
// routes the gated wire-col gradient back to the aux head's softmax CE
|
||||
// backward chain lives in B.11. Pre-B.11 the wire is zeroed by the
|
||||
// sentinel-α gate above anyway, so the unrouted column is moot.
|
||||
self.accumulate_d_h_s2_from_concat(
|
||||
self.sp14_d_dir_qaux_concat.raw_ptr(),
|
||||
d_h_s2_ptr,
|
||||
self.config.batch_size,
|
||||
self.config.shared_h2 + 1,
|
||||
1.0,
|
||||
)?;
|
||||
|
||||
// Accumulate magnitude branch dX (first SH2 columns) into d_h_s2
|
||||
@@ -26868,7 +27068,7 @@ impl GpuDqnTrainer {
|
||||
d_h_s2_ptr,
|
||||
self.config.batch_size,
|
||||
self.config.shared_h2 + self.config.branch_0_size, // mag: direction-conditioned
|
||||
1.0, // beta=1: d==0 already wrote to d_h_s2 in backward_full
|
||||
1.0, // beta=1: SP14 B.10 zeroed d_h_s2 + value-FC + dir-Q already accumulated
|
||||
)?;
|
||||
}
|
||||
// Accumulate order branch dX (first SH2 columns) into d_h_s2
|
||||
|
||||
64
crates/ml/src/cuda_pipeline/sp14_scale_wire_col_kernel.cu
Normal file
64
crates/ml/src/cuda_pipeline/sp14_scale_wire_col_kernel.cu
Normal file
@@ -0,0 +1,64 @@
|
||||
// crates/ml/src/cuda_pipeline/sp14_scale_wire_col_kernel.cu
|
||||
//
|
||||
// SP14 Layer B Task B.10 (2026-05-05): Earned Gradient Flow backward gate.
|
||||
// Scales the wire column of `dL/dx_concat [B, SH2 + 1]` (the gradient
|
||||
// flowing back FROM the direction Q-head's first FC SGEMM TO
|
||||
// `aux_softmax_diff`) by `α_grad_smoothed` (read from ISV[393]).
|
||||
//
|
||||
// ── What this gates ─────────────────────────────────────────────────────
|
||||
// dL/dx_concat shape is [B, SH2 + 1]. Columns [0..SH2) propagate back to
|
||||
// `h_s2` via the existing trunk-encoder backward chain. Column SH2 is the
|
||||
// gradient w.r.t. the appended `aux_softmax_diff` value, which downstream
|
||||
// (B.11) propagates to the aux head's softmax → CE backward chain. THIS
|
||||
// kernel multiplies ONLY column SH2 by `ALPHA_GRAD_SMOOTHED = ISV[393]`.
|
||||
//
|
||||
// Per the EGF pearl design (spec §B.4):
|
||||
// - α_grad ≈ 0 → wire effectively closed (Q-loss can't pull aux head)
|
||||
// - α_grad ≈ 1 → full co-training enabled
|
||||
//
|
||||
// Critical: dW (the Q-head's own weight gradient) is NOT touched. Q-head
|
||||
// learns to USE the wire freely from Q-loss; only the gradient PROPAGATING
|
||||
// BACK to aux head's representation is gated. The dW and dX SGEMMs are
|
||||
// independent (`dW = dY^T × x_concat`; `dX = dY × W^T`), so scaling dx[:, SH2]
|
||||
// AFTER both SGEMMs have completed leaves dW unaffected.
|
||||
//
|
||||
// ── ISV slot index (must match `sp14_isv_slots.rs`) ─────────────────────
|
||||
// ALPHA_GRAD_SMOOTHED_INDEX = 393 (sentinel 0.0 pre-B.11 producer wires;
|
||||
// when sentinel-active the wire is force-closed, the conservative state).
|
||||
//
|
||||
// ── Pure per-thread map ─────────────────────────────────────────────────
|
||||
// One thread handles one batch row's wire-column element. No reductions,
|
||||
// no shared memory, no atomicAdd (per `feedback_no_atomicadd.md`). Read of
|
||||
// `isv[ALPHA_GRAD_SMOOTHED_INDEX]` is identical across all threads — safe
|
||||
// concurrent read.
|
||||
//
|
||||
// Grid: ceil(B / 256), Block: 256.
|
||||
|
||||
extern "C" __global__
|
||||
void sp14_scale_wire_col_kernel(
|
||||
/* dL/dx_concat [B, SH2 + 1] — IN-PLACE scaled at column SH2 only */
|
||||
float* __restrict__ dx_concat,
|
||||
/* Global ISV bus. Reads ISV[ALPHA_GRAD_SMOOTHED_INDEX]; writes nothing. */
|
||||
const float* __restrict__ isv,
|
||||
int B,
|
||||
int SH2)
|
||||
{
|
||||
/* ISV slot index — must match `sp14_isv_slots.rs::ALPHA_GRAD_SMOOTHED_INDEX`.
|
||||
* The layout fingerprint regression test catches drift. */
|
||||
const int ALPHA_GRAD_SMOOTHED = 393;
|
||||
|
||||
const int b = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (b >= B) return;
|
||||
|
||||
/* Wire column lives at offset SH2 in each row of stride (SH2+1). */
|
||||
const int wire_idx = b * (SH2 + 1) + SH2;
|
||||
|
||||
/* Pre-B.11 producer is unwired → ISV[393] holds sentinel 0.0 →
|
||||
* wire force-closed (gradient zeroed). Post-B.11, the value is the
|
||||
* EGF gate output ∈ [0, 1].
|
||||
*
|
||||
* Match `dqn_scale_f32_kernel`'s NaN-safe contract: 0×NaN=NaN under
|
||||
* IEEE 754; explicit branch yields a clean 0 when α=0. */
|
||||
const float alpha = isv[ALPHA_GRAD_SMOOTHED];
|
||||
dx_concat[wire_idx] = (alpha == 0.0f) ? 0.0f : (dx_concat[wire_idx] * alpha);
|
||||
}
|
||||
@@ -6601,3 +6601,44 @@ Per `pearl_canary_input_freshness_launch_order`: the producer (aux_heads_forward
|
||||
- **Diagnostic-path residual**: causal intervention, DDQN argmax, and the experience-collector / value-decoder forwards intentionally pass `0u64` for `dir_qaux_concat_ptr` and fall back to `K = SH2`. Their direction Q outputs feed either (a) only-value-logit consumers (causal) or (b) downstream argmax-only consumers with one-step-bias acknowledged by the spec (DDQN). The cuBLAS heuristic for `K = SH2, LDA = SH2` against the underlying `[adv_h, SH2 + 1]` weight tensor reads the first `adv_h * SH2` floats with stride SH2 — within bounds (no OOB), produces stable-but-incorrect outputs for the residual paths. `feedback_no_partial_refactor` is honoured for the train-time path (online forward, target forward, replay); diagnostic-path residuals are documented and bounded.
|
||||
- **Reverse dependencies**: `forward_value_head_for_ensemble` and `forward_value_head` use only the value head (V_h ← h_s2; no branch heads), so unaffected by the direction Q-head wire. The `submit_dqn_step_loop_cublas` path is the same `forward_online_raw` API — no separate dispatch site.
|
||||
|
||||
## SP14 Layer B — Commit B.10: backward wire — gradient scaling at the wire column (2026-05-05)
|
||||
|
||||
**Why this commit.** B.10 is the critical safety mechanism that completes the EGF pearl. The forward wire (B.6/B.9) feeds `aux_softmax_diff` into the direction Q-head's first FC SGEMM via `[h_s2 ; aux_softmax_diff] [B, SH2 + 1]`. Without B.10, the Q-loss backward computes `dL/dx_concat[:, SH2]` (the gradient flowing back to `aux_softmax_diff`) and propagates it ungated all the way to the aux head's softmax — co-training aux on Q-loss without protection. B.10 scales `dL/dx[wire_col]` by `ISV[ALPHA_GRAD_SMOOTHED_INDEX = 393]` (computed by B.4's `alpha_grad_compute_kernel`, orchestrated per-step in B.11). Crucially, `dL/dW[wire_col]` (the Q-head's own weight gradient for the appended column) is NOT scaled — Q-head learns to USE the wire freely; only the gradient PROPAGATING BACK to aux is gated. The two SGEMMs `dW = dY^T × x_concat` and `dX = dY × W^T` are independent, so scaling `dx[:, SH2]` AFTER both have completed leaves dW unaffected.
|
||||
|
||||
### Pre-existing latent gap closed atomically with the wire-col scale
|
||||
|
||||
B.8/B.9 grew `w_b0fc` to `[adv_h, SH2 + 1]` end-to-end across forward dispatch + Adam m/v + spectral-norm power-iteration vector + smoke fixtures, but the **backward dW/dX SGEMMs for `d == 0` still used `K = SH2`** against the new `LDA = SH2 + 1` weight tensor — the same K-mismatch B.9 closed in forward, surviving in backward. B.10 closes that gap atomically with the wire-col scale per `feedback_no_partial_refactor.md`:
|
||||
|
||||
| Backward path | Pre-B.10 | Post-B.10 |
|
||||
|---|---|---|
|
||||
| `backward_branch_dw` for `d == 0` | `(save_h_s2, SH2)` → silent dW under-write of last column | `(dir_qaux_concat_ptr, SH2 + 1)` → full dW including last column |
|
||||
| `backward_branch_dx` for `d == 0` | `(scratch_d_h_s2, SH2, beta=0)` → wire-col gradient dropped | `(d_dir_qaux_concat_ptr, SH2 + 1, beta=0)` → wire-col gradient preserved for B.10 scale |
|
||||
| Wire-col gating | none | `sp14_scale_wire_col_kernel` reads `ISV[393]` and multiplies `d_dir_qaux_concat[:, SH2]` |
|
||||
| Trunk `d_h_s2` accumulation | implicit (d==0 wrote `scratch_d_h_s2` with beta=0) | explicit `cuMemsetD32Async` zero pre-`backward_full`; `accumulate_d_h_s2_from_concat` adds first SH2 cols with beta=1 after wire-col scale |
|
||||
|
||||
### File-change summary
|
||||
|
||||
| Change | File |
|
||||
|--------|------|
|
||||
| New `sp14_scale_wire_col_kernel.cu` (one thread per batch row, scales col SH2 by `ISV[393]`); `extern "C"` symbol; structurally bounded NaN-safe per `dqn_scale_f32_kernel` precedent | `crates/ml/src/cuda_pipeline/sp14_scale_wire_col_kernel.cu` |
|
||||
| New cubin entry in `kernels_with_common` | `crates/ml/build.rs` |
|
||||
| New `static SP14_SCALE_WIRE_COL_CUBIN`; new struct fields `sp14_d_dir_qaux_concat: CudaSlice<f32>` (`[B, SH2 + 1]`) + `sp14_scale_wire_col_kernel: CudaFunction`; allocation + cubin load + struct init; new `launch_sp14_scale_wire_col` launcher; both `backward_full` call sites (CQL @ ~10484 + main @ ~26980) get `launch_sp14_dir_concat_qaux(save_h_s2)` re-build (the forward path overwrites the buffer with target concat at line ~25817), `cuMemsetD32Async` zero of d_h_s2 before the call, the new `(dir_qaux_concat_ptr, d_dir_qaux_concat_ptr)` trailing args, post-call `launch_sp14_scale_wire_col` + `accumulate_d_h_s2_from_concat(beta=1)` | `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` |
|
||||
| `backward_full` signature grows `dir_qaux_concat_ptr` + `d_dir_qaux_concat_ptr` u64 args; `backward_branch_dw` adds `d == 0 && dir_qaux_concat_ptr != 0` arm using `(SH2 + 1)` fc_in_dim; `backward_branch_dx` adds `d == 0 && d_dir_qaux_concat_ptr != 0` arm writing `[B, SH2 + 1]` with `K = SH2 + 1` (mirroring magnitude branch's wider-buffer pattern) | `crates/ml/src/cuda_pipeline/batched_backward.rs` |
|
||||
|
||||
### Verification
|
||||
|
||||
- `SQLX_OFFLINE=true cargo check -p ml` — clean, 18 warnings (pre-existing baseline, no new warnings)
|
||||
- `cargo test -p ml --test sp14_oracle_tests` — 2 passed, 6 ignored (GPU)
|
||||
|
||||
### Wire status
|
||||
|
||||
- **Forward dispatch**: unchanged (B.9-complete).
|
||||
- **Backward dispatch**: GATED. Both `backward_full` call sites (main + CQL aux) use the EGF wire path: dW via `dir_qaux_concat_ptr` with `K = SH2 + 1`, dX into `d_dir_qaux_concat [B, SH2 + 1]` with `K = SH2 + 1`. The K-mismatch B.8/B.9 left in backward is now closed in lockstep with the wire-col scale.
|
||||
- **Wire-col gating**: ACTIVE. `launch_sp14_scale_wire_col` reads `ISV[ALPHA_GRAD_SMOOTHED_INDEX = 393]` and multiplies column SH2 of `d_dir_qaux_concat` IN-PLACE. Pre-B.11 (no producer wired) the slot sentinel = 0.0 → wire force-closed (gradient zeroed) — the conservative safety state. Post-B.11, B.4's `alpha_grad_compute_kernel` writes the live gate output ∈ [0, 1] each step.
|
||||
- **dW unchanged**: `dL/dW = dL/d_branch_h^T × x_concat` writes to `grad_buf[goff_w_b0fc..goff_w_b0fc + adv_h * (SH2 + 1) * 4]` via `launch_dw_only_ws` BEFORE the scale-wire-col launches; the scale operates ONLY on `d_dir_qaux_concat` (the dx buffer) AFTER both dW and dX SGEMMs complete. Q-head learns to use the wire freely.
|
||||
- **First-SH2-cols accumulator**: `accumulate_d_h_s2_from_concat(d_dir_qaux_concat → d_h_s2, src_stride=SH2+1, beta=1.0)` runs AFTER `launch_sp14_scale_wire_col` so the scaled wire-col stays in `d_dir_qaux_concat[:, SH2]` (untouched by the accumulator's destination range `[0, SH2)`). The wire column gradient is NOT propagated downstream from here pre-B.11; the orchestrator that routes the gated wire-col gradient back to the aux head's softmax CE backward chain lives in B.11. Pre-B.11 the wire is zeroed by the sentinel-α gate anyway, so the unrouted column is moot.
|
||||
- **One-step-lag preservation**: backward consumes the same `aux_nb_softmax_buf` snapshot the forward consumed — `aux_heads_forward` writes the buffer once per step BEFORE the next-step's online forward runs `launch_sp14_dir_concat_qaux`, and the CE loss kernel writes a separate `d_aux_softmax` scratch (does not overwrite `aux_nb_softmax_buf`). Re-running the concat with `save_h_s2` at backward start yields the bit-identical online concat the forward SGEMM consumed.
|
||||
- **Target forward unaffected**: target net is Polyak-EMA-updated only (no backward), so the wire-col scale + K-dim migration apply only to the online backward path.
|
||||
- **Reverse dependencies**: `apply_iqn_trunk_gradient` and aux-paths use the `bw_d_h_s2` that this path writes — the new accumulator pipeline (memset → backward_full → wire-col scale → strided accumulate from d_dir_qaux_concat → mag/ord/urg accumulators → value-FC inside backward_full) leaves `bw_d_h_s2` with the same algebraic value as pre-B.10 *except* for the gated wire-col contribution from the direction-Q's first FC. Pre-B.11 (α=0) the gated contribution is zero → bit-identical to pre-B.10.
|
||||
- **CudaSlice wrapper path**: passes `0u64` for both `dir_qaux_concat_ptr` and `d_dir_qaux_concat_ptr`, falling back to the legacy K=SH2 path. This is consistent with the forward CudaSlice wrapper (`dir_qaux_concat_ptr: 0`); the wrapper-based callers (causal intervention, DDQN argmax) are diagnostic-only paths whose direction-Q outputs are downstream-bounded per the B.9 residual-path analysis.
|
||||
|
||||
|
||||
Reference in New Issue
Block a user