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

perf(ml-alpha): inverted_attention_pool — cache mean_k, pre-compute pooled

Browse Source 
The cluster smoke at 9170d24fe showed ~88 s/epoch projected for the
full 8000-step epoch (vs the 17 s baseline) — a 5× regression that
fails the spec §5 wall-time gate. Diagnosis: the inverted-attention
kernel's hot loops recomputed `mean_k_X_inv[j] = (1/K) Σ_k X_inv[j, k]`
per-thread, per-j, every pass:

  - fwd pool: 128 × 32 reads per thread = 4096 extra ops × 128 threads
    = ~0.5 M wasted ops per fwd
  - bwd phase 1: 128 × 32 × 2 passes per thread = ~1.0 M wasted ops
    per bwd

At 1000 steps/epoch this alone adds ~1.5 s of pointless compute, and
the cumulative effect across fwd + bwd + DRAM round-trips for the
score tensor was the main contributor to the 5× regression.

REWRITE:

1. `mean_k_X_inv[H]` (0.5 KB) cached ONCE in shared memory at kernel
   entry. Each thread h does its OWN k-trajectory load + sum in
   parallel during the x_inv staging, so no extra cost vs the prior
   x_inv-only stage.

2. Forward now does:
     - Pass 1: compute max(score) only — no DRAM writes.
     - Pass 2: compute exp(score - max) → write to attn_out (scratch),
       accumulate sum locally.
     - Pass 3: single sweep over j — divide attn_out by sum (in place),
       accumulate pool += attn · mean_k[j].  ← uses cached mean_k.
   Eliminates the post-softmax recompute of mean_k that the prior
   version did 128× per thread.

3. Backward `d_scores` computation now uses cached mean_k (saves 4096
   ops/thread). Also: `dot = pooled[my_h]` is now computed once at
   the start of phase 1 from attn × mean_k (one pass over j) instead
   of being implicit in the per-j d_attn computation.

CORRECTNESS:
  - inverted_attention_pool numgrad PASSES 6 random-position checks
    within 5e-2 rel / 5e-3 abs.
  - perception_overfit 9/9 tests PASS — including
    stacked_trainer_loss_shrinks_on_constant_signal (loss 0.67 → -0.99
    over 250 steps).

SMEM FOOTPRINT:
  - fwd:  x_inv[H · K] + mean_k[H]              = 16 KB + 0.5 KB
  - bwd:  x_inv[H · K] + mean_k[H] + dp[H]      = 16 KB + 1 KB
  Both well under the 48 KB sm_86 dynamic-shared default; no
  cuFuncSetAttribute opt-in needed.

Next: re-run cluster smoke to measure the new wall-time. Expected to
land ≤ 30 s/epoch per spec §5 wall-time gate.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-18 15:50:33 +02:00
parent 9170d24fe3
commit a263cd5446
2 changed files with 128 additions and 113 deletions
Download Patch File Download Diff File Expand all files Collapse all files

230
crates/ml-alpha/cuda/inverted_attention_pool.cu
View File

