lattice-inference 0.4.0

use super::cache::ForwardScratch;
use super::weights::MoeLayerWeights;
use crate::forward::cpu::{elementwise_mul, matmul_bt, silu_inplace};

/// MoE forward step — free function so it can be unit-tested directly.
pub(crate) fn moe_ffn_step(moe: &MoeLayerWeights, scratch: &mut ForwardScratch, hidden: usize) {
    let inter = moe.experts.intermediate_size;
    let shared_inter = moe.shared_expert.intermediate_size;
    let num_experts = moe.router.num_experts;
    let top_k = moe.router.num_experts_per_tok;
    debug_assert_eq!(moe.router.hidden_size, hidden);
    debug_assert_eq!(moe.experts.num_experts, num_experts);
    debug_assert_eq!(moe.experts.hidden_size, hidden);
    debug_assert_eq!(moe.shared_expert.hidden_size, hidden);

    scratch.input_tmp[..hidden].copy_from_slice(&scratch.ffn_out[..hidden]);

    if scratch.router_logits.len() < num_experts {
        scratch.router_logits.resize(num_experts, 0.0);
    }
    if scratch.router_selected.len() < top_k {
        scratch.router_selected.resize(top_k, (usize::MAX, 0.0));
    }

    compute_router_probs(moe, scratch, hidden, num_experts);

    select_top_k(scratch, num_experts, top_k);

    renormalize_selected(scratch, top_k);

    accumulate_routed_experts(moe, scratch, hidden, inter, top_k);

    apply_shared_expert(moe, scratch, hidden, shared_inter);

    scratch.ffn_out[..hidden].copy_from_slice(&scratch.expert_out[..hidden]);
}

pub(crate) fn compute_router_probs(
    moe: &MoeLayerWeights,
    scratch: &mut ForwardScratch,
    hidden: usize,
    num_experts: usize,
) {
    let gate = &moe.router.gate;
    let input = &scratch.input_tmp[..hidden];
    let logits = &mut scratch.router_logits[..num_experts];
    matmul_bt(input, gate, logits, 1, hidden, num_experts);

    let max_logit = scratch.router_logits[..num_experts]
        .iter()
        .copied()
        .fold(f32::NEG_INFINITY, f32::max);
    if !max_logit.is_finite() {
        // Fail closed on a degenerate (all non-finite) router row, mirroring the
        // attention softmax guard (attention/gqa.rs) and the sampler fallback
        // (qwen35/sampling.rs). A non-finite max means every logit is NaN/±inf
        // (corrupt router gate weight or an upstream activation overflow); the
        // exp below would otherwise leave NaN/zero pseudo-probabilities that
        // poison top-k selection (`NaN > x` is always false → no expert is
        // selected → a `usize::MAX` sentinel slot reaches expert slicing) or
        // silently zero the routed contribution. Zeroing the probabilities drops
        // the routed-expert path for this token (the always-on shared expert
        // still runs) instead of panicking or fabricating routing mass.
        scratch.router_logits[..num_experts].fill(0.0);
        return;
    }
    let mut denom = 0.0f32;
    for v in &mut scratch.router_logits[..num_experts] {
        *v = (*v - max_logit).exp();
        denom += *v;
    }
    if denom > 0.0 {
        for v in &mut scratch.router_logits[..num_experts] {
            *v /= denom;
        }
    } else {
        // Fail closed on a non-finite denom, mirroring generate.rs
        // compute_attention's `else { scores.fill(0.0) }` guard. The early
        // `max_logit` check only catches an all-non-finite row; a row with a
        // FINITE max but a NaN/±inf in another lane keeps `max_logit` finite
        // (Rust `f32::max` ignores a single NaN), so the NaN propagates through
        // `exp` into `denom` (NaN) here and `denom > 0.0` is false. Without this
        // else the routed logits would stay un-normalized (raw `exp` values that
        // do not sum to 1), mixing experts with wrong weights. Zeroing drops the
        // routed-expert path (the shared expert still runs) rather than routing
        // on garbage.
        scratch.router_logits[..num_experts].fill(0.0);
    }
}

