From 11b96dac6ad5a003b7f1a9a81df8c9e96758aeca Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Mon, 18 May 2026 14:21:59 +0200 Subject: [PATCH] feat(ml-alpha): regime_moe_gate kernel + numgrad (v2 D) [V6] MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Top-1 Mixture-of-Experts gate per Switch Transformer. Three kernels in one .cu file: - regime_moe_gate_fwd: select top-1 expert from gate_logits, apply its [H, H] linear to fused_ctx → routed_ctx [B, N_H, H]. - regime_moe_gate_bwd: chain-rule grads through the SELECTED expert's W and bias (sparse-by-expert scratches), plus d_fused_ctx accumulator. Inactive experts get zero contribution. - regime_moe_gate_aux: softmax of gate_logits + load-balancing auxiliary loss (frac · prob_mean × N_EXPERTS). ARCHITECTURE: - N_EXPERTS = 4. Each expert is a [H, H] linear with bias. - Total expert params: 4 · 128 · 128 + 4 · 128 = 66 KB. Cheap. - STE on gate: gate logit grad is zero from the expert path (top-1 is non-differentiable); the load-balance aux loss provides the differentiable signal that pushes routing toward balanced usage. PERFORMANCE: - Grid (B, N_H, 1), block (HIDDEN_DIM). One block per (b, h). - Forward: each thread computes one output channel via a dot product over HIDDEN_DIM input dims (#pragma unroll 8). - Backward d_fused_ctx: each thread accumulates over d_out sequentially (HIDDEN_DIM iterations) since the weight matrix column is naturally aligned to the thread's d_in index. - Backward d_W/d_b scratches are sparse-by-expert; downstream reduce_axis0 collapses over (B, N_H). - Top-1 chosen by thread 0 per block, broadcast via shared mem. NUMGRAD VERIFICATION (RTX 3050 sm_86): forward_matches_host_reference_and_backward_matches_numgrad PASSES 11 checks (4 on d_W, 3 on d_b, 4 on d_fused_ctx) within 5e-2 rel / 5e-3 abs. Co-Authored-By: Claude Opus 4.7 --- crates/ml-alpha/build.rs | 1 + crates/ml-alpha/cuda/regime_moe_gate.cu | 245 ++++++++++++++ crates/ml-alpha/src/lib.rs | 1 + crates/ml-alpha/src/regime_moe_gate.rs | 140 ++++++++ .../ml-alpha/tests/regime_moe_gate_numgrad.rs | 300 ++++++++++++++++++ 5 files changed, 687 insertions(+) create mode 100644 crates/ml-alpha/cuda/regime_moe_gate.cu create mode 100644 crates/ml-alpha/src/regime_moe_gate.rs create mode 100644 crates/ml-alpha/tests/regime_moe_gate_numgrad.rs diff --git a/crates/ml-alpha/build.rs b/crates/ml-alpha/build.rs index 9022ddc73..028b40a4a 100644 --- a/crates/ml-alpha/build.rs +++ b/crates/ml-alpha/build.rs @@ -23,6 +23,7 @@ const KERNELS: &[&str] = &[ "attention_pool", // Phase 3: learned context summary at CfC k=0 "horizon_token_attention_pool", // v2-C: horizon-token K-prepend single-Q attention "inverted_attention_pool", // v2-E: iTransformer-style cross-variate attention + "regime_moe_gate", // v2-D: top-1 MoE gate + expert dispatch + aux loss "reduce_axis0", // Phase B: cross-batch param-grad reducer ]; diff --git a/crates/ml-alpha/cuda/regime_moe_gate.cu b/crates/ml-alpha/cuda/regime_moe_gate.cu new file mode 100644 index 000000000..2567b7346 --- /dev/null +++ b/crates/ml-alpha/cuda/regime_moe_gate.cu @@ -0,0 +1,245 @@ +// 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). +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, N_E, H, H] sparse-by-expert + float* __restrict__ d_b_scratch, // [B, N_H, N_E, H] sparse-by-expert + 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, e, 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, e, 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_N_EXPERTS + e) * 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_N_EXPERTS + e) * 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; + } +} + +// ── 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; +} diff --git a/crates/ml-alpha/src/lib.rs b/crates/ml-alpha/src/lib.rs index b5206705b..ef1599926 100644 --- a/crates/ml-alpha/src/lib.rs +++ b/crates/ml-alpha/src/lib.rs @@ -32,6 +32,7 @@ pub mod heads; pub mod horizon_token_attention_pool; pub mod inverted_attention_pool; pub mod isv; +pub mod regime_moe_gate; pub mod pinned; pub mod pinned_mem; pub mod trainer; diff --git a/crates/ml-alpha/src/regime_moe_gate.rs b/crates/ml-alpha/src/regime_moe_gate.rs new file mode 100644 index 000000000..0fec9e795 --- /dev/null +++ b/crates/ml-alpha/src/regime_moe_gate.rs @@ -0,0 +1,140 @@ +//! 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, + fwd_fn: CudaFunction, + bwd_fn: CudaFunction, + aux_fn: CudaFunction, + stream: Arc, +} + +impl RegimeMoeGate { + pub fn new(ctx: &Arc, stream: Arc) -> Result { + 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 aux_fn = module.load_function("regime_moe_gate_aux").context("aux")?; + Ok(Self { _module: module, fwd_fn, bwd_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, + fused_ctx: &CudaSlice, + experts_w: &CudaSlice, + experts_b: &CudaSlice, + n_batch: i32, + routed_out: &mut CudaSlice, + top_e_out: &mut CudaSlice, + ) -> 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, + experts_w: &CudaSlice, + d_routed: &CudaSlice, + top_e: &CudaSlice, + n_batch: i32, + d_w_scratch: &mut CudaSlice, + d_b_scratch: &mut CudaSlice, + d_fused_ctx_out: &mut CudaSlice, + ) -> 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(()) + } + + /// 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, + top_e: &CudaSlice, + n_batch: i32, + gate_probs_out: &mut CudaSlice, + aux_loss_out: &mut CudaSlice, + ) -> 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(()) + } +} diff --git a/crates/ml-alpha/tests/regime_moe_gate_numgrad.rs b/crates/ml-alpha/tests/regime_moe_gate_numgrad.rs new file mode 100644 index 000000000..d8146d57f --- /dev/null +++ b/crates/ml-alpha/tests/regime_moe_gate_numgrad.rs @@ -0,0 +1,300 @@ +//! 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, host: &[f32]) -> Result> { + 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::(n)?; + let nbytes = n * std::mem::size_of::(); + 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, host: &[i32]) -> Result> { + 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::(); + unsafe { + for i in 0..n { std::ptr::write_volatile(host_ptr.add(i), host[i]); } + } + let mut dst = stream.alloc_zeros::(n)?; + let nbytes = n * std::mem::size_of::(); + 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, src: &CudaSlice) -> Result> { + 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::(); + 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, src: &CudaSlice) -> Result> { + 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::(); + 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::(); + 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, + fused_ctx: Vec, + experts_w: Vec, + experts_b: Vec, + grad_routed: Vec, +} + +fn make_inputs(seed: u64) -> Inputs { + let mut r = rng(seed); + let gate_logits: Vec = (0..B * N_E).map(|_| r() * 1.5).collect(); + let fused_ctx: Vec = (0..B * N_H * H).map(|_| r() * 0.3).collect(); + let experts_w: Vec = (0..N_E * H * H).map(|_| r() * 0.05).collect(); + let experts_b: Vec = (0..N_E * H).map(|_| r() * 0.02).collect(); + let grad_routed: Vec = (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, + inputs: &Inputs, + fused_ctx_override: Option<&[f32]>, + experts_w_override: Option<&[f32]>, + experts_b_override: Option<&[f32]>, +) -> Result { + 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::(B * N_H * H)?; + let mut top_e_d = stream.alloc_zeros::(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::(B * N_H * H)?; + let mut top_e_d = stream.alloc_zeros::(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. + let grad_routed_d = upload(&stream, &inputs.grad_routed)?; + let mut d_w_scratch_d = stream.alloc_zeros::(B * N_H * N_E * H * H)?; + let mut d_b_scratch_d = stream.alloc_zeros::(B * N_H * N_E * H)?; + let mut d_fused_ctx_d = stream.alloc_zeros::(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, + )?; + stream.synchronize()?; + let d_w_scratch = download(&stream, &d_w_scratch_d)?; + let d_b_scratch = download(&stream, &d_b_scratch_d)?; + let d_fused_ctx = download(&stream, &d_fused_ctx_d)?; + + // Reduce scratches host-side to get d_W [N_E, H, H] and d_b [N_E, H]. + let mut d_w_full = vec![0.0_f32; N_E * H * H]; + let mut d_b_full = vec![0.0_f32; N_E * H]; + for b in 0..B { + for h in 0..N_H { + for e in 0..N_E { + for d_out in 0..H { + let bias_idx = ((b * N_H + h) * N_E + e) * H + d_out; + d_b_full[e * H + d_out] += d_b_scratch[bias_idx]; + for d_in in 0..H { + let w_idx = (((b * N_H + h) * N_E + e) * H + d_out) * H + d_in; + d_w_full[e * H * H + d_out * H + d_in] += d_w_scratch[w_idx]; + } + } + } + } + } + + // 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 = + 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(()) +}