From 3d17c2a93c4a44e22eb050c7408cdeccdc15e15a Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Mon, 27 Apr 2026 18:33:41 +0200 Subject: [PATCH] feat(moe): moe_mixture_forward kernel + Rust wrapper + unit test MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Single-thread-per-(b,c) kernel, no atomicAdd, capture-friendly. CPU-reference test verifies correctness for B=4, K=8, C=256 within 1e-6 tolerance. Test wrapper uses mapped pinned memory exclusively (feedback_no_htod_htoh_only_mapped_pinned.md): allocate MappedF32Buffer, write CPU-side via host_ptr, kernel reads via dev_ptr, output via host_ptr read after stream sync. NO memcpy_stod / memcpy_dtov anywhere. GpuMoeHead struct in crates/ml/src/cuda_pipeline/gpu_moe_head.rs follows existing cuda_pipeline head pattern (cubin loaded once, kernel handles cached). Subsequent kernels (moe_mixture_backward, moe_load_balance_loss, moe_expert_util_ema_update) extend the same struct. Spec: docs/superpowers/specs/2026-04-27-moe-regime-redesign-design.md §6.2. Plan: docs/superpowers/plans/2026-04-27-moe-regime-redesign.md Task 2.1. Co-Authored-By: Claude Opus 4.7 (1M context) --- crates/ml/build.rs | 6 + crates/ml/src/cuda_pipeline/gpu_moe_head.rs | 130 ++++++++++++++++++++ crates/ml/src/cuda_pipeline/mod.rs | 1 + crates/ml/src/cuda_pipeline/moe_kernels.cu | 33 +++++ crates/ml/tests/moe_kernels_test.rs | 62 ++++++++++ docs/dqn-wire-up-audit.md | 9 ++ 6 files changed, 241 insertions(+) create mode 100644 crates/ml/src/cuda_pipeline/gpu_moe_head.rs create mode 100644 crates/ml/src/cuda_pipeline/moe_kernels.cu create mode 100644 crates/ml/tests/moe_kernels_test.rs diff --git a/crates/ml/build.rs b/crates/ml/build.rs index f3ad295d6..e44ddcad3 100644 --- a/crates/ml/build.rs +++ b/crates/ml/build.rs @@ -187,6 +187,12 @@ fn main() { // integrate into `experience_action_select`. Test-only; no production // callers in this commit. "thompson_test_kernel.cu", + // MoE Regime Redesign Phase 2 Task 2.1: mixture-forward kernel. + // Single-thread-per-(b,c), no atomicAdd, capture-friendly. + // h_s2[b,c] = Σ_k g[b,k] · expert_outputs[k,b,c]. + // Subsequent tasks extend this cubin with moe_mixture_backward, + // moe_load_balance_loss, moe_expert_util_ema_update. + "moe_kernels.cu", ]; // ALL kernels get common header (BF16 types + wrappers) diff --git a/crates/ml/src/cuda_pipeline/gpu_moe_head.rs b/crates/ml/src/cuda_pipeline/gpu_moe_head.rs new file mode 100644 index 000000000..0d26a4816 --- /dev/null +++ b/crates/ml/src/cuda_pipeline/gpu_moe_head.rs @@ -0,0 +1,130 @@ +#![allow(unsafe_code)] // Required for CUDA kernel launches. + +//! Rust wrapper for MoE GPU kernels. +//! Loaded from cubin at runtime, like other cuda_pipeline heads (gpu_iqn_head, etc.). +//! +//! Per `feedback_no_htod_htoh_only_mapped_pinned.md`: all CPU<->GPU +//! communication uses mapped pinned memory. The test entry points +//! allocate MappedF32Buffers, write CPU data via host_ptr, launch the +//! kernel against dev_ptr, sync, and read GPU-written output via host_ptr. + +use std::sync::Arc; + +use cudarc::driver::{CudaFunction, CudaStream, LaunchConfig, PushKernelArg}; + +use crate::MLError; +use super::mapped_pinned::MappedF32Buffer; + +static MOE_CUBIN: &[u8] = + include_bytes!(concat!(env!("OUT_DIR"), "/moe_kernels.cubin")); + +/// 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. +#[allow(missing_debug_implementations)] +pub struct GpuMoeHead { + stream: Arc, + moe_mixture_forward: CudaFunction, +} + +impl GpuMoeHead { + /// Load MoE cubin and cache all kernel function handles. + /// Mirrors the loader pattern in `GpuAuxHeads::new` and `GrnBlock::new`. + pub fn new(stream: Arc) -> Result { + let context = stream.context(); + let module = context + .load_cubin(MOE_CUBIN.to_vec()) + .map_err(|e| MLError::ModelError(format!("MoE cubin load: {e}")))?; + let moe_mixture_forward = module + .load_function("moe_mixture_forward") + .map_err(|e| MLError::ModelError(format!("moe_mixture_forward load: {e}")))?; + Ok(Self { + stream, + moe_mixture_forward, + }) + } + + /// Test-only entry — stages CPU data via mapped pinned memory, runs the + /// kernel, reads mapped pinned output. NO HtoD/HtoH per + /// `feedback_no_htod_htoh_only_mapped_pinned.md`. + /// + /// Arguments: + /// `expert_outputs` — `[K, B, C]` row-major flat slice. + /// `gate` — `[B, K]` row-major flat slice (rows must sum to 1). + /// `b` — batch size. + /// `k` — number of experts. + /// `c` — hidden/channel dimension. + /// + /// Returns `h_s2` as a `[B, C]` flat `Vec`. + pub fn test_mixture_forward( + &self, + expert_outputs: &[f32], + gate: &[f32], + b: usize, + k: usize, + c: usize, + ) -> Result, MLError> { + assert_eq!( + expert_outputs.len(), + k * b * c, + "expert_outputs length mismatch: expected {} ({k}×{b}×{c}), got {}", + k * b * c, + expert_outputs.len() + ); + assert_eq!( + gate.len(), + b * k, + "gate length mismatch: expected {} ({b}×{k}), got {}", + b * k, + gate.len() + ); + + // Allocate mapped pinned buffers (CPU-and-GPU coherent). + 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 h_s2_buf = unsafe { MappedF32Buffer::new(b * c) } + .map_err(|e| MLError::ModelError(format!("h_s2 buf alloc: {e}")))?; + + // Direct host_ptr writes, no memcpy. + expert_buf.write_from_slice(expert_outputs); + gate_buf.write_from_slice(gate); + + let total = (b * c) as u32; + let block: u32 = 256; + let grid = total.div_ceil(block); + + let dev_expert = expert_buf.dev_ptr; + let dev_gate = gate_buf.dev_ptr; + let dev_h_s2 = h_s2_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_forward) + .arg(&dev_expert) + .arg(&dev_gate) + .arg(&dev_h_s2) + .arg(&b_i32) + .arg(&k_i32) + .arg(&c_i32) + .launch(LaunchConfig { + grid_dim: (grid, 1, 1), + block_dim: (block, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("moe_mixture_forward launch: {e}")))?; + } + self.stream + .synchronize() + .map_err(|e| MLError::ModelError(format!("synchronize: {e}")))?; + + // Read GPU-written result via host_ptr (mapped pinned coherence). + Ok(h_s2_buf.read_all()) + } +} diff --git a/crates/ml/src/cuda_pipeline/mod.rs b/crates/ml/src/cuda_pipeline/mod.rs index f600c6c48..6c7e2e8ca 100644 --- a/crates/ml/src/cuda_pipeline/mod.rs +++ b/crates/ml/src/cuda_pipeline/mod.rs @@ -43,6 +43,7 @@ pub mod gpu_attention; pub mod gpu_tlob; pub mod gpu_grn; pub mod gpu_aux_heads; +pub mod gpu_moe_head; pub mod decision_transformer; pub mod learning_health; pub mod q_snapshot; diff --git a/crates/ml/src/cuda_pipeline/moe_kernels.cu b/crates/ml/src/cuda_pipeline/moe_kernels.cu new file mode 100644 index 000000000..c81f69a65 --- /dev/null +++ b/crates/ml/src/cuda_pipeline/moe_kernels.cu @@ -0,0 +1,33 @@ +/* MoE mixture / load-balance / ISV-producer kernels for the Mixture-of- + * Experts redesign. Spec: + * docs/superpowers/specs/2026-04-27-moe-regime-redesign-design.md §6.2. + * + * GPU-resident, capture-friendly, no atomicAdd + * (per `feedback_no_atomicadd.md`). + */ + +extern "C" __global__ void moe_mixture_forward( + const float* __restrict__ expert_outputs, /* [K, B, C] */ + const float* __restrict__ gate, /* [B, K] */ + float* __restrict__ h_s2, /* [B, C] */ + int B, + int K, + int C +) { + /* One thread per (b, c) pair. */ + int total = B * C; + int tid = blockIdx.x * blockDim.x + threadIdx.x; + if (tid >= total) return; + + int b = tid / C; + int c = tid % C; + + float acc = 0.0f; + for (int k = 0; k < K; ++k) { + float g = gate[b * K + k]; + float ek = expert_outputs[k * (B * C) + b * C + c]; + acc += g * ek; + } + h_s2[tid] = acc; + __threadfence_system(); /* make write PCIe-visible to mapped pinned host_ptr */ +} diff --git a/crates/ml/tests/moe_kernels_test.rs b/crates/ml/tests/moe_kernels_test.rs new file mode 100644 index 000000000..7b30e2574 --- /dev/null +++ b/crates/ml/tests/moe_kernels_test.rs @@ -0,0 +1,62 @@ +//! Layer 1 unit tests for MoE kernels. +//! Spec: docs/superpowers/specs/2026-04-27-moe-regime-redesign-design.md §8. + +#![allow(unsafe_code)] // Required for CUDA + mapped pinned memory. + +use std::sync::Arc; +use cudarc::driver::CudaContext; +use ml::cuda_pipeline::gpu_moe_head::GpuMoeHead; + +#[test] +#[ignore] // Requires CUDA; run with --ignored. +fn moe_mixture_forward_matches_cpu_reference() { + let ctx = CudaContext::new(0).expect("CUDA context — GPU available?"); + let stream = ctx.default_stream(); + + const B: usize = 4; + const C: usize = 256; + const K: usize = 8; + + // Random expert outputs [K, B, C] and gate [B, K] + let expert_outputs_host: Vec = (0..K * B * C) + .map(|i| ((i * 7 + 3) % 100) as f32 / 100.0) + .collect(); + let mut gate_host: Vec = (0..B * K) + .map(|i| ((i * 11 + 5) % 100) as f32 / 100.0) + .collect(); + // Normalize gate rows to sum to 1 + for b in 0..B { + let row_sum: f32 = (0..K).map(|k| gate_host[b * K + k]).sum(); + for k in 0..K { + gate_host[b * K + k] /= row_sum; + } + } + + // CPU reference: h_s2[b, c] = sum_k g[b, k] * e_k[b, c] + let mut expected: Vec = vec![0.0; B * C]; + for b in 0..B { + for c in 0..C { + for k in 0..K { + expected[b * C + c] += + gate_host[b * K + k] * expert_outputs_host[k * B * C + b * C + c]; + } + } + } + + // GPU: stage via mapped pinned memory, run kernel, read mapped pinned output. + let head = GpuMoeHead::new(Arc::clone(&stream)).expect("MoE head init"); + let actual = head + .test_mixture_forward(&expert_outputs_host, &gate_host, B, K, C) + .expect("kernel run"); + + // Compare with tolerance 1e-6 + for i in 0..B * C { + assert!( + (expected[i] - actual[i]).abs() < 1e-6, + "Mismatch at index {}: expected {:.6}, got {:.6}", + i, + expected[i], + actual[i] + ); + } +} diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index a41d5fdd4..5b4091aee 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -787,3 +787,12 @@ all kernel test wrappers (Test 0.F + upcoming MoE kernel tests) share one implementation. Adds `write_from_slice` helper for direct host_ptr writes (no memcpy). Test 0.F bit-identical post-move. Per `feedback_no_htod_htoh_only_mapped_pinned.md`. + +MoE moe_mixture_forward kernel + Rust wrapper (2026-04-27): first MoE +CUDA kernel landed. Single-thread-per-(b,c) kernel computes h_s2[b,c] = +Σ_k g[b,k]·expert_outputs[k,b,c]. No atomicAdd, capture-friendly. Rust +wrapper `GpuMoeHead` in `crates/ml/src/cuda_pipeline/gpu_moe_head.rs` +loads cubin, exposes test_mixture_forward using mapped pinned buffers +exclusively (no HtoD/HtoH per `feedback_no_htod_htoh_only_mapped_pinned.md`). +Unit test in `crates/ml/tests/moe_kernels_test.rs` verifies CPU +reference match within 1e-6 for B=4, K=8, C=256.