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

feat(ml-alpha): regime_moe_gate kernel + numgrad (v2 D) [V6]

Browse Source 
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 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-18 14:21:59 +02:00
parent d94696620d
commit 11b96dac6a
5 changed files with 687 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

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

245
crates/ml-alpha/cuda/regime_moe_gate.cu Normal file
View File

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

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

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

140
crates/ml-alpha/src/regime_moe_gate.rs Normal file
View File

@@ -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<CudaModule>,
fwd_fn: CudaFunction,
bwd_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 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<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(())
}
/// 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(())
}
}

300
crates/ml-alpha/tests/regime_moe_gate_numgrad.rs Normal file
View File

@@ -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<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.
let grad_routed_d = upload(&stream, &inputs.grad_routed)?;
let mut d_w_scratch_d = stream.alloc_zeros::<f32>(B * N_H * N_E * H * H)?;
let mut d_b_scratch_d = stream.alloc_zeros::<f32>(B * N_H * N_E * 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,
)?;
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<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(())
}