From 4ae9a27f48e38409c02517485761df314f7760cf Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Mon, 18 May 2026 16:16:31 +0200 Subject: [PATCH] =?UTF-8?q?fix(ml-alpha):=20wire=20axis=20E=20into=20loss?= =?UTF-8?q?=20=E2=80=94=20real=20add=5Finv=5Fbroadcast=20kernel?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit CRITICAL ARCHITECTURAL FIX discovered while planning fused kernel: The previous Stage 2 `add_inv_broadcast` helper in `multi_horizon_attention.rs` was a STUB that returned Ok(()) without doing anything. Practical consequences: - Forward: `inv_pooled_d` (output of inverted_attention_pool.forward) was never added into `ctx_h_d`. Downstream MoE + heads never saw the inverted-attention signal. Axis E contributed ZERO to the forward output and the loss. - Backward: `inv_pool.backward` was being fed `grad_ctx_mean` as its "upstream gradient", but that's the gradient at the CHAIN TERMINUS — not the gradient w.r.t. inv_pooled_d (which is zero by construction since inv_pooled wasn't in the loss). The bwd was injecting incorrect noise into `grad_ln_out`. Net: paying inverted_attention compute for no gain, plus polluting ln_b's gradient. Two perf rewrite rounds earlier today showed no wall-time movement precisely because the slow path was wired into training while the optimized one was dead. FIX: (a) New kernel `cuda/inv_pooled_merge.cu`: fwd: ctx_h[b, h, d] += inv_pooled[b, d] (broadcast over h) bwd: grad_inv_pooled[b, d] = Σ_h grad_ctx_h[b, h, d] Tiny — single block-per-batch, no syncthreads, coalesced reads. (b) Host binding `src/inv_pooled_merge.rs` (InvPooledMerge). (c) `MultiHorizonAttention` adds: - `merge: InvPooledMerge` field. - `grad_inv_pooled_d [B, H]` buffer for the real upstream of inv_pool.bwd. (d) `MultiHorizonAttention.forward` now calls `merge.forward(...)` between inv_pool.forward and moe.forward. Axis E is now actually in the model's forward output. (e) `MultiHorizonAttention.backward` now calls `merge.backward` after moe.bwd writes `grad_ctx_h_d`, producing `grad_inv_pooled_d`. `inv_pool.backward` consumes the REAL upstream gradient (`grad_inv_pooled_d`) instead of the prior `grad_ctx_mean` fake. (f) Old stub `add_inv_broadcast` deleted. CORRECTNESS: - perception_overfit 9/9 PASS after the wiring. Loss still shrinks on the constant-signal test (0.67 → -0.99 over 250 steps), now with axis E actually contributing. - All numgrad kernel tests still pass (kernels themselves unchanged in this commit; only the wiring). PERF IMPACT (expected): - +2 tiny launches per step (merge fwd + bwd). Negligible. - The inverted-attention compute that was previously dead now actually feeds the loss → same wall-time, but it's earning the cost. This unblocks meaningful axis-E perf measurement on next smoke. NEXT: re-run cluster smoke to confirm wall-time + verify axis E gradient flow is healthy. After that, consider full MHA forward fusion as a follow-up. Co-Authored-By: Claude Opus 4.7 --- crates/ml-alpha/build.rs | 1 + crates/ml-alpha/cuda/inv_pooled_merge.cu | 64 +++++++++++++++ crates/ml-alpha/src/inv_pooled_merge.rs | 82 +++++++++++++++++++ crates/ml-alpha/src/lib.rs | 1 + .../src/trainer/multi_horizon_attention.rs | 55 ++++++------- 5 files changed, 172 insertions(+), 31 deletions(-) create mode 100644 crates/ml-alpha/cuda/inv_pooled_merge.cu create mode 100644 crates/ml-alpha/src/inv_pooled_merge.rs diff --git a/crates/ml-alpha/build.rs b/crates/ml-alpha/build.rs index c559216e4..1ba2aba13 100644 --- a/crates/ml-alpha/build.rs +++ b/crates/ml-alpha/build.rs @@ -24,6 +24,7 @@ const KERNELS: &[&str] = &[ "inverted_attention_pool", // Cross-variate (iTransformer-style) attention pool "regime_moe_gate", // Top-1 regime MoE gate + expert dispatch + aux loss "horizon_mean_collapse", // Mean over horizon axis (seeds CfC h_old at k=0) + "inv_pooled_merge", // Broadcast-add inv_pooled into ctx_h (wires axis E into loss) "anchor_l2", // L2 anchor regularization toward init "reduce_axis0", // Phase B: cross-batch param-grad reducer ]; diff --git a/crates/ml-alpha/cuda/inv_pooled_merge.cu b/crates/ml-alpha/cuda/inv_pooled_merge.cu new file mode 100644 index 000000000..994add3ef --- /dev/null +++ b/crates/ml-alpha/cuda/inv_pooled_merge.cu @@ -0,0 +1,64 @@ +// inv_pooled_merge.cu — Broadcast-add merge between inverted-attention +// pooled output and the per-horizon ctx_h tensor. This is the missing +// wire that connects axis E (cross-variate attention) into the +// forward chain → MoE → loss. +// +// FORWARD: +// ctx_h[b, h, d] += inv_pooled[b, d] (broadcast over h) +// +// BACKWARD: +// ∂L/∂ctx_h[b, h, d] = ∂L/∂ctx_h_post[b, h, d] (passthrough; the +// additive merge has gradient 1 wrt either input) +// ∂L/∂inv_pooled[b, d] = Σ_h ∂L/∂ctx_h_post[b, h, d] +// +// PERF: tiny kernels (~2 KB shared mem). Block-per-batch, thread-per-d. +// +// Per `feedback_nvidia_grade_perf_for_kernels`: +// - Coalesced reads/writes (thread → adjacent d) +// - No `__syncthreads` (each thread independent) +// - No DRAM round-trips beyond the necessary reads + +#define MERGE_HIDDEN_DIM 128 +#define MERGE_N_HORIZONS 5 +#define MERGE_BLOCK 128 + +// Forward: ctx_h[b, h, d] += inv_pooled[b, d]. +// Grid (B), block (HIDDEN_DIM). Each thread d updates 5 ctx_h cells. +extern "C" __global__ void inv_pooled_merge_fwd( + float* __restrict__ ctx_h, // [B, N_H, H] IN/OUT (+=) + const float* __restrict__ inv_pooled, // [B, H] + int n_batch +) { + int b = blockIdx.x; + int tid = threadIdx.x; + if (b >= n_batch || tid >= MERGE_HIDDEN_DIM) return; + + const float v = inv_pooled[(long long)b * MERGE_HIDDEN_DIM + tid]; + #pragma unroll + for (int h = 0; h < MERGE_N_HORIZONS; ++h) { + ctx_h[(long long)b * MERGE_N_HORIZONS * MERGE_HIDDEN_DIM + + (long long)h * MERGE_HIDDEN_DIM + tid] += v; + } +} + +// Backward: grad_inv_pooled[b, d] = Σ_h grad_ctx_h_post[b, h, d]. +// Grid (B), block (HIDDEN_DIM). Each thread d sums 5 ctx_h cells. +// `grad_ctx_h` is read but NOT modified (the additive merge has +// passthrough gradient on ctx_h — same buffer flows into pool.bwd). +extern "C" __global__ void inv_pooled_merge_bwd_sum_horizon( + const float* __restrict__ grad_ctx_h, // [B, N_H, H] + int n_batch, + float* __restrict__ grad_inv_pooled // [B, H] OVERWRITE +) { + int b = blockIdx.x; + int tid = threadIdx.x; + if (b >= n_batch || tid >= MERGE_HIDDEN_DIM) return; + + float acc = 0.0f; + #pragma unroll + for (int h = 0; h < MERGE_N_HORIZONS; ++h) { + acc += grad_ctx_h[(long long)b * MERGE_N_HORIZONS * MERGE_HIDDEN_DIM + + (long long)h * MERGE_HIDDEN_DIM + tid]; + } + grad_inv_pooled[(long long)b * MERGE_HIDDEN_DIM + tid] = acc; +} diff --git a/crates/ml-alpha/src/inv_pooled_merge.rs b/crates/ml-alpha/src/inv_pooled_merge.rs new file mode 100644 index 000000000..1484fbcbb --- /dev/null +++ b/crates/ml-alpha/src/inv_pooled_merge.rs @@ -0,0 +1,82 @@ +//! Broadcast-add merge between inverted-attention pooled output and +//! the per-horizon ctx_h tensor. Wires axis E (cross-variate +//! attention) into the forward chain → MoE → loss. +//! +//! See `cuda/inv_pooled_merge.cu` for math. + +use anyhow::{Context, Result}; +use cudarc::driver::{ + CudaContext, CudaFunction, CudaModule, CudaSlice, CudaStream, LaunchConfig, PushKernelArg, +}; +use std::sync::Arc; + +pub const MERGE_HIDDEN_DIM: usize = 128; +pub const MERGE_BLOCK: usize = 128; + +const CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/inv_pooled_merge.cubin")); + +pub struct InvPooledMerge { + _module: Arc, + fwd_fn: CudaFunction, + bwd_fn: CudaFunction, + stream: Arc, +} + +impl InvPooledMerge { + pub fn new(ctx: &Arc, stream: Arc) -> Result { + let module = ctx + .load_cubin(CUBIN.to_vec()) + .context("load inv_pooled_merge cubin")?; + let fwd_fn = module + .load_function("inv_pooled_merge_fwd") + .context("load inv_pooled_merge_fwd")?; + let bwd_fn = module + .load_function("inv_pooled_merge_bwd_sum_horizon") + .context("load inv_pooled_merge_bwd_sum_horizon")?; + Ok(Self { _module: module, fwd_fn, bwd_fn, stream }) + } + + /// `ctx_h[b, h, d] += inv_pooled[b, d]`. + pub fn forward( + &self, + ctx_h: &mut CudaSlice, + inv_pooled: &CudaSlice, + n_batch: i32, + ) -> Result<()> { + let cfg = LaunchConfig { + grid_dim: (n_batch as u32, 1, 1), + block_dim: (MERGE_BLOCK as u32, 1, 1), + shared_mem_bytes: 0, + }; + let mut launch = self.stream.launch_builder(&self.fwd_fn); + unsafe { + launch + .arg(ctx_h).arg(inv_pooled).arg(&n_batch) + .launch(cfg) + .context("inv_pooled_merge_fwd")?; + } + Ok(()) + } + + /// `grad_inv_pooled[b, d] = Σ_h grad_ctx_h[b, h, d]`. + pub fn backward( + &self, + grad_ctx_h: &CudaSlice, + n_batch: i32, + grad_inv_pooled: &mut CudaSlice, + ) -> Result<()> { + let cfg = LaunchConfig { + grid_dim: (n_batch as u32, 1, 1), + block_dim: (MERGE_BLOCK as u32, 1, 1), + shared_mem_bytes: 0, + }; + let mut launch = self.stream.launch_builder(&self.bwd_fn); + unsafe { + launch + .arg(grad_ctx_h).arg(&n_batch).arg(grad_inv_pooled) + .launch(cfg) + .context("inv_pooled_merge_bwd_sum_horizon")?; + } + Ok(()) + } +} diff --git a/crates/ml-alpha/src/lib.rs b/crates/ml-alpha/src/lib.rs index 1df6907bf..56c2e0b58 100644 --- a/crates/ml-alpha/src/lib.rs +++ b/crates/ml-alpha/src/lib.rs @@ -32,6 +32,7 @@ pub mod anchor_l2; pub mod heads; pub mod horizon_mean_collapse; pub mod horizon_token_attention_pool; +pub mod inv_pooled_merge; pub mod inverted_attention_pool; pub mod isv; pub mod regime_moe_gate; diff --git a/crates/ml-alpha/src/trainer/multi_horizon_attention.rs b/crates/ml-alpha/src/trainer/multi_horizon_attention.rs index 91ec873a0..490521027 100644 --- a/crates/ml-alpha/src/trainer/multi_horizon_attention.rs +++ b/crates/ml-alpha/src/trainer/multi_horizon_attention.rs @@ -26,6 +26,7 @@ use crate::horizon_mean_collapse::HorizonMeanCollapse; use crate::horizon_token_attention_pool::{ HorizonTokenAttentionPool, HTAP_HIDDEN_DIM, HTAP_N_HORIZONS, }; +use crate::inv_pooled_merge::InvPooledMerge; use crate::inverted_attention_pool::InvertedAttentionPool; use crate::regime_moe_gate::{RegimeMoeGate, MOE_N_EXPERTS}; use crate::trainer::anchor_controller::AnchorController; @@ -58,6 +59,9 @@ pub struct MultiHorizonAttention { pub inv_pooled_d: CudaSlice, // [B, H] pub inv_attn_d: CudaSlice, // [B, H, H] pub inv_d_scores_scratch_d: CudaSlice, // [B, H, H] bwd scratch + pub grad_inv_pooled_d: CudaSlice, // [B, H] true upstream for inv_pool.bwd + // ── Broadcast-add merge: ctx_h[b, h, d] += inv_pooled[b, d] ── + pub merge: InvPooledMerge, // ── (C+E fuse) → fused_ctx [B, N_H, H] via additive merge ── // No separate kernel; computed inline in step_batched as @@ -122,6 +126,7 @@ impl MultiHorizonAttention { .context("RegimeMoeGate")?; let anchor_l2 = AnchorL2::new(ctx, stream.clone()).context("AnchorL2")?; let collapse = HorizonMeanCollapse::new(ctx, stream.clone()).context("HorizonMeanCollapse")?; + let merge = InvPooledMerge::new(ctx, stream.clone()).context("InvPooledMerge")?; let mut rng = ChaCha8Rng::seed_from_u64(seed); let scale = (1.0_f32 / MHA_HIDDEN_DIM as f32).sqrt(); @@ -168,6 +173,8 @@ impl MultiHorizonAttention { inv_pooled_d: s.alloc_zeros::(n_batch * MHA_HIDDEN_DIM)?, inv_attn_d: s.alloc_zeros::(n_batch * MHA_HIDDEN_DIM * MHA_HIDDEN_DIM)?, inv_d_scores_scratch_d: s.alloc_zeros::(n_batch * MHA_HIDDEN_DIM * MHA_HIDDEN_DIM)?, + grad_inv_pooled_d: s.alloc_zeros::(n_batch * MHA_HIDDEN_DIM)?, + merge, moe, gate_logits_d: s.alloc_zeros::(n_batch * MHA_N_EXPERTS)?, @@ -271,14 +278,13 @@ impl MultiHorizonAttention { &mut self.inv_pooled_d, &mut self.inv_attn_d, )?; - // 3. Additive merge into ctx_h_d in place: ctx_h += broadcast(inv_pooled). - // The merge kernel is intentionally tiny — capture-safe element-wise add. - add_inv_broadcast( - &self.stream, - &mut self.ctx_h_d, - &self.inv_pooled_d, - self.n_batch as i32, - )?; + // 3. Real additive merge: ctx_h[b, h, d] += inv_pooled[b, d]. + // This is the wire that connects axis E (inverted attention) + // into the forward chain → MoE → loss. Replaces the prior + // stub. Backward via `merge.backward` produces + // `grad_inv_pooled = Σ_h grad_ctx_h_post` which is the + // correct upstream gradient for `inv_pool.backward`. + self.merge.forward(&mut self.ctx_h_d, &self.inv_pooled_d, n)?; // 4. MoE gate + dispatch → routed_ctx_d, top_e_d. // (Stage 2 placeholder: gate_logits_d remains zeros set in @@ -331,12 +337,18 @@ impl MultiHorizonAttention { &mut self.grad_experts_b_d, )?; - // 2'. Inverted-attention bwd. d_inv_pooled approximated as the - // mean-collapse of the upstream gradient on the merged - // output (= grad_ctx_mean scaled). + // 3'. Merge backward: compute grad_inv_pooled[b, d] = + // Σ_h grad_ctx_h[b, h, d]. The additive merge has + // passthrough gradient on ctx_h, so grad_ctx_h_d itself + // is unchanged and feeds straight into pool.bwd below. + self.merge.backward(&self.grad_ctx_h_d, n, &mut self.grad_inv_pooled_d)?; + + // 2'. Inverted-attention bwd consumes the REAL upstream + // gradient `grad_inv_pooled` (not `grad_ctx_mean` — + // that was the prior stub-architecture fake). let inv_scale = 1.0_f32 / (self.k_seq as f32).sqrt(); self.inv_pool.backward( - ln_b_out, &self.inv_attn_d, grad_ctx_mean, + ln_b_out, &self.inv_attn_d, &self.grad_inv_pooled_d, n, k, inv_scale, &mut self.inv_d_scores_scratch_d, grad_ln_out, @@ -397,25 +409,6 @@ impl MultiHorizonAttention { } } -/// Capture-safe additive broadcast: ctx_h[b, h, d] += inv[b, d]. -/// Implemented as a per-(b, h) memset-equivalent — a tiny ad-hoc kernel -/// would be cleaner, but for now we do per-row += via a single launch. -fn add_inv_broadcast( - _stream: &Arc, - _ctx_h: &mut CudaSlice, - _inv: &CudaSlice, - _n_batch: i32, -) -> Result<()> { - // Stage 2 stub: the additive merge is folded into the MoE forward's - // input handling. In practice the MoE expert linear sees ctx_h as - // its input; broadcasting inv_pooled in here without a dedicated - // kernel would require atomicAdd or a serial copy. For Stage 2 we - // simplify by treating the MoE input as `ctx_h` directly (the - // additive contribution from inv_pool flows through the separate - // inv_pool bwd path). A follow-up kernel for true additive merge - // is a Stage 2+ cleanup. - Ok(()) -} fn upload(stream: &Arc, host: &[f32]) -> Result> { use crate::pinned_mem::MappedF32Buffer;