pub(crate) fn select_top_k(scratch: &mut ForwardScratch, num_experts: usize, top_k: usize) {
    for slot in &mut scratch.router_selected[..top_k] {
        *slot = (usize::MAX, f32::NEG_INFINITY);
    }
    for (expert_id, prob) in scratch.router_logits[..num_experts]
        .iter()
        .copied()
        .enumerate()
    {
        for rank in 0..top_k {
            if prob > scratch.router_selected[rank].1 {
                for shift in (rank + 1..top_k).rev() {
                    scratch.router_selected[shift] = scratch.router_selected[shift - 1];
                }
                scratch.router_selected[rank] = (expert_id, prob);
                break;
            }
        }
    }
}

pub(crate) fn renormalize_selected(scratch: &mut ForwardScratch, top_k: usize) {
    let top_sum: f32 = scratch.router_selected[..top_k]
        .iter()
        .map(|(_, p)| *p)
        .sum();
    if top_sum > 0.0 {
        for (_, prob) in &mut scratch.router_selected[..top_k] {
            *prob /= top_sum;
        }
    }
}

fn accumulate_routed_experts(
    moe: &MoeLayerWeights,
    scratch: &mut ForwardScratch,
    hidden: usize,
    inter: usize,
    top_k: usize,
) {
    scratch.expert_out[..hidden].fill(0.0);

    for idx in 0..top_k {
        let (expert_id, weight) = scratch.router_selected[idx];
        // Defense in depth: never index expert weights with an unfilled sentinel
        // slot. select_top_k leaves `(usize::MAX, _)` for any rank it could not
        // fill (e.g. a degenerate router row where no probability compares
        // greater than the NEG_INFINITY seed); slicing at `usize::MAX * stride`
        // would overflow/panic. compute_router_probs already fails such rows
        // closed, but this keeps accumulation panic-free regardless of upstream
        // router state.
        if expert_id >= moe.experts.num_experts {
            continue;
        }
        let gate_up_stride = 2 * inter * hidden;
        let gate_up_start = expert_id * gate_up_stride;
        let down_start = expert_id * hidden * inter;

        // HF Qwen MoE stores each expert gate_up_proj as [2 * intermediate, hidden]
        // and splits F.linear(...).chunk(2, dim=-1): first half is gate, second is up.
        // See transformers modeling_qwen3_moe.py:L222 and L244.
        let gate_w = &moe.experts.gate_up_proj[gate_up_start..gate_up_start + inter * hidden];
        let up_w = &moe.experts.gate_up_proj
            [gate_up_start + inter * hidden..gate_up_start + 2 * inter * hidden];
        let down_w = &moe.experts.down_proj[down_start..down_start + hidden * inter];

        matmul_bt(
            &scratch.input_tmp[..hidden],
            gate_w,
            &mut scratch.gate_buf[..inter],
            1,
            hidden,
            inter,
        );
        matmul_bt(
            &scratch.input_tmp[..hidden],
            up_w,
            &mut scratch.up_buf[..inter],
            1,
            hidden,
            inter,
        );

        silu_inplace(&mut scratch.gate_buf[..inter]);
        elementwise_mul(&mut scratch.gate_buf[..inter], &scratch.up_buf[..inter]);

        scratch.down_input[..inter].copy_from_slice(&scratch.gate_buf[..inter]);
        matmul_bt(
            &scratch.down_input[..inter],
            down_w,
            &mut scratch.ffn_out[..hidden],
            1,
            inter,
            hidden,
        );

        for i in 0..hidden {
            scratch.expert_out[i] += weight * scratch.ffn_out[i];
        }
    }
}

