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

fix(ml-alpha): wire axis E into loss — real add_inv_broadcast kernel

Browse Source 
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 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-18 16:16:31 +02:00
parent 11b964359b
commit 4ae9a27f48
5 changed files with 172 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
crates/ml-alpha/build.rs
View File

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

64
crates/ml-alpha/cuda/inv_pooled_merge.cu Normal file
View File

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

82
crates/ml-alpha/src/inv_pooled_merge.rs Normal file
View File

@@ -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<CudaModule>,
fwd_fn: CudaFunction,
bwd_fn: CudaFunction,
stream: Arc<CudaStream>,
}
impl InvPooledMerge {
pub fn new(ctx: &Arc<CudaContext>, stream: Arc<CudaStream>) -> Result<Self> {
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<f32>,
inv_pooled: &CudaSlice<f32>,
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<f32>,
n_batch: i32,
grad_inv_pooled: &mut CudaSlice<f32>,
) -> 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(())
}
}

1
crates/ml-alpha/src/lib.rs
View File

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

55
crates/ml-alpha/src/trainer/multi_horizon_attention.rs
View File

@@ -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<f32>, // [B, H]
pub inv_attn_d: CudaSlice<f32>, // [B, H, H]
pub inv_d_scores_scratch_d: CudaSlice<f32>, // [B, H, H] bwd scratch
pub grad_inv_pooled_d: CudaSlice<f32>, // [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::<f32>(n_batch * MHA_HIDDEN_DIM)?,
inv_attn_d: s.alloc_zeros::<f32>(n_batch * MHA_HIDDEN_DIM * MHA_HIDDEN_DIM)?,
inv_d_scores_scratch_d: s.alloc_zeros::<f32>(n_batch * MHA_HIDDEN_DIM * MHA_HIDDEN_DIM)?,
grad_inv_pooled_d: s.alloc_zeros::<f32>(n_batch * MHA_HIDDEN_DIM)?,
merge,
moe,
gate_logits_d: s.alloc_zeros::<f32>(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<CudaStream>,
_ctx_h: &mut CudaSlice<f32>,
_inv: &CudaSlice<f32>,
_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<CudaStream>, host: &[f32]) -> Result<CudaSlice<f32>> {
use crate::pinned_mem::MappedF32Buffer;