From 95965dd11bc5dfb6bcf6e148c903c9429cef76c1 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Mon, 27 Apr 2026 18:47:34 +0200 Subject: [PATCH] feat(moe): moe_expert_util_ema_update ISV producer kernel + test MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Single-thread cold-path-cadence kernel writes 8 per-expert utilization EMAs + 1 gate-entropy EMA into ISV slots [118..127), α=0.05. Same shape as h_s2_rms_ema_update / aux_heads_loss_ema_update. GPU-resident, CPU read-only per pearl_cold_path_no_exception_to_gpu_drives.md. Co-Authored-By: Claude Opus 4.7 (1M context) --- crates/ml/src/cuda_pipeline/gpu_moe_head.rs | 67 ++++++++++++++++++++- crates/ml/src/cuda_pipeline/moe_kernels.cu | 35 +++++++++++ crates/ml/tests/moe_kernels_test.rs | 51 ++++++++++++++++ docs/dqn-wire-up-audit.md | 7 +++ 4 files changed, 157 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 fb0c4025d..de81d44b1 100644 --- a/crates/ml/src/cuda_pipeline/gpu_moe_head.rs +++ b/crates/ml/src/cuda_pipeline/gpu_moe_head.rs @@ -20,9 +20,8 @@ static MOE_CUBIN: &[u8] = /// GPU-accelerated MoE head. /// -/// Owns the kernel function handles for the MoE mixture forward/backward pass -/// and load-balance loss computation. -/// Subsequent tasks will add: moe_expert_util_ema_update. +/// Owns the kernel function handles for the MoE mixture forward/backward pass, +/// load-balance loss computation, and expert utilization EMA update. #[allow(missing_debug_implementations)] pub struct GpuMoeHead { stream: Arc, @@ -31,6 +30,7 @@ pub struct GpuMoeHead { moe_dgate_reduce: CudaFunction, moe_load_balance_loss: CudaFunction, moe_load_balance_reduce: CudaFunction, + moe_expert_util_ema_update: CudaFunction, } impl GpuMoeHead { @@ -56,6 +56,9 @@ impl GpuMoeHead { let moe_load_balance_reduce = module .load_function("moe_load_balance_reduce") .map_err(|e| MLError::ModelError(format!("moe_load_balance_reduce load: {e}")))?; + let moe_expert_util_ema_update = module + .load_function("moe_expert_util_ema_update") + .map_err(|e| MLError::ModelError(format!("moe_expert_util_ema_update load: {e}")))?; Ok(Self { stream, moe_mixture_forward, @@ -63,6 +66,7 @@ impl GpuMoeHead { moe_dgate_reduce, moe_load_balance_loss, moe_load_balance_reduce, + moe_expert_util_ema_update, }) } @@ -292,4 +296,61 @@ impl GpuMoeHead { Ok(loss_total_buf.read_all()[0]) } + + /// Test-only entry for the ISV EMA update kernel. + /// + /// Writes 8 per-expert utilization EMAs + 1 gate-entropy EMA into the + /// provided ISV buffer at `isv_util_base..isv_util_base+K` and + /// `isv_entropy_index`, with decay `alpha`. + /// + /// GPU-resident per `pearl_cold_path_no_exception_to_gpu_drives.md`. + /// Returns the full updated ISV as a `Vec`. + pub fn test_expert_util_ema_update( + &self, + isv: &[f32], + gate: &[f32], + b: usize, + k: usize, + alpha: f32, + isv_util_base: usize, + isv_entropy_index: usize, + ) -> Result, MLError> { + let isv_buf = unsafe { MappedF32Buffer::new(isv.len()) } + .map_err(|e| MLError::ModelError(format!("isv buf alloc: {e}")))?; + let gate_buf = unsafe { MappedF32Buffer::new(b * k) } + .map_err(|e| MLError::ModelError(format!("gate buf alloc: {e}")))?; + + isv_buf.write_from_slice(isv); + gate_buf.write_from_slice(gate); + + let dev_gate = gate_buf.dev_ptr; + let dev_isv = isv_buf.dev_ptr; + let b_i32 = b as i32; + let k_i32 = k as i32; + let util_base_i32 = isv_util_base as i32; + let entropy_idx_i32 = isv_entropy_index as i32; + + unsafe { + self.stream + .launch_builder(&self.moe_expert_util_ema_update) + .arg(&dev_gate) + .arg(&dev_isv) + .arg(&b_i32) + .arg(&k_i32) + .arg(&util_base_i32) + .arg(&entropy_idx_i32) + .arg(&alpha) + .launch(LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (1, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("moe_expert_util_ema_update launch: {e}")))?; + } + self.stream + .synchronize() + .map_err(|e| MLError::ModelError(format!("synchronize: {e}")))?; + + Ok(isv_buf.read_all()) + } } diff --git a/crates/ml/src/cuda_pipeline/moe_kernels.cu b/crates/ml/src/cuda_pipeline/moe_kernels.cu index d704bb483..4acf35fe1 100644 --- a/crates/ml/src/cuda_pipeline/moe_kernels.cu +++ b/crates/ml/src/cuda_pipeline/moe_kernels.cu @@ -136,3 +136,38 @@ extern "C" __global__ void moe_load_balance_reduce( loss_total[0] = acc; __threadfence_system(); } + +extern "C" __global__ void moe_expert_util_ema_update( + const float* __restrict__ gate, /* [B, K] */ + float* __restrict__ isv, /* [ISV_TOTAL_DIM] */ + int B, + int K, + int isv_util_base, + int isv_entropy_index, + float alpha +) { + /* Single-block, single-thread cold-path cadence kernel (matches + * h_s2_rms_ema_update / aux_heads_loss_ema_update shape). */ + if (blockIdx.x != 0 || threadIdx.x != 0) return; + + const float one_minus_alpha = 1.0f - alpha; + + for (int k = 0; k < K; ++k) { + float sum = 0.0f; + for (int b = 0; b < B; ++b) sum += gate[b * K + k]; + float col_mean = sum / (float)B; + float prev = isv[isv_util_base + k]; + isv[isv_util_base + k] = one_minus_alpha * prev + alpha * col_mean; + } + + float entropy = 0.0f; + for (int k = 0; k < K; ++k) { + float sum = 0.0f; + for (int b = 0; b < B; ++b) sum += gate[b * K + k]; + float col_mean = sum / (float)B; + if (col_mean > 1e-9f) entropy += -col_mean * logf(col_mean); + } + float prev_e = isv[isv_entropy_index]; + isv[isv_entropy_index] = one_minus_alpha * prev_e + alpha * entropy; + __threadfence_system(); +} diff --git a/crates/ml/tests/moe_kernels_test.rs b/crates/ml/tests/moe_kernels_test.rs index 8e07afcd0..ad29434ff 100644 --- a/crates/ml/tests/moe_kernels_test.rs +++ b/crates/ml/tests/moe_kernels_test.rs @@ -214,3 +214,54 @@ fn moe_load_balance_loss_uniform_minimum() { LAMBDA, loss, ); } + +#[test] +#[ignore] +fn moe_expert_util_ema_update_writes_correct_values() { + let ctx = CudaContext::new(0).unwrap(); + let stream = ctx.default_stream(); + + const B: usize = 16; + const K: usize = 8; + const ALPHA: f32 = 0.05; + + let mut isv = vec![0.0_f32; 127]; + let util_base = 118; + for k in 0..K { + isv[util_base + k] = 1.0 / K as f32; + } + isv[126] = (K as f32).ln(); + + let mut gate = vec![0.0_f32; B * K]; + for b in 0..B { + for k in 0..K { + gate[b * K + k] = if k == 3 { 0.6 } else { 0.057143 }; + } + } + + // CPU reference + let mut expected_util = vec![0.0_f32; K]; + for k in 0..K { + let cm: f32 = (0..B).map(|b| gate[b * K + k]).sum::() / B as f32; + expected_util[k] = (1.0 - ALPHA) * isv[util_base + k] + ALPHA * cm; + } + let avg: Vec = (0..K).map(|k| (0..B).map(|b| gate[b * K + k]).sum::() / B as f32).collect(); + let avg_ent: f32 = avg.iter().map(|p| if *p > 1e-9 { -p * p.ln() } else { 0.0 }).sum(); + let expected_ent = (1.0 - ALPHA) * isv[126] + ALPHA * avg_ent; + + let head = GpuMoeHead::new(Arc::clone(&stream)).unwrap(); + let new_isv = head + .test_expert_util_ema_update(&isv, &gate, B, K, ALPHA, util_base, 126) + .unwrap(); + + for k in 0..K { + assert!( + (new_isv[util_base + k] - expected_util[k]).abs() < 1e-5, + "util[{}]: expected {:.6}, got {:.6}", k, expected_util[k], new_isv[util_base + k], + ); + } + assert!( + (new_isv[126] - expected_ent).abs() < 1e-5, + "entropy: expected {:.6}, got {:.6}", expected_ent, new_isv[126], + ); +} diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index cba63e1d5..faa8984fc 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 expert util EMA T2.4 (2026-04-27): `moe_expert_util_ema_update` +single-thread cold-path-cadence kernel writes 8 per-expert utilization +EMAs (ISV[118..126)) + gate-entropy EMA (ISV[126]) with α=0.05. Same shape +as `h_s2_rms_ema_update`/`aux_heads_loss_ema_update`. GPU-resident per +`pearl_cold_path_no_exception_to_gpu_drives.md`. Test verifies EMA values +match CPU reference within 1e-5 using skewed gate (expert-3 = 0.6). + MoE load-balance loss T2.3 (2026-04-27): `moe_load_balance_loss` (one block per k, shmem reduction over B) computes `λ·K·(mean_b g[b,k])²` per expert without atomicAdd; `moe_load_balance_reduce` (single-thread scalar