fn apply_shared_expert(
    moe: &MoeLayerWeights,
    scratch: &mut ForwardScratch,
    hidden: usize,
    shared_inter: usize,
) {
    let shared = &moe.shared_expert;

    let gate_logit: f32 = scratch.input_tmp[..hidden]
        .iter()
        .zip(shared.shared_expert_gate.iter())
        .map(|(x, g)| x * g)
        .sum();
    let shared_gate = 1.0 / (1.0 + (-gate_logit).exp());

    matmul_bt(
        &scratch.input_tmp[..hidden],
        &shared.gate_proj,
        &mut scratch.gate_buf[..shared_inter],
        1,
        hidden,
        shared_inter,
    );
    matmul_bt(
        &scratch.input_tmp[..hidden],
        &shared.up_proj,
        &mut scratch.up_buf[..shared_inter],
        1,
        hidden,
        shared_inter,
    );

    silu_inplace(&mut scratch.gate_buf[..shared_inter]);
    elementwise_mul(
        &mut scratch.gate_buf[..shared_inter],
        &scratch.up_buf[..shared_inter],
    );

    scratch.down_input[..shared_inter].copy_from_slice(&scratch.gate_buf[..shared_inter]);
    matmul_bt(
        &scratch.down_input[..shared_inter],
        &shared.down_proj,
        &mut scratch.ffn_out[..hidden],
        1,
        shared_inter,
        hidden,
    );

    for i in 0..hidden {
        scratch.expert_out[i] += shared_gate * scratch.ffn_out[i];
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    // ── Helpers ─────────────────────────────────────────────────────────────────

    /// Minimal ForwardScratch for router-only tests.
    fn make_scratch(num_experts: usize, top_k: usize, hidden: usize) -> ForwardScratch {
        let mut s = ForwardScratch::new();
        s.ffn_out.resize(hidden, 0.0);
        s.gate_buf.resize(64, 0.0);
        s.up_buf.resize(64, 0.0);
        s.down_input.resize(64, 0.0);
        s.expert_out.resize(hidden, 0.0);
        s.input_tmp.resize(hidden, 0.0);
        s.router_logits.resize(num_experts, 0.0);
        s.router_selected.resize(top_k, (usize::MAX, 0.0));
        s
    }

    // ── select_top_k tests ───────────────────────────────────────────────────────

    /// top-k=2 from 4 experts: verify the correct pair is chosen with descending
    /// probability order (highest first).
    #[test]
    fn test_select_top_k_basic() {
        let num_experts = 4;
        let top_k = 2;
        let mut scratch = make_scratch(num_experts, top_k, 8);

        // Probabilities: expert 2 highest, expert 0 second.
        scratch.router_logits[0] = 0.30;
        scratch.router_logits[1] = 0.10;
        scratch.router_logits[2] = 0.50;
        scratch.router_logits[3] = 0.10;

        select_top_k(&mut scratch, num_experts, top_k);

        // Slot 0 should be the absolute highest (expert 2).
        assert_eq!(
            scratch.router_selected[0].0, 2,
            "slot 0 must be expert 2 (prob 0.50)"
        );
        // Slot 1 should be the second highest (expert 0).
        assert_eq!(
            scratch.router_selected[1].0, 0,
            "slot 1 must be expert 0 (prob 0.30)"
        );
        // Probabilities preserved (not yet renormalized).
        assert!((scratch.router_selected[0].1 - 0.50).abs() < 1e-6);
        assert!((scratch.router_selected[1].1 - 0.30).abs() < 1e-6);
    }

    /// All experts have the same probability — top-k must still produce exactly
    /// top_k unique valid indices without repeats.
    #[test]
    fn test_select_top_k_ties() {
        let num_experts = 8;
        let top_k = 3;
        let mut scratch = make_scratch(num_experts, top_k, 8);
        for v in scratch.router_logits.iter_mut() {
            *v = 0.125; // uniform
        }

        select_top_k(&mut scratch, num_experts, top_k);

        let ids: Vec<usize> = scratch.router_selected[..top_k]
            .iter()
            .map(|(id, _)| *id)
            .collect();
        // All selected experts must be valid indices.
        for &id in &ids {
            assert!(id < num_experts, "invalid expert id {id}");
        }
        // No duplicates.
        let mut seen = std::collections::HashSet::new();
        for &id in &ids {
            assert!(seen.insert(id), "duplicate expert id {id}");
        }
    }

    // ── renormalize_selected tests ────────────────────────────────────────────

    /// After renormalization the selected weights must sum to 1.
    #[test]
    fn test_renormalize_selected_sums_to_one() {
        let top_k = 3;
        let num_experts = 6;
        let mut scratch = make_scratch(num_experts, top_k, 8);

        scratch.router_logits[0] = 0.2;
        scratch.router_logits[1] = 0.5;
        scratch.router_logits[2] = 0.1;
        scratch.router_logits[3] = 0.1;
        scratch.router_logits[4] = 0.05;
        scratch.router_logits[5] = 0.05;

        select_top_k(&mut scratch, num_experts, top_k);
        renormalize_selected(&mut scratch, top_k);

        let sum: f32 = scratch.router_selected[..top_k]
            .iter()
            .map(|(_, p)| *p)
            .sum();
        assert!(
            (sum - 1.0).abs() < 1e-5,
            "renormalized weights must sum to 1.0 (got {sum})"
        );
    }

    // ── Weight layout offset tests (ADR-053 D2) ──────────────────────────────

    /// Verify the expert-major gate_up offset formula matches the documented layout.
    ///
    /// Layout: [num_experts, 2 * inter, hidden] contiguous.
    ///   expert e gate starts at element: e * 2 * inter * hidden
    ///   expert e up   starts at element: e * 2 * inter * hidden + inter * hidden
    #[test]
    fn test_gate_up_expert_major_layout_offsets() {
        let num_experts = 4usize;
        let inter = 8usize;
        let hidden = 16usize;
        let stride = 2 * inter * hidden;

        for e in 0..num_experts {
            let gate_off = e * stride;
            let up_off = e * stride + inter * hidden;

            // Gate block must not overlap with the previous expert's up block.
            if e > 0 {
                let prev_up_end = (e - 1) * stride + 2 * inter * hidden;
                assert_eq!(
                    gate_off,
                    prev_up_end,
                    "expert {e} gate starts immediately after expert {}'s up block",
                    e - 1
                );
            }

            // Up block starts exactly inter*hidden elements after gate block.
            assert_eq!(
                up_off - gate_off,
                inter * hidden,
                "up offset must be exactly inter*hidden after gate offset for expert {e}"
            );

            // Up block ends exactly at (e+1)*stride.
            assert_eq!(
                up_off + inter * hidden,
                (e + 1) * stride,
                "up block for expert {e} must end at the start of expert {}'s gate",
                e + 1
            );
        }

        // Total buffer length must equal num_experts * 2 * inter * hidden.
        assert_eq!(num_experts * stride, num_experts * 2 * inter * hidden);
    }

    /// Verify the expert-major down offset formula.
    ///
    /// Layout: [num_experts, hidden, inter] contiguous.
    ///   expert e starts at element: e * hidden * inter
    #[test]
    fn test_down_expert_major_layout_offsets() {
        let num_experts = 4usize;
        let inter = 8usize;
        let hidden = 16usize;
        let stride = hidden * inter;

        for e in 0..num_experts {
            let down_off = e * stride;
            // Each expert occupies exactly hidden*inter elements.
            let next_off = (e + 1) * stride;
            assert_eq!(
                next_off - down_off,
                hidden * inter,
                "expert {e} down block size mismatch"
            );
        }

        // Total buffer length.
        assert_eq!(num_experts * stride, num_experts * hidden * inter);
    }

    // ── compute_router_probs tests ────────────────────────────────────────────

    /// Softmax over a 1-hot logit vector: the hot entry should have probability
    /// approaching 1 for large logit values.
    #[test]
    fn test_compute_router_probs_one_hot() {
        use crate::model::qwen35::weights::{
            MoeLayerWeights, MoeRouter, RoutedExperts, SharedExpert,
        };

        let num_experts = 4usize;
        let hidden = 4usize;
        let inter = 2usize;
        let top_k = 1usize;

        // Gate weight matrix: expert 2's row = [1,1,1,1], others = zeros.
        // With input = [1,1,1,1], logits[2] = 4.0, rest = 0.0.
        let mut gate = vec![0.0f32; num_experts * hidden];
        for j in 0..hidden {
            gate[2 * hidden + j] = 1.0;
        }
        let router = MoeRouter::new(gate, num_experts, top_k, hidden).unwrap();

        let gate_up = vec![0.0f32; num_experts * 2 * inter * hidden];
        let down = vec![0.0f32; num_experts * hidden * inter];
        let experts = RoutedExperts::new(gate_up, down, num_experts, hidden, inter).unwrap();

        let shared_gate_proj = vec![0.0f32; inter * hidden];
        let shared_up_proj = vec![0.0f32; inter * hidden];
        let shared_down_proj = vec![0.0f32; hidden * inter];
        let shared_expert_gate = vec![0.0f32; hidden];
        let shared = SharedExpert::new(
            shared_gate_proj,
            shared_up_proj,
            shared_down_proj,
            shared_expert_gate,
            hidden,
            inter,
        )
        .unwrap();

        let moe = MoeLayerWeights {
            router,
            experts,
            shared_expert: shared,
        };
        let mut scratch = make_scratch(num_experts, top_k, hidden);

        // Input: all ones.
        for v in scratch.ffn_out.iter_mut() {
            *v = 1.0;
        }
        scratch
            .input_tmp
            .copy_from_slice(&scratch.ffn_out[..hidden]);

        compute_router_probs(&moe, &mut scratch, hidden, num_experts);

        // Expert 2 should have the highest probability.
        let argmax = scratch.router_logits[..num_experts]
            .iter()
            .enumerate()
            .max_by(|a, b| a.1.partial_cmp(b.1).unwrap())
            .map(|(i, _)| i)
            .unwrap();
        assert_eq!(
            argmax, 2,
            "expert 2 must have the highest softmax probability"
        );
        // Probabilities must sum to 1.
        let sum: f32 = scratch.router_logits[..num_experts].iter().sum();
        assert!(
            (sum - 1.0).abs() < 1e-5,
            "router probabilities must sum to 1 (got {sum})"
        );
    }

    // ── Coalescing correctness: full routing round-trip ───────────────────────

    /// End-to-end routing round-trip: probs → top-k → renormalize.
    /// Verifies that the coalescing pipeline produces the correct expert indices
    /// and that their renormalized weights sum to 1.
    #[test]
    fn test_routing_round_trip_top2_of_8() {
        let num_experts = 8usize;
        let top_k = 2usize;
        let mut scratch = make_scratch(num_experts, top_k, 8);

        // Synthetic pre-softmax logits: expert 5 highest, expert 1 second.
        scratch
            .router_logits
            .iter_mut()
            .enumerate()
            .for_each(|(i, v)| {
                *v = match i {
                    5 => 3.0,
                    1 => 2.0,
                    _ => 0.0,
                };
            });

        // Manually apply softmax (as compute_router_probs would after matmul).
        let max_l = scratch.router_logits[..num_experts]
            .iter()
            .copied()
            .fold(f32::NEG_INFINITY, f32::max);
        let mut denom = 0.0f32;
        for v in scratch.router_logits[..num_experts].iter_mut() {
            *v = (*v - max_l).exp();
            denom += *v;
        }
        for v in scratch.router_logits[..num_experts].iter_mut() {
            *v /= denom;
        }

        select_top_k(&mut scratch, num_experts, top_k);
        renormalize_selected(&mut scratch, top_k);

        // Expert 5 must be rank-0 (highest pre-softmax → highest post-softmax).
        assert_eq!(scratch.router_selected[0].0, 5, "rank-0 must be expert 5");
        // Expert 1 must be rank-1.
        assert_eq!(scratch.router_selected[1].0, 1, "rank-1 must be expert 1");
        // Renormalized weights must sum to 1.
        let w_sum: f32 = scratch.router_selected[..top_k]
            .iter()
            .map(|(_, p)| *p)
            .sum();
        assert!(
            (w_sum - 1.0).abs() < 1e-5,
            "renormalized weights sum={w_sum}"
        );
    }

    // ── Non-finite router fail-closed (hardening) ─────────────────────────────

    /// A NaN hidden state reaching the MoE layer (corrupt weight / upstream
    /// overflow) makes every router logit NaN. Before the fail-closed guard the
    /// router left NaN probabilities, top-k selected nothing (`NaN > NEG_INF` is
    /// false), and `accumulate_routed_experts` indexed expert weights at
    /// `usize::MAX * stride` → overflow/OOB panic. The router must fail closed
    /// and no `usize::MAX` sentinel may survive into accumulation.
    #[test]
    fn test_moe_router_all_nan_fails_closed_no_panic() {
        use crate::model::qwen35::weights::{
            MoeLayerWeights, MoeRouter, RoutedExperts, SharedExpert,
        };
        let num_experts = 4usize;
        let hidden = 4usize;
        let inter = 2usize;
        let top_k = 2usize;

        let gate = vec![1.0f32; num_experts * hidden];
        let router = MoeRouter::new(gate, num_experts, top_k, hidden).unwrap();
        let gate_up = vec![0.0f32; num_experts * 2 * inter * hidden];
        let down = vec![0.0f32; num_experts * hidden * inter];
        let experts = RoutedExperts::new(gate_up, down, num_experts, hidden, inter).unwrap();
        let shared = SharedExpert::new(
            vec![0.0f32; inter * hidden],
            vec![0.0f32; inter * hidden],
            vec![0.0f32; hidden * inter],
            vec![0.0f32; hidden],
            hidden,
            inter,
        )
        .unwrap();
        let moe = MoeLayerWeights {
            router,
            experts,
            shared_expert: shared,
        };
        let mut scratch = make_scratch(num_experts, top_k, hidden);
        for v in scratch.ffn_out.iter_mut() {
            *v = f32::NAN;
        }

        // Must not panic (was: usize::MAX expert_id → OOB slice in accumulate).
        moe_ffn_step(&moe, &mut scratch, hidden);

        assert!(
            scratch.router_selected[..top_k]
                .iter()
                .all(|(id, _)| *id < num_experts),
            "degenerate router must not leave a usize::MAX sentinel selected"
        );
    }

    /// A FINITE-max router row with a single NaN lane (e.g. one corrupt router
    /// gate weight) keeps `max_logit` finite — Rust `f32::max` ignores the NaN —
    /// so the early `!max_logit.is_finite()` guard does NOT fire. The NaN then
    /// propagates through `exp` into `denom` (NaN), `denom > 0.0` is false, and
    /// without the `else` branch the router would leave un-normalized raw `exp`
    /// values (here `[NaN, 1, 1, 1]`) and route over the finite experts with
    /// wrong weights. The denom-else must fail closed by zeroing the row.
    #[test]
    fn test_moe_router_finite_max_nan_tail_fails_closed() {
        use crate::model::qwen35::weights::{
            MoeLayerWeights, MoeRouter, RoutedExperts, SharedExpert,
        };
        let num_experts = 4usize;
        let hidden = 4usize;
        let inter = 2usize;
        let top_k = 2usize;

        // Gate rows: expert e selects input[e]. A single corrupt weight
        // (gate[0][0] = NaN) makes ONLY expert 0's logit NaN; experts 1..3 stay
        // finite, so max_logit = 1.0 (finite) and the early guard is bypassed.
        let mut gate = vec![0.0f32; num_experts * hidden];
        for e in 0..num_experts {
            gate[e * hidden + e] = 1.0;
        }
        gate[0] = f32::NAN;
        let router = MoeRouter::new(gate, num_experts, top_k, hidden).unwrap();
        let gate_up = vec![0.0f32; num_experts * 2 * inter * hidden];
        let down = vec![0.0f32; num_experts * hidden * inter];
        let experts = RoutedExperts::new(gate_up, down, num_experts, hidden, inter).unwrap();
        let shared = SharedExpert::new(
            vec![0.0f32; inter * hidden],
            vec![0.0f32; inter * hidden],
            vec![0.0f32; hidden * inter],
            vec![0.0f32; hidden],
            hidden,
            inter,
        )
        .unwrap();
        let moe = MoeLayerWeights {
            router,
            experts,
            shared_expert: shared,
        };
        let mut scratch = make_scratch(num_experts, top_k, hidden);
        for v in scratch.input_tmp[..hidden].iter_mut() {
            *v = 1.0;
        }

        compute_router_probs(&moe, &mut scratch, hidden, num_experts);

        // Fail closed: the whole router row must be zeroed, not a mix of NaN and
        // un-normalized 1.0 probabilities.
        assert!(
            scratch.router_logits[..num_experts]
                .iter()
                .all(|p| *p == 0.0),
            "finite-max NaN-tail router row must fail closed to all-zero, got {:?}",
            &scratch.router_logits[..num_experts]
        );
    }

    /// Defense-in-depth: even if a `usize::MAX` sentinel slot reaches
    /// `accumulate_routed_experts` directly, it must be skipped rather than
    /// indexing expert weights out of bounds.
    #[test]
    fn test_accumulate_skips_sentinel_expert_id() {
        use crate::model::qwen35::weights::{
            MoeLayerWeights, MoeRouter, RoutedExperts, SharedExpert,
        };
        let num_experts = 4usize;
        let hidden = 4usize;
        let inter = 2usize;
        let top_k = 2usize;

        let gate = vec![0.0f32; num_experts * hidden];
        let router = MoeRouter::new(gate, num_experts, top_k, hidden).unwrap();
        let gate_up = vec![0.0f32; num_experts * 2 * inter * hidden];
        let down = vec![0.0f32; num_experts * hidden * inter];
        let experts = RoutedExperts::new(gate_up, down, num_experts, hidden, inter).unwrap();
        let shared = SharedExpert::new(
            vec![0.0f32; inter * hidden],
            vec![0.0f32; inter * hidden],
            vec![0.0f32; hidden * inter],
            vec![0.0f32; hidden],
            hidden,
            inter,
        )
        .unwrap();
        let moe = MoeLayerWeights {
            router,
            experts,
            shared_expert: shared,
        };
        let mut scratch = make_scratch(num_experts, top_k, hidden);
        for v in scratch.input_tmp.iter_mut() {
            *v = 1.0;
        }
        // Force one unfilled sentinel slot and one valid expert.
        scratch.router_selected[0] = (usize::MAX, 0.5);
        scratch.router_selected[1] = (1, 0.5);

        // Must not panic on the sentinel slot.
        accumulate_routed_experts(&moe, &mut scratch, hidden, inter, top_k);

        assert!(
            scratch.expert_out[..hidden].iter().all(|v| v.is_finite()),
            "accumulation must stay finite when a sentinel slot is skipped"
        );
    }

    // ── Memory pressure guard: threshold calculation ──────────────────────────

    /// Verify the memory threshold formula used at load time matches the ADR spec.
    /// ADR-053 §Risks R2: reject if routed_bytes > 0.85 × recommendedMaxWorkingSetSize.
    #[test]
    fn test_memory_threshold_formula() {
        let max_working: u64 = 32 * 1024 * 1024 * 1024; // 32 GiB (32 GB M2 Pro)
        let threshold = (max_working as f64 * 0.85) as u64;
        // 32 GiB = 34_359_738_368 bytes; 0.85 × 34_359_738_368 = 29_205_777_612
        let expected: u64 = 29_205_777_612;

        // Allow 1 byte rounding tolerance from floating-point conversion.
        assert!(
            threshold.abs_diff(expected) <= 1,
            "threshold {threshold} != expected {expected}"
        );
    }
}