From 9eea7f2977a7441152006e32f16adb47cf0f760e Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Mon, 27 Apr 2026 18:42:29 +0200 Subject: [PATCH] feat(moe): moe_mixture_backward + moe_dgate_reduce kernels + FD test MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two-stage backward: moe_mixture_backward computes de_k = g · dh_s2 in one kernel; moe_dgate_reduce computes dg via shmem reduction over c. FD test verifies analytic backward matches numerical for B=2, K=4, C=32 within 1e-3 tolerance (perturbation 1e-3). All test data flows via mapped pinned buffers (no HtoD/HtoH). Co-Authored-By: Claude Opus 4.7 (1M context) --- crates/ml/src/cuda_pipeline/gpu_moe_head.rs | 97 ++++++++++++++++++++- crates/ml/src/cuda_pipeline/moe_kernels.cu | 57 ++++++++++++ crates/ml/tests/moe_kernels_test.rs | 97 +++++++++++++++++++++ docs/dqn-wire-up-audit.md | 7 ++ 4 files changed, 255 insertions(+), 3 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/gpu_moe_head.rs b/crates/ml/src/cuda_pipeline/gpu_moe_head.rs index 0d26a4816..560c16770 100644 --- a/crates/ml/src/cuda_pipeline/gpu_moe_head.rs +++ b/crates/ml/src/cuda_pipeline/gpu_moe_head.rs @@ -20,13 +20,15 @@ static MOE_CUBIN: &[u8] = /// GPU-accelerated MoE head. /// -/// Owns the kernel function handles for the MoE forward mixture pass. -/// Subsequent tasks will add: moe_mixture_backward, moe_dgate_reduce, -/// moe_load_balance_loss, moe_load_balance_reduce, moe_expert_util_ema_update. +/// Owns the kernel function handles for the MoE mixture forward/backward pass. +/// Subsequent tasks will add: moe_load_balance_loss, moe_load_balance_reduce, +/// moe_expert_util_ema_update. #[allow(missing_debug_implementations)] pub struct GpuMoeHead { stream: Arc, moe_mixture_forward: CudaFunction, + moe_mixture_backward: CudaFunction, + moe_dgate_reduce: CudaFunction, } impl GpuMoeHead { @@ -40,9 +42,17 @@ impl GpuMoeHead { let moe_mixture_forward = module .load_function("moe_mixture_forward") .map_err(|e| MLError::ModelError(format!("moe_mixture_forward load: {e}")))?; + let moe_mixture_backward = module + .load_function("moe_mixture_backward") + .map_err(|e| MLError::ModelError(format!("moe_mixture_backward load: {e}")))?; + let moe_dgate_reduce = module + .load_function("moe_dgate_reduce") + .map_err(|e| MLError::ModelError(format!("moe_dgate_reduce load: {e}")))?; Ok(Self { stream, moe_mixture_forward, + moe_mixture_backward, + moe_dgate_reduce, }) } @@ -127,4 +137,85 @@ impl GpuMoeHead { // Read GPU-written result via host_ptr (mapped pinned coherence). Ok(h_s2_buf.read_all()) } + + /// Test-only entry for the backward pass. + /// + /// Launches `moe_mixture_backward` (de_k = g · dh_s2) and + /// `moe_dgate_reduce` (dg = Σ_c e_k · dh_s2, shmem reduction over c). + /// + /// Returns `(de_k [K*B*C], dg [B*K])`. + pub fn test_mixture_backward( + &self, + expert_outputs: &[f32], + gate: &[f32], + dh_s2: &[f32], + b: usize, + k: usize, + c: usize, + ) -> Result<(Vec, Vec), MLError> { + let expert_buf = unsafe { MappedF32Buffer::new(k * b * c) } + .map_err(|e| MLError::ModelError(format!("expert buf alloc: {e}")))?; + let gate_buf = unsafe { MappedF32Buffer::new(b * k) } + .map_err(|e| MLError::ModelError(format!("gate buf alloc: {e}")))?; + let dh_s2_buf = unsafe { MappedF32Buffer::new(b * c) } + .map_err(|e| MLError::ModelError(format!("dh_s2 buf alloc: {e}")))?; + let de_k_buf = unsafe { MappedF32Buffer::new(k * b * c) } + .map_err(|e| MLError::ModelError(format!("de_k buf alloc: {e}")))?; + let dg_buf = unsafe { MappedF32Buffer::new(b * k) } + .map_err(|e| MLError::ModelError(format!("dg buf alloc: {e}")))?; + + expert_buf.write_from_slice(expert_outputs); + gate_buf.write_from_slice(gate); + dh_s2_buf.write_from_slice(dh_s2); + + let block: u32 = 256; + let total_kbc = (k * b * c) as u32; + let grid_kbc = total_kbc.div_ceil(block); + + let dev_expert = expert_buf.dev_ptr; + let dev_gate = gate_buf.dev_ptr; + let dev_dh_s2 = dh_s2_buf.dev_ptr; + let dev_de_k = de_k_buf.dev_ptr; + let dev_dg = dg_buf.dev_ptr; + let b_i32 = b as i32; + let k_i32 = k as i32; + let c_i32 = c as i32; + + unsafe { + self.stream + .launch_builder(&self.moe_mixture_backward) + .arg(&dev_dh_s2) + .arg(&dev_gate) + .arg(&dev_de_k) + .arg(&b_i32) + .arg(&k_i32) + .arg(&c_i32) + .launch(LaunchConfig { + grid_dim: (grid_kbc, 1, 1), + block_dim: (block, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("moe_mixture_backward launch: {e}")))?; + + self.stream + .launch_builder(&self.moe_dgate_reduce) + .arg(&dev_dh_s2) + .arg(&dev_expert) + .arg(&dev_dg) + .arg(&b_i32) + .arg(&k_i32) + .arg(&c_i32) + .launch(LaunchConfig { + grid_dim: (b as u32, k as u32, 1), + block_dim: (block, 1, 1), + shared_mem_bytes: block * 4, // block * sizeof(f32) + }) + .map_err(|e| MLError::ModelError(format!("moe_dgate_reduce launch: {e}")))?; + } + self.stream + .synchronize() + .map_err(|e| MLError::ModelError(format!("synchronize: {e}")))?; + + Ok((de_k_buf.read_all(), dg_buf.read_all())) + } } diff --git a/crates/ml/src/cuda_pipeline/moe_kernels.cu b/crates/ml/src/cuda_pipeline/moe_kernels.cu index c81f69a65..5a8a3f787 100644 --- a/crates/ml/src/cuda_pipeline/moe_kernels.cu +++ b/crates/ml/src/cuda_pipeline/moe_kernels.cu @@ -31,3 +31,60 @@ extern "C" __global__ void moe_mixture_forward( h_s2[tid] = acc; __threadfence_system(); /* make write PCIe-visible to mapped pinned host_ptr */ } + +extern "C" __global__ void moe_mixture_backward( + const float* __restrict__ dh_s2, /* [B, C] */ + const float* __restrict__ gate, /* [B, K] */ + float* __restrict__ de_k, /* [K, B, C] */ + int B, + int K, + int C +) { + /* One thread per (k, b, c). de_k[k,b,c] = g[b,k] * dh_s2[b,c]. */ + int total = K * B * C; + int tid = blockIdx.x * blockDim.x + threadIdx.x; + if (tid >= total) return; + + int k = tid / (B * C); + int rem = tid - k * (B * C); + int b = rem / C; + int c = rem - b * C; + + de_k[tid] = gate[b * K + k] * dh_s2[b * C + c]; + __threadfence_system(); +} + +extern "C" __global__ void moe_dgate_reduce( + const float* __restrict__ dh_s2, /* [B, C] */ + const float* __restrict__ expert_outputs, /* [K, B, C] */ + float* __restrict__ dg, /* [B, K] */ + int B, + int K, + int C +) { + /* Grid: (B, K, 1); block: 256 threads (or chosen by host). */ + int b = blockIdx.x; + int k = blockIdx.y; + if (b >= B || k >= K) return; + + extern __shared__ float smem[]; + int tid = threadIdx.x; + int block = blockDim.x; + + float local = 0.0f; + for (int c = tid; c < C; c += block) { + local += dh_s2[b * C + c] * expert_outputs[k * (B * C) + b * C + c]; + } + smem[tid] = local; + __syncthreads(); + + for (int s = block / 2; s > 0; s >>= 1) { + if (tid < s) smem[tid] += smem[tid + s]; + __syncthreads(); + } + + if (tid == 0) { + dg[b * K + k] = smem[0]; + __threadfence_system(); + } +} diff --git a/crates/ml/tests/moe_kernels_test.rs b/crates/ml/tests/moe_kernels_test.rs index 7b30e2574..ef1b6a091 100644 --- a/crates/ml/tests/moe_kernels_test.rs +++ b/crates/ml/tests/moe_kernels_test.rs @@ -60,3 +60,100 @@ fn moe_mixture_forward_matches_cpu_reference() { ); } } + +#[test] +#[ignore] +fn moe_mixture_backward_finite_differences() { + let ctx = CudaContext::new(0).expect("CUDA context"); + let stream = ctx.default_stream(); + + const B: usize = 2; + const C: usize = 32; + const K: usize = 4; + const EPS: f32 = 1e-3; + const TOL: f32 = 1e-3; + + let expert_outputs: Vec = (0..K * B * C) + .map(|i| ((i * 13 + 1) % 200) as f32 / 200.0 - 0.25) + .collect(); + let mut gate: Vec = (0..B * K) + .map(|i| ((i * 17 + 7) % 100) as f32 / 100.0) + .collect(); + for b in 0..B { + let s: f32 = (0..K).map(|k| gate[b * K + k]).sum(); + for k in 0..K { + gate[b * K + k] /= s; + } + } + + let dh_s2: Vec = (0..B * C) + .map(|i| ((i * 19 + 11) % 100) as f32 / 100.0 - 0.5) + .collect(); + + let head = GpuMoeHead::new(Arc::clone(&stream)).unwrap(); + let (de_analytic, dg_analytic) = head + .test_mixture_backward(&expert_outputs, &gate, &dh_s2, B, K, C) + .unwrap(); + + // Numerical de_k via finite differences over h_s2 = sum_k g_k * e_k. + // d(loss)/d(e_k0[b0,c0]) = dh_s2[b0,c0] * gate[b0,k0] (analytic). + // Numerically: reuse the cpu_mixture_forward helper. + for k0 in 0..K { + for b0 in 0..B { + for c0 in 0..C { + let idx = k0 * B * C + b0 * C + c0; + let mut e_plus = expert_outputs.clone(); + e_plus[idx] += EPS; + let h_plus = cpu_mixture_forward(&e_plus, &gate, B, K, C); + let mut e_minus = expert_outputs.clone(); + e_minus[idx] -= EPS; + let h_minus = cpu_mixture_forward(&e_minus, &gate, B, K, C); + let mut numerical: f32 = 0.0; + for i in 0..B * C { + numerical += dh_s2[i] * (h_plus[i] - h_minus[i]) / (2.0 * EPS); + } + assert!( + (de_analytic[idx] - numerical).abs() < TOL, + "de_k[{},{},{}]: analytic {:.6}, numerical {:.6}", + k0, b0, c0, de_analytic[idx], numerical, + ); + } + } + } + + // Numerical dg over gate (note: perturbing g without renormalizing is OK + // for the local gradient check since we're testing the kernel's stage + // before any softmax/gate-softmax backward — that lives in cuBLAS). + for b0 in 0..B { + for k0 in 0..K { + let idx = b0 * K + k0; + let mut g_plus = gate.clone(); + g_plus[idx] += EPS; + let h_plus = cpu_mixture_forward(&expert_outputs, &g_plus, B, K, C); + let mut g_minus = gate.clone(); + g_minus[idx] -= EPS; + let h_minus = cpu_mixture_forward(&expert_outputs, &g_minus, B, K, C); + let mut numerical: f32 = 0.0; + for i in 0..B * C { + numerical += dh_s2[i] * (h_plus[i] - h_minus[i]) / (2.0 * EPS); + } + assert!( + (dg_analytic[idx] - numerical).abs() < TOL, + "dg[{},{}]: analytic {:.6}, numerical {:.6}", + b0, k0, dg_analytic[idx], numerical, + ); + } + } +} + +fn cpu_mixture_forward(expert_outputs: &[f32], gate: &[f32], b: usize, k: usize, c: usize) -> Vec { + let mut out = vec![0.0_f32; b * c]; + for bi in 0..b { + for ci in 0..c { + for ki in 0..k { + out[bi * c + ci] += gate[bi * k + ki] * expert_outputs[ki * b * c + bi * c + ci]; + } + } + } + out +} diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index 5b4091aee..1cea6ad48 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -2,6 +2,13 @@ **Status:** Populated during Plan 1 Task 6 (A.5 orphan audit). Updated on every commit per Invariant 7. +MoE backward kernels T2.2 (2026-04-27): `moe_mixture_backward` computes +`de_k[k,b,c] = g[b,k] · dh_s2[b,c]` (one thread per (k,b,c)); +`moe_dgate_reduce` computes `dg[b,k] = Σ_c e_k[k,b,c] · dh_s2[b,c]` via +shmem-tree reduction over c (one block per (b,k)). Finite-differences test +verifies analytic backward matches numerical for B=2, K=4, C=32 within 1e-3. +All data flows via mapped pinned buffers. Consumers land in Phase 3 wire-up. + MoE ISV slot reset registration (2026-04-27): registered fold-boundary reset entries for the 8 `MOE_EXPERT_UTIL_EMA` slots (reset to 1/K=0.125) and the `MOE_GATE_ENTROPY_EMA_INDEX` slot (reset to ln(8)). Producers