Mixture of Experts (MoE)¶

pub const MoEConfig = struct {
    n_layers: u32,
    d_model: u32,
    n_heads: u32,
    n_kv_heads: u32,
    n_experts: u32,               // N: total number of experts
    n_experts_active: u32,        // k: experts active per token
    d_expert_ff: u32,             // FFN hidden dim per expert
    vocab_size: u32,
    max_seq_len: u32,
    rope_theta: f32 = 10000.0,
    norm_eps: f32 = 1e-5,
    router_aux_loss_coef: f32 = 0.01,

    pub fn headDim(self: MoEConfig) u32 {
        return self.d_model / self.n_heads;
    }

    pub fn totalParams(self: MoEConfig) u64 {
        const expert_params = @as(u64, self.n_experts)
            * 3 * self.d_model * self.d_expert_ff;
        // ... add attention, embedding, norm params
        return expert_params;
    }

    pub fn activeParams(self: MoEConfig) u64 {
        const active_expert_params = @as(u64, self.n_experts_active)
            * 3 * self.d_model * self.d_expert_ff;
        return active_expert_params;
    }
};

pub const ExpertRouter = struct {
    gate_weight: Tensor(f32),    // [d_model, n_experts]
    n_experts: u32,
    n_active: u32,

    pub const RoutingResult = struct {
        expert_indices: []u32,    // top-k expert indices
        gate_values: []f32,       // normalized gate scores
    };

    pub fn route(self: *ExpertRouter, x: []f32) !RoutingResult {
        // Compute gate scores: x @ W_g
        var scores = try linearForward(x, self.gate_weight);

        // Softmax over all experts
        softmaxInPlace(scores);

        // Select top-k
        var result = try topK(scores, self.n_active);

        // Re-normalize gate values over selected experts
        var gate_sum: f32 = 0;
        for (result.gate_values) |g| gate_sum += g;
        for (result.gate_values) |*g| g.* /= gate_sum;

        return result;
    }
};

pub const MoELayer = struct {
    router: ExpertRouter,
    experts: []SwiGLU_FFN,        // N independent FFN experts
    config: MoEConfig,

    pub fn forward(self: *MoELayer, x: []f32) ![]f32 {
        // Route: determine which experts process this token
        const routing = try self.router.route(x);

        // Compute weighted sum of selected expert outputs
        var output = try allocator.alloc(f32, x.len);
        @memset(output, 0);

        for (routing.expert_indices, routing.gate_values) |idx, gate| {
            const expert_out = try self.experts[idx].forward(x);
            defer allocator.free(expert_out);

            // Accumulate: output += gate * expert_output
            for (0..output.len) |i| {
                output[i] += gate * expert_out[i];
            }
        }

        return output;
    }
};

pub fn computeAuxLoss(
    routing_decisions: []const ExpertRouter.RoutingResult,
    n_experts: u32,
    n_tokens: u32,
    coef: f32,
) f32 {
    var expert_counts = try allocator.alloc(f32, n_experts);
    var expert_probs = try allocator.alloc(f32, n_experts);
    @memset(expert_counts, 0);
    @memset(expert_probs, 0);

    for (routing_decisions) |decision| {
        for (decision.expert_indices) |idx| {
            expert_counts[idx] += 1.0;
        }
        // Accumulate full probability distribution
        // ...
    }

    // Normalize
    const n = @as(f32, @floatFromInt(n_tokens));
    for (0..n_experts) |i| {
        expert_counts[i] /= n;    // f_i
        expert_probs[i] /= n;     // P_i
    }

    // L_aux = alpha * N * sum(f_i * P_i)
    var loss: f32 = 0;
    for (0..n_experts) |i| {
        loss += expert_counts[i] * expert_probs[i];
    }
    return coef * @as(f32, @floatFromInt(n_experts)) * loss;
}