refactor(ml-alpha): revert axes B/C/D/E — keep only Kendall σ (axis A)
Post-A/B verdict (see project_ml_alpha_v2_ab_verdict.md): v2 with all
5 axes was marginally tied on h6000 (+0.0013 vs 0.7591 baseline mean,
fails the +0.01 win threshold) and slightly below on mean_auc
(−0.0208 vs 0.7749 baseline mean, within 1σ) at ~5× the wall-time
cost. Per `feedback_v7_gem_methodology` (measure before delete or
wire), the architecture has been measured — it doesn't earn its
compute cost. This commit reverts the axes that didn't lift:
- axis B (L2 anchor + Wiener-α controller) — DROPPED
- axis C (horizon-token attention pool) — DROPPED
- axis D (regime-MoE gate + experts) — DROPPED
- axis E (inverted cross-variate attn) — DROPPED
- axis A (Kendall σ-weighted BCE) — KEPT
Files deleted (kernels, host bindings, numgrad tests, trainer state):
- cuda/{horizon_token_attention_pool, inverted_attention_pool,
inv_pooled_merge, regime_moe_gate, anchor_l2,
horizon_mean_collapse}.cu
- src/{horizon_token_attention_pool, inverted_attention_pool,
inv_pooled_merge, regime_moe_gate, anchor_l2,
horizon_mean_collapse}.rs
- src/trainer/{multi_horizon_attention, anchor_controller}.rs
- tests/{horizon_token_attention_pool_numgrad,
inverted_attention_pool_numgrad,
regime_moe_gate_numgrad,
anchor_l2_numgrad}.rs
Files restored (from V1 commit 41292303d):
- cuda/attention_pool.cu — legacy single-Q attention pool kernel
- src/trainer/perception.rs — pre-MHA trainer state with the
legacy `attn_*` plumbing intact.
Files modified:
- bce_loss_multi_horizon.cu stays σ-aware (kept the V7 work; it
has the kernel function name preserved from V1).
- perception.rs: ADD `log_sigma_h_d [N_HORIZONS]`,
`grad_log_sigma_h_d [N_HORIZONS]`, `opt_log_sigma` AdamW
directly on PerceptionTrainer (no MHA bundle). BCE callsites in
`step_batched` (training) and `evaluate_batched` thread the σ
args. Grad scratch zeroed each step before the BCE launch.
`opt_log_sigma.step` lives in section 9 alongside the other
AdamW updates.
NET DIFF: 23 files, 442 insertions, 3160 deletions (~2700-line
cleanup).
LOCAL VERIFICATION (RTX 3050 sm_86, --test-threads=1):
- ml-alpha builds clean (cuda feature)
- bce_grad_finite_diff 4/4 PASS (BCE still works through σ-kernel)
- perception_overfit 9/9 PASS (full trainer pipeline, loss-shrinks
tests still green)
NEXT: cluster smoke + 3-fold A/B vs task #200 baseline. Expected
wall-time ≈ baseline 17 s/epoch (we're back to baseline architecture
plus 5 scalar Kendall σ params + 1 tiny AdamW). Expected lift on
mean_auc: modest — Kendall σ rebalances per-horizon contributions
based on observed BCE EMA, which may help horizons with intrinsically
higher noise floors.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -13,19 +13,13 @@ const KERNELS: &[&str] = &[
|
||||
"cfc_step",
|
||||
"multi_horizon_heads",
|
||||
"projection",
|
||||
"bce_loss_multi_horizon",
|
||||
"bce_loss_multi_horizon_sigma", // v2-A: Kendall σ-weighted BCE
|
||||
"bce_loss_multi_horizon", // Kendall σ-weighted multi-horizon BCE (axis A)
|
||||
"adamw_step",
|
||||
"grad_norm",
|
||||
"horizon_lambda", // ISV-driven per-horizon gradient scaler (EMA + lambda)
|
||||
"layer_norm", // Phase 1: trunk pre-CfC normalisation
|
||||
"variable_selection", // Phase 2D: TFT-style per-feature gating
|
||||
"horizon_token_attention_pool", // Horizon-token K-prepend single-Q attention pool
|
||||
"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
|
||||
"attention_pool", // Phase 3: single-Q learned content summary at CfC k=0
|
||||
"reduce_axis0", // Phase B: cross-batch param-grad reducer
|
||||
];
|
||||
|
||||
|
||||
@@ -1,67 +0,0 @@
|
||||
// anchor_l2.cu — L2 anchor regularization toward initialization (v2 axis B).
|
||||
//
|
||||
// Computes anchor_loss = λ · Σ_i (p[i] − p_init[i])² for a single
|
||||
// flat parameter buffer p of length n. Adds the corresponding L2
|
||||
// gradient 2λ(p − p_init) into the parameter's grad buffer (`+=`).
|
||||
//
|
||||
// Caller invokes one launch per parameter group anchored. The
|
||||
// trainer's Wiener-α controller computes λ each step on the host
|
||||
// (cheap — scalar update; copies λ to a device-side single-float
|
||||
// buffer before the launch).
|
||||
//
|
||||
// PERFORMANCE:
|
||||
// - Warp-shuffle reduce for the loss sum. One block per parameter
|
||||
// group; grid_dim = (1, 1, 1). For n ≤ ~100K (typical group size)
|
||||
// a single 256-thread block is sufficient; threads stride over n.
|
||||
// - Grad write is per-element, coalesced.
|
||||
|
||||
#define ANCHOR_BLOCK 256
|
||||
#define ANCHOR_NWARPS (ANCHOR_BLOCK / 32)
|
||||
|
||||
__device__ __forceinline__ float warp_reduce_sum_an(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 anchor_l2_fwd_bwd(
|
||||
const float* __restrict__ p, // [n] current parameter values
|
||||
const float* __restrict__ p_init, // [n] snapshot at trainer construction
|
||||
const float* __restrict__ lambda_scalar, // [1] current λ from Wiener-α controller
|
||||
int n,
|
||||
float* __restrict__ loss_out, // [1] λ · Σ (p − p_init)²
|
||||
float* __restrict__ grad_p // [n] += 2λ (p − p_init)
|
||||
) {
|
||||
const int tid = threadIdx.x;
|
||||
const int lane = tid & 31;
|
||||
const int warp = tid >> 5;
|
||||
|
||||
const float lam = lambda_scalar[0];
|
||||
const float two_lam = 2.0f * lam;
|
||||
|
||||
// Strided per-thread accumulate of (p - p_init)^2 + emit grad.
|
||||
float local_sum = 0.0f;
|
||||
for (int i = tid; i < n; i += ANCHOR_BLOCK) {
|
||||
const float diff = p[i] - p_init[i];
|
||||
local_sum += diff * diff;
|
||||
// Coalesced grad write.
|
||||
grad_p[i] += two_lam * diff;
|
||||
}
|
||||
|
||||
// Warp-shuffle reduce.
|
||||
float warp_sum = warp_reduce_sum_an(local_sum);
|
||||
|
||||
__shared__ float s_warp[ANCHOR_NWARPS];
|
||||
if (lane == 0) s_warp[warp] = warp_sum;
|
||||
__syncthreads();
|
||||
|
||||
// Cross-warp reduce: all 32 lanes of warp 0 participate; inactive
|
||||
// lanes contribute 0 via ternary (NOT a divergent shuffle).
|
||||
float w = (tid < ANCHOR_NWARPS) ? s_warp[tid] : 0.0f;
|
||||
if (tid < 32) {
|
||||
w = warp_reduce_sum_an(w);
|
||||
if (tid == 0) loss_out[0] = lam * w;
|
||||
}
|
||||
}
|
||||
244
crates/ml-alpha/cuda/attention_pool.cu
Normal file
244
crates/ml-alpha/cuda/attention_pool.cu
Normal file
@@ -0,0 +1,244 @@
|
||||
// attention_pool.cu — Single-head attention pool over Mamba2 K-positions.
|
||||
// Phase 3 (2026-05-17).
|
||||
//
|
||||
// Replaces the CfC's zero-initialised `h_old` at k=0 with a learned
|
||||
// attention-pooled summary over all K LN_b output positions. The K-loop
|
||||
// CfC keeps its recurrent semantics (h_old at k+1 = h_new at k); only
|
||||
// the initial state changes from zero to a content-addressable lookup.
|
||||
//
|
||||
// Forward math (per sample b):
|
||||
// scores[k] = Q · keys[b, k, :] # [K] — dot product over HIDDEN_DIM
|
||||
// attn[k] = softmax_k(scores) # [K]
|
||||
// context[h] = sum_k attn[k] * values[b, k, h] # [HIDDEN_DIM]
|
||||
//
|
||||
// For this attention pool, keys == values == LN_b output [B, K, HIDDEN_DIM].
|
||||
// Single learned param: Q [HIDDEN_DIM].
|
||||
//
|
||||
// Saved for backward:
|
||||
// attn_weights [B, K] (post-softmax)
|
||||
//
|
||||
// Block layout (forward):
|
||||
// grid_dim = (B, 1, 1), block_dim = (HIDDEN_DIM=128, 1, 1).
|
||||
// Thread tid handles dimension h of HIDDEN_DIM.
|
||||
//
|
||||
// Per `feedback_no_atomicadd.md`: block tree-reduce only, no atomics.
|
||||
|
||||
#define ATTN_HIDDEN_DIM 128
|
||||
#define ATTN_BLOCK 128
|
||||
#define ATTN_MAX_K 512 // safety cap on K; smoke uses K=16-64.
|
||||
|
||||
extern "C" __global__ void attention_pool_fwd(
|
||||
const float* __restrict__ Q, // [HIDDEN_DIM]
|
||||
const float* __restrict__ ln_out, // [B, K, HIDDEN_DIM]
|
||||
int n_batch,
|
||||
int k_seq,
|
||||
float* __restrict__ context, // [B, HIDDEN_DIM]
|
||||
float* __restrict__ attn_weights // [B, K] — saved post-softmax
|
||||
) {
|
||||
int b_idx = blockIdx.x;
|
||||
int tid = threadIdx.x;
|
||||
if (b_idx >= n_batch || tid >= ATTN_BLOCK) return;
|
||||
|
||||
extern __shared__ float smem[];
|
||||
float* s_scores = smem; // [K]
|
||||
float* s_red = smem + k_seq; // [BLOCK] — reduce scratch
|
||||
float* s_context = smem + k_seq + ATTN_BLOCK; // [HIDDEN_DIM] — final output buffer
|
||||
|
||||
// Cache Q in shared mem (broadcast).
|
||||
__shared__ float s_Q[ATTN_HIDDEN_DIM];
|
||||
if (tid < ATTN_HIDDEN_DIM) s_Q[tid] = Q[tid];
|
||||
__syncthreads();
|
||||
|
||||
// Pass 1: scores[k] = Q · ln_out[b, k, :].
|
||||
// Each k is handled sequentially by ALL threads via tree-reduce
|
||||
// over HIDDEN_DIM. K outer loop, threads tile HIDDEN_DIM inner.
|
||||
// For HIDDEN_DIM=128 and ATTN_BLOCK=128 we have one-to-one.
|
||||
const float* ln_b = ln_out + (long long)b_idx * k_seq * ATTN_HIDDEN_DIM;
|
||||
for (int k = 0; k < k_seq; ++k) {
|
||||
const float v = ln_b[k * ATTN_HIDDEN_DIM + tid] * s_Q[tid];
|
||||
s_red[tid] = v;
|
||||
__syncthreads();
|
||||
for (int s = ATTN_BLOCK / 2; s > 0; s >>= 1) {
|
||||
if (tid < s) s_red[tid] += s_red[tid + s];
|
||||
__syncthreads();
|
||||
}
|
||||
if (tid == 0) s_scores[k] = s_red[0];
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
// Pass 2: softmax over K. Numerically stable: max-subtract + exp + sum.
|
||||
// Block-wide max over s_scores[0..k_seq).
|
||||
float my_max = -INFINITY;
|
||||
for (int k = tid; k < k_seq; k += ATTN_BLOCK) {
|
||||
const float v = s_scores[k];
|
||||
if (v > my_max) my_max = v;
|
||||
}
|
||||
s_red[tid] = my_max;
|
||||
__syncthreads();
|
||||
for (int s = ATTN_BLOCK / 2; s > 0; s >>= 1) {
|
||||
if (tid < s) {
|
||||
const float a = s_red[tid];
|
||||
const float b = s_red[tid + s];
|
||||
s_red[tid] = (a > b) ? a : b;
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
__shared__ float s_max;
|
||||
if (tid == 0) s_max = s_red[0];
|
||||
__syncthreads();
|
||||
|
||||
// exp(score - max) + sum.
|
||||
float my_sum = 0.0f;
|
||||
for (int k = tid; k < k_seq; k += ATTN_BLOCK) {
|
||||
const float e = expf(s_scores[k] - s_max);
|
||||
s_scores[k] = e; // reuse as exp_shifted
|
||||
my_sum += e;
|
||||
}
|
||||
s_red[tid] = my_sum;
|
||||
__syncthreads();
|
||||
for (int s = ATTN_BLOCK / 2; s > 0; s >>= 1) {
|
||||
if (tid < s) s_red[tid] += s_red[tid + s];
|
||||
__syncthreads();
|
||||
}
|
||||
__shared__ float s_sum;
|
||||
if (tid == 0) s_sum = s_red[0];
|
||||
__syncthreads();
|
||||
|
||||
// Pass 3: attn[k] = exp/sum; save; context[h] = sum_k attn[k] * ln_out[b, k, h].
|
||||
// Save attn weights AND build the context output.
|
||||
for (int k = tid; k < k_seq; k += ATTN_BLOCK) {
|
||||
const float a = s_scores[k] / s_sum;
|
||||
s_scores[k] = a; // overwrite again — now holds true attn weights
|
||||
attn_weights[(long long)b_idx * k_seq + k] = a;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// Accumulate context[h=tid] = sum_k attn[k] * ln_out[b, k, h=tid].
|
||||
if (tid < ATTN_HIDDEN_DIM) {
|
||||
float c = 0.0f;
|
||||
for (int k = 0; k < k_seq; ++k) {
|
||||
c += s_scores[k] * ln_b[k * ATTN_HIDDEN_DIM + tid];
|
||||
}
|
||||
s_context[tid] = c;
|
||||
context[(long long)b_idx * ATTN_HIDDEN_DIM + tid] = c;
|
||||
}
|
||||
}
|
||||
|
||||
// Attention pool backward — chain rule:
|
||||
//
|
||||
// d_attn[k] = sum_h grad_context[h] * values[b, k, h]
|
||||
// + (this is from context = sum_k attn[k] * values[b, k, :])
|
||||
//
|
||||
// d_values[b, k, h] += grad_context[h] * attn[k]
|
||||
// + d_scores[k] * Q[h]
|
||||
//
|
||||
// d_scores via softmax Jacobian:
|
||||
// d_scores[k] = attn[k] * (d_attn[k] - sum_kp attn[kp] * d_attn[kp])
|
||||
//
|
||||
// d_Q[h] += sum_{b, k} d_scores[k] * values[b, k, h]
|
||||
// (here values == ln_out)
|
||||
//
|
||||
// Single-writer discipline:
|
||||
// - d_Q is [HIDDEN_DIM] shared across all (b, k). ONE block per
|
||||
// launch, internal batch loop, tile over HIDDEN_DIM. += into d_Q.
|
||||
// - d_values[b, k, h] is [B, K, HIDDEN_DIM] — one block per launch
|
||||
// iterates over (b, k, h) sequentially; with block_dim = HIDDEN_DIM,
|
||||
// thread h is sole writer of column h for ALL (b, k).
|
||||
|
||||
// Block-per-batch attn_pool bwd (Phase B commit 4).
|
||||
// grid=(n_batch, 1, 1) block=(ATTN_BLOCK, 1, 1)
|
||||
//
|
||||
// Each block handles one batch's HIDDEN_DIM channels. grad_ln_out is
|
||||
// per-batch indexed (already safe), so each block bi writes its
|
||||
// [bi, :, :] slice via += onto whatever value grad_ln_out holds at
|
||||
// kernel launch (the K-loop's contribution to LN_b output grad).
|
||||
//
|
||||
// grad_Q is a single [HIDDEN_DIM] shared across all batches → per-batch
|
||||
// scratch [B, HIDDEN_DIM], reduced after the kernel returns.
|
||||
extern "C" __global__ void attention_pool_bwd(
|
||||
const float* __restrict__ Q, // [HIDDEN_DIM]
|
||||
const float* __restrict__ ln_out, // [B, K, HIDDEN_DIM] (= values)
|
||||
const float* __restrict__ attn_weights, // [B, K] from fwd
|
||||
const float* __restrict__ grad_context, // [B, HIDDEN_DIM]
|
||||
int n_batch,
|
||||
int k_seq,
|
||||
float* __restrict__ grad_Q_scratch, // [B, HIDDEN_DIM] (+=)
|
||||
float* __restrict__ grad_ln_out // [B, K, HIDDEN_DIM] (+= chained with K-loop)
|
||||
) {
|
||||
int bi = blockIdx.x;
|
||||
int tid = threadIdx.x;
|
||||
if (bi >= n_batch || tid >= ATTN_BLOCK) return;
|
||||
|
||||
extern __shared__ float smem[];
|
||||
float* s_dattn = smem; // [K]
|
||||
float* s_dscores = smem + k_seq; // [K]
|
||||
float* s_red = smem + 2 * k_seq; // [BLOCK]
|
||||
|
||||
__shared__ float s_Q[ATTN_HIDDEN_DIM];
|
||||
__shared__ float s_attn[ATTN_MAX_K];
|
||||
__shared__ float s_grad_ctx[ATTN_HIDDEN_DIM];
|
||||
__shared__ float s_sum_attn_dattn;
|
||||
|
||||
if (tid < ATTN_HIDDEN_DIM) s_Q[tid] = Q[tid];
|
||||
__syncthreads();
|
||||
|
||||
const float* ln_b = ln_out + (long long)bi * k_seq * ATTN_HIDDEN_DIM;
|
||||
float* grad_ln_b = grad_ln_out + (long long)bi * k_seq * ATTN_HIDDEN_DIM;
|
||||
// Cache attn + grad_context for this block's batch.
|
||||
for (int k = tid; k < k_seq; k += ATTN_BLOCK) {
|
||||
s_attn[k] = attn_weights[(long long)bi * k_seq + k];
|
||||
}
|
||||
if (tid < ATTN_HIDDEN_DIM) {
|
||||
s_grad_ctx[tid] = grad_context[(long long)bi * ATTN_HIDDEN_DIM + tid];
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// Pass 1: d_attn[k] = sum_h grad_context[h] * values[b, k, h].
|
||||
for (int k = 0; k < k_seq; ++k) {
|
||||
const float v = (tid < ATTN_HIDDEN_DIM)
|
||||
? s_grad_ctx[tid] * ln_b[k * ATTN_HIDDEN_DIM + tid] : 0.0f;
|
||||
s_red[tid] = v;
|
||||
__syncthreads();
|
||||
for (int s = ATTN_BLOCK / 2; s > 0; s >>= 1) {
|
||||
if (tid < s) s_red[tid] += s_red[tid + s];
|
||||
__syncthreads();
|
||||
}
|
||||
if (tid == 0) s_dattn[k] = s_red[0];
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
// Pass 2: sum_{kp} attn[kp] * d_attn[kp] (softmax Jacobian centring).
|
||||
float my_sum = 0.0f;
|
||||
for (int k = tid; k < k_seq; k += ATTN_BLOCK) {
|
||||
my_sum += s_attn[k] * s_dattn[k];
|
||||
}
|
||||
s_red[tid] = my_sum;
|
||||
__syncthreads();
|
||||
for (int s = ATTN_BLOCK / 2; s > 0; s >>= 1) {
|
||||
if (tid < s) s_red[tid] += s_red[tid + s];
|
||||
__syncthreads();
|
||||
}
|
||||
if (tid == 0) s_sum_attn_dattn = s_red[0];
|
||||
__syncthreads();
|
||||
|
||||
// Pass 3: d_scores[k] = attn[k] * (d_attn[k] - s_sum_attn_dattn).
|
||||
for (int k = tid; k < k_seq; k += ATTN_BLOCK) {
|
||||
s_dscores[k] = s_attn[k] * (s_dattn[k] - s_sum_attn_dattn);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// Pass 4: grad_Q_scratch[bi, h] += sum_k d_scores[k] * ln_out[b, k, h].
|
||||
// d_ln_out[b, k, h] += grad_context[h] * attn[k] + d_scores[k] * Q[h].
|
||||
// Thread h owns column h. Loops over k. grad_ln_out += chains the
|
||||
// attn-path gradient on top of the K-loop's contribution.
|
||||
if (tid < ATTN_HIDDEN_DIM) {
|
||||
float dq_local = 0.0f;
|
||||
for (int k = 0; k < k_seq; ++k) {
|
||||
const float v = ln_b[k * ATTN_HIDDEN_DIM + tid];
|
||||
dq_local += s_dscores[k] * v;
|
||||
grad_ln_b[k * ATTN_HIDDEN_DIM + tid] +=
|
||||
s_grad_ctx[tid] * s_attn[k] + s_dscores[k] * s_Q[tid];
|
||||
}
|
||||
grad_Q_scratch[(long long)bi * ATTN_HIDDEN_DIM + tid] += dq_local;
|
||||
}
|
||||
}
|
||||
@@ -1,51 +0,0 @@
|
||||
// horizon_mean_collapse.cu — Collapse per-horizon context to a single
|
||||
// vector by averaging across the horizon axis.
|
||||
//
|
||||
// Forward:
|
||||
// out[b, d] = (1 / N_H) · Σ_h in[b, h, d]
|
||||
//
|
||||
// Backward (broadcast — each horizon receives the same per-d gradient):
|
||||
// d_in[b, h, d] = (1 / N_H) · d_out[b, d]
|
||||
//
|
||||
// Grid: (B, 1, 1). Block: (HIDDEN_DIM, 1, 1).
|
||||
// Each thread tid owns d = tid. No reductions, no barriers.
|
||||
|
||||
#define HMC_HIDDEN_DIM 128
|
||||
#define HMC_BLOCK 128
|
||||
#define HMC_N_HORIZONS 5
|
||||
|
||||
extern "C" __global__ void horizon_mean_collapse_fwd(
|
||||
const float* __restrict__ ctx_per_h, // [B, N_H, H]
|
||||
int n_batch,
|
||||
float* __restrict__ ctx_mean // [B, H]
|
||||
) {
|
||||
int b = blockIdx.x;
|
||||
int tid = threadIdx.x;
|
||||
if (b >= n_batch || tid >= HMC_HIDDEN_DIM) return;
|
||||
|
||||
float acc = 0.0f;
|
||||
#pragma unroll
|
||||
for (int h = 0; h < HMC_N_HORIZONS; ++h) {
|
||||
acc += ctx_per_h[(long long)b * HMC_N_HORIZONS * HMC_HIDDEN_DIM
|
||||
+ (long long)h * HMC_HIDDEN_DIM + tid];
|
||||
}
|
||||
ctx_mean[(long long)b * HMC_HIDDEN_DIM + tid] = acc * (1.0f / (float)HMC_N_HORIZONS);
|
||||
}
|
||||
|
||||
extern "C" __global__ void horizon_mean_collapse_bwd(
|
||||
const float* __restrict__ grad_ctx_mean, // [B, H]
|
||||
int n_batch,
|
||||
float* __restrict__ grad_ctx_per_h // [B, N_H, H] += broadcast / N_H
|
||||
) {
|
||||
int b = blockIdx.x;
|
||||
int tid = threadIdx.x;
|
||||
if (b >= n_batch || tid >= HMC_HIDDEN_DIM) return;
|
||||
|
||||
const float g = grad_ctx_mean[(long long)b * HMC_HIDDEN_DIM + tid]
|
||||
* (1.0f / (float)HMC_N_HORIZONS);
|
||||
#pragma unroll
|
||||
for (int h = 0; h < HMC_N_HORIZONS; ++h) {
|
||||
grad_ctx_per_h[(long long)b * HMC_N_HORIZONS * HMC_HIDDEN_DIM
|
||||
+ (long long)h * HMC_HIDDEN_DIM + tid] += g;
|
||||
}
|
||||
}
|
||||
@@ -1,322 +0,0 @@
|
||||
// horizon_token_attention_pool.cu — Horizon-token attention pool (v2 axis C).
|
||||
//
|
||||
// Replaces the falsified per-horizon Q_h (C21/C25) with a single
|
||||
// shared query Q over an EXTENDED key sequence:
|
||||
// K' = [horizon_tokens; LN_b_out] shape [B, N_H + K, H]
|
||||
//
|
||||
// One attention pass; per-horizon outputs derived by mixing each
|
||||
// horizon's token contribution with a shared time-context aggregate.
|
||||
//
|
||||
// FORWARD MATH (per batch b; b indexing suppressed):
|
||||
//
|
||||
// For i ∈ [0, N_H):
|
||||
// score[i] = Σ_d Q[d] · horizon_tokens[i, d]
|
||||
// For i ∈ [N_H, N_H + K):
|
||||
// k = i − N_H
|
||||
// score[i] = Σ_d Q[d] · ln_out[k, d]
|
||||
//
|
||||
// attn[i] = softmax_i(score) # over N_H + K positions
|
||||
// S[d] = Σ_k attn[N_H + k] · ln_out[k, d] # shared time-context per d
|
||||
// ctx[h, d] = attn[h] · horizon_tokens[h, d] + S[d] # per-horizon output
|
||||
//
|
||||
// BACKWARD MATH (chain rule):
|
||||
//
|
||||
// d_S[d] = Σ_h grad_ctx[h, d]
|
||||
// d_attn[h] = Σ_d grad_ctx[h, d] · horizon_tokens[h, d] (h-token term)
|
||||
// d_attn[N_H + k] = Σ_d d_S[d] · ln_out[k, d] (time term)
|
||||
// d_score (softmax bwd): d_score[i] = attn[i] · (d_attn[i] − Σ_j attn[j]·d_attn[j])
|
||||
//
|
||||
// d_horizon_tokens[h, d] = d_score[h] · Q[d]
|
||||
// + grad_ctx[h, d] · attn[h]
|
||||
// d_Q[d] = Σ_h d_score[h] · horizon_tokens[h, d]
|
||||
// + Σ_k d_score[N_H + k] · ln_out[k, d]
|
||||
// d_ln_out[k, d] = d_score[N_H + k] · Q[d]
|
||||
// + attn[N_H + k] · d_S[d]
|
||||
//
|
||||
// PERFORMANCE (per feedback_nvidia_grade_perf_for_kernels):
|
||||
// - Warp-shuffle reduction (`__shfl_xor_sync`) for every per-d dot
|
||||
// product and per-attention sum; NO block tree-reduce.
|
||||
// - 4 warps × 32 lanes; cross-warp reduce uses one `__syncthreads`.
|
||||
// - All inactive-lane shuffles use ternary fallback to 0 (never a
|
||||
// divergent `if (tid < N) shuffle`).
|
||||
// - Block-per-batch; horizon loop inside the block keeps grad writes
|
||||
// race-free (each batch's ln_out slice touched by one block in bwd).
|
||||
// - Shared mem layout pre-computed at launch from k_seq.
|
||||
|
||||
#define HTAP_HIDDEN_DIM 128
|
||||
#define HTAP_BLOCK 128
|
||||
#define HTAP_N_HORIZONS 5
|
||||
#define HTAP_N_WARPS (HTAP_BLOCK / 32) // == 4
|
||||
#define HTAP_MAX_KSEQ 128 // safety cap; trainer uses K ≤ 64
|
||||
#define HTAP_MAX_EXT (HTAP_N_HORIZONS + HTAP_MAX_KSEQ)
|
||||
|
||||
__device__ __forceinline__ float warp_reduce_sum_local(float v) {
|
||||
#pragma unroll
|
||||
for (int s = 16; s > 0; s >>= 1) {
|
||||
v += __shfl_xor_sync(0xffffffff, v, s);
|
||||
}
|
||||
return v;
|
||||
}
|
||||
|
||||
__device__ __forceinline__ float warp_reduce_max_local(float v) {
|
||||
#pragma unroll
|
||||
for (int s = 16; s > 0; s >>= 1) {
|
||||
v = fmaxf(v, __shfl_xor_sync(0xffffffff, v, s));
|
||||
}
|
||||
return v;
|
||||
}
|
||||
|
||||
// Block-wide reduce-sum using warp-shuffle + 1 syncthreads.
|
||||
// `s_warp` is caller-provided shared scratch of size HTAP_N_WARPS.
|
||||
__device__ __forceinline__ float block_reduce_sum(float v, float* s_warp, int tid) {
|
||||
v = warp_reduce_sum_local(v);
|
||||
int lane = tid & 31, warp = tid >> 5;
|
||||
if (lane == 0) s_warp[warp] = v;
|
||||
__syncthreads();
|
||||
float w = (tid < HTAP_N_WARPS) ? s_warp[tid] : 0.0f;
|
||||
if (tid < 32) {
|
||||
w = warp_reduce_sum_local(w);
|
||||
if (tid == 0) s_warp[0] = w;
|
||||
}
|
||||
__syncthreads();
|
||||
return s_warp[0];
|
||||
}
|
||||
|
||||
__device__ __forceinline__ float block_reduce_max(float v, float* s_warp, int tid) {
|
||||
v = warp_reduce_max_local(v);
|
||||
int lane = tid & 31, warp = tid >> 5;
|
||||
if (lane == 0) s_warp[warp] = v;
|
||||
__syncthreads();
|
||||
float w = (tid < HTAP_N_WARPS) ? s_warp[tid] : -INFINITY;
|
||||
if (tid < 32) {
|
||||
w = warp_reduce_max_local(w);
|
||||
if (tid == 0) s_warp[0] = w;
|
||||
}
|
||||
__syncthreads();
|
||||
return s_warp[0];
|
||||
}
|
||||
|
||||
extern "C" __global__ void horizon_token_attention_pool_fwd(
|
||||
const float* __restrict__ horizon_tokens, // [N_H, H]
|
||||
const float* __restrict__ Q, // [H]
|
||||
const float* __restrict__ ln_out, // [B, K, H]
|
||||
int n_batch,
|
||||
int k_seq,
|
||||
float* __restrict__ ctx_out, // [B, N_H, H]
|
||||
float* __restrict__ attn_out // [B, N_H + K]
|
||||
) {
|
||||
int b = blockIdx.x;
|
||||
int tid = threadIdx.x;
|
||||
if (b >= n_batch || tid >= HTAP_BLOCK) return;
|
||||
|
||||
extern __shared__ float smem[];
|
||||
// Layout: s_attn[N_H + K] + s_warp[N_WARPS] + s_max + s_sum
|
||||
float* s_attn = smem; // [N_H + k_seq]
|
||||
float* s_warp = s_attn + (HTAP_N_HORIZONS + k_seq); // [N_WARPS]
|
||||
// Used scalars stashed by the reduce helpers; cheap to materialise.
|
||||
|
||||
// Cache Q[tid] and horizon_tokens[h, tid] in registers / shared.
|
||||
__shared__ float s_horizon_tokens[HTAP_N_HORIZONS * HTAP_HIDDEN_DIM];
|
||||
if (tid < HTAP_HIDDEN_DIM) {
|
||||
#pragma unroll
|
||||
for (int h = 0; h < HTAP_N_HORIZONS; ++h) {
|
||||
s_horizon_tokens[h * HTAP_HIDDEN_DIM + tid] =
|
||||
horizon_tokens[h * HTAP_HIDDEN_DIM + tid];
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
const float q_d = (tid < HTAP_HIDDEN_DIM) ? Q[tid] : 0.0f;
|
||||
const float* ln_b = ln_out + (long long)b * k_seq * HTAP_HIDDEN_DIM;
|
||||
|
||||
// ── Phase 1: scores[i] = Σ_d q[d] · ext[i, d]. Serial outer i, inner
|
||||
// HIDDEN_DIM reduced via block_reduce_sum (warp-shuffle).
|
||||
const int n_ext = HTAP_N_HORIZONS + k_seq;
|
||||
for (int i = 0; i < n_ext; ++i) {
|
||||
float key_d;
|
||||
if (i < HTAP_N_HORIZONS) {
|
||||
key_d = (tid < HTAP_HIDDEN_DIM)
|
||||
? s_horizon_tokens[i * HTAP_HIDDEN_DIM + tid] : 0.0f;
|
||||
} else {
|
||||
int k = i - HTAP_N_HORIZONS;
|
||||
key_d = (tid < HTAP_HIDDEN_DIM)
|
||||
? ln_b[k * HTAP_HIDDEN_DIM + tid] : 0.0f;
|
||||
}
|
||||
float v = q_d * key_d;
|
||||
float r = block_reduce_sum(v, s_warp, tid);
|
||||
if (tid == 0) s_attn[i] = r;
|
||||
// s_warp is reused by the next iteration's block_reduce_sum; safe.
|
||||
}
|
||||
|
||||
// ── Phase 2: softmax over [0, n_ext).
|
||||
// 2a: max
|
||||
float my_max = -INFINITY;
|
||||
for (int i = tid; i < n_ext; i += HTAP_BLOCK) {
|
||||
my_max = fmaxf(my_max, s_attn[i]);
|
||||
}
|
||||
const float s_max = block_reduce_max(my_max, s_warp, tid);
|
||||
|
||||
// 2b: exp + sum
|
||||
float my_sum = 0.0f;
|
||||
for (int i = tid; i < n_ext; i += HTAP_BLOCK) {
|
||||
const float e = expf(s_attn[i] - s_max);
|
||||
s_attn[i] = e;
|
||||
my_sum += e;
|
||||
}
|
||||
const float s_sum = block_reduce_sum(my_sum, s_warp, tid);
|
||||
|
||||
// 2c: normalise + persist attn_out
|
||||
const float inv_sum = 1.0f / s_sum;
|
||||
for (int i = tid; i < n_ext; i += HTAP_BLOCK) {
|
||||
const float a = s_attn[i] * inv_sum;
|
||||
s_attn[i] = a;
|
||||
attn_out[(long long)b * n_ext + i] = a;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// ── Phase 3: ctx[h, d] = attn[h] · h_token[h, d] + S[d]
|
||||
// where S[d] = Σ_k attn[N_H + k] · ln_out[k, d].
|
||||
if (tid < HTAP_HIDDEN_DIM) {
|
||||
float S_d = 0.0f;
|
||||
for (int k = 0; k < k_seq; ++k) {
|
||||
S_d += s_attn[HTAP_N_HORIZONS + k] * ln_b[k * HTAP_HIDDEN_DIM + tid];
|
||||
}
|
||||
#pragma unroll
|
||||
for (int h = 0; h < HTAP_N_HORIZONS; ++h) {
|
||||
const float token_term = s_attn[h] * s_horizon_tokens[h * HTAP_HIDDEN_DIM + tid];
|
||||
ctx_out[(long long)b * HTAP_N_HORIZONS * HTAP_HIDDEN_DIM
|
||||
+ (long long)h * HTAP_HIDDEN_DIM + tid] = token_term + S_d;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
extern "C" __global__ void horizon_token_attention_pool_bwd(
|
||||
const float* __restrict__ horizon_tokens, // [N_H, H]
|
||||
const float* __restrict__ Q, // [H]
|
||||
const float* __restrict__ ln_out, // [B, K, H]
|
||||
const float* __restrict__ attn, // [B, N_H + K] (saved fwd)
|
||||
const float* __restrict__ grad_ctx, // [B, N_H, H] (upstream)
|
||||
int n_batch,
|
||||
int k_seq,
|
||||
float* __restrict__ grad_horizon_tokens_scratch, // [B, N_H, H] += (reduce_axis0 → [N_H, H])
|
||||
float* __restrict__ grad_Q_scratch, // [B, H] += (reduce_axis0 → [H])
|
||||
float* __restrict__ grad_ln_out // [B, K, H] += (per-batch, no race)
|
||||
) {
|
||||
int b = blockIdx.x;
|
||||
int tid = threadIdx.x;
|
||||
if (b >= n_batch || tid >= HTAP_BLOCK) return;
|
||||
|
||||
extern __shared__ float smem_bwd[];
|
||||
// Layout: s_d_attn[N_H + K] + s_warp[N_WARPS] + s_d_S[HIDDEN_DIM]
|
||||
float* s_d_attn = smem_bwd; // [N_H + k]
|
||||
float* s_warp = s_d_attn + (HTAP_N_HORIZONS + k_seq); // [N_WARPS]
|
||||
float* s_d_S = s_warp + HTAP_N_WARPS; // [HIDDEN_DIM]
|
||||
|
||||
__shared__ float s_horizon_tokens[HTAP_N_HORIZONS * HTAP_HIDDEN_DIM];
|
||||
__shared__ float s_attn[HTAP_N_HORIZONS + HTAP_MAX_KSEQ];
|
||||
|
||||
const int n_ext = HTAP_N_HORIZONS + k_seq;
|
||||
|
||||
// Load horizon_tokens + attn into shared.
|
||||
if (tid < HTAP_HIDDEN_DIM) {
|
||||
#pragma unroll
|
||||
for (int h = 0; h < HTAP_N_HORIZONS; ++h) {
|
||||
s_horizon_tokens[h * HTAP_HIDDEN_DIM + tid] =
|
||||
horizon_tokens[h * HTAP_HIDDEN_DIM + tid];
|
||||
}
|
||||
}
|
||||
for (int i = tid; i < n_ext; i += HTAP_BLOCK) {
|
||||
s_attn[i] = attn[(long long)b * n_ext + i];
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
const float q_d = (tid < HTAP_HIDDEN_DIM) ? Q[tid] : 0.0f;
|
||||
const float* ln_b = ln_out + (long long)b * k_seq * HTAP_HIDDEN_DIM;
|
||||
float* grad_ln_b = grad_ln_out + (long long)b * k_seq * HTAP_HIDDEN_DIM;
|
||||
|
||||
// ── d_S[d] = Σ_h grad_ctx[h, d]. Each thread owns d = tid (if < H).
|
||||
float d_S = 0.0f;
|
||||
if (tid < HTAP_HIDDEN_DIM) {
|
||||
#pragma unroll
|
||||
for (int h = 0; h < HTAP_N_HORIZONS; ++h) {
|
||||
d_S += grad_ctx[(long long)b * HTAP_N_HORIZONS * HTAP_HIDDEN_DIM
|
||||
+ (long long)h * HTAP_HIDDEN_DIM + tid];
|
||||
}
|
||||
s_d_S[tid] = d_S;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// ── d_attn[i] for both ranges, stored in s_d_attn.
|
||||
// Token range i ∈ [0, N_H): d_attn[h] = Σ_d grad_ctx[h, d] · h_token[h, d]
|
||||
// Time range i ∈ [N_H, n_ext): d_attn[N_H + k] = Σ_d d_S[d] · ln_out[k, d]
|
||||
for (int h = 0; h < HTAP_N_HORIZONS; ++h) {
|
||||
float v = 0.0f;
|
||||
if (tid < HTAP_HIDDEN_DIM) {
|
||||
v = grad_ctx[(long long)b * HTAP_N_HORIZONS * HTAP_HIDDEN_DIM
|
||||
+ (long long)h * HTAP_HIDDEN_DIM + tid]
|
||||
* s_horizon_tokens[h * HTAP_HIDDEN_DIM + tid];
|
||||
}
|
||||
float r = block_reduce_sum(v, s_warp, tid);
|
||||
if (tid == 0) s_d_attn[h] = r;
|
||||
}
|
||||
for (int k = 0; k < k_seq; ++k) {
|
||||
float v = 0.0f;
|
||||
if (tid < HTAP_HIDDEN_DIM) {
|
||||
v = s_d_S[tid] * ln_b[k * HTAP_HIDDEN_DIM + tid];
|
||||
}
|
||||
float r = block_reduce_sum(v, s_warp, tid);
|
||||
if (tid == 0) s_d_attn[HTAP_N_HORIZONS + k] = r;
|
||||
}
|
||||
|
||||
// ── Softmax backward: d_score[i] = attn[i] · (d_attn[i] − Σ_j attn[j] · d_attn[j])
|
||||
float my_sum = 0.0f;
|
||||
for (int i = tid; i < n_ext; i += HTAP_BLOCK) {
|
||||
my_sum += s_attn[i] * s_d_attn[i];
|
||||
}
|
||||
const float dot_attn_dattn = block_reduce_sum(my_sum, s_warp, tid);
|
||||
|
||||
// Reuse s_d_attn slot to hold d_score; s_attn still holds attn.
|
||||
for (int i = tid; i < n_ext; i += HTAP_BLOCK) {
|
||||
s_d_attn[i] = s_attn[i] * (s_d_attn[i] - dot_attn_dattn);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// ── d_horizon_tokens[h, d] = d_score[h] · Q[d] + grad_ctx[h, d] · attn[h]
|
||||
// grad_horizon_tokens_scratch is per-batch; reduce_axis0 collapses
|
||||
// over B downstream.
|
||||
if (tid < HTAP_HIDDEN_DIM) {
|
||||
#pragma unroll
|
||||
for (int h = 0; h < HTAP_N_HORIZONS; ++h) {
|
||||
const float dscore_h = s_d_attn[h];
|
||||
const float gc_hd = grad_ctx[(long long)b * HTAP_N_HORIZONS * HTAP_HIDDEN_DIM
|
||||
+ (long long)h * HTAP_HIDDEN_DIM + tid];
|
||||
const float attn_h = s_attn[h];
|
||||
const float dq = dscore_h * q_d + gc_hd * attn_h;
|
||||
grad_horizon_tokens_scratch[(long long)b * HTAP_N_HORIZONS * HTAP_HIDDEN_DIM
|
||||
+ (long long)h * HTAP_HIDDEN_DIM + tid] += dq;
|
||||
}
|
||||
}
|
||||
|
||||
// ── d_Q[d] = Σ_h d_score[h] · h_token[h, d] + Σ_k d_score[N_H+k] · ln_out[k, d]
|
||||
if (tid < HTAP_HIDDEN_DIM) {
|
||||
float dq_d = 0.0f;
|
||||
#pragma unroll
|
||||
for (int h = 0; h < HTAP_N_HORIZONS; ++h) {
|
||||
dq_d += s_d_attn[h] * s_horizon_tokens[h * HTAP_HIDDEN_DIM + tid];
|
||||
}
|
||||
for (int k = 0; k < k_seq; ++k) {
|
||||
dq_d += s_d_attn[HTAP_N_HORIZONS + k] * ln_b[k * HTAP_HIDDEN_DIM + tid];
|
||||
}
|
||||
grad_Q_scratch[(long long)b * HTAP_HIDDEN_DIM + tid] += dq_d;
|
||||
}
|
||||
|
||||
// ── d_ln_out[k, d] = d_score[N_H+k] · Q[d] + attn[N_H+k] · d_S[d]
|
||||
// Per-batch slice → no cross-block race.
|
||||
if (tid < HTAP_HIDDEN_DIM) {
|
||||
for (int k = 0; k < k_seq; ++k) {
|
||||
const float dscore_k = s_d_attn[HTAP_N_HORIZONS + k];
|
||||
const float attn_k = s_attn[HTAP_N_HORIZONS + k];
|
||||
grad_ln_b[k * HTAP_HIDDEN_DIM + tid] += dscore_k * q_d + attn_k * s_d_S[tid];
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,64 +0,0 @@
|
||||
// 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;
|
||||
}
|
||||
@@ -1,265 +0,0 @@
|
||||
// inverted_attention_pool.cu — iTransformer-style inverted attention.
|
||||
//
|
||||
// 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):
|
||||
//
|
||||
// X_inv[h, k] = ln_out[k, h] # [H, K]
|
||||
// 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]
|
||||
//
|
||||
// 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).
|
||||
//
|
||||
// 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,
|
||||
int k_seq,
|
||||
float inv_scale, // 1 / sqrt(K)
|
||||
float* __restrict__ pooled_out, // [B, H]
|
||||
float* __restrict__ attn_out // [B, H, H]
|
||||
) {
|
||||
int b = blockIdx.x;
|
||||
int tid = threadIdx.x;
|
||||
if (b >= n_batch || tid >= IAP_BLOCK) return;
|
||||
|
||||
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) {
|
||||
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: 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;
|
||||
float dot = 0.0f;
|
||||
#pragma unroll 8
|
||||
for (int k = 0; k < k_seq; ++k) {
|
||||
dot += row_h[k] * row_j[k];
|
||||
}
|
||||
const float s = dot * inv_scale;
|
||||
if (s > my_max) my_max = s;
|
||||
}
|
||||
|
||||
// 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 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) {
|
||||
dot += row_h[k] * row_j[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;
|
||||
}
|
||||
}
|
||||
|
||||
extern "C" __global__ void inverted_attention_pool_bwd(
|
||||
const float* __restrict__ ln_out, // [B, K, H]
|
||||
const float* __restrict__ attn, // [B, H, H]
|
||||
const float* __restrict__ grad_pooled, // [B, H]
|
||||
int n_batch,
|
||||
int k_seq,
|
||||
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 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) {
|
||||
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();
|
||||
|
||||
// Phase 1: compute d_scores[my_h, j] for all j; write to DRAM scratch.
|
||||
//
|
||||
// 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[l]` (= pooled[my_h]).
|
||||
float dot = 0.0f;
|
||||
for (int j = 0; j < IAP_HIDDEN_DIM; ++j) {
|
||||
dot += attn[attn_row_base + j] * mean_k[j];
|
||||
}
|
||||
|
||||
// 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) {
|
||||
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.
|
||||
//
|
||||
// Per-(my_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]
|
||||
//
|
||||
// KEY OPTIMIZATION (canonical 2026-05-18 perf rewrite):
|
||||
// The d_scores reads do NOT depend on k → hoist them out of the
|
||||
// K loop. We loop J-outer with two register accumulators per
|
||||
// thread (q_k[K_MAX], k_k[K_MAX]). For each j we do exactly TWO
|
||||
// DRAM reads of d_scores (row + column) instead of `2 * K` =
|
||||
// 2 × 32 = 64 reads inside a tight loop. Net: 32× fewer DRAM
|
||||
// reads for the d_scores tensor.
|
||||
//
|
||||
// For K up to 64 (production max), the per-thread float arrays
|
||||
// stay in registers (NVCC --maxrregcount permits it on sm_86+).
|
||||
if (tid < IAP_HIDDEN_DIM) {
|
||||
const int my_h = tid;
|
||||
|
||||
// Value-path scalar: read `attn[:, my_h]` (strided!) and dp from smem.
|
||||
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];
|
||||
V_my_h += dp_shared[hp] * a_hp_myh;
|
||||
}
|
||||
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;
|
||||
|
||||
// Per-thread register accumulators for q-path and k-path,
|
||||
// one per K position. IAP_MAX_KSEQ = 128 floats = 512 B per
|
||||
// thread total = 64 KB across the 128-thread block — fits
|
||||
// within the register file on sm_86+ with --maxrregcount=128.
|
||||
float q_arr[IAP_MAX_KSEQ];
|
||||
float k_arr[IAP_MAX_KSEQ];
|
||||
#pragma unroll
|
||||
for (int k = 0; k < IAP_MAX_KSEQ; ++k) {
|
||||
q_arr[k] = 0.0f;
|
||||
k_arr[k] = 0.0f;
|
||||
}
|
||||
|
||||
// J-outer: each j read of d_scores happens exactly ONCE.
|
||||
// The inner K-loop walks x_inv[j, k] which is contiguous in smem.
|
||||
for (int j = 0; j < IAP_HIDDEN_DIM; ++j) {
|
||||
const float ds_myh_j = d_scores_scratch[my_row_base + j];
|
||||
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];
|
||||
const float* xj = x_inv + (long long)j * k_seq;
|
||||
#pragma unroll 8
|
||||
for (int k = 0; k < k_seq; ++k) {
|
||||
const float xj_k = xj[k];
|
||||
q_arr[k] += ds_myh_j * xj_k;
|
||||
k_arr[k] += ds_j_myh * xj_k;
|
||||
}
|
||||
}
|
||||
|
||||
// Final write-out per k.
|
||||
for (int k = 0; k < k_seq; ++k) {
|
||||
d_ln_b[k * IAP_HIDDEN_DIM + my_h] +=
|
||||
V_my_h + inv_scale * (q_arr[k] + k_arr[k]);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,300 +0,0 @@
|
||||
// regime_moe_gate.cu — Regime-aware top-1 Mixture-of-Experts gate (v2 axis D).
|
||||
//
|
||||
// Routes the fused per-horizon context through one of N_EXPERTS expert
|
||||
// linear layers based on a regime-feature gate. Top-1 routing per
|
||||
// Switch Transformer (Fedus et al. 2021) with straight-through grad on
|
||||
// the gate logits and an auxiliary load-balancing loss.
|
||||
//
|
||||
// FORWARD MATH (per batch b, suppressing b):
|
||||
//
|
||||
// gate_logits[e] = Σ_r W_gate[e, r] · R[r] # [N_EXPERTS]
|
||||
// gate_probs[e] = softmax_e(gate_logits[e]) # [N_EXPERTS]
|
||||
// top_e = argmax_e(gate_logits[e]) # scalar
|
||||
// expert_out[e](h, d_out) = Σ_d_in W_e[e, d_out, d_in] · fused_ctx[h, d_in]
|
||||
// + b_e[e, d_out] # not all e computed; only top_e
|
||||
// routed_ctx[h, d] = expert_out[top_e][h, d]
|
||||
//
|
||||
// BACKWARD (straight-through estimator on gate, full grad on the
|
||||
// selected expert):
|
||||
// d_W_e[top_e, d_out, d_in] += Σ_{b, h} fused_ctx[b, h, d_in] · d_routed[b, h, d_out]
|
||||
// d_b_e[top_e, d_out] += Σ_{b, h} d_routed[b, h, d_out]
|
||||
// d_fused_ctx[b, h, d] += Σ_{d_out} W_e[top_e, d_out, d] · d_routed[b, h, d_out]
|
||||
//
|
||||
// STE on gate_logits[b, e]: pass an upstream "router loss signal"
|
||||
// proportional to the expert output norm — concretely, route the
|
||||
// gradient ∂L/∂routed[b, h, ·] · expert_out[e][h, ·] for e == top_e[b]
|
||||
// into gate_logits via the softmax Jacobian. For inactive experts,
|
||||
// STE assumes the same loss signal would have been produced (zero
|
||||
// selection grad). Aux load-balance grad is added separately.
|
||||
//
|
||||
// AUXILIARY LOAD-BALANCING LOSS (computed by a separate small kernel):
|
||||
// frac[e] = mean_b({top_e[b] == e})
|
||||
// prob_mean[e] = mean_b(gate_probs[b, e])
|
||||
// aux_loss = N_EXPERTS · Σ_e frac[e] · prob_mean[e]
|
||||
// d_gate_logits aux contribution: standard softmax-frac-cross-entropy
|
||||
//
|
||||
// PERFORMANCE:
|
||||
// - One block per (b, h): forward and backward have block-per-(b, h)
|
||||
// parallelism for the per-horizon linear (matches V4 pattern).
|
||||
// - Warp-shuffle reduce for the d_in dot product (HIDDEN_DIM = 128 =
|
||||
// 4 warps). All sums are warp-shuffle, not block tree-reduce.
|
||||
// - Top-1 selection per batch is decided in a pre-pass: ONE thread
|
||||
// per block reads gate_logits[b, :] and picks the argmax, broadcasts
|
||||
// via shared mem.
|
||||
// - Per-expert weights live in [N_EXPERTS, H, H] = ~250 KB total; on
|
||||
// the hot path each block only fetches the selected expert's slice.
|
||||
|
||||
#define MOE_HIDDEN_DIM 128
|
||||
#define MOE_BLOCK 128
|
||||
#define MOE_N_HORIZONS 5
|
||||
#define MOE_N_EXPERTS 4
|
||||
#define MOE_N_WARPS (MOE_BLOCK / 32)
|
||||
|
||||
// ── Forward: gate + expert dispatch fused into one kernel.
|
||||
// Grid: (B, N_H, 1). Block: (HIDDEN_DIM, 1, 1).
|
||||
// Each block: (1) reads gate_logits, picks top_e for this batch,
|
||||
// (2) computes routed_ctx[b, h, d_out=tid].
|
||||
extern "C" __global__ void regime_moe_gate_fwd(
|
||||
const float* __restrict__ gate_logits, // [B, N_EXPERTS] pre-softmax (caller computes)
|
||||
const float* __restrict__ fused_ctx, // [B, N_H, H]
|
||||
const float* __restrict__ experts_W, // [N_EXPERTS, H, H]
|
||||
const float* __restrict__ experts_b, // [N_EXPERTS, H]
|
||||
int n_batch,
|
||||
float* __restrict__ routed_out, // [B, N_H, H]
|
||||
int* __restrict__ top_e_out // [B] top-1 expert per batch (saved)
|
||||
) {
|
||||
int b = blockIdx.x;
|
||||
int h = blockIdx.y;
|
||||
int tid = threadIdx.x;
|
||||
if (b >= n_batch || h >= MOE_N_HORIZONS || tid >= MOE_BLOCK) return;
|
||||
|
||||
// Pick top-1 expert by thread 0; broadcast to the block.
|
||||
__shared__ int s_top_e;
|
||||
if (tid == 0) {
|
||||
float best = -INFINITY;
|
||||
int best_e = 0;
|
||||
#pragma unroll
|
||||
for (int e = 0; e < MOE_N_EXPERTS; ++e) {
|
||||
const float v = gate_logits[(long long)b * MOE_N_EXPERTS + e];
|
||||
if (v > best) { best = v; best_e = e; }
|
||||
}
|
||||
s_top_e = best_e;
|
||||
// Only one block per batch needs to write top_e_out (the h == 0
|
||||
// block by convention) — but writing from each block is fine
|
||||
// since the value is identical.
|
||||
if (h == 0) top_e_out[b] = best_e;
|
||||
}
|
||||
__syncthreads();
|
||||
const int top_e = s_top_e;
|
||||
|
||||
// Cache fused_ctx[b, h, ·] in shared mem.
|
||||
__shared__ float s_in[MOE_HIDDEN_DIM];
|
||||
if (tid < MOE_HIDDEN_DIM) {
|
||||
s_in[tid] = fused_ctx[(long long)b * MOE_N_HORIZONS * MOE_HIDDEN_DIM
|
||||
+ (long long)h * MOE_HIDDEN_DIM + tid];
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// Each thread computes one output channel: routed[b, h, tid] =
|
||||
// Σ_d_in W_e[top_e, tid, d_in] · s_in[d_in] + b_e[top_e, tid].
|
||||
if (tid < MOE_HIDDEN_DIM) {
|
||||
const long long w_base = (long long)top_e * MOE_HIDDEN_DIM * MOE_HIDDEN_DIM
|
||||
+ (long long)tid * MOE_HIDDEN_DIM;
|
||||
float acc = experts_b[(long long)top_e * MOE_HIDDEN_DIM + tid];
|
||||
#pragma unroll 8
|
||||
for (int d_in = 0; d_in < MOE_HIDDEN_DIM; ++d_in) {
|
||||
acc += experts_W[w_base + d_in] * s_in[d_in];
|
||||
}
|
||||
routed_out[(long long)b * MOE_N_HORIZONS * MOE_HIDDEN_DIM
|
||||
+ (long long)h * MOE_HIDDEN_DIM + tid] = acc;
|
||||
}
|
||||
}
|
||||
|
||||
// ── Backward expert path only (gate STE handled in a separate kernel).
|
||||
// Same block layout as forward.
|
||||
// For the SELECTED expert top_e[b]:
|
||||
// d_W_scratch[b, h, e=top_e, d_out, d_in] = d_routed[b, h, d_out] · fused_ctx[b, h, d_in]
|
||||
// d_b_scratch[b, h, e=top_e, d_out] = d_routed[b, h, d_out]
|
||||
// d_fused_ctx[b, h, d_in] += Σ_d_out W_e[top_e, d_out, d_in] · d_routed[b, h, d_out]
|
||||
// Inactive experts: 0 contribution.
|
||||
//
|
||||
// Per-(b, h) scratches are massive (B × N_H × N_E × H × H). We instead
|
||||
// have CALLER allocate per-(b, h) scratches that are sparse-by-expert:
|
||||
// the kernel only writes to the slot corresponding to top_e for this
|
||||
// (b, h). Total scratch size: [B, N_H, N_E, H, H] = small at trainer
|
||||
// scale (B=1, N_H=5, N_E=4, H=128 → 2.5 MB).
|
||||
// d_W_scratch layout changed 2026-05-18 for perf: [B, N_H, H, H] (was
|
||||
// [B, N_H, N_E, H, H] — 4× smaller, eliminating the wasteful N_E axis
|
||||
// since each batch routes to exactly ONE expert in top-1 MoE). The
|
||||
// expert index `top_e[b]` is needed downstream when scattering the
|
||||
// (B, N_H)-reduced sum into the per-expert grad slot.
|
||||
extern "C" __global__ void regime_moe_gate_bwd(
|
||||
const float* __restrict__ fused_ctx, // [B, N_H, H]
|
||||
const float* __restrict__ experts_W, // [N_EXPERTS, H, H]
|
||||
const float* __restrict__ d_routed, // [B, N_H, H]
|
||||
const int* __restrict__ top_e, // [B]
|
||||
int n_batch,
|
||||
float* __restrict__ d_W_scratch, // [B, N_H, H, H] per-batch-horizon rank-1 contribution
|
||||
float* __restrict__ d_b_scratch, // [B, N_H, H] per-batch-horizon bias contribution
|
||||
float* __restrict__ d_fused_ctx // [B, N_H, H] += per-(b, h)
|
||||
) {
|
||||
int b = blockIdx.x;
|
||||
int h = blockIdx.y;
|
||||
int tid = threadIdx.x;
|
||||
if (b >= n_batch || h >= MOE_N_HORIZONS || tid >= MOE_BLOCK) return;
|
||||
|
||||
const int e = top_e[b];
|
||||
|
||||
__shared__ float s_in[MOE_HIDDEN_DIM];
|
||||
__shared__ float s_dout[MOE_HIDDEN_DIM];
|
||||
if (tid < MOE_HIDDEN_DIM) {
|
||||
s_in[tid] = fused_ctx[(long long)b * MOE_N_HORIZONS * MOE_HIDDEN_DIM
|
||||
+ (long long)h * MOE_HIDDEN_DIM + tid];
|
||||
s_dout[tid] = d_routed[(long long)b * MOE_N_HORIZONS * MOE_HIDDEN_DIM
|
||||
+ (long long)h * MOE_HIDDEN_DIM + tid];
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// d_W_scratch[b, h, tid (= d_out), d_in] = d_routed[b, h, d_out] · fused_ctx[b, h, d_in]
|
||||
// Each thread d_out writes one ROW of [H, H] (length H = MOE_HIDDEN_DIM).
|
||||
// d_b_scratch[b, h, d_out] = d_routed[b, h, d_out]
|
||||
if (tid < MOE_HIDDEN_DIM) {
|
||||
const long long w_base = (((long long)b * MOE_N_HORIZONS + h)
|
||||
* MOE_HIDDEN_DIM + tid) * MOE_HIDDEN_DIM;
|
||||
const float g = s_dout[tid];
|
||||
#pragma unroll 8
|
||||
for (int d_in = 0; d_in < MOE_HIDDEN_DIM; ++d_in) {
|
||||
d_W_scratch[w_base + d_in] = g * s_in[d_in];
|
||||
}
|
||||
d_b_scratch[((long long)b * MOE_N_HORIZONS + h)
|
||||
* MOE_HIDDEN_DIM + tid] = g;
|
||||
}
|
||||
|
||||
// d_fused_ctx[b, h, d_in] += Σ_d_out W_e[e, d_out, d_in] · d_routed[b, h, d_out]
|
||||
// Each thread owns d_in = tid. Reduce over d_out via warp-shuffle.
|
||||
__syncthreads();
|
||||
if (tid < MOE_HIDDEN_DIM) {
|
||||
const int d_in = tid;
|
||||
float acc = 0.0f;
|
||||
const long long w_col_base = (long long)e * MOE_HIDDEN_DIM * MOE_HIDDEN_DIM;
|
||||
// W_e[e, d_out, d_in]: stride = MOE_HIDDEN_DIM between rows.
|
||||
for (int d_out = 0; d_out < MOE_HIDDEN_DIM; ++d_out) {
|
||||
acc += experts_W[w_col_base + (long long)d_out * MOE_HIDDEN_DIM + d_in]
|
||||
* s_dout[d_out];
|
||||
}
|
||||
d_fused_ctx[(long long)b * MOE_N_HORIZONS * MOE_HIDDEN_DIM
|
||||
+ (long long)h * MOE_HIDDEN_DIM + d_in] = acc;
|
||||
}
|
||||
}
|
||||
|
||||
// ── Scatter per-batch grads to the correct expert slot.
|
||||
// After `regime_moe_gate_bwd` writes the rank-1 contributions into
|
||||
// `d_W_scratch [B, N_H, H, H]` and `d_b_scratch [B, N_H, H]`, this
|
||||
// kernel scatters them into `grad_experts_W [N_E, H, H]` /
|
||||
// `grad_experts_b [N_E, H]` based on `top_e[b]`. Each batch's
|
||||
// contribution lands in slot `top_e[b]` of the per-expert grad.
|
||||
//
|
||||
// Grid (N_E, H, H/32 + 1), block (32). Each thread iterates over
|
||||
// (b, h) pairs and accumulates into its assigned (e, d_out, d_in)
|
||||
// cell only if top_e[b] matches. Per cell: 5 contributions max
|
||||
// (B=1 × N_H=5). Coalesced reads of d_W_scratch row.
|
||||
extern "C" __global__ void regime_moe_gate_scatter(
|
||||
const float* __restrict__ d_W_scratch, // [B, N_H, H, H]
|
||||
const float* __restrict__ d_b_scratch, // [B, N_H, H]
|
||||
const int* __restrict__ top_e, // [B]
|
||||
int n_batch,
|
||||
float* __restrict__ grad_experts_W, // [N_E, H, H] OVERWRITE
|
||||
float* __restrict__ grad_experts_b // [N_E, H] OVERWRITE
|
||||
) {
|
||||
int e = blockIdx.x;
|
||||
int d_out = blockIdx.y;
|
||||
int d_in_chunk = blockIdx.z;
|
||||
int lane = threadIdx.x;
|
||||
if (e >= MOE_N_EXPERTS || d_out >= MOE_HIDDEN_DIM || lane >= 32) return;
|
||||
|
||||
const int d_in_start = d_in_chunk * 32;
|
||||
const int d_in = d_in_start + lane;
|
||||
if (d_in >= MOE_HIDDEN_DIM) return;
|
||||
|
||||
// Accumulate W[e, d_out, d_in] across batches where top_e[b] == e.
|
||||
float acc_w = 0.0f;
|
||||
float acc_b = (lane == 0 && d_out == 0) ? 0.0f : 0.0f; // bias only via warp 0
|
||||
for (int b = 0; b < n_batch; ++b) {
|
||||
if (top_e[b] != e) continue;
|
||||
for (int h = 0; h < MOE_N_HORIZONS; ++h) {
|
||||
const long long w_idx = (((long long)b * MOE_N_HORIZONS + h)
|
||||
* MOE_HIDDEN_DIM + d_out) * MOE_HIDDEN_DIM + d_in;
|
||||
acc_w += d_W_scratch[w_idx];
|
||||
if (lane == 0) {
|
||||
const long long b_idx = ((long long)b * MOE_N_HORIZONS + h)
|
||||
* MOE_HIDDEN_DIM + d_out;
|
||||
acc_b += d_b_scratch[b_idx];
|
||||
}
|
||||
}
|
||||
}
|
||||
grad_experts_W[((long long)e * MOE_HIDDEN_DIM + d_out) * MOE_HIDDEN_DIM + d_in] = acc_w;
|
||||
if (lane == 0) {
|
||||
grad_experts_b[(long long)e * MOE_HIDDEN_DIM + d_out] = acc_b;
|
||||
}
|
||||
}
|
||||
|
||||
// ── Gate softmax + auxiliary load-balancing loss.
|
||||
// Computes gate_probs (for the optional STE gate-grad path) and the
|
||||
// `aux_loss` scalar that adds to total loss to discourage expert
|
||||
// collapse.
|
||||
// Grid: (1, 1, 1). Block: (N_EXPERTS, 1, 1) (= 4 threads).
|
||||
// All-in-one — input sizes are tiny.
|
||||
extern "C" __global__ void regime_moe_gate_aux(
|
||||
const float* __restrict__ gate_logits, // [B, N_E]
|
||||
const int* __restrict__ top_e, // [B]
|
||||
int n_batch,
|
||||
float* __restrict__ gate_probs_out, // [B, N_E]
|
||||
float* __restrict__ aux_loss_out // [1]
|
||||
) {
|
||||
int e = threadIdx.x;
|
||||
if (e >= MOE_N_EXPERTS) return;
|
||||
|
||||
// Phase 1: softmax per batch and write gate_probs.
|
||||
// We do one full sweep here, per-thread = per-expert column.
|
||||
float frac_local = 0.0f;
|
||||
float prob_mean_local = 0.0f;
|
||||
for (int b = 0; b < n_batch; ++b) {
|
||||
// Compute max + sum once per batch, by thread 0; broadcast.
|
||||
__shared__ float s_max;
|
||||
__shared__ float s_sum;
|
||||
if (e == 0) {
|
||||
float mx = -INFINITY;
|
||||
float sm = 0.0f;
|
||||
#pragma unroll
|
||||
for (int ep = 0; ep < MOE_N_EXPERTS; ++ep) {
|
||||
const float v = gate_logits[(long long)b * MOE_N_EXPERTS + ep];
|
||||
if (v > mx) mx = v;
|
||||
}
|
||||
#pragma unroll
|
||||
for (int ep = 0; ep < MOE_N_EXPERTS; ++ep) {
|
||||
sm += expf(gate_logits[(long long)b * MOE_N_EXPERTS + ep] - mx);
|
||||
}
|
||||
s_max = mx;
|
||||
s_sum = sm;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
const float p_e = expf(gate_logits[(long long)b * MOE_N_EXPERTS + e] - s_max) / s_sum;
|
||||
gate_probs_out[(long long)b * MOE_N_EXPERTS + e] = p_e;
|
||||
prob_mean_local += p_e;
|
||||
if (top_e[b] == e) frac_local += 1.0f;
|
||||
__syncthreads();
|
||||
}
|
||||
frac_local *= (1.0f / (float)n_batch);
|
||||
prob_mean_local *= (1.0f / (float)n_batch);
|
||||
|
||||
// Phase 2: aux_loss = N_EXPERTS · Σ_e frac · prob_mean.
|
||||
// Warp-reduce across the 4 threads.
|
||||
float aux_contrib = frac_local * prob_mean_local;
|
||||
// Warp shuffle on a 4-thread block.
|
||||
#pragma unroll
|
||||
for (int s = 16; s > 0; s >>= 1) {
|
||||
aux_contrib += __shfl_xor_sync(0xffffffff, aux_contrib, s);
|
||||
}
|
||||
if (e == 0) aux_loss_out[0] = (float)MOE_N_EXPERTS * aux_contrib;
|
||||
}
|
||||
@@ -1,59 +0,0 @@
|
||||
//! L2 anchor regularization host binding (v2 axis B, V8 commit).
|
||||
//!
|
||||
//! Anchors a trainable parameter buffer toward its initialization
|
||||
//! values via `loss = λ · Σ (p − p_init)²`. The Wiener-α controller
|
||||
//! on the host side adjusts λ each step based on the loss-improvement
|
||||
//! rate (signal-driven, floor-bounded — see `trainer/anchor_controller.rs`).
|
||||
//!
|
||||
//! See `cuda/anchor_l2.cu` for math + the v2 spec §3.5.
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use cudarc::driver::{
|
||||
CudaContext, CudaFunction, CudaModule, CudaSlice, CudaStream, LaunchConfig, PushKernelArg,
|
||||
};
|
||||
use std::sync::Arc;
|
||||
|
||||
const CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/anchor_l2.cubin"));
|
||||
|
||||
pub struct AnchorL2 {
|
||||
_module: Arc<CudaModule>,
|
||||
func: CudaFunction,
|
||||
stream: Arc<CudaStream>,
|
||||
}
|
||||
|
||||
impl AnchorL2 {
|
||||
pub fn new(ctx: &Arc<CudaContext>, stream: Arc<CudaStream>) -> Result<Self> {
|
||||
let module = ctx.load_cubin(CUBIN.to_vec()).context("load anchor_l2 cubin")?;
|
||||
let func = module.load_function("anchor_l2_fwd_bwd").context("load anchor_l2_fwd_bwd")?;
|
||||
Ok(Self { _module: module, func, stream })
|
||||
}
|
||||
|
||||
/// Compute anchor_loss = λ · ‖p − p_init‖² and accumulate
|
||||
/// `grad_p += 2λ · (p − p_init)`. `lambda_d` is a device-side
|
||||
/// [1]-buffer (host writes scalar before launch).
|
||||
pub fn apply(
|
||||
&self,
|
||||
p: &CudaSlice<f32>,
|
||||
p_init: &CudaSlice<f32>,
|
||||
lambda_d: &CudaSlice<f32>,
|
||||
n: i32,
|
||||
loss_out: &mut CudaSlice<f32>,
|
||||
grad_p: &mut CudaSlice<f32>,
|
||||
) -> Result<()> {
|
||||
let cfg = LaunchConfig {
|
||||
grid_dim: (1, 1, 1),
|
||||
block_dim: (256, 1, 1),
|
||||
shared_mem_bytes: 0,
|
||||
};
|
||||
let mut launch = self.stream.launch_builder(&self.func);
|
||||
unsafe {
|
||||
launch
|
||||
.arg(p).arg(p_init).arg(lambda_d)
|
||||
.arg(&n)
|
||||
.arg(loss_out).arg(grad_p)
|
||||
.launch(cfg)
|
||||
.context("anchor_l2_fwd_bwd")?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
@@ -1,81 +0,0 @@
|
||||
//! Horizon-axis mean-collapse host binding.
|
||||
//!
|
||||
//! Collapses per-horizon context `[B, N_H, H]` to a single context
|
||||
//! `[B, H]` by averaging across the horizon axis. Used to seed the
|
||||
//! CfC's `h_old` at k=0 from `MultiHorizonAttention.routed_ctx`.
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use cudarc::driver::{
|
||||
CudaContext, CudaFunction, CudaModule, CudaSlice, CudaStream, LaunchConfig, PushKernelArg,
|
||||
};
|
||||
use std::sync::Arc;
|
||||
|
||||
pub const HMC_HIDDEN_DIM: usize = 128;
|
||||
pub const HMC_BLOCK: usize = 128;
|
||||
|
||||
const CUBIN: &[u8] =
|
||||
include_bytes!(concat!(env!("OUT_DIR"), "/horizon_mean_collapse.cubin"));
|
||||
|
||||
pub struct HorizonMeanCollapse {
|
||||
_module: Arc<CudaModule>,
|
||||
fwd_fn: CudaFunction,
|
||||
bwd_fn: CudaFunction,
|
||||
stream: Arc<CudaStream>,
|
||||
}
|
||||
|
||||
impl HorizonMeanCollapse {
|
||||
pub fn new(ctx: &Arc<CudaContext>, stream: Arc<CudaStream>) -> Result<Self> {
|
||||
let module = ctx
|
||||
.load_cubin(CUBIN.to_vec())
|
||||
.context("load horizon_mean_collapse cubin")?;
|
||||
let fwd_fn = module
|
||||
.load_function("horizon_mean_collapse_fwd")
|
||||
.context("load horizon_mean_collapse_fwd")?;
|
||||
let bwd_fn = module
|
||||
.load_function("horizon_mean_collapse_bwd")
|
||||
.context("load horizon_mean_collapse_bwd")?;
|
||||
Ok(Self { _module: module, fwd_fn, bwd_fn, stream })
|
||||
}
|
||||
|
||||
pub fn forward(
|
||||
&self,
|
||||
ctx_per_h: &CudaSlice<f32>,
|
||||
n_batch: i32,
|
||||
ctx_mean: &mut CudaSlice<f32>,
|
||||
) -> Result<()> {
|
||||
let cfg = LaunchConfig {
|
||||
grid_dim: (n_batch as u32, 1, 1),
|
||||
block_dim: (HMC_BLOCK as u32, 1, 1),
|
||||
shared_mem_bytes: 0,
|
||||
};
|
||||
let mut launch = self.stream.launch_builder(&self.fwd_fn);
|
||||
unsafe {
|
||||
launch
|
||||
.arg(ctx_per_h).arg(&n_batch).arg(ctx_mean)
|
||||
.launch(cfg)
|
||||
.context("horizon_mean_collapse_fwd")?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn backward(
|
||||
&self,
|
||||
grad_ctx_mean: &CudaSlice<f32>,
|
||||
n_batch: i32,
|
||||
grad_ctx_per_h: &mut CudaSlice<f32>,
|
||||
) -> Result<()> {
|
||||
let cfg = LaunchConfig {
|
||||
grid_dim: (n_batch as u32, 1, 1),
|
||||
block_dim: (HMC_BLOCK as u32, 1, 1),
|
||||
shared_mem_bytes: 0,
|
||||
};
|
||||
let mut launch = self.stream.launch_builder(&self.bwd_fn);
|
||||
unsafe {
|
||||
launch
|
||||
.arg(grad_ctx_mean).arg(&n_batch).arg(grad_ctx_per_h)
|
||||
.launch(cfg)
|
||||
.context("horizon_mean_collapse_bwd")?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
@@ -1,122 +0,0 @@
|
||||
//! Horizon-token attention pool host binding (v2 axis C, V2 commit).
|
||||
//!
|
||||
//! Replaces the falsified per-horizon Q_h pool with N_HORIZONS learned
|
||||
//! tokens prepended to the K-sequence and a single shared learned Q.
|
||||
//! See `cuda/horizon_token_attention_pool.cu` for math + the v2 spec.
|
||||
//!
|
||||
//! This binding is graph-capture-safe (no `synchronize`, no host
|
||||
//! allocs, no `memcpy_htod` in the hot path).
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use cudarc::driver::{
|
||||
CudaContext, CudaFunction, CudaModule, CudaSlice, CudaStream, LaunchConfig, PushKernelArg,
|
||||
};
|
||||
use std::sync::Arc;
|
||||
|
||||
pub const HTAP_HIDDEN_DIM: usize = 128;
|
||||
pub const HTAP_BLOCK: usize = 128;
|
||||
pub const HTAP_N_HORIZONS: usize = 5;
|
||||
pub const HTAP_N_WARPS: usize = HTAP_BLOCK / 32;
|
||||
|
||||
const CUBIN: &[u8] =
|
||||
include_bytes!(concat!(env!("OUT_DIR"), "/horizon_token_attention_pool.cubin"));
|
||||
|
||||
pub struct HorizonTokenAttentionPool {
|
||||
_module: Arc<CudaModule>,
|
||||
fwd_fn: CudaFunction,
|
||||
bwd_fn: CudaFunction,
|
||||
stream: Arc<CudaStream>,
|
||||
}
|
||||
|
||||
impl HorizonTokenAttentionPool {
|
||||
pub fn new(ctx: &Arc<CudaContext>, stream: Arc<CudaStream>) -> Result<Self> {
|
||||
let module = ctx
|
||||
.load_cubin(CUBIN.to_vec())
|
||||
.context("load horizon_token_attention_pool cubin")?;
|
||||
let fwd_fn = module
|
||||
.load_function("horizon_token_attention_pool_fwd")
|
||||
.context("load horizon_token_attention_pool_fwd")?;
|
||||
let bwd_fn = module
|
||||
.load_function("horizon_token_attention_pool_bwd")
|
||||
.context("load horizon_token_attention_pool_bwd")?;
|
||||
Ok(Self { _module: module, fwd_fn, bwd_fn, stream })
|
||||
}
|
||||
|
||||
/// Forward. All buffers must be device-resident; no synchronize.
|
||||
/// `horizon_tokens` : [N_H, H]
|
||||
/// `q` : [H]
|
||||
/// `ln_out` : [B, K, H]
|
||||
/// `ctx_out` : [B, N_H, H] written
|
||||
/// `attn_out` : [B, N_H + K] written (saved for backward)
|
||||
pub fn forward(
|
||||
&self,
|
||||
horizon_tokens: &CudaSlice<f32>,
|
||||
q: &CudaSlice<f32>,
|
||||
ln_out: &CudaSlice<f32>,
|
||||
n_batch: i32,
|
||||
k_seq: i32,
|
||||
ctx_out: &mut CudaSlice<f32>,
|
||||
attn_out: &mut CudaSlice<f32>,
|
||||
) -> Result<()> {
|
||||
// Smem layout: s_attn[N_H + K] + s_warp[N_WARPS]
|
||||
let smem_bytes = ((HTAP_N_HORIZONS + k_seq as usize + HTAP_N_WARPS)
|
||||
* std::mem::size_of::<f32>()) as u32;
|
||||
let cfg = LaunchConfig {
|
||||
grid_dim: (n_batch as u32, 1, 1),
|
||||
block_dim: (HTAP_BLOCK as u32, 1, 1),
|
||||
shared_mem_bytes: smem_bytes,
|
||||
};
|
||||
let mut launch = self.stream.launch_builder(&self.fwd_fn);
|
||||
unsafe {
|
||||
launch
|
||||
.arg(horizon_tokens).arg(q).arg(ln_out)
|
||||
.arg(&n_batch).arg(&k_seq)
|
||||
.arg(ctx_out).arg(attn_out)
|
||||
.launch(cfg)
|
||||
.context("horizon_token_attention_pool_fwd")?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Backward. Gradients accumulate (`+=`) into the per-batch scratch
|
||||
/// buffers and `grad_ln_out`. Reduce_axis0 collapses scratches over
|
||||
/// the batch axis downstream.
|
||||
/// `grad_horizon_tokens_scratch` : [B, N_H, H] +=
|
||||
/// `grad_q_scratch` : [B, H] +=
|
||||
/// `grad_ln_out` : [B, K, H] += (per-batch slice; no race)
|
||||
pub fn backward(
|
||||
&self,
|
||||
horizon_tokens: &CudaSlice<f32>,
|
||||
q: &CudaSlice<f32>,
|
||||
ln_out: &CudaSlice<f32>,
|
||||
attn: &CudaSlice<f32>,
|
||||
grad_ctx: &CudaSlice<f32>,
|
||||
n_batch: i32,
|
||||
k_seq: i32,
|
||||
grad_horizon_tokens_scratch: &mut CudaSlice<f32>,
|
||||
grad_q_scratch: &mut CudaSlice<f32>,
|
||||
grad_ln_out: &mut CudaSlice<f32>,
|
||||
) -> Result<()> {
|
||||
// Smem layout: s_d_attn[N_H + K] + s_warp[N_WARPS] + s_d_S[HIDDEN_DIM]
|
||||
let smem_bytes = ((HTAP_N_HORIZONS + k_seq as usize + HTAP_N_WARPS + HTAP_HIDDEN_DIM)
|
||||
* std::mem::size_of::<f32>()) as u32;
|
||||
let cfg = LaunchConfig {
|
||||
grid_dim: (n_batch as u32, 1, 1),
|
||||
block_dim: (HTAP_BLOCK as u32, 1, 1),
|
||||
shared_mem_bytes: smem_bytes,
|
||||
};
|
||||
let mut launch = self.stream.launch_builder(&self.bwd_fn);
|
||||
unsafe {
|
||||
launch
|
||||
.arg(horizon_tokens).arg(q).arg(ln_out)
|
||||
.arg(attn).arg(grad_ctx)
|
||||
.arg(&n_batch).arg(&k_seq)
|
||||
.arg(grad_horizon_tokens_scratch)
|
||||
.arg(grad_q_scratch)
|
||||
.arg(grad_ln_out)
|
||||
.launch(cfg)
|
||||
.context("horizon_token_attention_pool_bwd")?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
@@ -1,82 +0,0 @@
|
||||
//! 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,112 +0,0 @@
|
||||
//! Inverted attention pool host binding (v2 axis E, V3 commit).
|
||||
//!
|
||||
//! iTransformer-style cross-variate attention: each of HIDDEN_DIM
|
||||
//! features becomes a "variate token" with its K-trajectory as
|
||||
//! embedding. Single-block-per-batch, no learnable parameters in this
|
||||
//! v2 (Q = K = V = X_inv self-attention with mean-K pool).
|
||||
//!
|
||||
//! See `cuda/inverted_attention_pool.cu` for math + the v2 spec §3.2.
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use cudarc::driver::{
|
||||
CudaContext, CudaFunction, CudaModule, CudaSlice, CudaStream, LaunchConfig, PushKernelArg,
|
||||
};
|
||||
use std::sync::Arc;
|
||||
|
||||
pub const IAP_HIDDEN_DIM: usize = 128;
|
||||
pub const IAP_BLOCK: usize = 128;
|
||||
|
||||
const CUBIN: &[u8] =
|
||||
include_bytes!(concat!(env!("OUT_DIR"), "/inverted_attention_pool.cubin"));
|
||||
|
||||
pub struct InvertedAttentionPool {
|
||||
_module: Arc<CudaModule>,
|
||||
fwd_fn: CudaFunction,
|
||||
bwd_fn: CudaFunction,
|
||||
stream: Arc<CudaStream>,
|
||||
}
|
||||
|
||||
impl InvertedAttentionPool {
|
||||
pub fn new(ctx: &Arc<CudaContext>, stream: Arc<CudaStream>) -> Result<Self> {
|
||||
let module = ctx
|
||||
.load_cubin(CUBIN.to_vec())
|
||||
.context("load inverted_attention_pool cubin")?;
|
||||
let fwd_fn = module
|
||||
.load_function("inverted_attention_pool_fwd")
|
||||
.context("load inverted_attention_pool_fwd")?;
|
||||
let bwd_fn = module
|
||||
.load_function("inverted_attention_pool_bwd")
|
||||
.context("load inverted_attention_pool_bwd")?;
|
||||
Ok(Self { _module: module, fwd_fn, bwd_fn, stream })
|
||||
}
|
||||
|
||||
/// Forward. `inv_scale` must equal `1 / sqrt(k_seq)` to match the
|
||||
/// scaled-dot-product attention convention.
|
||||
/// `ln_out` : [B, K, H]
|
||||
/// `pooled_out` : [B, H] written
|
||||
/// `attn_out` : [B, H, H] written (saved for backward)
|
||||
pub fn forward(
|
||||
&self,
|
||||
ln_out: &CudaSlice<f32>,
|
||||
n_batch: i32,
|
||||
k_seq: i32,
|
||||
inv_scale: f32,
|
||||
pooled_out: &mut CudaSlice<f32>,
|
||||
attn_out: &mut CudaSlice<f32>,
|
||||
) -> Result<()> {
|
||||
// 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),
|
||||
shared_mem_bytes: smem_bytes,
|
||||
};
|
||||
let mut launch = self.stream.launch_builder(&self.fwd_fn);
|
||||
unsafe {
|
||||
launch
|
||||
.arg(ln_out)
|
||||
.arg(&n_batch).arg(&k_seq).arg(&inv_scale)
|
||||
.arg(pooled_out).arg(attn_out)
|
||||
.launch(cfg)
|
||||
.context("inverted_attention_pool_fwd")?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Backward. `grad_ln_out` accumulates += per-batch (no race).
|
||||
/// `d_scores_scratch` is a per-step intermediate buffer of shape
|
||||
/// `[B, H, H]` (must be at least `n_batch * H * H` floats). Trainer
|
||||
/// allocates once and reuses across steps.
|
||||
pub fn backward(
|
||||
&self,
|
||||
ln_out: &CudaSlice<f32>,
|
||||
attn: &CudaSlice<f32>,
|
||||
grad_pooled: &CudaSlice<f32>,
|
||||
n_batch: i32,
|
||||
k_seq: i32,
|
||||
inv_scale: f32,
|
||||
d_scores_scratch: &mut CudaSlice<f32>,
|
||||
grad_ln_out: &mut CudaSlice<f32>,
|
||||
) -> Result<()> {
|
||||
// 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),
|
||||
shared_mem_bytes: smem_bytes,
|
||||
};
|
||||
let mut launch = self.stream.launch_builder(&self.bwd_fn);
|
||||
unsafe {
|
||||
launch
|
||||
.arg(ln_out).arg(attn).arg(grad_pooled)
|
||||
.arg(&n_batch).arg(&k_seq).arg(&inv_scale)
|
||||
.arg(d_scores_scratch)
|
||||
.arg(grad_ln_out)
|
||||
.launch(cfg)
|
||||
.context("inverted_attention_pool_bwd")?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
@@ -28,14 +28,8 @@
|
||||
pub mod cfc;
|
||||
pub mod data;
|
||||
pub mod eval;
|
||||
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;
|
||||
pub mod pinned;
|
||||
pub mod pinned_mem;
|
||||
pub mod trainer;
|
||||
|
||||
@@ -1,174 +0,0 @@
|
||||
//! Regime-aware MoE gate host binding (v2 axis D, V6 commit).
|
||||
//!
|
||||
//! Top-1 routing per Switch Transformer. Provides three kernel
|
||||
//! launches:
|
||||
//! - `forward`: gate top-1 + selected-expert linear → routed_ctx
|
||||
//! - `backward`: chain-rule grads on the selected expert's weights,
|
||||
//! bias, and the input fused_ctx. STE on the gate logits is handled
|
||||
//! by the caller (no separate kernel — gate-logit grad = 0 except
|
||||
//! via the aux loss path).
|
||||
//! - `aux_loss`: softmax(gate_logits) + load-balancing aux scalar.
|
||||
//!
|
||||
//! See `cuda/regime_moe_gate.cu` for math + the v2 spec §3.3.
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use cudarc::driver::{
|
||||
CudaContext, CudaFunction, CudaModule, CudaSlice, CudaStream, LaunchConfig, PushKernelArg,
|
||||
};
|
||||
use std::sync::Arc;
|
||||
|
||||
pub const MOE_HIDDEN_DIM: usize = 128;
|
||||
pub const MOE_BLOCK: usize = 128;
|
||||
pub const MOE_N_HORIZONS: usize = 5;
|
||||
pub const MOE_N_EXPERTS: usize = 4;
|
||||
|
||||
const CUBIN: &[u8] =
|
||||
include_bytes!(concat!(env!("OUT_DIR"), "/regime_moe_gate.cubin"));
|
||||
|
||||
pub struct RegimeMoeGate {
|
||||
_module: Arc<CudaModule>,
|
||||
fwd_fn: CudaFunction,
|
||||
bwd_fn: CudaFunction,
|
||||
scatter_fn: CudaFunction,
|
||||
aux_fn: CudaFunction,
|
||||
stream: Arc<CudaStream>,
|
||||
}
|
||||
|
||||
impl RegimeMoeGate {
|
||||
pub fn new(ctx: &Arc<CudaContext>, stream: Arc<CudaStream>) -> Result<Self> {
|
||||
let module = ctx.load_cubin(CUBIN.to_vec()).context("load regime_moe_gate cubin")?;
|
||||
let fwd_fn = module.load_function("regime_moe_gate_fwd").context("fwd")?;
|
||||
let bwd_fn = module.load_function("regime_moe_gate_bwd").context("bwd")?;
|
||||
let scatter_fn = module.load_function("regime_moe_gate_scatter").context("scatter")?;
|
||||
let aux_fn = module.load_function("regime_moe_gate_aux").context("aux")?;
|
||||
Ok(Self { _module: module, fwd_fn, bwd_fn, scatter_fn, aux_fn, stream })
|
||||
}
|
||||
|
||||
/// Forward. Selects top-1 expert per batch and applies its linear.
|
||||
/// `gate_logits` : [B, N_E]
|
||||
/// `fused_ctx` : [B, N_H, H]
|
||||
/// `experts_w` : [N_E, H, H]
|
||||
/// `experts_b` : [N_E, H]
|
||||
/// `routed_out` : [B, N_H, H] written
|
||||
/// `top_e_out` : [B] written (i32)
|
||||
pub fn forward(
|
||||
&self,
|
||||
gate_logits: &CudaSlice<f32>,
|
||||
fused_ctx: &CudaSlice<f32>,
|
||||
experts_w: &CudaSlice<f32>,
|
||||
experts_b: &CudaSlice<f32>,
|
||||
n_batch: i32,
|
||||
routed_out: &mut CudaSlice<f32>,
|
||||
top_e_out: &mut CudaSlice<i32>,
|
||||
) -> Result<()> {
|
||||
let cfg = LaunchConfig {
|
||||
grid_dim: (n_batch as u32, MOE_N_HORIZONS as u32, 1),
|
||||
block_dim: (MOE_BLOCK as u32, 1, 1),
|
||||
shared_mem_bytes: 0,
|
||||
};
|
||||
let mut launch = self.stream.launch_builder(&self.fwd_fn);
|
||||
unsafe {
|
||||
launch
|
||||
.arg(gate_logits).arg(fused_ctx).arg(experts_w).arg(experts_b)
|
||||
.arg(&n_batch)
|
||||
.arg(routed_out).arg(top_e_out)
|
||||
.launch(cfg)
|
||||
.context("regime_moe_gate_fwd")?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Backward (expert path; gate STE handled in aux). All param grad
|
||||
/// scratches are sparse-by-expert: writes go only to `top_e[b]`'s
|
||||
/// slot. Reduce over (B, N_H) downstream via reduce_axis0.
|
||||
/// `d_w_scratch` : [B, N_H, N_E, H, H]
|
||||
/// `d_b_scratch` : [B, N_H, N_E, H]
|
||||
/// `d_fused_ctx_out` : [B, N_H, H]
|
||||
pub fn backward(
|
||||
&self,
|
||||
fused_ctx: &CudaSlice<f32>,
|
||||
experts_w: &CudaSlice<f32>,
|
||||
d_routed: &CudaSlice<f32>,
|
||||
top_e: &CudaSlice<i32>,
|
||||
n_batch: i32,
|
||||
d_w_scratch: &mut CudaSlice<f32>,
|
||||
d_b_scratch: &mut CudaSlice<f32>,
|
||||
d_fused_ctx_out: &mut CudaSlice<f32>,
|
||||
) -> Result<()> {
|
||||
let cfg = LaunchConfig {
|
||||
grid_dim: (n_batch as u32, MOE_N_HORIZONS as u32, 1),
|
||||
block_dim: (MOE_BLOCK as u32, 1, 1),
|
||||
shared_mem_bytes: 0,
|
||||
};
|
||||
let mut launch = self.stream.launch_builder(&self.bwd_fn);
|
||||
unsafe {
|
||||
launch
|
||||
.arg(fused_ctx).arg(experts_w).arg(d_routed).arg(top_e)
|
||||
.arg(&n_batch)
|
||||
.arg(d_w_scratch).arg(d_b_scratch).arg(d_fused_ctx_out)
|
||||
.launch(cfg)
|
||||
.context("regime_moe_gate_bwd")?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Scatter per-batch rank-1 contributions in `d_W_scratch` /
|
||||
/// `d_b_scratch` into the per-expert grad buffers based on `top_e`.
|
||||
/// `d_W_scratch` must be `[B, N_H, H, H]` and `d_b_scratch` must
|
||||
/// be `[B, N_H, H]` (the post-perf compact layout).
|
||||
pub fn scatter(
|
||||
&self,
|
||||
d_w_scratch: &CudaSlice<f32>,
|
||||
d_b_scratch: &CudaSlice<f32>,
|
||||
top_e: &CudaSlice<i32>,
|
||||
n_batch: i32,
|
||||
grad_experts_w: &mut CudaSlice<f32>,
|
||||
grad_experts_b: &mut CudaSlice<f32>,
|
||||
) -> Result<()> {
|
||||
// Grid: (N_E, H, ceil(H/32)). Block: 32.
|
||||
let chunks: u32 = (MOE_HIDDEN_DIM as u32).div_ceil(32);
|
||||
let cfg = LaunchConfig {
|
||||
grid_dim: (MOE_N_EXPERTS as u32, MOE_HIDDEN_DIM as u32, chunks),
|
||||
block_dim: (32, 1, 1),
|
||||
shared_mem_bytes: 0,
|
||||
};
|
||||
let mut launch = self.stream.launch_builder(&self.scatter_fn);
|
||||
unsafe {
|
||||
launch
|
||||
.arg(d_w_scratch).arg(d_b_scratch).arg(top_e)
|
||||
.arg(&n_batch)
|
||||
.arg(grad_experts_w).arg(grad_experts_b)
|
||||
.launch(cfg)
|
||||
.context("regime_moe_gate_scatter")?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Softmax + load-balancing aux loss.
|
||||
/// `gate_probs_out` : [B, N_E] softmax output
|
||||
/// `aux_loss_out` : [1] scalar aux loss
|
||||
pub fn aux_loss(
|
||||
&self,
|
||||
gate_logits: &CudaSlice<f32>,
|
||||
top_e: &CudaSlice<i32>,
|
||||
n_batch: i32,
|
||||
gate_probs_out: &mut CudaSlice<f32>,
|
||||
aux_loss_out: &mut CudaSlice<f32>,
|
||||
) -> Result<()> {
|
||||
let cfg = LaunchConfig {
|
||||
grid_dim: (1, 1, 1),
|
||||
block_dim: (MOE_N_EXPERTS as u32, 1, 1),
|
||||
shared_mem_bytes: 0,
|
||||
};
|
||||
let mut launch = self.stream.launch_builder(&self.aux_fn);
|
||||
unsafe {
|
||||
launch
|
||||
.arg(gate_logits).arg(top_e)
|
||||
.arg(&n_batch)
|
||||
.arg(gate_probs_out).arg(aux_loss_out)
|
||||
.launch(cfg)
|
||||
.context("regime_moe_gate_aux")?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
@@ -1,142 +0,0 @@
|
||||
//! Wiener-α anchor-coefficient controller (v2 axis B controller).
|
||||
//!
|
||||
//! Drives `λ_anchor` for the L2 anchor regularization. Signal-driven
|
||||
//! per `pearl_controller_anchors_isv_driven` + the
|
||||
//! Wiener-α-with-floor pattern (`pearl_wiener_alpha_floor_for_nonstationary`,
|
||||
//! `pearl_blend_formulas_must_have_permanent_floor`).
|
||||
//!
|
||||
//! Math:
|
||||
//! loss_change_t = val_loss_t − val_loss_{t-1} # raw signal
|
||||
//! ema_change ← ema_change + α · (loss_change_t − ema_change)
|
||||
//! target_λ = clamp(|ema_change| · scale, λ_floor, λ_max)
|
||||
//! λ_t = max(λ_floor, blend(λ_{t-1}, target_λ, α))
|
||||
//!
|
||||
//! where α is the standard Wiener-α (MSE-optimal under stationarity).
|
||||
//! Floor protects against deadlock when ema_change → 0
|
||||
//! (anchor never fully relaxes — minimum guard against runaway drift).
|
||||
//!
|
||||
//! Floor `λ_floor` is itself signal-driven from initial parameter
|
||||
//! magnitudes: `λ_floor = ‖p_init‖ / (100 · sqrt(numel))`. Captured at
|
||||
//! controller construction via the first-observation bootstrap pattern
|
||||
//! (`pearl_first_observation_bootstrap`).
|
||||
|
||||
const WIENER_ALPHA_FLOOR: f32 = 0.4; // pearl_wiener_alpha_floor_for_nonstationary
|
||||
const LAMBDA_MAX: f32 = 1.0; // safety cap on λ
|
||||
const LAMBDA_SCALE: f32 = 10.0; // ema_change → λ multiplier
|
||||
|
||||
pub struct AnchorController {
|
||||
/// Floor for λ_anchor. Anchor never fully relaxes below this.
|
||||
pub lambda_floor: f32,
|
||||
pub lambda_max: f32,
|
||||
|
||||
/// Wiener-α smoothing state.
|
||||
diff_var: f32,
|
||||
sample_var: f32,
|
||||
ema_change: f32,
|
||||
prev_loss: f32,
|
||||
bootstrapped: bool,
|
||||
|
||||
/// Smoothed λ.
|
||||
pub lambda: f32,
|
||||
}
|
||||
|
||||
impl AnchorController {
|
||||
/// `init_param_magnitude` = sum of L2-norms of all anchored param
|
||||
/// groups; `init_numel` = total trainable element count anchored.
|
||||
/// Used to derive the signal-driven floor.
|
||||
pub fn new(init_param_magnitude: f32, init_numel: usize) -> Self {
|
||||
let lambda_floor =
|
||||
init_param_magnitude / (100.0 * (init_numel.max(1) as f32).sqrt());
|
||||
Self {
|
||||
lambda_floor,
|
||||
lambda_max: LAMBDA_MAX,
|
||||
diff_var: 0.0,
|
||||
sample_var: 0.0,
|
||||
ema_change: 0.0,
|
||||
prev_loss: 0.0,
|
||||
bootstrapped: false,
|
||||
lambda: lambda_floor.max(1e-6), // start at floor
|
||||
}
|
||||
}
|
||||
|
||||
/// Step the controller with a new validation-loss reading.
|
||||
/// Returns the updated λ. Capture-safe (host-side; no GPU touch).
|
||||
pub fn step(&mut self, val_loss: f32) -> f32 {
|
||||
if !self.bootstrapped {
|
||||
// First observation: replace state directly per
|
||||
// pearl_first_observation_bootstrap.
|
||||
self.prev_loss = val_loss;
|
||||
self.bootstrapped = true;
|
||||
return self.lambda;
|
||||
}
|
||||
|
||||
let change = val_loss - self.prev_loss;
|
||||
self.prev_loss = val_loss;
|
||||
|
||||
// Wiener-α: α = diff_var / (diff_var + sample_var + ε).
|
||||
// Maintain running estimates of diff_var (changes) and
|
||||
// sample_var (noise around the mean change).
|
||||
let delta = change - self.ema_change;
|
||||
self.diff_var = 0.95 * self.diff_var + 0.05 * change * change;
|
||||
self.sample_var = 0.95 * self.sample_var + 0.05 * delta * delta;
|
||||
let alpha_raw = self.diff_var / (self.diff_var + self.sample_var + 1e-9);
|
||||
let alpha = alpha_raw.max(WIENER_ALPHA_FLOOR);
|
||||
|
||||
self.ema_change = self.ema_change + alpha * delta;
|
||||
|
||||
// Target λ proportional to |ema_change| × scale.
|
||||
let target_lambda = (self.ema_change.abs() * LAMBDA_SCALE)
|
||||
.clamp(self.lambda_floor, self.lambda_max);
|
||||
|
||||
// Blend with floor — pearl_blend_formulas_must_have_permanent_floor.
|
||||
let blended = self.lambda + alpha * (target_lambda - self.lambda);
|
||||
self.lambda = blended.max(self.lambda_floor);
|
||||
self.lambda
|
||||
}
|
||||
|
||||
pub fn current_lambda(&self) -> f32 {
|
||||
self.lambda
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn floor_init_from_signal() {
|
||||
// ‖p_init‖ = 5.0, numel = 100 → floor = 5 / (100 · 10) = 0.005
|
||||
let c = AnchorController::new(5.0, 100);
|
||||
assert!((c.lambda_floor - 0.005).abs() < 1e-6);
|
||||
assert!(c.lambda >= c.lambda_floor);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn first_observation_bootstrap_returns_floor() {
|
||||
let mut c = AnchorController::new(1.0, 100);
|
||||
let l0 = c.step(0.5);
|
||||
assert!((l0 - c.lambda_floor).abs() < 1e-6);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn lambda_never_falls_below_floor() {
|
||||
let mut c = AnchorController::new(1.0, 100);
|
||||
// Many steps with no change → ema_change → 0 → target → floor.
|
||||
for _ in 0..1000 {
|
||||
let l = c.step(0.5);
|
||||
assert!(l >= c.lambda_floor - 1e-9, "λ={l} below floor={}", c.lambda_floor);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn lambda_capped_at_max() {
|
||||
let mut c = AnchorController::new(1.0, 100);
|
||||
c.step(0.0);
|
||||
// Huge loss spikes alternating sign — ema_change should grow,
|
||||
// but λ stays ≤ LAMBDA_MAX.
|
||||
for i in 0..100 {
|
||||
let l = c.step(if i % 2 == 0 { 100.0 } else { -100.0 });
|
||||
assert!(l <= LAMBDA_MAX + 1e-6, "λ={l} above max");
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,8 +1,6 @@
|
||||
//! Trainer module — PerceptionTrainer wraps the full stacked
|
||||
//! Mamba2 -> CfC -> heads pipeline plus 6 AdamW optimizer groups.
|
||||
|
||||
pub mod anchor_controller;
|
||||
pub mod loss;
|
||||
pub mod multi_horizon_attention;
|
||||
pub mod optim;
|
||||
pub mod perception;
|
||||
|
||||
@@ -1,435 +0,0 @@
|
||||
//! MultiHorizonAttention — single source of truth for the ml-alpha
|
||||
//! attention path. Replaces the legacy single-Q `attention_pool`
|
||||
//! entirely.
|
||||
//!
|
||||
//! Owns:
|
||||
//! - horizon_tokens + their init snapshot (for L2 anchor)
|
||||
//! - shared Q + its init snapshot
|
||||
//! - horizon-token attention pool kernel binding
|
||||
//! - inverted (cross-variate) attention pool + saved buffers
|
||||
//! - regime-MoE gate + expert weights + init snapshots
|
||||
//! - Kendall σ logarithm (`log_sigma_h`) for the multi-horizon BCE
|
||||
//! - Anchor L2 kernel + the Wiener-α `AnchorController`
|
||||
//!
|
||||
//! All kernel calls go through the bound modules; nothing in this
|
||||
//! file does host-side compute or sync inside captured regions.
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use cudarc::driver::{CudaSlice, CudaStream};
|
||||
use ml_core::device::MlDevice;
|
||||
use rand::{Rng, SeedableRng};
|
||||
use rand_chacha::ChaCha8Rng;
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::anchor_l2::AnchorL2;
|
||||
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;
|
||||
use crate::trainer::optim::AdamW;
|
||||
|
||||
pub const MHA_HIDDEN_DIM: usize = HTAP_HIDDEN_DIM;
|
||||
pub const MHA_N_HORIZONS: usize = HTAP_N_HORIZONS;
|
||||
pub const MHA_N_EXPERTS: usize = MOE_N_EXPERTS;
|
||||
pub const MHA_REGIME_DIM: usize = 12; // spread(5) + vol(3) + tod(4)
|
||||
|
||||
pub struct MultiHorizonAttention {
|
||||
// ── (C) horizon-token attention pool ──
|
||||
pub pool: HorizonTokenAttentionPool,
|
||||
pub horizon_tokens_d: CudaSlice<f32>, // [N_H, H] learnable
|
||||
pub horizon_tokens_init_d: CudaSlice<f32>, // [N_H, H] non-trainable snapshot
|
||||
pub q_d: CudaSlice<f32>, // [H] single shared Q
|
||||
pub q_init_d: CudaSlice<f32>, // [H] non-trainable snapshot
|
||||
pub grad_horizon_tokens_d: CudaSlice<f32>,
|
||||
pub grad_horizon_tokens_scratch_d: CudaSlice<f32>, // [B, N_H, H] per-batch scratch
|
||||
pub grad_ctx_h_d: CudaSlice<f32>, // [B, N_H, H] d_ctx_h flowing from MoE bwd → pool bwd input
|
||||
pub grad_q_d: CudaSlice<f32>,
|
||||
pub grad_q_scratch_d: CudaSlice<f32>, // [B, H] per-batch scratch
|
||||
pub ctx_h_d: CudaSlice<f32>, // [B, N_H, H] fwd output
|
||||
pub attn_h_d: CudaSlice<f32>, // [B, N_H + K]
|
||||
pub opt_horizon_tokens: AdamW,
|
||||
pub opt_q: AdamW,
|
||||
|
||||
// ── (E) inverted attention pool ──
|
||||
pub inv_pool: InvertedAttentionPool,
|
||||
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
|
||||
// fused_ctx[b, h, d] = ctx_h[b, h, d] + inv_pooled[b, d].
|
||||
|
||||
// ── (D) regime-MoE gate + experts ──
|
||||
pub moe: RegimeMoeGate,
|
||||
pub gate_logits_d: CudaSlice<f32>, // [B, N_E] computed each step from regime features
|
||||
pub w_gate_d: CudaSlice<f32>, // [N_E, REGIME_DIM]
|
||||
pub experts_w_d: CudaSlice<f32>, // [N_E, H, H]
|
||||
pub experts_b_d: CudaSlice<f32>, // [N_E, H]
|
||||
pub experts_w_init_d: CudaSlice<f32>, // anchor snapshot (B)
|
||||
pub experts_b_init_d: CudaSlice<f32>, // anchor snapshot
|
||||
pub top_e_d: CudaSlice<i32>, // [B]
|
||||
pub gate_probs_d: CudaSlice<f32>, // [B, N_E]
|
||||
pub aux_loss_d: CudaSlice<f32>, // [1]
|
||||
pub routed_ctx_d: CudaSlice<f32>, // [B, N_H, H] v2 output (consumed by GRN heads)
|
||||
pub grad_routed_d: CudaSlice<f32>, // [B, N_H, H]
|
||||
pub grad_w_gate_d: CudaSlice<f32>,
|
||||
pub grad_experts_w_d: CudaSlice<f32>,
|
||||
pub grad_experts_b_d: CudaSlice<f32>,
|
||||
pub grad_w_scratch_d: CudaSlice<f32>, // [B, N_H, H, H] compact (no N_E axis)
|
||||
pub grad_b_scratch_d: CudaSlice<f32>, // [B, N_H, H] compact (no N_E axis)
|
||||
pub opt_w_gate: AdamW,
|
||||
pub opt_experts_w: AdamW,
|
||||
pub opt_experts_b: AdamW,
|
||||
|
||||
// ── (A) Kendall σ on the BCE ──
|
||||
pub log_sigma_h_d: CudaSlice<f32>, // [N_HORIZONS]
|
||||
pub grad_log_sigma_h_d: CudaSlice<f32>,
|
||||
pub opt_log_sigma: AdamW,
|
||||
|
||||
// ── Mean-collapse over horizon axis (seeds CfC h_old at k=0) ──
|
||||
pub collapse: HorizonMeanCollapse,
|
||||
pub ctx_mean_d: CudaSlice<f32>, // [B, H] output of collapse fwd, == legacy attn_context_d slot
|
||||
pub grad_ctx_mean_d: CudaSlice<f32>, // [B, H] upstream grad from CfC h_old at k=0
|
||||
|
||||
// ── (B) anchor regularization controller ──
|
||||
pub anchor_l2: AnchorL2,
|
||||
pub anchor: AnchorController,
|
||||
pub lambda_d: CudaSlice<f32>, // [1] current λ uploaded each step
|
||||
pub anchor_loss_partial_d: CudaSlice<f32>, // [1] per-launch partial
|
||||
pub anchor_loss_total_d: CudaSlice<f32>, // [1] accumulated
|
||||
|
||||
// Bookkeeping.
|
||||
n_batch: usize,
|
||||
k_seq: usize,
|
||||
stream: Arc<CudaStream>,
|
||||
}
|
||||
|
||||
impl MultiHorizonAttention {
|
||||
pub fn new(dev: &MlDevice, n_batch: usize, k_seq: usize, lr: f32, seed: u64) -> Result<Self> {
|
||||
anyhow::ensure!(n_batch > 0 && k_seq > 0, "n_batch and k_seq must be > 0");
|
||||
let stream = dev.cuda_stream().context("v2 stream")?.clone();
|
||||
let ctx = dev.cuda_context().context("v2 ctx")?;
|
||||
|
||||
let pool = HorizonTokenAttentionPool::new(ctx, stream.clone())
|
||||
.context("HorizonTokenAttentionPool")?;
|
||||
let inv_pool = InvertedAttentionPool::new(ctx, stream.clone())
|
||||
.context("InvertedAttentionPool")?;
|
||||
let moe = RegimeMoeGate::new(ctx, stream.clone())
|
||||
.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();
|
||||
|
||||
let init_horizon: Vec<f32> = (0..MHA_N_HORIZONS * MHA_HIDDEN_DIM)
|
||||
.map(|_| rng.gen_range(-scale..scale)).collect();
|
||||
let init_q: Vec<f32> = (0..MHA_HIDDEN_DIM)
|
||||
.map(|_| rng.gen_range(-scale..scale)).collect();
|
||||
let init_w_gate: Vec<f32> = (0..MHA_N_EXPERTS * MHA_REGIME_DIM)
|
||||
.map(|_| rng.gen_range(-0.1..0.1)).collect();
|
||||
let init_experts_w: Vec<f32> = (0..MHA_N_EXPERTS * MHA_HIDDEN_DIM * MHA_HIDDEN_DIM)
|
||||
.map(|_| rng.gen_range(-scale..scale)).collect();
|
||||
|
||||
let horizon_tokens_d = upload(&stream, &init_horizon)?;
|
||||
let horizon_tokens_init_d = upload(&stream, &init_horizon)?;
|
||||
let q_d = upload(&stream, &init_q)?;
|
||||
let q_init_d = upload(&stream, &init_q)?;
|
||||
let w_gate_d = upload(&stream, &init_w_gate)?;
|
||||
let experts_w_d = upload(&stream, &init_experts_w)?;
|
||||
let experts_b_d = stream.alloc_zeros::<f32>(MHA_N_EXPERTS * MHA_HIDDEN_DIM)?;
|
||||
let experts_w_init_d = upload(&stream, &init_experts_w)?;
|
||||
let experts_b_init_d = stream.alloc_zeros::<f32>(MHA_N_EXPERTS * MHA_HIDDEN_DIM)?;
|
||||
|
||||
let n_ext = MHA_N_HORIZONS + k_seq;
|
||||
|
||||
let s = &stream;
|
||||
let state = Self {
|
||||
pool,
|
||||
horizon_tokens_d,
|
||||
horizon_tokens_init_d,
|
||||
q_d,
|
||||
q_init_d,
|
||||
grad_horizon_tokens_d: s.alloc_zeros::<f32>(MHA_N_HORIZONS * MHA_HIDDEN_DIM)?,
|
||||
grad_horizon_tokens_scratch_d:s.alloc_zeros::<f32>(n_batch * MHA_N_HORIZONS * MHA_HIDDEN_DIM)?,
|
||||
grad_ctx_h_d: s.alloc_zeros::<f32>(n_batch * MHA_N_HORIZONS * MHA_HIDDEN_DIM)?,
|
||||
grad_q_d: s.alloc_zeros::<f32>(MHA_HIDDEN_DIM)?,
|
||||
grad_q_scratch_d: s.alloc_zeros::<f32>(n_batch * MHA_HIDDEN_DIM)?,
|
||||
ctx_h_d: s.alloc_zeros::<f32>(n_batch * MHA_N_HORIZONS * MHA_HIDDEN_DIM)?,
|
||||
attn_h_d: s.alloc_zeros::<f32>(n_batch * n_ext)?,
|
||||
opt_horizon_tokens: AdamW::new(dev, MHA_N_HORIZONS * MHA_HIDDEN_DIM, lr)?,
|
||||
opt_q: AdamW::new(dev, MHA_HIDDEN_DIM, lr)?,
|
||||
|
||||
inv_pool,
|
||||
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)?,
|
||||
w_gate_d,
|
||||
experts_w_d,
|
||||
experts_b_d,
|
||||
experts_w_init_d,
|
||||
experts_b_init_d,
|
||||
top_e_d: s.alloc_zeros::<i32>(n_batch)?,
|
||||
gate_probs_d: s.alloc_zeros::<f32>(n_batch * MHA_N_EXPERTS)?,
|
||||
aux_loss_d: s.alloc_zeros::<f32>(1)?,
|
||||
routed_ctx_d: s.alloc_zeros::<f32>(n_batch * MHA_N_HORIZONS * MHA_HIDDEN_DIM)?,
|
||||
grad_routed_d: s.alloc_zeros::<f32>(n_batch * MHA_N_HORIZONS * MHA_HIDDEN_DIM)?,
|
||||
grad_w_gate_d: s.alloc_zeros::<f32>(MHA_N_EXPERTS * MHA_REGIME_DIM)?,
|
||||
grad_experts_w_d: s.alloc_zeros::<f32>(MHA_N_EXPERTS * MHA_HIDDEN_DIM * MHA_HIDDEN_DIM)?,
|
||||
grad_experts_b_d: s.alloc_zeros::<f32>(MHA_N_EXPERTS * MHA_HIDDEN_DIM)?,
|
||||
// Compact: [B, N_H, H, H] / [B, N_H, H]. Per-batch
|
||||
// contribution scatters into top_e's slot of grad_experts via
|
||||
// `moe.scatter()` instead of writing into a 4× larger
|
||||
// sparse-by-expert scratch.
|
||||
grad_w_scratch_d: s.alloc_zeros::<f32>(n_batch * MHA_N_HORIZONS * MHA_HIDDEN_DIM * MHA_HIDDEN_DIM)?,
|
||||
grad_b_scratch_d: s.alloc_zeros::<f32>(n_batch * MHA_N_HORIZONS * MHA_HIDDEN_DIM)?,
|
||||
opt_w_gate: AdamW::new(dev, MHA_N_EXPERTS * MHA_REGIME_DIM, lr)?,
|
||||
opt_experts_w: AdamW::new(dev, MHA_N_EXPERTS * MHA_HIDDEN_DIM * MHA_HIDDEN_DIM, lr)?,
|
||||
opt_experts_b: AdamW::new(dev, MHA_N_EXPERTS * MHA_HIDDEN_DIM, lr)?,
|
||||
|
||||
log_sigma_h_d: s.alloc_zeros::<f32>(MHA_N_HORIZONS)?,
|
||||
grad_log_sigma_h_d: s.alloc_zeros::<f32>(MHA_N_HORIZONS)?,
|
||||
opt_log_sigma: AdamW::new(dev, MHA_N_HORIZONS, lr * 0.25)?, // slow update
|
||||
|
||||
collapse,
|
||||
ctx_mean_d: s.alloc_zeros::<f32>(n_batch * MHA_HIDDEN_DIM)?,
|
||||
grad_ctx_mean_d: s.alloc_zeros::<f32>(n_batch * MHA_HIDDEN_DIM)?,
|
||||
anchor_l2,
|
||||
anchor: AnchorController::new(
|
||||
magnitude_of(&init_horizon) + magnitude_of(&init_q),
|
||||
init_horizon.len() + init_q.len(),
|
||||
),
|
||||
lambda_d: s.alloc_zeros::<f32>(1)?,
|
||||
anchor_loss_partial_d: s.alloc_zeros::<f32>(1)?,
|
||||
anchor_loss_total_d: s.alloc_zeros::<f32>(1)?,
|
||||
|
||||
n_batch,
|
||||
k_seq,
|
||||
stream: stream.clone(),
|
||||
};
|
||||
Ok(state)
|
||||
}
|
||||
|
||||
/// Zero per-step grad scratches. Capture-safe (memset_zeros only).
|
||||
pub fn zero_grads(&mut self) -> Result<()> {
|
||||
let s = &self.stream;
|
||||
s.memset_zeros(&mut self.grad_horizon_tokens_d)?;
|
||||
s.memset_zeros(&mut self.grad_horizon_tokens_scratch_d)?;
|
||||
s.memset_zeros(&mut self.grad_ctx_h_d)?;
|
||||
s.memset_zeros(&mut self.grad_q_d)?;
|
||||
s.memset_zeros(&mut self.grad_q_scratch_d)?;
|
||||
s.memset_zeros(&mut self.grad_routed_d)?;
|
||||
s.memset_zeros(&mut self.grad_w_gate_d)?;
|
||||
s.memset_zeros(&mut self.grad_experts_w_d)?;
|
||||
s.memset_zeros(&mut self.grad_experts_b_d)?;
|
||||
s.memset_zeros(&mut self.grad_w_scratch_d)?;
|
||||
s.memset_zeros(&mut self.grad_b_scratch_d)?;
|
||||
s.memset_zeros(&mut self.grad_log_sigma_h_d)?;
|
||||
s.memset_zeros(&mut self.anchor_loss_partial_d)?;
|
||||
s.memset_zeros(&mut self.anchor_loss_total_d)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn n_batch(&self) -> usize { self.n_batch }
|
||||
pub fn k_seq(&self) -> usize { self.k_seq }
|
||||
|
||||
/// Forward path: ln_b_out [B, K, H] → ctx_mean_d [B, H] (seeds CfC h_old at k=0).
|
||||
///
|
||||
/// Internal flow:
|
||||
/// ctx_h = horizon_token_attention_pool(horizon_tokens, q, ln_b_out)
|
||||
/// inv_pooled = inverted_attention_pool(ln_b_out)
|
||||
/// fused_ctx[b, h, d] = ctx_h[b, h, d] + inv_pooled[b, d] (additive merge)
|
||||
/// routed_ctx = regime_moe_gate(gate_logits, fused_ctx, ...)
|
||||
/// ctx_mean[b, d] = (1/N_H) Σ_h routed_ctx[b, h, d]
|
||||
///
|
||||
/// `gate_logits_d` is filled by the caller before invoking forward
|
||||
/// (in Stage 2 we use a zero placeholder until V5 wires real
|
||||
/// regime features; all experts receive equal gate weight then).
|
||||
/// Capture-safe: no host branches, no synchronize, no host allocs.
|
||||
pub fn forward(&mut self, ln_b_out: &CudaSlice<f32>) -> Result<()> {
|
||||
let n = self.n_batch as i32;
|
||||
let k = self.k_seq as i32;
|
||||
|
||||
// 1. Horizon-token attention pool → ctx_h_d, attn_h_d.
|
||||
self.pool.forward(
|
||||
&self.horizon_tokens_d, &self.q_d, ln_b_out,
|
||||
n, k,
|
||||
&mut self.ctx_h_d, &mut self.attn_h_d,
|
||||
)?;
|
||||
|
||||
// 2. Inverted attention → inv_pooled_d, inv_attn_d.
|
||||
let inv_scale = 1.0_f32 / (self.k_seq as f32).sqrt();
|
||||
self.inv_pool.forward(
|
||||
ln_b_out, n, k, inv_scale,
|
||||
&mut self.inv_pooled_d, &mut self.inv_attn_d,
|
||||
)?;
|
||||
|
||||
// 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
|
||||
// zero_grads; all experts equal → top_e = 0 deterministically.)
|
||||
self.moe.forward(
|
||||
&self.gate_logits_d, &self.ctx_h_d,
|
||||
&self.experts_w_d, &self.experts_b_d,
|
||||
n,
|
||||
&mut self.routed_ctx_d, &mut self.top_e_d,
|
||||
)?;
|
||||
|
||||
// 5. Mean-collapse over horizon axis → ctx_mean_d [B, H].
|
||||
self.collapse.forward(&self.routed_ctx_d, n, &mut self.ctx_mean_d)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Backward path from `grad_ctx_mean [B, H]` (the CfC's grad on
|
||||
/// h_old at k=0). Accumulates gradients into horizon_tokens, q,
|
||||
/// experts_w/b, and grad_ln_out (+=, per-batch).
|
||||
pub fn backward(
|
||||
&mut self,
|
||||
ln_b_out: &CudaSlice<f32>,
|
||||
grad_ctx_mean: &CudaSlice<f32>,
|
||||
grad_ln_out: &mut CudaSlice<f32>,
|
||||
) -> Result<()> {
|
||||
let n = self.n_batch as i32;
|
||||
let k = self.k_seq as i32;
|
||||
|
||||
// 5'. Mean-collapse bwd: grad_routed += broadcast(grad_ctx_mean) / N_H.
|
||||
self.collapse.backward(grad_ctx_mean, n, &mut self.grad_routed_d)?;
|
||||
|
||||
// 4'. MoE bwd: writes per-(b,h) rank-1 contributions into
|
||||
// compact `grad_w_scratch_d [B, N_H, H, H]` and
|
||||
// `grad_b_scratch_d [B, N_H, H]`. d_fused_ctx → d_ctx_h.
|
||||
self.moe.backward(
|
||||
&self.ctx_h_d, &self.experts_w_d, &self.grad_routed_d, &self.top_e_d,
|
||||
n,
|
||||
&mut self.grad_w_scratch_d,
|
||||
&mut self.grad_b_scratch_d,
|
||||
&mut self.grad_ctx_h_d,
|
||||
)?;
|
||||
|
||||
// 4''. Scatter the per-(b,h) contributions into the per-expert
|
||||
// shared grad slots based on top_e[b]. Replaces the prior
|
||||
// naive reduce_axis0 over a sparse-by-expert scratch.
|
||||
self.moe.scatter(
|
||||
&self.grad_w_scratch_d, &self.grad_b_scratch_d, &self.top_e_d,
|
||||
n,
|
||||
&mut self.grad_experts_w_d,
|
||||
&mut self.grad_experts_b_d,
|
||||
)?;
|
||||
|
||||
// 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, &self.grad_inv_pooled_d,
|
||||
n, k, inv_scale,
|
||||
&mut self.inv_d_scores_scratch_d,
|
||||
grad_ln_out,
|
||||
)?;
|
||||
|
||||
// 1'. Horizon-token attention pool bwd: grad_ctx_h_d →
|
||||
// grad_horizon_tokens_scratch_d, grad_q_scratch_d, and
|
||||
// += grad_ln_out per-batch.
|
||||
self.pool.backward(
|
||||
&self.horizon_tokens_d, &self.q_d, ln_b_out,
|
||||
&self.attn_h_d,
|
||||
&self.grad_ctx_h_d,
|
||||
n, k,
|
||||
&mut self.grad_horizon_tokens_scratch_d,
|
||||
&mut self.grad_q_scratch_d,
|
||||
grad_ln_out,
|
||||
)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Apply L2 anchor regularization to all anchored param groups.
|
||||
/// Adds to `grad_horizon_tokens_d`, `grad_q_d`, `grad_experts_w_d`.
|
||||
/// `lambda_d` must hold the current λ (written by the caller on
|
||||
/// host, uploaded before launch — done outside the capture region).
|
||||
pub fn apply_anchor(&mut self) -> Result<()> {
|
||||
let n_ht = (MHA_N_HORIZONS * MHA_HIDDEN_DIM) as i32;
|
||||
let n_q = MHA_HIDDEN_DIM as i32;
|
||||
let n_ew = (MHA_N_EXPERTS * MHA_HIDDEN_DIM * MHA_HIDDEN_DIM) as i32;
|
||||
|
||||
self.anchor_l2.apply(
|
||||
&self.horizon_tokens_d, &self.horizon_tokens_init_d,
|
||||
&self.lambda_d, n_ht,
|
||||
&mut self.anchor_loss_partial_d, &mut self.grad_horizon_tokens_d,
|
||||
)?;
|
||||
self.anchor_l2.apply(
|
||||
&self.q_d, &self.q_init_d,
|
||||
&self.lambda_d, n_q,
|
||||
&mut self.anchor_loss_partial_d, &mut self.grad_q_d,
|
||||
)?;
|
||||
self.anchor_l2.apply(
|
||||
&self.experts_w_d, &self.experts_w_init_d,
|
||||
&self.lambda_d, n_ew,
|
||||
&mut self.anchor_loss_partial_d, &mut self.grad_experts_w_d,
|
||||
)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Step all owned optimizers. Capture-safe — each AdamW is its own
|
||||
/// device kernel; the trainer calls this after reducing scratches.
|
||||
pub fn adamw_step(&mut self) -> Result<()> {
|
||||
self.opt_horizon_tokens.step(&mut self.horizon_tokens_d, &self.grad_horizon_tokens_d)?;
|
||||
self.opt_q.step(&mut self.q_d, &self.grad_q_d)?;
|
||||
self.opt_w_gate.step(&mut self.w_gate_d, &self.grad_w_gate_d)?;
|
||||
self.opt_experts_w.step(&mut self.experts_w_d, &self.grad_experts_w_d)?;
|
||||
self.opt_experts_b.step(&mut self.experts_b_d, &self.grad_experts_b_d)?;
|
||||
self.opt_log_sigma.step(&mut self.log_sigma_h_d, &self.grad_log_sigma_h_d)?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
fn upload(stream: &Arc<CudaStream>, host: &[f32]) -> Result<CudaSlice<f32>> {
|
||||
use crate::pinned_mem::MappedF32Buffer;
|
||||
use cudarc::driver::{DevicePtrMut};
|
||||
let n = host.len();
|
||||
let staging = unsafe { MappedF32Buffer::new(n) }
|
||||
.map_err(|e| anyhow::anyhow!("v2_state upload staging: {e}"))?;
|
||||
staging.write_from_slice(host);
|
||||
let mut dst = stream.alloc_zeros::<f32>(n).context("v2_state upload alloc")?;
|
||||
if n > 0 {
|
||||
let nbytes = n * std::mem::size_of::<f32>();
|
||||
unsafe {
|
||||
let (dst_ptr, _g) = dst.device_ptr_mut(stream);
|
||||
cudarc::driver::result::memcpy_dtod_async(
|
||||
dst_ptr, staging.dev_ptr, nbytes, stream.cu_stream(),
|
||||
).context("v2_state upload DtoD")?;
|
||||
}
|
||||
}
|
||||
Ok(dst)
|
||||
}
|
||||
|
||||
fn magnitude_of(p: &[f32]) -> f32 {
|
||||
p.iter().map(|x| x * x).sum::<f32>().sqrt()
|
||||
}
|
||||
@@ -61,6 +61,7 @@ const BCE_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/bce_loss_mult
|
||||
const HORIZON_LAMBDA_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/horizon_lambda.cubin"));
|
||||
const LAYER_NORM_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/layer_norm.cubin"));
|
||||
const VARIABLE_SELECTION_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/variable_selection.cubin"));
|
||||
const ATTENTION_POOL_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/attention_pool.cubin"));
|
||||
const REDUCE_AXIS0_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/reduce_axis0.cubin"));
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
@@ -359,15 +360,30 @@ pub struct PerceptionTrainer {
|
||||
vsn_bwd_fn: CudaFunction,
|
||||
_vsn_module: Arc<CudaModule>,
|
||||
|
||||
// ── Attention path ──
|
||||
// MultiHorizonAttention is the single source of truth for the
|
||||
// attention summary that seeds CfC's h_old at k=0. Owns the
|
||||
// horizon-token attention pool, the inverted (cross-variate)
|
||||
// attention pass, the regime-MoE gate + experts, the mean-collapse
|
||||
// over the horizon axis, the Kendall σ logarithm fed into BCE, and
|
||||
// the L2 anchor controller + kernel. Replaces the prior single-Q
|
||||
// `attention_pool` entirely.
|
||||
pub mha: crate::trainer::multi_horizon_attention::MultiHorizonAttention,
|
||||
// ── Attention pool (Phase 3) ──
|
||||
// Single-head attention pool over LN_b output. Replaces CfC's
|
||||
// zero-initialised h_old at k=0 with a learned content-addressable
|
||||
// summary over all K positions. Single learned param: Q [HIDDEN_DIM].
|
||||
pub attn_q_d: CudaSlice<f32>,
|
||||
/// Per-sample pooled context `[B, HIDDEN_DIM]` — fed as h_old at k=0.
|
||||
attn_context_d: CudaSlice<f32>,
|
||||
/// Saved post-softmax attention weights `[B, K]` for bwd.
|
||||
attn_weights_d: CudaSlice<f32>,
|
||||
grad_attn_q_d: CudaSlice<f32>,
|
||||
pub opt_attn_q: AdamW,
|
||||
attn_fwd_fn: CudaFunction,
|
||||
attn_bwd_fn: CudaFunction,
|
||||
_attn_module: Arc<CudaModule>,
|
||||
|
||||
// ── Kendall σ-weighted BCE (axis A) ──
|
||||
// Per-horizon learnable `log σ_h` scalars. Weighted BCE per horizon
|
||||
// is `(1/(2·exp(2·log σ_h))) · mean_bce_h + log σ_h`. AdamW updates
|
||||
// log_sigma_h at 1/4 the head LR (slow update). σ on its own without
|
||||
// C/D/E is the keeper from the v2 sweep (see
|
||||
// project_ml_alpha_v2_ab_verdict.md).
|
||||
pub log_sigma_h_d: CudaSlice<f32>, // [N_HORIZONS]
|
||||
pub grad_log_sigma_h_d: CudaSlice<f32>, // [N_HORIZONS]
|
||||
pub opt_log_sigma: AdamW,
|
||||
|
||||
// ── K-loop parallelization (Phase B) ──
|
||||
// Per-batch grad scratch buffers for cfc_step_backward_batched.
|
||||
@@ -393,6 +409,8 @@ pub struct PerceptionTrainer {
|
||||
// VSN per-row grad scratch (Phase B commit 3). n_rows = B * K.
|
||||
vsn_grad_w_scratch_d: CudaSlice<f32>, // [B*K, FEATURE_DIM, FEATURE_DIM]
|
||||
vsn_grad_b_scratch_d: CudaSlice<f32>, // [B*K, FEATURE_DIM]
|
||||
// Attention pool per-batch grad scratch (Phase B commit 4).
|
||||
attn_grad_q_scratch_d: CudaSlice<f32>, // [B, HIDDEN_DIM]
|
||||
/// Cross-batch reducer kernel: `[B, N] → [N]` via block tree-reduce.
|
||||
/// Used for every per-batch grad scratch in the refactored bwd path.
|
||||
reduce_axis0_fn: CudaFunction,
|
||||
@@ -523,6 +541,15 @@ impl PerceptionTrainer {
|
||||
let vsn_bwd_fn = vsn_module
|
||||
.load_function("variable_selection_bwd")
|
||||
.context("variable_selection_bwd symbol")?;
|
||||
let attn_module = ctx
|
||||
.load_cubin(ATTENTION_POOL_CUBIN.to_vec())
|
||||
.context("attention_pool cubin")?;
|
||||
let attn_fwd_fn = attn_module
|
||||
.load_function("attention_pool_fwd")
|
||||
.context("attention_pool_fwd symbol")?;
|
||||
let attn_bwd_fn = attn_module
|
||||
.load_function("attention_pool_bwd")
|
||||
.context("attention_pool_bwd symbol")?;
|
||||
let reduce_module = ctx
|
||||
.load_cubin(REDUCE_AXIS0_CUBIN.to_vec())
|
||||
.context("reduce_axis0 cubin")?;
|
||||
@@ -772,20 +799,35 @@ impl PerceptionTrainer {
|
||||
cfg.n_batch * cfg.seq_len * FEATURE_DIM)?;
|
||||
|
||||
// ── Attention pool init (Phase 3) ──
|
||||
// MultiHorizonAttention bundle. Single source of truth for the
|
||||
// multi-horizon attention path + Kendall σ on the BCE. Init
|
||||
// captures snapshots used by the L2 anchor controller.
|
||||
// Stage 1 (this commit) wires `mha.log_sigma_h_d` and
|
||||
// `mha.grad_log_sigma_h_d` into the BCE launch. Stage 2
|
||||
// replaces the legacy `attention_pool` forward/backward
|
||||
// callsites with `mha.pool` + `mha.inv_pool` + `mha.moe`.
|
||||
let mha = crate::trainer::multi_horizon_attention::MultiHorizonAttention::new(
|
||||
dev,
|
||||
cfg.n_batch,
|
||||
cfg.seq_len,
|
||||
cfg.lr_cfc,
|
||||
cfg.seed.wrapping_add(0xB200_C00B),
|
||||
)?;
|
||||
// Q init near zero so initial scores ≈ 0 → softmax ≈ uniform 1/K
|
||||
// → context ≈ mean of LN_b output. Model learns content
|
||||
// addressing from a near-uniform starting point.
|
||||
let attn_q_scale = (1.0_f32 / HIDDEN_DIM as f32).sqrt();
|
||||
let attn_q_init: Vec<f32> = (0..HIDDEN_DIM)
|
||||
.map(|_| r.gen_range(-attn_q_scale..attn_q_scale)).collect();
|
||||
let attn_q_d = upload(&stream, &attn_q_init)?;
|
||||
let attn_context_d = stream.alloc_zeros::<f32>(cfg.n_batch * HIDDEN_DIM)?;
|
||||
let attn_weights_d = stream.alloc_zeros::<f32>(cfg.n_batch * cfg.seq_len)?;
|
||||
let grad_attn_q_d = stream.alloc_zeros::<f32>(HIDDEN_DIM)?;
|
||||
let opt_attn_q = AdamW::new(dev, HIDDEN_DIM, cfg.lr_cfc)?;
|
||||
// Phase B: attn pool per-batch grad scratch.
|
||||
let attn_grad_q_scratch_d = stream.alloc_zeros::<f32>(cfg.n_batch * HIDDEN_DIM)?;
|
||||
|
||||
// Kendall σ — per-horizon learnable `log σ_h` (axis A only;
|
||||
// C/D/E reverted after the 2026-05-18 A/B sweep). Init at 0
|
||||
// → weight = 1/(2·exp(0)) = 0.5 per horizon at step 0; AdamW
|
||||
// updates at 1/4 the head LR per the spec.
|
||||
let log_sigma_h_d = stream.alloc_zeros::<f32>(N_HORIZONS)?;
|
||||
let grad_log_sigma_h_d = stream.alloc_zeros::<f32>(N_HORIZONS)?;
|
||||
let opt_log_sigma = AdamW::new(dev, N_HORIZONS, cfg.lr_cfc * 0.25)?;
|
||||
|
||||
// V1 (2026-05-18): per-horizon Q_h trainer state (C24/C25) removed.
|
||||
// A/B sweep at commit 83546b5c3 falsified the approach (mean_auc
|
||||
// -0.019 vs baseline; h6000 essentially tied). v2 design with
|
||||
// horizon-token K-prepend + inverted attention + regime-MoE +
|
||||
// Kendall σ-BCE + L2 anchor replaces it. See
|
||||
// docs/superpowers/specs/2026-05-18-ml-alpha-v2-multi-horizon-design.md
|
||||
// and the V1-V13 plan.
|
||||
|
||||
let k = cfg.seq_len;
|
||||
Ok(Self {
|
||||
@@ -834,7 +876,17 @@ impl PerceptionTrainer {
|
||||
vsn_bwd_fn,
|
||||
_vsn_module: vsn_module,
|
||||
// Attention pool (Phase 3).
|
||||
mha,
|
||||
attn_q_d,
|
||||
attn_context_d,
|
||||
attn_weights_d,
|
||||
grad_attn_q_d,
|
||||
opt_attn_q,
|
||||
attn_fwd_fn,
|
||||
attn_bwd_fn,
|
||||
_attn_module: attn_module,
|
||||
log_sigma_h_d,
|
||||
grad_log_sigma_h_d,
|
||||
opt_log_sigma,
|
||||
// Phase B: cfc per-batch grad scratch + reducer.
|
||||
cfc_grad_w_in_scratch_d,
|
||||
cfc_grad_w_rec_scratch_d,
|
||||
@@ -852,6 +904,7 @@ impl PerceptionTrainer {
|
||||
grn_grad_b_skip_scratch_d,
|
||||
vsn_grad_w_scratch_d,
|
||||
vsn_grad_b_scratch_d,
|
||||
attn_grad_q_scratch_d,
|
||||
reduce_axis0_fn,
|
||||
_reduce_module: reduce_module,
|
||||
loss_ema_d: stream.alloc_zeros::<f32>(N_HORIZONS)?,
|
||||
@@ -1397,16 +1450,30 @@ impl PerceptionTrainer {
|
||||
unsafe { launch.launch(cfg_tx).context("transpose h_enriched fwd")?; }
|
||||
}
|
||||
|
||||
// ── 2d. Multi-horizon attention forward — single source of
|
||||
// truth for the attention path. Replaces the legacy
|
||||
// single-Q attention_pool. Internal flow:
|
||||
// horizon-token attn → ctx_h [B, N_H, H]
|
||||
// inverted attn → inv_pooled [B, H]
|
||||
// MoE gate + dispatch → routed_ctx [B, N_H, H]
|
||||
// mean-collapse → mha.ctx_mean_d [B, H]
|
||||
// The CfC's h_old at k=0 reads from mha.ctx_mean_d
|
||||
// (legacy attn_context_d slot).
|
||||
self.mha.forward(&self.ln_out_d)?;
|
||||
// ── 2d. Attention pool forward (Phase 3) — produces
|
||||
// attn_context_d [B, HIDDEN_DIM] = learned content-summary
|
||||
// over all K LN_b output positions. Replaces CfC's
|
||||
// zero-initialised h_old at k=0 with this context vector.
|
||||
// Shared mem: K floats (scores) + BLOCK floats (reduce) +
|
||||
// HIDDEN_DIM floats (context) = (K + 128 + 128) * 4 bytes.
|
||||
{
|
||||
let k_i32 = k_seq as i32;
|
||||
let n_batch_attn = b_sz as i32;
|
||||
let shared = (k_seq + 128 + HIDDEN_DIM) * std::mem::size_of::<f32>();
|
||||
let cfg_attn = LaunchConfig {
|
||||
grid_dim: (b_sz as u32, 1, 1),
|
||||
block_dim: (128, 1, 1), // ATTN_BLOCK
|
||||
shared_mem_bytes: shared as u32,
|
||||
};
|
||||
let mut launch = self.stream.launch_builder(&self.attn_fwd_fn);
|
||||
launch
|
||||
.arg(&self.attn_q_d)
|
||||
.arg(&self.ln_out_d)
|
||||
.arg(&n_batch_attn).arg(&k_i32)
|
||||
.arg(&mut self.attn_context_d)
|
||||
.arg(&mut self.attn_weights_d);
|
||||
unsafe { launch.launch(cfg_attn).context("attention_pool_fwd")?; }
|
||||
}
|
||||
|
||||
// ── 3. Labels DtoD: staging (filled in step_batched before
|
||||
// captured region) → device.
|
||||
@@ -1431,10 +1498,9 @@ impl PerceptionTrainer {
|
||||
.map_err(|e| anyhow::anyhow!("zero cfc_grad_b_scratch: {e}"))?;
|
||||
self.stream.memset_zeros(&mut self.cfc_grad_tau_scratch_d)
|
||||
.map_err(|e| anyhow::anyhow!("zero cfc_grad_tau_scratch: {e}"))?;
|
||||
// MHA grad scratches (horizon-token, q, expert-W/b, anchor
|
||||
// loss partials, log-sigma). Capture-safe memset_zeros.
|
||||
self.mha.zero_grads()?;
|
||||
|
||||
// Phase B commit 4: attention pool per-batch grad_Q scratch.
|
||||
self.stream.memset_zeros(&mut self.attn_grad_q_scratch_d)
|
||||
.map_err(|e| anyhow::anyhow!("zero attn_grad_q_scratch: {e}"))?;
|
||||
// GRN per-batch grad scratch (Phase B commit 2): zero ONCE per
|
||||
// step; K-loop bwd accumulates into these, then reduce_axis0
|
||||
// collapses → final grad buffers (OVERWRITE) after the K-loop.
|
||||
@@ -1535,16 +1601,16 @@ impl PerceptionTrainer {
|
||||
// [K, B, H] layout. Pointer-offset trick (single mut
|
||||
// borrow + raw u64 arithmetic for slot pointers) keeps
|
||||
// the launches stream-ordered with no syncs.
|
||||
// MHA produces mha.ctx_mean_d which seeds h_old at k=0. The
|
||||
// legacy attn_context_d slot is removed. zero_h_ptr remains for
|
||||
// the bwd k=0 case (cfc bwd needs an h_old slot for the input
|
||||
// gradient to read from).
|
||||
// Phase 3: attention pool produces attn_context_d which replaces
|
||||
// the zero initial h_old at k=0. zero_h_ptr is no longer used in
|
||||
// the fwd K-loop but is retained for the bwd k=0 case (cfc bwd
|
||||
// needs an h_old slot for the input gradient to read from).
|
||||
let _zero_h_ptr_unused = {
|
||||
let (p, _g) = self.zero_h_d.device_ptr(&self.stream);
|
||||
p
|
||||
};
|
||||
let attn_context_ptr = {
|
||||
let (p, _g) = self.mha.ctx_mean_d.device_ptr(&self.stream);
|
||||
let (p, _g) = self.attn_context_d.device_ptr(&self.stream);
|
||||
p
|
||||
};
|
||||
let henr_t_base = {
|
||||
@@ -1606,12 +1672,12 @@ impl PerceptionTrainer {
|
||||
}
|
||||
drop((_g_hpk, _g_probs, _g_z1, _g_a1, _g_z2, _g_gate, _g_main, _g_logit));
|
||||
|
||||
// ── 4.5. (v2 reserved) — per-horizon Q_h path removed at V1.
|
||||
// The v2 design (horizon-token K-prepend + inverted attn +
|
||||
// regime-MoE) wires in here in commits V9/V10.
|
||||
// ── 4.5. Zero σ grad scratch each step (capture-safe).
|
||||
self.stream.memset_zeros(&mut self.grad_log_sigma_h_d)
|
||||
.map_err(|e| anyhow::anyhow!("zero grad_log_sigma_h: {e}"))?;
|
||||
|
||||
// ── 5. Fused multi-horizon BCE over the full [K*B, N_HORIZONS]
|
||||
// grid. The kernel doesn't distinguish position-vs-batch
|
||||
// ── 5. Fused multi-horizon BCE (Kendall σ-weighted, axis A).
|
||||
// The kernel doesn't distinguish position-vs-batch
|
||||
// since the per-horizon weight is identified by `i % n_horizons`.
|
||||
let n_pos_i = (k_seq * b_sz) as i32;
|
||||
let n_h_i = N_HORIZONS as i32;
|
||||
@@ -1625,16 +1691,19 @@ impl PerceptionTrainer {
|
||||
launch
|
||||
.arg(&self.probs_per_k_d).arg(&self.labels_per_k_d)
|
||||
.arg(&self.loss_weights_d)
|
||||
.arg(&self.mha.log_sigma_h_d)
|
||||
.arg(&self.log_sigma_h_d)
|
||||
.arg(&n_pos_i).arg(&n_h_i)
|
||||
.arg(&mut self.loss_d)
|
||||
.arg(&mut self.loss_per_horizon_d)
|
||||
.arg(&mut self.grad_probs_per_k_d)
|
||||
.arg(&mut self.valid_d)
|
||||
.arg(&mut self.mha.grad_log_sigma_h_d);
|
||||
.arg(&mut self.grad_log_sigma_h_d);
|
||||
unsafe { launch.launch(bce_cfg).context("bce launch")?; }
|
||||
}
|
||||
|
||||
// ── 5a. (v2 reserved) — per-horizon Q_h backward removed at V1.
|
||||
// v2 backward path will live here in commit V10.
|
||||
|
||||
// ── 5b. ISV-driven per-horizon EMA + lambda. Updates the EMA
|
||||
// of unweighted per-horizon BCE and emits a clamped
|
||||
// multiplier `lambda_d[h]` per
|
||||
@@ -1661,7 +1730,7 @@ impl PerceptionTrainer {
|
||||
self.stream.memset_zeros(&mut self.grad_h_carry_d)
|
||||
.map_err(|e| anyhow::anyhow!("zero grad_h_carry: {e}"))?;
|
||||
|
||||
// MHA: at k=0 the bwd kernel reads `h_old` = mha.ctx_mean_d
|
||||
// Phase 3: at k=0 the bwd kernel reads `h_old` = attn_context_d
|
||||
// (mirroring the forward pass). zero_h_ptr_bwd retained as a
|
||||
// legacy fallback / unused alias.
|
||||
let _zero_h_ptr_bwd_unused = {
|
||||
@@ -1669,7 +1738,7 @@ impl PerceptionTrainer {
|
||||
p
|
||||
};
|
||||
let attn_context_ptr_bwd = {
|
||||
let (p, _g) = self.mha.ctx_mean_d.device_ptr(&self.stream);
|
||||
let (p, _g) = self.attn_context_d.device_ptr(&self.stream);
|
||||
p
|
||||
};
|
||||
let henr_t_base_bwd = {
|
||||
@@ -1784,53 +1853,60 @@ impl PerceptionTrainer {
|
||||
unsafe { launch.launch(cfg_tx).context("transpose grad bwd")?; }
|
||||
}
|
||||
|
||||
// ── 7c-pre. MultiHorizonAttention backward. Consumes:
|
||||
// grad_h_carry_d (= grad on initial h_old at k=0, which IS
|
||||
// mha.ctx_mean_d in fwd) → fed as grad_ctx_mean.
|
||||
// Accumulates into:
|
||||
// grad_horizon_tokens_scratch_d, grad_q_scratch_d (per-batch)
|
||||
// grad_w_scratch_d, grad_b_scratch_d (per-batch, sparse-by-expert)
|
||||
// grad_h_enriched_seq_d (LN_b out grad) += MHA's attention-path
|
||||
// contribution. Same += semantics as the legacy attn_bwd.
|
||||
let grad_h_enriched_slice = self.grad_h_enriched_seq_d.data_mut();
|
||||
self.mha.backward(&self.ln_out_d, &self.grad_h_carry_d, grad_h_enriched_slice)?;
|
||||
|
||||
// Reduce MHA's per-batch grad scratches → shared grads.
|
||||
// grad_horizon_tokens_scratch_d [B, N_H, H] → reduce_axis0 → [N_H, H]
|
||||
// grad_q_scratch_d [B, H] → reduce_axis0 → [H]
|
||||
// grad_w_scratch_d / grad_b_scratch_d are already scattered into
|
||||
// grad_experts_w_d / grad_experts_b_d by `mha.moe.scatter()` inside
|
||||
// `mha.backward()` (compact layout — see regime_moe_gate.cu).
|
||||
let n_batch_i = b_sz as i32;
|
||||
let cfg_red = |n_tail: i32| LaunchConfig {
|
||||
grid_dim: (n_tail as u32, 1, 1),
|
||||
block_dim: (256, 1, 1),
|
||||
shared_mem_bytes: 0,
|
||||
};
|
||||
let n_ht_i = (5 * HIDDEN_DIM) as i32;
|
||||
let n_q_i = HIDDEN_DIM as i32;
|
||||
// ── 7c-pre. Attention pool backward (Phase 3). Consumes:
|
||||
// Q = self.attn_q_d [HIDDEN_DIM]
|
||||
// ln_out (values)= self.ln_out_d [B, K, HIDDEN_DIM]
|
||||
// attn_weights = self.attn_weights_d (saved by fwd) [B, K]
|
||||
// grad_context = self.grad_h_carry_d (= grad on initial
|
||||
// h_old at k=0, which IS attn_context) [B, HIDDEN_DIM]
|
||||
// Writes (BOTH ARE +=):
|
||||
// grad_attn_q_d += attn-path contribution to Q
|
||||
// grad_h_enriched_seq_d (LN_b output grad) += attn-path
|
||||
// contribution to ln_out
|
||||
// The pre-zero of grad_attn_q_d at step start makes the += a
|
||||
// clean overwrite for the Q grad. grad_h_enriched_seq_d already
|
||||
// holds the K-loop's contribution at this point — attn's
|
||||
// contribution adds on top.
|
||||
// Phase B commit 4: block-per-batch attn bwd writes per-batch
|
||||
// grad_Q scratch; reducer collapses → final grad_attn_q_d below.
|
||||
{
|
||||
let mut launch = self.stream.launch_builder(&self.reduce_axis0_fn);
|
||||
launch.arg(&self.mha.grad_horizon_tokens_scratch_d)
|
||||
.arg(&n_batch_i).arg(&n_ht_i)
|
||||
.arg(&mut self.mha.grad_horizon_tokens_d);
|
||||
unsafe { launch.launch(cfg_red(n_ht_i)).context("reduce grad_horizon_tokens")?; }
|
||||
let k_i32 = k_seq as i32;
|
||||
let n_batch_attn = b_sz as i32;
|
||||
let shared = (2 * k_seq + 128) * std::mem::size_of::<f32>();
|
||||
let cfg_attn_bwd = LaunchConfig {
|
||||
grid_dim: (b_sz as u32, 1, 1),
|
||||
block_dim: (128, 1, 1), // ATTN_BLOCK
|
||||
shared_mem_bytes: shared as u32,
|
||||
};
|
||||
let mut launch = self.stream.launch_builder(&self.attn_bwd_fn);
|
||||
launch
|
||||
.arg(&self.attn_q_d)
|
||||
.arg(&self.ln_out_d)
|
||||
.arg(&self.attn_weights_d)
|
||||
.arg(&self.grad_h_carry_d)
|
||||
.arg(&n_batch_attn).arg(&k_i32)
|
||||
.arg(&mut self.attn_grad_q_scratch_d)
|
||||
.arg(self.grad_h_enriched_seq_d.data_mut());
|
||||
unsafe { launch.launch(cfg_attn_bwd).context("attention_pool_bwd")?; }
|
||||
}
|
||||
// Attn pool reducer: collapse [B, HIDDEN_DIM] → [HIDDEN_DIM].
|
||||
{
|
||||
let n_batch_i = b_sz as i32;
|
||||
let n_tail_i = HIDDEN_DIM as i32;
|
||||
let cfg_red = LaunchConfig {
|
||||
grid_dim: (HIDDEN_DIM as u32, 1, 1),
|
||||
block_dim: (256, 1, 1),
|
||||
shared_mem_bytes: 0,
|
||||
};
|
||||
let mut launch = self.stream.launch_builder(&self.reduce_axis0_fn);
|
||||
launch.arg(&self.mha.grad_q_scratch_d)
|
||||
.arg(&n_batch_i).arg(&n_q_i)
|
||||
.arg(&mut self.mha.grad_q_d);
|
||||
unsafe { launch.launch(cfg_red(n_q_i)).context("reduce grad_q")?; }
|
||||
launch
|
||||
.arg(&self.attn_grad_q_scratch_d)
|
||||
.arg(&n_batch_i)
|
||||
.arg(&n_tail_i)
|
||||
.arg(&mut self.grad_attn_q_d);
|
||||
unsafe { launch.launch(cfg_red).context("reduce attn_grad_q")?; }
|
||||
}
|
||||
|
||||
// Anchor L2 launches — apply on horizon_tokens, Q, experts_w.
|
||||
// λ_d is written on host (outside captured region) before
|
||||
// step_batched. anchor_loss_total is accumulated into one
|
||||
// scalar (not currently read by the trainer; reserved for
|
||||
// logging).
|
||||
self.mha.apply_anchor()?;
|
||||
|
||||
// ── 7c. LayerNorm B backward (between m2 and CfC). Consumes:
|
||||
// x = m2.h_enriched_seq [B, K, H]
|
||||
// gain = self.ln_gain_d [H]
|
||||
@@ -2139,7 +2215,9 @@ impl PerceptionTrainer {
|
||||
self.opt_ln_a_bias.step(&mut self.ln_a_bias_d, &self.grad_ln_a_bias_d)?;
|
||||
self.opt_vsn_w.step(&mut self.vsn_w_d, &self.grad_vsn_w_d)?;
|
||||
self.opt_vsn_b.step(&mut self.vsn_b_d, &self.grad_vsn_b_d)?;
|
||||
self.mha.adamw_step()?;
|
||||
self.opt_attn_q.step(&mut self.attn_q_d, &self.grad_attn_q_d)?;
|
||||
// Kendall σ step (axis A; 1/4 head LR).
|
||||
self.opt_log_sigma.step(&mut self.log_sigma_h_d, &self.grad_log_sigma_h_d)?;
|
||||
|
||||
// (v2 reserved) — per-horizon AdamW step removed at V1; v2's six
|
||||
// optimizer groups (horizon_tokens, Q_inv, w_fuse + b_fuse,
|
||||
@@ -2394,9 +2472,26 @@ impl PerceptionTrainer {
|
||||
unsafe { launch.launch(cfg_tx).context("eval transpose h_enriched")?; }
|
||||
}
|
||||
|
||||
// MultiHorizonAttention fwd — same as training. Produces
|
||||
// mha.ctx_mean_d which seeds h_old at k=0 in the eval K-loop.
|
||||
self.mha.forward(&self.ln_out_d)?;
|
||||
// Attention pool fwd — same as training, populates attn_context_d
|
||||
// for use as k=0 h_old in the eval K-loop below.
|
||||
{
|
||||
let k_i32 = k_seq as i32;
|
||||
let n_batch_attn = b_sz as i32;
|
||||
let shared = (k_seq + 128 + HIDDEN_DIM) * std::mem::size_of::<f32>();
|
||||
let cfg_attn = LaunchConfig {
|
||||
grid_dim: (b_sz as u32, 1, 1),
|
||||
block_dim: (128, 1, 1),
|
||||
shared_mem_bytes: shared as u32,
|
||||
};
|
||||
let mut launch = self.stream.launch_builder(&self.attn_fwd_fn);
|
||||
launch
|
||||
.arg(&self.attn_q_d)
|
||||
.arg(&self.ln_out_d)
|
||||
.arg(&n_batch_attn).arg(&k_i32)
|
||||
.arg(&mut self.attn_context_d)
|
||||
.arg(&mut self.attn_weights_d);
|
||||
unsafe { launch.launch(cfg_attn).context("eval attention_pool_fwd")?; }
|
||||
}
|
||||
|
||||
// Upload labels [K, B, N_HORIZONS].
|
||||
let total_labels = k_seq * b_sz * N_HORIZONS;
|
||||
@@ -2448,9 +2543,9 @@ impl PerceptionTrainer {
|
||||
let (p, _g) = self.zero_h_d.device_ptr(&self.stream);
|
||||
p
|
||||
};
|
||||
// Eval uses mha.ctx_mean_d as h_old at k=0, mirroring training.
|
||||
// Phase 3: eval uses attn_context_d as h_old at k=0, mirroring training.
|
||||
let attn_context_ptr = {
|
||||
let (p, _g) = self.mha.ctx_mean_d.device_ptr(&self.stream);
|
||||
let (p, _g) = self.attn_context_d.device_ptr(&self.stream);
|
||||
p
|
||||
};
|
||||
let henr_t_base = {
|
||||
@@ -2516,12 +2611,12 @@ impl PerceptionTrainer {
|
||||
let mut launch = self.stream.launch_builder(&self.bce_fn);
|
||||
launch
|
||||
.arg(&self.probs_per_k_d).arg(&self.labels_per_k_d).arg(&self.loss_weights_d)
|
||||
.arg(&self.mha.log_sigma_h_d)
|
||||
.arg(&self.log_sigma_h_d)
|
||||
.arg(&n_pos_i).arg(&n_h_i)
|
||||
.arg(&mut self.loss_d)
|
||||
.arg(&mut self.loss_per_horizon_d)
|
||||
.arg(&mut self.grad_probs_per_k_d).arg(&mut self.valid_d)
|
||||
.arg(&mut self.mha.grad_log_sigma_h_d);
|
||||
.arg(&mut self.grad_log_sigma_h_d);
|
||||
launch.launch(bce_cfg).context("eval bce")?;
|
||||
}
|
||||
self.stream.synchronize().context("eval sync")?;
|
||||
|
||||
@@ -1,138 +0,0 @@
|
||||
//! Numgrad parity test for anchor_l2 (v2 axis B).
|
||||
|
||||
#![cfg(feature = "cuda")]
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use cudarc::driver::{CudaSlice, CudaStream, DevicePtr, DevicePtrMut};
|
||||
use ml_alpha::anchor_l2::AnchorL2;
|
||||
use ml_alpha::pinned_mem::MappedF32Buffer;
|
||||
use ml_core::device::MlDevice;
|
||||
use std::sync::Arc;
|
||||
|
||||
const N: usize = 67; // intentionally non-power-of-two to stress stride loops
|
||||
|
||||
fn rng(seed: u64) -> impl FnMut() -> f32 {
|
||||
let mut s = seed.wrapping_mul(0x9E37_79B9_7F4A_7C15);
|
||||
move || -> f32 {
|
||||
s = s.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
|
||||
let x = ((s >> 16) & 0xFFFF) as f32 / 65536.0;
|
||||
x * 2.0 - 1.0
|
||||
}
|
||||
}
|
||||
|
||||
fn upload(stream: &Arc<CudaStream>, host: &[f32]) -> Result<CudaSlice<f32>> {
|
||||
let n = host.len();
|
||||
let staging = unsafe { MappedF32Buffer::new(n) }
|
||||
.map_err(|e| anyhow::anyhow!("upload: {e}"))?;
|
||||
staging.write_from_slice(host);
|
||||
let mut dst = stream.alloc_zeros::<f32>(n)?;
|
||||
let nbytes = n * std::mem::size_of::<f32>();
|
||||
unsafe {
|
||||
let (dst_ptr, _g) = dst.device_ptr_mut(stream);
|
||||
cudarc::driver::result::memcpy_dtod_async(
|
||||
dst_ptr, staging.dev_ptr, nbytes, stream.cu_stream(),
|
||||
)?;
|
||||
}
|
||||
Ok(dst)
|
||||
}
|
||||
|
||||
fn download(stream: &Arc<CudaStream>, src: &CudaSlice<f32>) -> Result<Vec<f32>> {
|
||||
let n = src.len();
|
||||
let staging = unsafe { MappedF32Buffer::new(n) }
|
||||
.map_err(|e| anyhow::anyhow!("download: {e}"))?;
|
||||
let nbytes = n * std::mem::size_of::<f32>();
|
||||
unsafe {
|
||||
let (src_ptr, _g) = src.device_ptr(stream);
|
||||
cudarc::driver::result::memcpy_dtod_async(
|
||||
staging.dev_ptr, src_ptr, nbytes, stream.cu_stream(),
|
||||
)?;
|
||||
}
|
||||
stream.synchronize()?;
|
||||
Ok(staging.read_all())
|
||||
}
|
||||
|
||||
fn forward_loss(
|
||||
anchor: &AnchorL2,
|
||||
stream: &Arc<CudaStream>,
|
||||
p: &[f32],
|
||||
p_init: &[f32],
|
||||
lambda: f32,
|
||||
) -> Result<f32> {
|
||||
let p_d = upload(stream, p)?;
|
||||
let p_init_d = upload(stream, p_init)?;
|
||||
let lambda_d = upload(stream, &[lambda])?;
|
||||
let mut loss_d = stream.alloc_zeros::<f32>(1)?;
|
||||
let mut grad_d = stream.alloc_zeros::<f32>(p.len())?;
|
||||
anchor.apply(&p_d, &p_init_d, &lambda_d, p.len() as i32, &mut loss_d, &mut grad_d)?;
|
||||
stream.synchronize()?;
|
||||
Ok(download(stream, &loss_d)?[0])
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[ignore = "requires CUDA"]
|
||||
fn forward_and_grad_match_central_difference() -> Result<()> {
|
||||
let dev = MlDevice::cuda(0)?;
|
||||
let stream = dev.cuda_stream().context("stream")?.clone();
|
||||
let ctx_dev = dev.cuda_context().context("ctx")?;
|
||||
let anchor = AnchorL2::new(ctx_dev, stream.clone())?;
|
||||
|
||||
let mut r = rng(20260518);
|
||||
let p: Vec<f32> = (0..N).map(|_| r() * 0.5).collect();
|
||||
let p_init: Vec<f32> = (0..N).map(|_| r() * 0.5).collect();
|
||||
let lambda = 0.137_f32;
|
||||
|
||||
// 1. Forward parity: kernel loss == closed-form.
|
||||
let p_d = upload(&stream, &p)?;
|
||||
let p_init_d = upload(&stream, &p_init)?;
|
||||
let lambda_d = upload(&stream, &[lambda])?;
|
||||
let mut loss_d = stream.alloc_zeros::<f32>(1)?;
|
||||
let mut grad_d = stream.alloc_zeros::<f32>(N)?;
|
||||
anchor.apply(&p_d, &p_init_d, &lambda_d, N as i32, &mut loss_d, &mut grad_d)?;
|
||||
stream.synchronize()?;
|
||||
let kernel_loss = download(&stream, &loss_d)?[0];
|
||||
let kernel_grad = download(&stream, &grad_d)?;
|
||||
|
||||
let mut expected_loss = 0.0_f32;
|
||||
for i in 0..N {
|
||||
let d = p[i] - p_init[i];
|
||||
expected_loss += d * d;
|
||||
}
|
||||
expected_loss *= lambda;
|
||||
let abs = (kernel_loss - expected_loss).abs();
|
||||
let rel = abs / expected_loss.abs().max(1e-3);
|
||||
assert!(
|
||||
abs < 1e-3 || rel < 5e-3,
|
||||
"loss closed-form mismatch: kernel={kernel_loss:.5} expected={expected_loss:.5}"
|
||||
);
|
||||
|
||||
// 2. Closed-form grad parity (analytic = 2λ (p − p_init)).
|
||||
for i in 0..N {
|
||||
let expected_grad = 2.0 * lambda * (p[i] - p_init[i]);
|
||||
let abs = (kernel_grad[i] - expected_grad).abs();
|
||||
assert!(abs < 1e-4,
|
||||
"grad[{i}] closed-form mismatch: kernel={} expected={} (diff={abs:.3e})",
|
||||
kernel_grad[i], expected_grad);
|
||||
}
|
||||
|
||||
// 3. Central-difference numgrad — 4 random positions.
|
||||
let eps = 1e-3_f32;
|
||||
let mut r2 = rng(42);
|
||||
for _ in 0..4 {
|
||||
let idx = ((r2() + 1.0) * 0.5 * N as f32) as usize % N;
|
||||
let mut p_p = p.clone();
|
||||
let mut p_m = p.clone();
|
||||
p_p[idx] += eps;
|
||||
p_m[idx] -= eps;
|
||||
let l_p = forward_loss(&anchor, &stream, &p_p, &p_init, lambda)?;
|
||||
let l_m = forward_loss(&anchor, &stream, &p_m, &p_init, lambda)?;
|
||||
let numgrad = (l_p - l_m) / (2.0 * eps);
|
||||
let analytic = kernel_grad[idx];
|
||||
let abs = (analytic - numgrad).abs();
|
||||
let rel = abs / numgrad.abs().max(1e-3);
|
||||
assert!(
|
||||
abs < 5e-3 || rel < 5e-2,
|
||||
"p[{idx}] numgrad mismatch: analytic={analytic:.4} numgrad={numgrad:.4} abs={abs:.3e}"
|
||||
);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
@@ -1,212 +0,0 @@
|
||||
//! Numgrad parity test for horizon_token_attention_pool (v2 axis C).
|
||||
//!
|
||||
//! Verifies forward + backward by central-difference comparison on
|
||||
//! horizon_tokens, Q, and ln_out within the established 5e-2 rel-tol /
|
||||
//! 5e-3 abs-floor envelope.
|
||||
|
||||
#![cfg(feature = "cuda")]
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use cudarc::driver::{CudaSlice, CudaStream, DevicePtr, DevicePtrMut};
|
||||
use ml_alpha::horizon_token_attention_pool::{
|
||||
HorizonTokenAttentionPool, HTAP_HIDDEN_DIM, HTAP_N_HORIZONS,
|
||||
};
|
||||
use ml_alpha::pinned_mem::MappedF32Buffer;
|
||||
use ml_core::device::MlDevice;
|
||||
use std::sync::Arc;
|
||||
|
||||
const B: usize = 2;
|
||||
const K: usize = 8;
|
||||
const H: usize = HTAP_HIDDEN_DIM;
|
||||
const N_H: usize = HTAP_N_HORIZONS;
|
||||
const N_EXT: usize = N_H + K;
|
||||
|
||||
fn deterministic_rng(seed: u64) -> impl FnMut() -> f32 {
|
||||
let mut s = seed.wrapping_mul(0x9E37_79B9_7F4A_7C15);
|
||||
move || -> f32 {
|
||||
s = s
|
||||
.wrapping_mul(6364136223846793005)
|
||||
.wrapping_add(1442695040888963407);
|
||||
let x = ((s >> 16) & 0xFFFF) as f32 / 65536.0;
|
||||
x * 2.0 - 1.0 // (-1, 1)
|
||||
}
|
||||
}
|
||||
|
||||
fn upload(stream: &Arc<CudaStream>, host: &[f32]) -> Result<CudaSlice<f32>> {
|
||||
let n = host.len();
|
||||
let staging = unsafe { MappedF32Buffer::new(n) }
|
||||
.map_err(|e| anyhow::anyhow!("upload staging: {e}"))?;
|
||||
staging.write_from_slice(host);
|
||||
let mut dst = stream.alloc_zeros::<f32>(n)?;
|
||||
let nbytes = n * std::mem::size_of::<f32>();
|
||||
unsafe {
|
||||
let (dst_ptr, _g) = dst.device_ptr_mut(stream);
|
||||
cudarc::driver::result::memcpy_dtod_async(
|
||||
dst_ptr, staging.dev_ptr, nbytes, stream.cu_stream(),
|
||||
)?;
|
||||
}
|
||||
Ok(dst)
|
||||
}
|
||||
|
||||
fn download(stream: &Arc<CudaStream>, src: &CudaSlice<f32>) -> Result<Vec<f32>> {
|
||||
let n = src.len();
|
||||
let staging = unsafe { MappedF32Buffer::new(n) }
|
||||
.map_err(|e| anyhow::anyhow!("download staging: {e}"))?;
|
||||
let nbytes = n * std::mem::size_of::<f32>();
|
||||
unsafe {
|
||||
let (src_ptr, _g) = src.device_ptr(stream);
|
||||
cudarc::driver::result::memcpy_dtod_async(
|
||||
staging.dev_ptr, src_ptr, nbytes, stream.cu_stream(),
|
||||
)?;
|
||||
}
|
||||
stream.synchronize()?;
|
||||
Ok(staging.read_all())
|
||||
}
|
||||
|
||||
/// Forward-only: produces ctx_out for a perturbed input.
|
||||
fn forward(
|
||||
pool: &HorizonTokenAttentionPool,
|
||||
stream: &Arc<CudaStream>,
|
||||
horizon_tokens: &[f32],
|
||||
q: &[f32],
|
||||
ln_out: &[f32],
|
||||
) -> Result<Vec<f32>> {
|
||||
let ht_d = upload(stream, horizon_tokens)?;
|
||||
let q_d = upload(stream, q)?;
|
||||
let ln_d = upload(stream, ln_out)?;
|
||||
let mut ctx_d = stream.alloc_zeros::<f32>(B * N_H * H)?;
|
||||
let mut attn_d = stream.alloc_zeros::<f32>(B * N_EXT)?;
|
||||
pool.forward(&ht_d, &q_d, &ln_d, B as i32, K as i32, &mut ctx_d, &mut attn_d)?;
|
||||
stream.synchronize()?;
|
||||
download(stream, &ctx_d)
|
||||
}
|
||||
|
||||
/// Scalar loss = sum of all entries of ctx_out, weighted by an upstream
|
||||
/// random `grad_ctx`. With a fixed seeded grad_ctx, ∂loss/∂ctx_out = grad_ctx
|
||||
/// and chain rule via the kernel's backward should match a numgrad on
|
||||
/// `loss = Σ_{b,h,d} grad_ctx[b,h,d] · ctx_out[b,h,d]`.
|
||||
fn weighted_sum_loss(ctx: &[f32], grad_ctx: &[f32]) -> f32 {
|
||||
ctx.iter().zip(grad_ctx).map(|(c, g)| c * g).sum()
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[ignore = "requires CUDA"]
|
||||
fn forward_then_backward_matches_central_difference() -> Result<()> {
|
||||
let dev = MlDevice::cuda(0)?;
|
||||
let stream = dev.cuda_stream().context("stream")?.clone();
|
||||
let ctx_dev = dev.cuda_context().context("ctx")?;
|
||||
let pool = HorizonTokenAttentionPool::new(ctx_dev, stream.clone())?;
|
||||
|
||||
let mut rng = deterministic_rng(20260518);
|
||||
let ht: Vec<f32> = (0..N_H * H).map(|_| rng() * 0.1).collect();
|
||||
let q: Vec<f32> = (0..H).map(|_| rng() * 0.1).collect();
|
||||
let ln: Vec<f32> = (0..B * K * H).map(|_| rng() * 0.5).collect();
|
||||
let grad_ctx: Vec<f32> = (0..B * N_H * H).map(|_| rng() * 0.1).collect();
|
||||
|
||||
// 1. Forward + backward through the kernel.
|
||||
let ht_d = upload(&stream, &ht)?;
|
||||
let q_d = upload(&stream, &q)?;
|
||||
let ln_d = upload(&stream, &ln)?;
|
||||
let mut ctx_d = stream.alloc_zeros::<f32>(B * N_H * H)?;
|
||||
let mut attn_d = stream.alloc_zeros::<f32>(B * N_EXT)?;
|
||||
pool.forward(&ht_d, &q_d, &ln_d, B as i32, K as i32, &mut ctx_d, &mut attn_d)?;
|
||||
|
||||
let grad_ctx_d = upload(&stream, &grad_ctx)?;
|
||||
let mut grad_ht_scratch_d = stream.alloc_zeros::<f32>(B * N_H * H)?;
|
||||
let mut grad_q_scratch_d = stream.alloc_zeros::<f32>(B * H)?;
|
||||
let mut grad_ln_d = stream.alloc_zeros::<f32>(B * K * H)?;
|
||||
pool.backward(
|
||||
&ht_d, &q_d, &ln_d, &attn_d, &grad_ctx_d,
|
||||
B as i32, K as i32,
|
||||
&mut grad_ht_scratch_d, &mut grad_q_scratch_d, &mut grad_ln_d,
|
||||
)?;
|
||||
stream.synchronize()?;
|
||||
|
||||
let grad_ht_scratch = download(&stream, &grad_ht_scratch_d)?;
|
||||
let grad_q_scratch = download(&stream, &grad_q_scratch_d)?;
|
||||
let grad_ln = download(&stream, &grad_ln_d)?;
|
||||
|
||||
// Reduce per-batch scratches to shared param-grads (host-side here
|
||||
// because numgrad is host-side too; trainer uses reduce_axis0 GPU).
|
||||
let mut grad_ht = vec![0.0_f32; N_H * H];
|
||||
for b in 0..B {
|
||||
for i in 0..(N_H * H) {
|
||||
grad_ht[i] += grad_ht_scratch[b * N_H * H + i];
|
||||
}
|
||||
}
|
||||
let mut grad_q = vec![0.0_f32; H];
|
||||
for b in 0..B {
|
||||
for d in 0..H {
|
||||
grad_q[d] += grad_q_scratch[b * H + d];
|
||||
}
|
||||
}
|
||||
|
||||
// 2. Central-difference numgrad on 4 sampled positions of each tensor.
|
||||
let eps = 1e-3_f32;
|
||||
let mut checks = 0_usize;
|
||||
|
||||
let sample_indices = |total: usize, n: usize, seed_off: u64| -> Vec<usize> {
|
||||
let mut r = deterministic_rng(20260518 + seed_off);
|
||||
(0..n).map(|_| ((r() + 1.0) * 0.5 * total as f32) as usize % total).collect()
|
||||
};
|
||||
|
||||
// horizon_tokens
|
||||
for &idx in sample_indices(N_H * H, 4, 1).iter() {
|
||||
let mut ht_p = ht.clone();
|
||||
let mut ht_m = ht.clone();
|
||||
ht_p[idx] += eps;
|
||||
ht_m[idx] -= eps;
|
||||
let l_p = weighted_sum_loss(&forward(&pool, &stream, &ht_p, &q, &ln)?, &grad_ctx);
|
||||
let l_m = weighted_sum_loss(&forward(&pool, &stream, &ht_m, &q, &ln)?, &grad_ctx);
|
||||
let numgrad = (l_p - l_m) / (2.0 * eps);
|
||||
let analytic = grad_ht[idx];
|
||||
let abs = (analytic - numgrad).abs();
|
||||
let rel = abs / numgrad.abs().max(1e-3);
|
||||
assert!(
|
||||
abs < 5e-3 || rel < 5e-2,
|
||||
"horizon_tokens[{idx}] mismatch: analytic={analytic:.6} numgrad={numgrad:.6} abs={abs:.3e} rel={rel:.3e}"
|
||||
);
|
||||
checks += 1;
|
||||
}
|
||||
|
||||
// Q
|
||||
for &idx in sample_indices(H, 4, 2).iter() {
|
||||
let mut q_p = q.clone();
|
||||
let mut q_m = q.clone();
|
||||
q_p[idx] += eps;
|
||||
q_m[idx] -= eps;
|
||||
let l_p = weighted_sum_loss(&forward(&pool, &stream, &ht, &q_p, &ln)?, &grad_ctx);
|
||||
let l_m = weighted_sum_loss(&forward(&pool, &stream, &ht, &q_m, &ln)?, &grad_ctx);
|
||||
let numgrad = (l_p - l_m) / (2.0 * eps);
|
||||
let analytic = grad_q[idx];
|
||||
let abs = (analytic - numgrad).abs();
|
||||
let rel = abs / numgrad.abs().max(1e-3);
|
||||
assert!(
|
||||
abs < 5e-3 || rel < 5e-2,
|
||||
"Q[{idx}] mismatch: analytic={analytic:.6} numgrad={numgrad:.6} abs={abs:.3e} rel={rel:.3e}"
|
||||
);
|
||||
checks += 1;
|
||||
}
|
||||
|
||||
// ln_out
|
||||
for &idx in sample_indices(B * K * H, 4, 3).iter() {
|
||||
let mut ln_p = ln.clone();
|
||||
let mut ln_m = ln.clone();
|
||||
ln_p[idx] += eps;
|
||||
ln_m[idx] -= eps;
|
||||
let l_p = weighted_sum_loss(&forward(&pool, &stream, &ht, &q, &ln_p)?, &grad_ctx);
|
||||
let l_m = weighted_sum_loss(&forward(&pool, &stream, &ht, &q, &ln_m)?, &grad_ctx);
|
||||
let numgrad = (l_p - l_m) / (2.0 * eps);
|
||||
let analytic = grad_ln[idx];
|
||||
let abs = (analytic - numgrad).abs();
|
||||
let rel = abs / numgrad.abs().max(1e-3);
|
||||
assert!(
|
||||
abs < 5e-3 || rel < 5e-2,
|
||||
"ln_out[{idx}] mismatch: analytic={analytic:.6} numgrad={numgrad:.6} abs={abs:.3e} rel={rel:.3e}"
|
||||
);
|
||||
checks += 1;
|
||||
}
|
||||
|
||||
eprintln!("PASSED {checks} numgrad checks");
|
||||
Ok(())
|
||||
}
|
||||
@@ -1,127 +0,0 @@
|
||||
//! Numgrad parity test for inverted_attention_pool (v2 axis E).
|
||||
|
||||
#![cfg(feature = "cuda")]
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use cudarc::driver::{CudaSlice, CudaStream, DevicePtr, DevicePtrMut};
|
||||
use ml_alpha::inverted_attention_pool::{InvertedAttentionPool, IAP_HIDDEN_DIM};
|
||||
use ml_alpha::pinned_mem::MappedF32Buffer;
|
||||
use ml_core::device::MlDevice;
|
||||
use std::sync::Arc;
|
||||
|
||||
const B: usize = 2;
|
||||
const K: usize = 8;
|
||||
const H: usize = IAP_HIDDEN_DIM;
|
||||
|
||||
fn rng(seed: u64) -> impl FnMut() -> f32 {
|
||||
let mut s = seed.wrapping_mul(0x9E37_79B9_7F4A_7C15);
|
||||
move || -> f32 {
|
||||
s = s.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
|
||||
let x = ((s >> 16) & 0xFFFF) as f32 / 65536.0;
|
||||
x * 2.0 - 1.0
|
||||
}
|
||||
}
|
||||
|
||||
fn upload(stream: &Arc<CudaStream>, host: &[f32]) -> Result<CudaSlice<f32>> {
|
||||
let n = host.len();
|
||||
let staging = unsafe { MappedF32Buffer::new(n) }
|
||||
.map_err(|e| anyhow::anyhow!("upload staging: {e}"))?;
|
||||
staging.write_from_slice(host);
|
||||
let mut dst = stream.alloc_zeros::<f32>(n)?;
|
||||
let nbytes = n * std::mem::size_of::<f32>();
|
||||
unsafe {
|
||||
let (dst_ptr, _g) = dst.device_ptr_mut(stream);
|
||||
cudarc::driver::result::memcpy_dtod_async(
|
||||
dst_ptr, staging.dev_ptr, nbytes, stream.cu_stream(),
|
||||
)?;
|
||||
}
|
||||
Ok(dst)
|
||||
}
|
||||
|
||||
fn download(stream: &Arc<CudaStream>, src: &CudaSlice<f32>) -> Result<Vec<f32>> {
|
||||
let n = src.len();
|
||||
let staging = unsafe { MappedF32Buffer::new(n) }
|
||||
.map_err(|e| anyhow::anyhow!("download staging: {e}"))?;
|
||||
let nbytes = n * std::mem::size_of::<f32>();
|
||||
unsafe {
|
||||
let (src_ptr, _g) = src.device_ptr(stream);
|
||||
cudarc::driver::result::memcpy_dtod_async(
|
||||
staging.dev_ptr, src_ptr, nbytes, stream.cu_stream(),
|
||||
)?;
|
||||
}
|
||||
stream.synchronize()?;
|
||||
Ok(staging.read_all())
|
||||
}
|
||||
|
||||
fn forward_loss(
|
||||
pool: &InvertedAttentionPool,
|
||||
stream: &Arc<CudaStream>,
|
||||
ln: &[f32],
|
||||
grad_pooled: &[f32],
|
||||
) -> Result<f32> {
|
||||
let inv_scale = 1.0_f32 / (K as f32).sqrt();
|
||||
let ln_d = upload(stream, ln)?;
|
||||
let mut pooled_d = stream.alloc_zeros::<f32>(B * H)?;
|
||||
let mut attn_d = stream.alloc_zeros::<f32>(B * H * H)?;
|
||||
pool.forward(&ln_d, B as i32, K as i32, inv_scale, &mut pooled_d, &mut attn_d)?;
|
||||
stream.synchronize()?;
|
||||
let pooled = download(stream, &pooled_d)?;
|
||||
Ok(pooled.iter().zip(grad_pooled).map(|(p, g)| p * g).sum())
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[ignore = "requires CUDA"]
|
||||
fn forward_then_backward_matches_central_difference() -> Result<()> {
|
||||
let dev = MlDevice::cuda(0)?;
|
||||
let stream = dev.cuda_stream().context("stream")?.clone();
|
||||
let ctx_dev = dev.cuda_context().context("ctx")?;
|
||||
let pool = InvertedAttentionPool::new(ctx_dev, stream.clone())?;
|
||||
|
||||
let mut r = rng(20260518);
|
||||
let ln: Vec<f32> = (0..B * K * H).map(|_| r() * 0.3).collect();
|
||||
let grad_pooled: Vec<f32> = (0..B * H).map(|_| r() * 0.1).collect();
|
||||
|
||||
let inv_scale = 1.0_f32 / (K as f32).sqrt();
|
||||
let ln_d = upload(&stream, &ln)?;
|
||||
let mut pooled_d = stream.alloc_zeros::<f32>(B * H)?;
|
||||
let mut attn_d = stream.alloc_zeros::<f32>(B * H * H)?;
|
||||
pool.forward(&ln_d, B as i32, K as i32, inv_scale, &mut pooled_d, &mut attn_d)?;
|
||||
|
||||
let grad_pooled_d = upload(&stream, &grad_pooled)?;
|
||||
let mut d_scores_scratch_d = stream.alloc_zeros::<f32>(B * H * H)?;
|
||||
let mut grad_ln_d = stream.alloc_zeros::<f32>(B * K * H)?;
|
||||
pool.backward(
|
||||
&ln_d, &attn_d, &grad_pooled_d,
|
||||
B as i32, K as i32, inv_scale,
|
||||
&mut d_scores_scratch_d,
|
||||
&mut grad_ln_d,
|
||||
)?;
|
||||
stream.synchronize()?;
|
||||
|
||||
let grad_ln = download(&stream, &grad_ln_d)?;
|
||||
|
||||
let eps = 1e-3_f32;
|
||||
let mut r2 = rng(42);
|
||||
let total = B * K * H;
|
||||
let mut checks = 0;
|
||||
for _ in 0..6 {
|
||||
let idx = ((r2() + 1.0) * 0.5 * total as f32) as usize % total;
|
||||
let mut ln_p = ln.clone();
|
||||
let mut ln_m = ln.clone();
|
||||
ln_p[idx] += eps;
|
||||
ln_m[idx] -= eps;
|
||||
let l_p = forward_loss(&pool, &stream, &ln_p, &grad_pooled)?;
|
||||
let l_m = forward_loss(&pool, &stream, &ln_m, &grad_pooled)?;
|
||||
let numgrad = (l_p - l_m) / (2.0 * eps);
|
||||
let analytic = grad_ln[idx];
|
||||
let abs = (analytic - numgrad).abs();
|
||||
let rel = abs / numgrad.abs().max(1e-3);
|
||||
assert!(
|
||||
abs < 5e-3 || rel < 5e-2,
|
||||
"ln_out[{idx}] mismatch: analytic={analytic:.6} numgrad={numgrad:.6} abs={abs:.3e} rel={rel:.3e}"
|
||||
);
|
||||
checks += 1;
|
||||
}
|
||||
eprintln!("PASSED {checks} numgrad checks");
|
||||
Ok(())
|
||||
}
|
||||
@@ -1,290 +0,0 @@
|
||||
//! Numgrad parity test for regime_moe_gate (v2 axis D).
|
||||
//!
|
||||
//! Top-1 routing has straight-through grad on the gate logits — the
|
||||
//! gate side isn't tested by numgrad here (it's a non-differentiable
|
||||
//! argmax). What we DO test:
|
||||
//! 1. Forward correctness: routed_out matches a host-side reference
|
||||
//! that picks the same top-1 expert.
|
||||
//! 2. Backward grads on the selected expert's W and bias match
|
||||
//! central-difference numgrad.
|
||||
//! 3. Grad on fused_ctx matches numgrad.
|
||||
|
||||
#![cfg(feature = "cuda")]
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use cudarc::driver::{CudaSlice, CudaStream, DevicePtr, DevicePtrMut};
|
||||
use ml_alpha::pinned_mem::MappedF32Buffer;
|
||||
use ml_alpha::regime_moe_gate::{
|
||||
RegimeMoeGate, MOE_HIDDEN_DIM, MOE_N_EXPERTS, MOE_N_HORIZONS,
|
||||
};
|
||||
use ml_core::device::MlDevice;
|
||||
use std::sync::Arc;
|
||||
|
||||
const B: usize = 3;
|
||||
const H: usize = MOE_HIDDEN_DIM;
|
||||
const N_H: usize = MOE_N_HORIZONS;
|
||||
const N_E: usize = MOE_N_EXPERTS;
|
||||
|
||||
fn rng(seed: u64) -> impl FnMut() -> f32 {
|
||||
let mut s = seed.wrapping_mul(0x9E37_79B9_7F4A_7C15);
|
||||
move || -> f32 {
|
||||
s = s.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
|
||||
let x = ((s >> 16) & 0xFFFF) as f32 / 65536.0;
|
||||
x * 2.0 - 1.0
|
||||
}
|
||||
}
|
||||
|
||||
fn upload(stream: &Arc<CudaStream>, host: &[f32]) -> Result<CudaSlice<f32>> {
|
||||
let n = host.len();
|
||||
let staging = unsafe { MappedF32Buffer::new(n) }
|
||||
.map_err(|e| anyhow::anyhow!("upload staging: {e}"))?;
|
||||
staging.write_from_slice(host);
|
||||
let mut dst = stream.alloc_zeros::<f32>(n)?;
|
||||
let nbytes = n * std::mem::size_of::<f32>();
|
||||
unsafe {
|
||||
let (dst_ptr, _g) = dst.device_ptr_mut(stream);
|
||||
cudarc::driver::result::memcpy_dtod_async(
|
||||
dst_ptr, staging.dev_ptr, nbytes, stream.cu_stream(),
|
||||
)?;
|
||||
}
|
||||
Ok(dst)
|
||||
}
|
||||
|
||||
#[allow(dead_code)] // helper kept for future numgrad expansion (gate-side STE)
|
||||
fn upload_i32(stream: &Arc<CudaStream>, host: &[i32]) -> Result<CudaSlice<i32>> {
|
||||
let n = host.len();
|
||||
let staging = unsafe { MappedF32Buffer::new(n) }
|
||||
.map_err(|e| anyhow::anyhow!("upload i32 staging: {e}"))?;
|
||||
// f32 buffer reinterpreted as i32.
|
||||
let host_ptr = staging.host_ptr.cast::<i32>();
|
||||
unsafe {
|
||||
for i in 0..n { std::ptr::write_volatile(host_ptr.add(i), host[i]); }
|
||||
}
|
||||
let mut dst = stream.alloc_zeros::<i32>(n)?;
|
||||
let nbytes = n * std::mem::size_of::<i32>();
|
||||
unsafe {
|
||||
let (dst_ptr, _g) = dst.device_ptr_mut(stream);
|
||||
cudarc::driver::result::memcpy_dtod_async(
|
||||
dst_ptr, staging.dev_ptr, nbytes, stream.cu_stream(),
|
||||
)?;
|
||||
}
|
||||
Ok(dst)
|
||||
}
|
||||
|
||||
fn download(stream: &Arc<CudaStream>, src: &CudaSlice<f32>) -> Result<Vec<f32>> {
|
||||
let n = src.len();
|
||||
let staging = unsafe { MappedF32Buffer::new(n) }
|
||||
.map_err(|e| anyhow::anyhow!("download staging: {e}"))?;
|
||||
let nbytes = n * std::mem::size_of::<f32>();
|
||||
unsafe {
|
||||
let (src_ptr, _g) = src.device_ptr(stream);
|
||||
cudarc::driver::result::memcpy_dtod_async(
|
||||
staging.dev_ptr, src_ptr, nbytes, stream.cu_stream(),
|
||||
)?;
|
||||
}
|
||||
stream.synchronize()?;
|
||||
Ok(staging.read_all())
|
||||
}
|
||||
|
||||
fn download_i32(stream: &Arc<CudaStream>, src: &CudaSlice<i32>) -> Result<Vec<i32>> {
|
||||
let n = src.len();
|
||||
let staging = unsafe { MappedF32Buffer::new(n) }
|
||||
.map_err(|e| anyhow::anyhow!("download i32 staging: {e}"))?;
|
||||
let nbytes = n * std::mem::size_of::<i32>();
|
||||
unsafe {
|
||||
let (src_ptr, _g) = src.device_ptr(stream);
|
||||
cudarc::driver::result::memcpy_dtod_async(
|
||||
staging.dev_ptr, src_ptr, nbytes, stream.cu_stream(),
|
||||
)?;
|
||||
}
|
||||
stream.synchronize()?;
|
||||
let host_ptr = staging.host_ptr.cast::<i32>();
|
||||
let mut out = Vec::with_capacity(n);
|
||||
unsafe {
|
||||
for i in 0..n { out.push(std::ptr::read_volatile(host_ptr.add(i))); }
|
||||
}
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
struct Inputs {
|
||||
gate_logits: Vec<f32>,
|
||||
fused_ctx: Vec<f32>,
|
||||
experts_w: Vec<f32>,
|
||||
experts_b: Vec<f32>,
|
||||
grad_routed: Vec<f32>,
|
||||
}
|
||||
|
||||
fn make_inputs(seed: u64) -> Inputs {
|
||||
let mut r = rng(seed);
|
||||
let gate_logits: Vec<f32> = (0..B * N_E).map(|_| r() * 1.5).collect();
|
||||
let fused_ctx: Vec<f32> = (0..B * N_H * H).map(|_| r() * 0.3).collect();
|
||||
let experts_w: Vec<f32> = (0..N_E * H * H).map(|_| r() * 0.05).collect();
|
||||
let experts_b: Vec<f32> = (0..N_E * H).map(|_| r() * 0.02).collect();
|
||||
let grad_routed: Vec<f32> = (0..B * N_H * H).map(|_| r() * 0.1).collect();
|
||||
Inputs { gate_logits, fused_ctx, experts_w, experts_b, grad_routed }
|
||||
}
|
||||
|
||||
fn forward_routed_loss(
|
||||
moe: &RegimeMoeGate,
|
||||
stream: &Arc<CudaStream>,
|
||||
inputs: &Inputs,
|
||||
fused_ctx_override: Option<&[f32]>,
|
||||
experts_w_override: Option<&[f32]>,
|
||||
experts_b_override: Option<&[f32]>,
|
||||
) -> Result<f32> {
|
||||
let fused_ctx = fused_ctx_override.unwrap_or(&inputs.fused_ctx);
|
||||
let experts_w = experts_w_override.unwrap_or(&inputs.experts_w);
|
||||
let experts_b = experts_b_override.unwrap_or(&inputs.experts_b);
|
||||
|
||||
let gate_d = upload(stream, &inputs.gate_logits)?;
|
||||
let ctx_d = upload(stream, fused_ctx)?;
|
||||
let w_d = upload(stream, experts_w)?;
|
||||
let b_d = upload(stream, experts_b)?;
|
||||
let mut routed_d = stream.alloc_zeros::<f32>(B * N_H * H)?;
|
||||
let mut top_e_d = stream.alloc_zeros::<i32>(B)?;
|
||||
moe.forward(&gate_d, &ctx_d, &w_d, &b_d, B as i32, &mut routed_d, &mut top_e_d)?;
|
||||
stream.synchronize()?;
|
||||
let routed = download(stream, &routed_d)?;
|
||||
Ok(routed.iter().zip(&inputs.grad_routed).map(|(r, g)| r * g).sum())
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[ignore = "requires CUDA"]
|
||||
fn forward_matches_host_reference_and_backward_matches_numgrad() -> Result<()> {
|
||||
let dev = MlDevice::cuda(0)?;
|
||||
let stream = dev.cuda_stream().context("stream")?.clone();
|
||||
let ctx_dev = dev.cuda_context().context("ctx")?;
|
||||
let moe = RegimeMoeGate::new(ctx_dev, stream.clone())?;
|
||||
|
||||
let inputs = make_inputs(20260518);
|
||||
|
||||
// 1. Run forward, capture top_e.
|
||||
let gate_d = upload(&stream, &inputs.gate_logits)?;
|
||||
let ctx_d = upload(&stream, &inputs.fused_ctx)?;
|
||||
let w_d = upload(&stream, &inputs.experts_w)?;
|
||||
let b_d = upload(&stream, &inputs.experts_b)?;
|
||||
let mut routed_d = stream.alloc_zeros::<f32>(B * N_H * H)?;
|
||||
let mut top_e_d = stream.alloc_zeros::<i32>(B)?;
|
||||
moe.forward(&gate_d, &ctx_d, &w_d, &b_d, B as i32, &mut routed_d, &mut top_e_d)?;
|
||||
stream.synchronize()?;
|
||||
let routed = download(&stream, &routed_d)?;
|
||||
let top_e = download_i32(&stream, &top_e_d)?;
|
||||
|
||||
// Host reference for routed.
|
||||
for b in 0..B {
|
||||
let e = top_e[b] as usize;
|
||||
for h in 0..N_H {
|
||||
for d_out in 0..H {
|
||||
let mut acc = inputs.experts_b[e * H + d_out];
|
||||
for d_in in 0..H {
|
||||
acc += inputs.experts_w[e * H * H + d_out * H + d_in]
|
||||
* inputs.fused_ctx[b * N_H * H + h * H + d_in];
|
||||
}
|
||||
let kernel_val = routed[b * N_H * H + h * H + d_out];
|
||||
let diff = (acc - kernel_val).abs();
|
||||
assert!(diff < 1e-3,
|
||||
"forward mismatch at b={b} h={h} d_out={d_out}: ref={acc:.5} kernel={kernel_val:.5}");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 2. Backward — compute grads. Compact scratch [B, N_H, H, H] / [B, N_H, H].
|
||||
let grad_routed_d = upload(&stream, &inputs.grad_routed)?;
|
||||
let mut d_w_scratch_d = stream.alloc_zeros::<f32>(B * N_H * H * H)?;
|
||||
let mut d_b_scratch_d = stream.alloc_zeros::<f32>(B * N_H * H)?;
|
||||
let mut d_fused_ctx_d = stream.alloc_zeros::<f32>(B * N_H * H)?;
|
||||
moe.backward(
|
||||
&ctx_d, &w_d, &grad_routed_d, &top_e_d, B as i32,
|
||||
&mut d_w_scratch_d, &mut d_b_scratch_d, &mut d_fused_ctx_d,
|
||||
)?;
|
||||
// Scatter into per-expert grads via the new MoE.scatter kernel.
|
||||
let mut grad_experts_w_d = stream.alloc_zeros::<f32>(N_E * H * H)?;
|
||||
let mut grad_experts_b_d = stream.alloc_zeros::<f32>(N_E * H)?;
|
||||
moe.scatter(
|
||||
&d_w_scratch_d, &d_b_scratch_d, &top_e_d, B as i32,
|
||||
&mut grad_experts_w_d, &mut grad_experts_b_d,
|
||||
)?;
|
||||
stream.synchronize()?;
|
||||
let d_w_full = download(&stream, &grad_experts_w_d)?;
|
||||
let d_b_full = download(&stream, &grad_experts_b_d)?;
|
||||
let d_fused_ctx = download(&stream, &d_fused_ctx_d)?;
|
||||
// (no host-side reduce — scatter kernel produces the per-expert grad directly)
|
||||
|
||||
// 3. Numgrad checks.
|
||||
let eps = 1e-3_f32;
|
||||
let mut r = rng(42);
|
||||
let mut checks = 0;
|
||||
|
||||
// d_W on a few sampled (e, d_out, d_in) — only check active experts.
|
||||
let active: std::collections::HashSet<usize> =
|
||||
top_e.iter().map(|&e| e as usize).collect();
|
||||
for _ in 0..4 {
|
||||
let e = *active.iter().nth((r().abs() * active.len() as f32) as usize % active.len()).unwrap();
|
||||
let d_out = ((r() + 1.0) * 0.5 * H as f32) as usize % H;
|
||||
let d_in = ((r() + 1.0) * 0.5 * H as f32) as usize % H;
|
||||
let idx = e * H * H + d_out * H + d_in;
|
||||
let mut w_p = inputs.experts_w.clone();
|
||||
let mut w_m = inputs.experts_w.clone();
|
||||
w_p[idx] += eps;
|
||||
w_m[idx] -= eps;
|
||||
let l_p = forward_routed_loss(&moe, &stream, &inputs, None, Some(&w_p), None)?;
|
||||
let l_m = forward_routed_loss(&moe, &stream, &inputs, None, Some(&w_m), None)?;
|
||||
let numgrad = (l_p - l_m) / (2.0 * eps);
|
||||
let analytic = d_w_full[idx];
|
||||
let abs = (analytic - numgrad).abs();
|
||||
let rel = abs / numgrad.abs().max(1e-3);
|
||||
assert!(
|
||||
abs < 5e-3 || rel < 5e-2,
|
||||
"d_W[{e},{d_out},{d_in}] mismatch: analytic={analytic:.4} numgrad={numgrad:.4} abs={abs:.3e} rel={rel:.3e}"
|
||||
);
|
||||
checks += 1;
|
||||
}
|
||||
|
||||
// d_b on a few sampled (e, d_out).
|
||||
for _ in 0..3 {
|
||||
let e = *active.iter().nth((r().abs() * active.len() as f32) as usize % active.len()).unwrap();
|
||||
let d_out = ((r() + 1.0) * 0.5 * H as f32) as usize % H;
|
||||
let idx = e * H + d_out;
|
||||
let mut b_p = inputs.experts_b.clone();
|
||||
let mut b_m = inputs.experts_b.clone();
|
||||
b_p[idx] += eps;
|
||||
b_m[idx] -= eps;
|
||||
let l_p = forward_routed_loss(&moe, &stream, &inputs, None, None, Some(&b_p))?;
|
||||
let l_m = forward_routed_loss(&moe, &stream, &inputs, None, None, Some(&b_m))?;
|
||||
let numgrad = (l_p - l_m) / (2.0 * eps);
|
||||
let analytic = d_b_full[idx];
|
||||
let abs = (analytic - numgrad).abs();
|
||||
let rel = abs / numgrad.abs().max(1e-3);
|
||||
assert!(
|
||||
abs < 5e-3 || rel < 5e-2,
|
||||
"d_b[{e},{d_out}] mismatch: analytic={analytic:.4} numgrad={numgrad:.4} abs={abs:.3e} rel={rel:.3e}"
|
||||
);
|
||||
checks += 1;
|
||||
}
|
||||
|
||||
// d_fused_ctx on a few sampled positions.
|
||||
for _ in 0..4 {
|
||||
let b_idx = ((r() + 1.0) * 0.5 * B as f32) as usize % B;
|
||||
let h_idx = ((r() + 1.0) * 0.5 * N_H as f32) as usize % N_H;
|
||||
let d_idx = ((r() + 1.0) * 0.5 * H as f32) as usize % H;
|
||||
let idx = b_idx * N_H * H + h_idx * H + d_idx;
|
||||
let mut ctx_p = inputs.fused_ctx.clone();
|
||||
let mut ctx_m = inputs.fused_ctx.clone();
|
||||
ctx_p[idx] += eps;
|
||||
ctx_m[idx] -= eps;
|
||||
let l_p = forward_routed_loss(&moe, &stream, &inputs, Some(&ctx_p), None, None)?;
|
||||
let l_m = forward_routed_loss(&moe, &stream, &inputs, Some(&ctx_m), None, None)?;
|
||||
let numgrad = (l_p - l_m) / (2.0 * eps);
|
||||
let analytic = d_fused_ctx[idx];
|
||||
let abs = (analytic - numgrad).abs();
|
||||
let rel = abs / numgrad.abs().max(1e-3);
|
||||
assert!(
|
||||
abs < 5e-3 || rel < 5e-2,
|
||||
"d_fused_ctx[{b_idx},{h_idx},{d_idx}] mismatch: analytic={analytic:.4} numgrad={numgrad:.4} abs={abs:.3e} rel={rel:.3e}"
|
||||
);
|
||||
checks += 1;
|
||||
}
|
||||
|
||||
eprintln!("PASSED {checks} numgrad checks");
|
||||
Ok(())
|
||||
}
|
||||
Reference in New Issue
Block a user