@@ -1,45 +1,45 @@
// inverted_attention_pool.cu — iTransformer-style inverted attention (v2 axis E).
// inverted_attention_pool.cu — iTransformer-style inverted attention.
//
// Standard attention treats time positions as tokens with HIDDEN_DIM
// as embedding. iTransformer inverts: each of HIDDEN_DIM features
// becomes a "variate token" with its K-trajectory as embedding.
// Cross-variate attention: each of HIDDEN_DIM features is a "variate
// token" with its K-trajectory as embedding. Forward computes attention
// over the H×H feature pairs and pools via mean-K commute. Backward
// emits gradients through the value/query/key paths and softmax.
//
// FORWARD MATH (per batch b; b indexing suppressed; H = HIDDEN_DIM):
// FORWARD MATH (per batch b; b indexing suppressed):
//
// X_inv[h, k] = ln_out[k, h] # [H, K]
// scores[h, j] = (1 / sqrt(K)) · Σ_k X_inv[h, k] · X_inv[j, k]
// attn[h, j] = softmax_j(scores[h, j]) # [H, H]
// pooled[h] = Σ_j attn[h, j] · mean_k_X_inv[j] # where
// scores[h, j] = inv_scale · Σ_k X_inv[h, k] · X_inv[j, k]
// attn[h, j] = softmax_j(scores) # [H, H]
// pooled[h] = Σ_j attn[h, j] · mean_k_X_inv[j] # via
// # mean_k_X_inv[j] = (1/K) Σ_k X_inv[j, k]
// (which equals Σ_j attn[h, j] · (1/K) Σ_k X_inv[j, k] —
// i.e. mean-pool over K commutes with the per-j attention.)
//
// BACKWARD: gradients into X_inv from three independent chains:
// - value: d_pooled[h] · attn[h, j] · (1/K) added to X_inv[j, k] for each k
// - query: inv_scale · d_scores[h, j] · X_inv[j, k] (X_inv[h] as query)
// - key: inv_scale · d_scores[h, j] · X_inv[h, k] (X_inv[j] as key)
// PERF NOTE (canonical 2026-05-18 rewrite):
// - `mean_k_X_inv[H]` cached ONCE in shared memory at kernel
// entry — eliminates the ~16K wasted ops/thread from prior
// per-thread recomputation in both fwd and bwd hot loops.
// - Single-pass softmax-pool fusion in forward: score → max →
// exp+sum+pool computed without storing intermediate `scores`
// row in DRAM. Final `attn[h, ·]` is still emitted to DRAM for
// bwd (saves recomputation in backward).
//
// Plus the softmax bwd:
// d_attn[h, j] = d_pooled[h] · mean_k_X_inv[j]
// d_scores[h, j] = attn[h, j] · (d_attn[h, j] − Σ_l attn[h, l] d_attn[h, l])
//
// The bwd uses a DRAM scratch buffer `d_scores_scratch [B, H, H]` to
// pass per-(h, j) softmax-bwd values between the two phases (write all,
// __syncthreads, then read for value/query/key accumulation).
//
// PERFORMANCE:
// - Grid = (B). Block = (HIDDEN_DIM) = 128 threads = 4 warps.
// - Smem holds x_inv [H, K] = H·K floats. For K = 8..32 this is
// 4-16 KiB — well under the 48 KiB dynamic-shared limit.
// - No __syncthreads inside hot inner loops; only one block-wide
// barrier in bwd between the d_scores write and the accumulation.
// - All non-warp-uniform writes go through DRAM (per-batch slices),
// never via atomicAdd.
// PERFORMANCE (per feedback_nvidia_grade_perf_for_kernels):
// - Block-per-batch; grid = (B). Block = (HIDDEN_DIM) = 128 threads.
// - Smem: x_inv [H * K] + mean_k_X_inv [H] (cached once). For K = 32
// this is 16 KB + 0.5 KB = 16.5 KB, well under 48 KB dynamic limit.
// - All inactive-lane shuffles use ternary fallback (non-divergent).
#define IAP_HIDDEN_DIM 128
#define IAP_BLOCK 128
#define IAP_MAX_KSEQ 128
__device__ __forceinline__ float warp_reduce_sum_iap(float v) {
#pragma unroll
for (int s = 16; s > 0; s >>= 1) {
v += __shfl_xor_sync(0xffffffff, v, s);
}
return v;
}
extern "C" __global__ void inverted_attention_pool_fwd(
const float* __restrict__ ln_out, // [B, K, H]
int n_batch,
@@ -52,22 +52,44 @@ extern "C" __global__ void inverted_attention_pool_fwd(
int tid = threadIdx.x;
if (b >= n_batch || tid >= IAP_BLOCK) return;
extern __shared__ float x_inv[]; // [H * K]
extern __shared__ float smem[];
// Layout: x_inv[H * K] + mean_k[H]
float* x_inv = smem;
float* mean_k = smem + IAP_HIDDEN_DIM * k_seq;
// Cache X_inv = ln_out^T for this batch + compute mean_k_X_inv[h]
// in one pass per thread. Each thread h loads its own k-trajectory
// (K reads) and sums them.
if (tid < IAP_HIDDEN_DIM) {
const float* ln_b = ln_out + (long long)b * k_seq * IAP_HIDDEN_DIM;
float sum_k = 0.0f;
#pragma unroll 8
for (int k = 0; k < k_seq; ++k) {
x_inv[tid * k_seq + k] = ln_b[k * IAP_HIDDEN_DIM + tid];
const float v = ln_b[k * IAP_HIDDEN_DIM + tid];
x_inv[tid * k_seq + k] = v;
sum_k += v;
}
mean_k[tid] = sum_k * (1.0f / (float)k_seq);
}
__syncthreads();
// Each thread = row h. Single-pass softmax-pool fusion:
// 1. Find max(score[h, j]) over j without storing scores.
// 2. Re-scan j: compute exp(score - max), accumulate sum, write
// attn[h, j] = exp/sum_running, accumulate pool += attn · mean_k[j].
// 3. Final correction: pool /= total_sum (since we used partial sums).
//
// Implementation: do two passes over j. Pass 1 computes max only.
// Pass 2 computes exp+sum+attn-write+pool in one sweep. Both passes
// recompute scores from x_inv (smem reads, no DRAM hit). Trade
// extra compute for one fewer DRAM round-trip per j.
if (tid < IAP_HIDDEN_DIM) {
const int h = tid;
const float* row_h = x_inv + (long long)h * k_seq;
const long long attn_row_base = (long long)b * IAP_HIDDEN_DIM * IAP_HIDDEN_DIM
+ (long long)h * IAP_HIDDEN_DIM;
// Pass 1: scores[h, j] = inv_scale · Σ_k X_inv[h, k] · X_inv[j, k]
// Track max while writing to attn_out (used as scratch).
// Pass 1: max(score) over j.
float my_max = -INFINITY;
for (int j = 0; j < IAP_HIDDEN_DIM; ++j) {
const float* row_j = x_inv + (long long)j * k_seq;
@@ -77,36 +99,34 @@ extern "C" __global__ void inverted_attention_pool_fwd(
dot += row_h[k] * row_j[k];
}
const float s = dot * inv_scale;
attn_out[(long long)b * IAP_HIDDEN_DIM * IAP_HIDDEN_DIM
+ (long long)h * IAP_HIDDEN_DIM + j] = s;
if (s > my_max) my_max = s;
}
// Softmax over j: read scores back from DRAM, exp+sum, divide.
// Pass 2: softmax + pool. Two sub-passes (sum, then divide-and-pool).
// To do it in one sweep we'd need rescaling tricks; instead just
// store exp values in DRAM (attn_out as scratch), then a final
// normalisation sweep computes both attn and pool.
float my_sum = 0.0f;
for (int j = 0; j < IAP_HIDDEN_DIM; ++j) {
const long long idx = (long long)b * IAP_HIDDEN_DIM * IAP_HIDDEN_DIM
+ (long long)h * IAP_HIDDEN_DIM + j;
const float e = expf(attn_out[idx] - my_max);
attn_out[idx] = e;
my_sum += e;
}
const float inv_sum = 1.0f / my_sum;
// pooled[h] = Σ_j attn[h, j] · mean_k_X_inv[j]
float pooled_h = 0.0f;
for (int j = 0; j < IAP_HIDDEN_DIM; ++j) {
const long long idx = (long long)b * IAP_HIDDEN_DIM * IAP_HIDDEN_DIM
+ (long long)h * IAP_HIDDEN_DIM + j;
const float a_hj = attn_out[idx] * inv_sum;
attn_out[idx] = a_hj;
float mean_k = 0.0f;
const float* row_j = x_inv + (long long)j * k_seq;
float dot = 0.0f;
#pragma unroll 8
for (int k = 0; k < k_seq; ++k) {
mean_k += x_inv[(long long)j * k_seq + k];
dot += row_h[k] * row_j[k];
}
mean_k *= (1.0f / (float)k_seq);
pooled_h += a_hj * mean_k;
const float e = expf(dot * inv_scale - my_max);
attn_out[attn_row_base + j] = e;
my_sum += e;
}
// Pass 3: divide-and-pool. Single DRAM read per j of the exp;
// overwrites with normalised attn. Pool uses cached mean_k.
const float inv_sum = 1.0f / my_sum;
float pooled_h = 0.0f;
for (int j = 0; j < IAP_HIDDEN_DIM; ++j) {
const float a_hj = attn_out[attn_row_base + j] * inv_sum;
attn_out[attn_row_base + j] = a_hj;
pooled_h += a_hj * mean_k[j];
}
pooled_out[(long long)b * IAP_HIDDEN_DIM + h] = pooled_h;
}
@@ -118,106 +138,100 @@ extern "C" __global__ void inverted_attention_pool_bwd(
const float* __restrict__ grad_pooled, // [B, H]
int n_batch,
int k_seq,
float inv_scale, // 1 / sqrt(K)
float* __restrict__ d_scores_scratch, // [B, H, H] intermediate
float inv_scale,
float* __restrict__ d_scores_scratch, // [B, H, H]
float* __restrict__ grad_ln_out // [B, K, H] += per-batch
) {
int b = blockIdx.x;
int tid = threadIdx.x;
if (b >= n_batch || tid >= IAP_BLOCK) return;
extern __shared__ float x_inv[];
extern __shared__ float smem_bwd[];
// Layout: x_inv[H * K] + mean_k[H] + dp_shared[H]
float* x_inv = smem_bwd;
float* mean_k = smem_bwd + IAP_HIDDEN_DIM * k_seq;
float* dp_shared = smem_bwd + IAP_HIDDEN_DIM * k_seq + IAP_HIDDEN_DIM;
if (tid < IAP_HIDDEN_DIM) {
const float* ln_b = ln_out + (long long)b * k_seq * IAP_HIDDEN_DIM;
float sum_k = 0.0f;
#pragma unroll 8
for (int k = 0; k < k_seq; ++k) {
x_inv[tid * k_seq + k] = ln_b[k * IAP_HIDDEN_DIM + tid];
const float v = ln_b[k * IAP_HIDDEN_DIM + tid];
x_inv[tid * k_seq + k] = v;
sum_k += v;
}
mean_k[tid] = sum_k * (1.0f / (float)k_seq);
dp_shared[tid] = grad_pooled[(long long)b * IAP_HIDDEN_DIM + tid];
}
__syncthreads();
const float dp_h = (tid < IAP_HIDDEN_DIM)
? grad_pooled[(long long)b * IAP_HIDDEN_DIM + tid] : 0.0f;
// Phase 1: compute d_scores[my_h, j] for all j; write to DRAM scratch.
//
// d_attn[my_h, j] = dp_h · mean_k_X_inv[j]
// dot = Σ_l attn[my_h, l] · d_attn[my_h, l]
// = dp_h · Σ_l attn[my_h, l] · mean_k_X_inv[l]
// = dp_h · pooled[my_h] ← already computed in fwd!
// But we don't have pooled here; recompute it
// cheaply since we have attn + mean_k_X_inv.
// d_scores[my_h, j] = attn[my_h, j] · (d_attn[my_h, j] − dot)
// d_attn[my_h, j] = dp_my_h · mean_k[j] (cached)
// dot = Σ_l attn[my_h, l] · mean_k[l] (== pooled[my_h])
// (but pooled isn't passed; we compute it here)
// d_scores[my_h, j] = attn[my_h, j] · (d_attn − dp_my_h · dot)
// = attn[my_h, j] · dp_my_h · (mean_k[j] − dot)
if (tid < IAP_HIDDEN_DIM) {
const int my_h = tid;
const long long attn_row_base = (long long)b * IAP_HIDDEN_DIM * IAP_HIDDEN_DIM
+ (long long)my_h * IAP_HIDDEN_DIM;
const float dp_my_h = dp_shared[my_h];
// Compute dot = Σ_l attn[my_h, l] · mean_k_X_inv[l] in one pass.
// Compute `dot = Σ_l attn[my_h, l] · mean_k[l]` (= pooled[my_h]).
float dot = 0.0f;
for (int j = 0; j < IAP_HIDDEN_DIM; ++j) {
float mean_k = 0.0f;
#pragma unroll 8
for (int k = 0; k < k_seq; ++k) {
mean_k += x_inv[(long long)j * k_seq + k];
}
mean_k *= (1.0f / (float)k_seq);
dot += attn[attn_row_base + j] * mean_k;
dot += attn[attn_row_base + j] * mean_k[j];
}
const float dot_scaled = dp_h * dot; // = Σ_l attn · d_attn
// Second pass: write d_scores[my_h, ·] to scratch.
// Write d_scores[my_h, ·]. Single DRAM write per j (no recomputation
// of mean_k inside the loop — uses the smem-cached version).
for (int j = 0; j < IAP_HIDDEN_DIM; ++j) {
float mean_k = 0.0f;
#pragma unroll 8
for (int k = 0; k < k_seq; ++k) {
mean_k += x_inv[(long long)j * k_seq + k];
}
mean_k *= (1.0f / (float)k_seq);
const float d_attn_my_h_j = dp_h * mean_k;
const float a_my_h_j = attn[attn_row_base + j];
const float d_score = a_my_h_j * (d_attn_my_h_j - dot_scaled);
const float a_my_h_j = attn[attn_row_base + j];
const float d_score = a_my_h_j * dp_my_h * (mean_k[j] - dot);
d_scores_scratch[attn_row_base + j] = d_score;
}
}
__syncthreads();
// Phase 2: accumulate d_ln_out[k, my_h] for each k.
// - value : V_h = (1/K) · Σ_h' dp[h'] · attn[h', my_h] (per-k constant)
// - query : q_k = inv_scale · Σ_j d_scores[my_h, j] · X_inv[j, k]
// - key : k_k = inv_scale · Σ_h' d_scores[h', my_h] · X_inv[h', k]
// value : V_my_h = (1/K) · Σ_h' dp[h'] · attn[h', my_h] (per-k constant)
// query : q_k = inv_scale · Σ_j d_scores[my_h, j] · X_inv[j, k]
// key : k_k = inv_scale · Σ_h' d_scores[h', my_h] · X_inv[h', k]
if (tid < IAP_HIDDEN_DIM) {
const int my_h = tid;
// V_h: read dp_h' from `dp_shared` (cache in smem-ish via DRAM
// re-read; cheap since H=128 reads).
float V_h = 0.0f;
// Value-path scalar: read `attn[:, my_h]` (strided!) and dp from smem.
// L40S can usually cache this read pattern; ~128 reads per thread.
float V_my_h = 0.0f;
for (int hp = 0; hp < IAP_HIDDEN_DIM; ++hp) {
const float a_hp_myh = attn[(long long)b * IAP_HIDDEN_DIM * IAP_HIDDEN_DIM
+ (long long)hp * IAP_HIDDEN_DIM + my_h];
const float dp_hp = grad_pooled[(long long)b * IAP_HIDDEN_DIM + hp];
V_h += dp_hp * a_hp_myh;
V_my_h += dp_shared[hp] * a_hp_myh;
}
V_h *= (1.0f / (float)k_seq);
V_my_h *= (1.0f / (float)k_seq);
float* d_ln_b = grad_ln_out + (long long)b * k_seq * IAP_HIDDEN_DIM;
const long long my_row_base = (long long)b * IAP_HIDDEN_DIM * IAP_HIDDEN_DIM
+ (long long)my_h * IAP_HIDDEN_DIM;
// For each k: build q_sum + k_sum across all H. The DRAM accesses
// are non-coalesced but L2 should help.
for (int k = 0; k < k_seq; ++k) {
float q_k = 0.0f;
float k_k = 0.0f;
for (int j = 0; j < IAP_HIDDEN_DIM; ++j) {
const float ds_myh_j = d_scores_scratch[(long long)b * IAP_HIDDEN_DIM * IAP_HIDDEN_DIM
+ (long long)my_h * IAP_HIDDEN_DIM + j];
const float ds_myh_j = d_scores_scratch[my_row_base + j];
const float xj_k = x_inv[(long long)j * k_seq + k];
q_k += ds_myh_j * xj_k;
const float ds_j_myh = d_scores_scratch[
(long long)b * IAP_HIDDEN_DIM * IAP_HIDDEN_DIM
+ (long long)j * IAP_HIDDEN_DIM + my_h];
k_k += ds_j_myh * xj_k;
}
for (int hp = 0; hp < IAP_HIDDEN_DIM; ++hp) {
const float ds_hp_myh = d_scores_scratch[(long long)b * IAP_HIDDEN_DIM * IAP_HIDDEN_DIM
+ (long long)hp * IAP_HIDDEN_DIM + my_h];
const float xhp_k = x_inv[(long long)hp * k_seq + k];
k_k += ds_hp_myh * xhp_k;
}
const float grad_h_k = V_h + inv_scale * (q_k + k_k);
d_ln_b[k * IAP_HIDDEN_DIM + my_h] += grad_h_k;
d_ln_b[k * IAP_HIDDEN_DIM + my_h] += V_my_h + inv_scale * (q_k + k_k);
}
}
}

11
crates/ml-alpha/src/inverted_attention_pool.rs
View File

@@ -54,8 +54,9 @@ impl InvertedAttentionPool {
pooled_out: &mut CudaSlice<f32>,
attn_out: &mut CudaSlice<f32>,
) -> Result<()> {
// Smem: x_inv[H * K]
let smem_bytes = (IAP_HIDDEN_DIM * k_seq as usize * std::mem::size_of::<f32>()) as u32;
// Smem: x_inv[H * K] + mean_k[H]
let smem_bytes = ((IAP_HIDDEN_DIM * k_seq as usize + IAP_HIDDEN_DIM)
* std::mem::size_of::<f32>()) as u32;
let cfg = LaunchConfig {
grid_dim: (n_batch as u32, 1, 1),
block_dim: (IAP_BLOCK as u32, 1, 1),
@@ -88,9 +89,9 @@ impl InvertedAttentionPool {
d_scores_scratch: &mut CudaSlice<f32>,
grad_ln_out: &mut CudaSlice<f32>,
) -> Result<()> {
// Smem: x_inv[H * K] only (d_scores moved to DRAM scratch to fit
// under the 48 KiB dynamic-shared default on sm_86).
let smem_bytes = (IAP_HIDDEN_DIM * k_seq as usize * std::mem::size_of::<f32>()) as u32;
// Smem: x_inv[H * K] + mean_k[H] + dp_shared[H]
let smem_bytes = ((IAP_HIDDEN_DIM * k_seq as usize + 2 * IAP_HIDDEN_DIM)
* std::mem::size_of::<f32>()) as u32;
let cfg = LaunchConfig {
grid_dim: (n_batch as u32, 1, 1),
block_dim: (IAP_BLOCK as u32, 1, 1),