foxhunt/crates/ml/tests/moe_kernels_test.rs

//! 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<f32> = (0..K * B * C)
        .map(|i| ((i * 7 + 3) % 100) as f32 / 100.0)
        .collect();
    let mut gate_host: Vec<f32> = (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<f32> = 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]
        );
    }
}

#[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<f32> = (0..K * B * C)
        .map(|i| ((i * 13 + 1) % 200) as f32 / 200.0 - 0.25)
        .collect();
    let mut gate: Vec<f32> = (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<f32> = (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<f32> {
    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
}

#[test]
#[ignore]
fn moe_load_balance_loss_correctness() {
    let ctx = CudaContext::new(0).unwrap();
    let stream = ctx.default_stream();

    const B: usize = 8;
    const K: usize = 8;
    const LAMBDA: f32 = 0.01;

    let mut gate: Vec<f32> = (0..B * K)
        .map(|i| ((i * 23 + 5) % 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 mut col_means = vec![0.0_f32; K];
    for b in 0..B {
        for k in 0..K {
            col_means[k] += gate[b * K + k] / B as f32;
        }
    }
    let expected: f32 = LAMBDA * K as f32 * col_means.iter().map(|m| m * m).sum::<f32>();

    let head = GpuMoeHead::new(Arc::clone(&stream)).unwrap();
    let actual = head.test_load_balance_loss(&gate, B, K, LAMBDA).unwrap();

    assert!(
        (actual - expected).abs() < 1e-5,
        "loss expected {:.6}, got {:.6}", expected, actual,
    );
}

#[test]
#[ignore]
fn moe_load_balance_loss_uniform_minimum() {
    let ctx = CudaContext::new(0).unwrap();
    let stream = ctx.default_stream();

    const B: usize = 8;
    const K: usize = 8;
    const LAMBDA: f32 = 0.01;

    let gate: Vec<f32> = vec![1.0 / K as f32; B * K];
    let head = GpuMoeHead::new(Arc::clone(&stream)).unwrap();
    let loss = head.test_load_balance_loss(&gate, B, K, LAMBDA).unwrap();
    assert!(
        (loss - LAMBDA).abs() < 1e-5,
        "uniform loss = λ ({}), got {}",
        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::<f32>() / B as f32;
        expected_util[k] = (1.0 - ALPHA) * isv[util_base + k] + ALPHA * cm;
    }
    let avg: Vec<f32> = (0..K).map(|k| (0..B).map(|b| gate[b * K + k]).sum::<f32>() / 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],
    );
}

/// Adaptive load-balance λ controller (2026-04-27).
///
/// Verifies the formula:
///   target  = entropy_target_frac × ln(K=8)
///   deficit = clamp((target - ent_ema) / target, 0, 1)
///   λ_eff   = λ_floor + λ_max_extra × deficit
///
/// Three regimes:
///   (a) ent_ema ≥ target  → deficit = 0, λ_eff = floor
///   (b) ent_ema = 0       → deficit = 1, λ_eff = floor + max_extra
///   (c) half-collapse    → deficit ≈ 0.5, λ_eff ≈ floor + 0.5×max_extra
#[test]
#[ignore]
fn moe_lambda_eff_update_writes_correct_values() {
    use std::f32::consts::E;

    let ctx = CudaContext::new(0).unwrap();
    let stream = ctx.default_stream();
    let head = GpuMoeHead::new(Arc::clone(&stream)).unwrap();

    const LAMBDA_FLOOR: f32 = 0.01;
    const MAX_EXTRA: f32 = 0.09;
    const TARGET_FRAC: f32 = 0.7;
    const LAMBDA_EFF_IDX: usize = 128;
    const ENTROPY_IDX: usize = 126;

    let ln_k = (8.0_f32).ln(); // 2.07944
    let target = TARGET_FRAC * ln_k; // ≈ 1.4556
    let _ = E; // keep std::f32::consts::E import warning quiet

    // Helper: stage ISV with a chosen entropy value, run the kernel,
    // return the resulting λ_eff at slot 128.
    let run = |ent_ema: f32| -> f32 {
        let mut isv = vec![0.0_f32; 129];
        isv[ENTROPY_IDX] = ent_ema;
        // Pre-fill λ_eff with garbage to confirm the kernel overwrites it.
        isv[LAMBDA_EFF_IDX] = -42.0;
        let new_isv = head
            .test_lambda_eff_update(
                &isv,
                LAMBDA_EFF_IDX,
                ENTROPY_IDX,
                LAMBDA_FLOOR,
                MAX_EXTRA,
                TARGET_FRAC,
            )
            .unwrap();
        new_isv[LAMBDA_EFF_IDX]
    };

    // (a) Diverse gate: ent_ema = 1.6 (above target ≈ 1.456) → λ_eff = floor.
    let lambda_a = run(1.6);
    assert!(
        (lambda_a - LAMBDA_FLOOR).abs() < 1e-6,
        "ent=1.6 → λ_eff expected {:.6}, got {:.6}",
        LAMBDA_FLOOR, lambda_a,
    );

    // (b) Full collapse: ent_ema = 0 → λ_eff = floor + max_extra.
    let lambda_b = run(0.0);
    let expected_b = LAMBDA_FLOOR + MAX_EXTRA;
    assert!(
        (lambda_b - expected_b).abs() < 1e-6,
        "ent=0 → λ_eff expected {:.6}, got {:.6}",
        expected_b, lambda_b,
    );

    // (c) Half collapse: ent_ema = 0.5 × target → deficit = 0.5 → λ_eff = floor + 0.5×max_extra.
    let half_target = 0.5 * target;
    let lambda_c = run(half_target);
    let expected_c = LAMBDA_FLOOR + 0.5 * MAX_EXTRA;
    assert!(
        (lambda_c - expected_c).abs() < 1e-5,
        "ent=0.5×target ({:.4}) → λ_eff expected {:.6}, got {:.6}",
        half_target, expected_c, lambda_c,
    );

    // (d) Exactly at target: deficit = 0, λ_eff = floor.
    let lambda_d = run(target);
    assert!(
        (lambda_d - LAMBDA_FLOOR).abs() < 1e-6,
        "ent=target ({:.4}) → λ_eff expected {:.6}, got {:.6}",
        target, LAMBDA_FLOOR, lambda_d,
    );

    // (e) Above-target (deficit clamped to 0): λ_eff = floor.
    let lambda_e = run(target + 0.5);
    assert!(
        (lambda_e - LAMBDA_FLOOR).abs() < 1e-6,
        "ent=target+0.5 → λ_eff expected {:.6}, got {:.6}",
        LAMBDA_FLOOR, lambda_e,
    );
}