memra-engine 0.131.0

//! Two-card whole-expert residency and exact FP8 dispatch/slot return.
//! Attention/shared experts stay on their existing owner; routed GEMVs use both cards.
use crate::dsv4_ffi as k;
use cudarc::driver::{CudaEvent, CudaSlice, CudaStream, DevicePtr, DevicePtrMut, DeviceRepr};
use memra_runtime::Gpu;
use std::{
    ffi::c_void,
    sync::{
        Arc, OnceLock,
        atomic::{AtomicU64, Ordering},
    },
};

type Res<T> = Result<T, String>;

/// Gate-only route-load receipt for the whole-expert EP experiment. The matrix
/// executor owns one contiguous expert half per rank; these counters record the
/// actual selected-slot split before any future row/column-split rewrite.
#[derive(Clone, Copy, Debug, Default)]
pub struct EpRouteStats {
    pub calls: u64,
    pub observed_calls: u64,
    pub unobserved_calls: u64,
    pub local_slots: u64,
    pub peer_slots: u64,
    pub busier_slots: u64,
    pub one_row_calls: u64,
    pub local_hist: [u64; 9],
    pub busier_hist: [u64; 9],
}

static ROUTE_STATS_ENABLED: OnceLock<bool> = OnceLock::new();
static ROUTE_STATS_CALLS: AtomicU64 = AtomicU64::new(0);
static ROUTE_STATS_OBSERVED_CALLS: AtomicU64 = AtomicU64::new(0);
static ROUTE_STATS_UNOBSERVED_CALLS: AtomicU64 = AtomicU64::new(0);
static ROUTE_STATS_LOCAL_SLOTS: AtomicU64 = AtomicU64::new(0);
static ROUTE_STATS_PEER_SLOTS: AtomicU64 = AtomicU64::new(0);
static ROUTE_STATS_BUSIER_SLOTS: AtomicU64 = AtomicU64::new(0);
static ROUTE_STATS_ONE_ROW_CALLS: AtomicU64 = AtomicU64::new(0);
static ROUTE_STATS_LOCAL_HIST: OnceLock<[AtomicU64; 9]> = OnceLock::new();
static ROUTE_STATS_BUSIER_HIST: OnceLock<[AtomicU64; 9]> = OnceLock::new();

fn route_stats_enabled() -> bool {
    *ROUTE_STATS_ENABLED.get_or_init(|| {
        matches!(
            std::env::var("MEMRA_DSV4_EP_ROUTE_STATS").as_deref(),
            Ok("1")
        )
    })
}

fn route_stats_hist() -> (&'static [AtomicU64; 9], &'static [AtomicU64; 9]) {
    (
        ROUTE_STATS_LOCAL_HIST.get_or_init(|| std::array::from_fn(|_| AtomicU64::new(0))),
        ROUTE_STATS_BUSIER_HIST.get_or_init(|| std::array::from_fn(|_| AtomicU64::new(0))),
    )
}

fn route_stats_call_delta(
    observed_split: Option<(usize, usize)>,
    rows: usize,
    topk: usize,
) -> EpRouteStats {
    let mut delta = EpRouteStats {
        calls: 1,
        ..EpRouteStats::default()
    };
    let Some((local_slots, peer_slots)) = observed_split else {
        delta.unobserved_calls = 1;
        return delta;
    };
    delta.observed_calls = 1;
    delta.local_slots = local_slots as u64;
    delta.peer_slots = peer_slots as u64;
    let busier = local_slots.max(peer_slots);
    delta.busier_slots = busier as u64;
    if rows == 1 && topk < 9 {
        delta.one_row_calls = 1;
        delta.local_hist[local_slots.min(8)] = 1;
        delta.busier_hist[busier.min(8)] = 1;
    }
    delta
}

fn record_route_stats(observed_split: Option<(usize, usize)>, rows: usize, topk: usize) {
    if !route_stats_enabled() {
        return;
    }
    let delta = route_stats_call_delta(observed_split, rows, topk);
    ROUTE_STATS_CALLS.fetch_add(delta.calls, Ordering::Relaxed);
    ROUTE_STATS_OBSERVED_CALLS.fetch_add(delta.observed_calls, Ordering::Relaxed);
    ROUTE_STATS_UNOBSERVED_CALLS.fetch_add(delta.unobserved_calls, Ordering::Relaxed);
    ROUTE_STATS_LOCAL_SLOTS.fetch_add(delta.local_slots, Ordering::Relaxed);
    ROUTE_STATS_PEER_SLOTS.fetch_add(delta.peer_slots, Ordering::Relaxed);
    ROUTE_STATS_BUSIER_SLOTS.fetch_add(delta.busier_slots, Ordering::Relaxed);
    ROUTE_STATS_ONE_ROW_CALLS.fetch_add(delta.one_row_calls, Ordering::Relaxed);
    let (local, busy) = route_stats_hist();
    for (i, count) in delta.local_hist.into_iter().enumerate() {
        local[i].fetch_add(count, Ordering::Relaxed);
    }
    for (i, count) in delta.busier_hist.into_iter().enumerate() {
        busy[i].fetch_add(count, Ordering::Relaxed);
    }
}

pub fn route_stats_snapshot() -> EpRouteStats {
    let (local, busy) = route_stats_hist();
    EpRouteStats {
        calls: ROUTE_STATS_CALLS.load(Ordering::Relaxed),
        observed_calls: ROUTE_STATS_OBSERVED_CALLS.load(Ordering::Relaxed),
        unobserved_calls: ROUTE_STATS_UNOBSERVED_CALLS.load(Ordering::Relaxed),
        local_slots: ROUTE_STATS_LOCAL_SLOTS.load(Ordering::Relaxed),
        peer_slots: ROUTE_STATS_PEER_SLOTS.load(Ordering::Relaxed),
        busier_slots: ROUTE_STATS_BUSIER_SLOTS.load(Ordering::Relaxed),
        one_row_calls: ROUTE_STATS_ONE_ROW_CALLS.load(Ordering::Relaxed),
        local_hist: std::array::from_fn(|i| local[i].load(Ordering::Relaxed)),
        busier_hist: std::array::from_fn(|i| busy[i].load(Ordering::Relaxed)),
    }
}

impl EpRouteStats {
    pub fn delta(self, before: Self) -> Self {
        Self {
            calls: self.calls.saturating_sub(before.calls),
            observed_calls: self.observed_calls.saturating_sub(before.observed_calls),
            unobserved_calls: self
                .unobserved_calls
                .saturating_sub(before.unobserved_calls),
            local_slots: self.local_slots.saturating_sub(before.local_slots),
            peer_slots: self.peer_slots.saturating_sub(before.peer_slots),
            busier_slots: self.busier_slots.saturating_sub(before.busier_slots),
            one_row_calls: self.one_row_calls.saturating_sub(before.one_row_calls),
            local_hist: std::array::from_fn(|i| {
                self.local_hist[i].saturating_sub(before.local_hist[i])
            }),
            busier_hist: std::array::from_fn(|i| {
                self.busier_hist[i].saturating_sub(before.busier_hist[i])
            }),
        }
    }
}

struct PeerFailureDrain {
    owner: Arc<CudaStream>,
    peer: Arc<CudaStream>,
    complete: bool,
}
impl Drop for PeerFailureDrain {
    fn drop(&mut self) {
        if !self.complete {
            // Drain before borrowed peer destinations can be released on error.
            let _ = self.owner.synchronize();
            let _ = self.peer.synchronize();
        }
    }
}

pub(crate) struct EpLayer {
    pub peer_stage: usize,
    pub local_first: usize,
    pub count: usize,
    pub peer_first: usize,
    pub peer_w: CudaSlice<u8>,
    pub peer_sc: CudaSlice<u8>,
    /// Tiny global-id metadata, replicated; expert code/scale banks are not replicated.
    pub peer_s2: CudaSlice<f32>,
    pub peer_table: Option<CudaSlice<u64>>,
}

pub(crate) struct EpScratch {
    pub xq: CudaSlice<u8>,
    pub xs: CudaSlice<f32>,
    pub ids: CudaSlice<i32>,
    pub weights: CudaSlice<f32>,
    pub g1: CudaSlice<f32>,
    pub g3: CudaSlice<f32>,
    pub h: CudaSlice<f32>,
    pub hq: CudaSlice<u8>,
    pub hs: CudaSlice<f32>,
    pub contribution: CudaSlice<f32>,
    pub returned: CudaSlice<f32>,
    pub grouped: Option<crate::dsv4_grouped::GroupedWork>,
    tx_done: CudaEvent,
    rx_done: CudaEvent,
    pub owner_bytes: u64,
    pub peer_bytes: u64,
}

pub(crate) struct EpCompute<'a> {
    pub xq: &'a CudaSlice<u8>,
    pub xs: &'a CudaSlice<f32>,
    pub ids: &'a CudaSlice<i32>,
    pub weights: &'a CudaSlice<f32>,
    pub g1: &'a mut CudaSlice<f32>,
    pub g3: &'a mut CudaSlice<f32>,
    pub h: &'a mut CudaSlice<f32>,
    pub hq: &'a mut CudaSlice<u8>,
    pub hs: &'a mut CudaSlice<f32>,
    pub contribution: &'a mut CudaSlice<f32>,
}

impl EpScratch {
    #[allow(clippy::too_many_arguments)]
    pub(crate) fn new(
        owner: &Gpu,
        peer: &Gpu,
        rows: usize,
        topk: usize,
        hidden: usize,
        inter: usize,
        matrix_partition: Option<(usize, usize, usize)>,
    ) -> Res<Self> {
        if rows == 0
            || rows > 512
            || topk == 0
            || !hidden.is_multiple_of(128)
            || !inter.is_multiple_of(128)
        {
            return Err("invalid EP workspace shape".into());
        }
        let slots = rows * topk;
        let ps = peer.stream();
        let grouped = match matrix_partition {
            Some((global, first, count)) => Some(crate::dsv4_grouped::GroupedWork::new_partition(
                &ps, global, first, count, slots, hidden, inter,
            )?),
            None => None,
        };
        let xq = ps.alloc_zeros(rows * hidden).map_err(|e| e.to_string())?;
        let xs = ps
            .alloc_zeros(rows * hidden / 128)
            .map_err(|e| e.to_string())?;
        let ids = ps.alloc_zeros(slots).map_err(|e| e.to_string())?;
        let weights = ps.alloc_zeros(slots).map_err(|e| e.to_string())?;
        let g1 = ps.alloc_zeros(slots * inter).map_err(|e| e.to_string())?;
        let g3 = ps.alloc_zeros(slots * inter).map_err(|e| e.to_string())?;
        let h = ps.alloc_zeros(slots * inter).map_err(|e| e.to_string())?;
        let hq = ps.alloc_zeros(slots * inter).map_err(|e| e.to_string())?;
        let hs = ps
            .alloc_zeros(slots * inter / 128)
            .map_err(|e| e.to_string())?;
        let contribution = ps.alloc_zeros(slots * hidden).map_err(|e| e.to_string())?;
        let rx_done = peer.ctx.new_event(None).map_err(|e| e.to_string())?;
        let os = owner.stream();
        let returned = os.alloc_zeros(slots * hidden).map_err(|e| e.to_string())?;
        let tx_done = owner.ctx.new_event(None).map_err(|e| e.to_string())?;
        // Explicit ordering is mandatory with per-buffer event tracking off.
        ps.synchronize().map_err(|e| e.to_string())?;
        os.synchronize().map_err(|e| e.to_string())?;
        let owner_bytes = (slots * hidden * 4) as u64;
        let peer_bytes = (rows * (hidden + hidden / 128 * 4)
            + slots * (8 + inter * 13 + inter / 128 * 4 + hidden * 4))
            as u64
            + grouped.as_ref().map_or(0, |work| work.bytes);
        Ok(Self {
            xq,
            xs,
            ids,
            weights,
            g1,
            g3,
            h,
            hq,
            hs,
            contribution,
            returned,
            grouped,
            tx_done,
            rx_done,
            owner_bytes,
            peer_bytes,
        })
    }
}

pub(crate) fn split_slab(
    owner: &Gpu,
    peer: &Gpu,
    source: &CudaSlice<u8>,
    local_first: usize,
    peer_first: usize,
    count: usize,
    expert_stride: usize,
) -> Res<(CudaSlice<u8>, CudaSlice<u8>)> {
    let bytes = count
        .checked_mul(expert_stride)
        .ok_or("EP slab size overflow")?;
    let local_start = local_first
        .checked_mul(expert_stride)
        .ok_or("EP slab offset overflow")?;
    let peer_start = peer_first
        .checked_mul(expert_stride)
        .ok_or("EP slab offset overflow")?;
    if local_start
        .checked_add(bytes)
        .is_none_or(|n| n > source.len())
        || peer_start
            .checked_add(bytes)
            .is_none_or(|n| n > source.len())
    {
        return Err("EP shard outside original bank".into());
    }
    let os = owner.stream();
    let ps = peer.stream();
    let mut local = os.alloc_zeros::<u8>(bytes).map_err(|e| e.to_string())?;
    let mut remote = ps.alloc_zeros::<u8>(bytes).map_err(|e| e.to_string())?;
    let mut drain = PeerFailureDrain {
        owner: os.clone(),
        peer: ps.clone(),
        complete: false,
    };
    ps.synchronize().map_err(|e| e.to_string())?;
    os.memcpy_dtod(&source.slice(local_start..local_start + bytes), &mut local)
        .map_err(|e| e.to_string())?;
    // Use the already-probed 64 MiB copy class rather than one untested GiB copy.
    for offset in (0..bytes).step_by(64 << 20) {
        let end = (offset + (64 << 20)).min(bytes);
        peer_copy(
            &os,
            &ps,
            &source.slice(peer_start + offset..peer_start + end),
            &mut remote.slice_mut(offset..end),
            end - offset,
        )?;
    }
    os.synchronize().map_err(|e| e.to_string())?;
    // Every retained code/scale byte is checked before the full bank is retired.
    for offset in (0..bytes).step_by(64 << 20) {
        let end = (offset + (64 << 20)).min(bytes);
        let expected_local = os
            .clone_dtoh(&source.slice(local_start + offset..local_start + end))
            .map_err(|e| e.to_string())?;
        let actual_local = os
            .clone_dtoh(&local.slice(offset..end))
            .map_err(|e| e.to_string())?;
        if expected_local != actual_local {
            return Err(format!("EP local shard mismatch at chunk {offset}"));
        }
        let expected_peer = os
            .clone_dtoh(&source.slice(peer_start + offset..peer_start + end))
            .map_err(|e| e.to_string())?;
        let actual_peer = ps
            .clone_dtoh(&remote.slice(offset..end))
            .map_err(|e| e.to_string())?;
        if expected_peer != actual_peer {
            return Err(format!("EP peer shard mismatch at chunk {offset}"));
        }
    }
    drain.complete = true;
    Ok((local, remote))
}

pub(crate) fn peer_copy<T: DeviceRepr, S: DevicePtr<T>, D: DevicePtrMut<T>>(
    source: &Arc<CudaStream>,
    destination: &Arc<CudaStream>,
    src: &S,
    dst: &mut D,
    n: usize,
) -> Res<()> {
    if n > src.len() || n > dst.len() {
        return Err("EP peer copy exceeds workspace".into());
    }
    source
        .context()
        .bind_to_thread()
        .map_err(|e| e.to_string())?;
    let (sp, _read) = src.device_ptr(source);
    let (dp, _write) = dst.device_ptr_mut(source);
    unsafe {
        cudarc::driver::result::memcpy_peer_async(
            destination.context().cu_ctx(),
            dp,
            source.context().cu_ctx(),
            sp,
            n * std::mem::size_of::<T>(),
            source.cu_stream(),
        )
        .map_err(|e| format!("EP peer copy: {e}"))
    }
}

#[allow(clippy::too_many_arguments)]
fn chain(
    gpu: &Gpu,
    w: &CudaSlice<u8>,
    sc: &CudaSlice<u8>,
    s2: &CudaSlice<f32>,
    first: usize,
    count: usize,
    ws: &mut EpCompute<'_>,
    rows: usize,
    topk: usize,
    hidden: usize,
    inter: usize,
    limit: f32,
    reduction: i32,
) -> Res<()> {
    gpu.ctx.bind_to_thread().map_err(|e| e.to_string())?;
    let s = gpu.stream();
    let sv = s.cu_stream().cast::<c_void>();
    let slots = rows * topk;
    let wstride = (inter * hidden / 2) as i64;
    let sstride = (inter * hidden / 16) as i64;
    let wp = w.device_ptr(&s).0 as *const c_void;
    let sp = sc.device_ptr(&s).0 as *const c_void;
    let s2p = s2.device_ptr(&s).0 as *const f32;
    let selected = ws.ids.device_ptr(&s).0 as *const i32;
    unsafe {
        for (projection, out) in [(0, &mut ws.g1), (2, &mut ws.g3)] {
            k::ck(
                "EP w1/w3",
                k::memra_dsv4_fp4_gemm_sel_ep(
                    ws.xq.device_ptr(&s).0 as *const c_void,
                    ws.xs.device_ptr(&s).0 as *const f32,
                    wp,
                    sp,
                    s2p,
                    selected,
                    projection,
                    0,
                    0,
                    out.device_ptr_mut(&s).0 as *mut f32,
                    slots as i32,
                    inter as i32,
                    hidden as i32,
                    wstride,
                    sstride,
                    topk as i32,
                    reduction,
                    first as i32,
                    count as i32,
                    sv,
                ),
            )?;
        }
        k::ck(
            "EP swiglu",
            k::memra_dsv4_swiglu(
                ws.g1.device_ptr(&s).0 as *const f32,
                ws.g3.device_ptr(&s).0 as *const f32,
                ws.h.device_ptr_mut(&s).0 as *mut f32,
                slots as i32,
                inter as i32,
                limit,
                ws.weights.device_ptr(&s).0 as *const f32,
                sv,
            ),
        )?;
        k::ck(
            "EP FP8 h",
            k::memra_dsv4_act_quant_fp8(
                ws.h.device_ptr(&s).0 as *const f32,
                ws.hq.device_ptr_mut(&s).0 as *mut c_void,
                ws.hs.device_ptr_mut(&s).0 as *mut f32,
                slots as i32,
                inter as i32,
                sv,
            ),
        )?;
        k::ck(
            "EP w2",
            k::memra_dsv4_fp4_gemm_sel_ep(
                ws.hq.device_ptr(&s).0 as *const c_void,
                ws.hs.device_ptr(&s).0 as *const f32,
                wp,
                sp,
                s2p,
                selected,
                1,
                1,
                0,
                ws.contribution.device_ptr_mut(&s).0 as *mut f32,
                slots as i32,
                hidden as i32,
                inter as i32,
                wstride,
                sstride,
                0,
                reduction,
                first as i32,
                count as i32,
                sv,
            ),
        )?;
    }
    Ok(())
}

#[allow(clippy::too_many_arguments)]
pub(crate) fn execute(
    owner: &Gpu,
    peer: &Gpu,
    bank: &EpLayer,
    w: &CudaSlice<u8>,
    sc: &CudaSlice<u8>,
    s2: &CudaSlice<f32>,
    local: &mut EpCompute<'_>,
    remote: &mut EpScratch,
    rows: usize,
    topk: usize,
    hidden: usize,
    inter: usize,
    limit: f32,
    reduction: i32,
    serial_control: bool,
) -> Res<(u64, u64)> {
    let os = owner.stream();
    let ps = peer.stream();
    let slots = rows * topk;
    let mut drain = PeerFailureDrain {
        owner: os.clone(),
        peer: ps.clone(),
        complete: false,
    };
    // Quantize once on the owner. The wire carries the exact FP8 codes/scales,
    // global router ids and routing weights, never a substitute activation program.
    peer_copy(&os, &ps, local.xq, &mut remote.xq, rows * hidden)?;
    peer_copy(&os, &ps, local.xs, &mut remote.xs, rows * hidden / 128)?;
    peer_copy(&os, &ps, local.ids, &mut remote.ids, slots)?;
    peer_copy(&os, &ps, local.weights, &mut remote.weights, slots)?;
    remote.tx_done.record(&os).map_err(|e| e.to_string())?;
    chain(
        owner,
        w,
        sc,
        s2,
        bank.local_first,
        bank.count,
        local,
        rows,
        topk,
        hidden,
        inter,
        limit,
        reduction,
    )?;
    if serial_control {
        os.synchronize().map_err(|e| e.to_string())?;
    }
    peer.ctx.bind_to_thread().map_err(|e| e.to_string())?;
    ps.wait(&remote.tx_done).map_err(|e| e.to_string())?;
    chain(
        peer,
        &bank.peer_w,
        &bank.peer_sc,
        &bank.peer_s2,
        bank.peer_first,
        bank.count,
        &mut EpCompute {
            xq: &remote.xq,
            xs: &remote.xs,
            ids: &remote.ids,
            weights: &remote.weights,
            g1: &mut remote.g1,
            g3: &mut remote.g3,
            h: &mut remote.h,
            hq: &mut remote.hq,
            hs: &mut remote.hs,
            contribution: &mut remote.contribution,
        },
        rows,
        topk,
        hidden,
        inter,
        limit,
        reduction,
    )?;
    peer_copy(
        &ps,
        &os,
        &remote.contribution,
        &mut remote.returned,
        slots * hidden,
    )?;
    remote.rx_done.record(&ps).map_err(|e| e.to_string())?;
    owner.ctx.bind_to_thread().map_err(|e| e.to_string())?;
    os.wait(&remote.rx_done).map_err(|e| e.to_string())?;
    // Overwrite peer-owned slots rather than adding partial vectors. This keeps
    // signed-zero bits and the existing ascending-expert combine order unchanged.
    unsafe {
        k::ck(
            "EP original-slot return",
            k::memra_dsv4_ep_merge_slots(
                local.contribution.device_ptr_mut(&os).0 as *mut f32,
                remote.returned.device_ptr(&os).0 as *const f32,
                local.ids.device_ptr(&os).0 as *const i32,
                slots as i32,
                hidden as i32,
                bank.peer_first as i32,
                bank.count as i32,
                os.cu_stream().cast(),
            ),
        )?;
    }
    drain.complete = true;
    Ok((
        (rows * (hidden + hidden / 128 * 4) + slots * 8) as u64,
        (slots * hidden * 4) as u64,
    ))
}

/// Same matrix expert program on both ranks, staged around explicit dependencies.
/// Both gate/up chains are queued before either intermediate mirror validation.
#[allow(clippy::too_many_arguments)]
pub(crate) fn execute_matrix(
    owner: &Gpu,
    peer: &Gpu,
    bank: &EpLayer,
    table: &CudaSlice<u64>,
    scale2: &CudaSlice<f32>,
    scale2_host: &[f32],
    local: &mut EpCompute<'_>,
    local_work: &mut crate::dsv4_grouped::GroupedWork,
    remote: &mut EpScratch,
    rows: usize,
    topk: usize,
    hidden: usize,
    limit: f32,
    allow_gu_fuse: bool,
    serial_control: bool,
) -> Res<u64> {
    let peer_table = bank
        .peer_table
        .as_ref()
        .ok_or("EP matrix peer table missing")?;
    let global = bank
        .count
        .checked_mul(2)
        .ok_or("EP matrix expert count overflow")?;
    if !((bank.local_first == 0 && bank.peer_first == bank.count)
        || (bank.peer_first == 0 && bank.local_first == bank.count))
        || !local_work
            .routes
            .matches_partition(global, bank.local_first, bank.count)
        || !remote.grouped.as_ref().is_some_and(|work| {
            work.routes
                .matches_partition(global, bank.peer_first, bank.count)
        })
    {
        return Err("EP matrix bank/workspace ownership mismatch".into());
    }
    let os = owner.stream();
    let ps = peer.stream();
    let slots = rows * topk;
    let mut drain = PeerFailureDrain {
        owner: os.clone(),
        peer: ps.clone(),
        complete: false,
    };
    peer_copy(&os, &ps, local.xq, &mut remote.xq, rows * hidden)?;
    peer_copy(&os, &ps, local.xs, &mut remote.xs, rows * hidden / 128)?;
    peer_copy(&os, &ps, local.ids, &mut remote.ids, slots)?;
    peer_copy(&os, &ps, local.weights, &mut remote.weights, slots)?;
    remote.tx_done.record(&os).map_err(|e| e.to_string())?;
    peer.ctx.bind_to_thread().map_err(|e| e.to_string())?;
    ps.wait(&remote.tx_done).map_err(|e| e.to_string())?;
    let peer_work = remote
        .grouped
        .as_mut()
        .ok_or("EP matrix peer workspace missing")?;
    let mut peer_compute = EpCompute {
        xq: &remote.xq,
        xs: &remote.xs,
        ids: &remote.ids,
        weights: &remote.weights,
        g1: &mut remote.g1,
        g3: &mut remote.g3,
        h: &mut remote.h,
        hq: &mut remote.hq,
        hs: &mut remote.hs,
        contribution: &mut remote.contribution,
    };
    let mut route_calls =
        u64::from(local_work.prepare(owner, local, scale2, scale2_host, rows, topk, true)?);
    route_calls += u64::from(peer_work.prepare(
        peer,
        &peer_compute,
        &bank.peer_s2,
        scale2_host,
        rows,
        topk,
        true,
    )?);
    local_work.set_gu_fuse_for_plain(allow_gu_fuse);
    peer_work.set_gu_fuse_for_plain(allow_gu_fuse);
    if crate::dsv4_grouped::route_validation_enabled()
        && local_work.routes.live_slots + peer_work.routes.live_slots != slots
    {
        return Err("EP matrix partitions did not cover every selected slot".into());
    }
    record_route_stats(
        if local_work.routes.live_slots_observed && peer_work.routes.live_slots_observed {
            Some((local_work.routes.live_slots, peer_work.routes.live_slots))
        } else {
            None
        },
        rows,
        topk,
    );
    local_work.gate_up(owner, table, local, limit)?;
    if serial_control {
        local_work.down(owner, table, local)?;
        os.synchronize().map_err(|e| e.to_string())?;
        peer_work.gate_up(peer, peer_table, &mut peer_compute, limit)?;
    } else {
        peer_work.gate_up(peer, peer_table, &mut peer_compute, limit)?;
        local_work.down(owner, table, local)?;
    }
    peer_work.down(peer, peer_table, &mut peer_compute)?;
    peer_copy(
        &ps,
        &os,
        &remote.contribution,
        &mut remote.returned,
        slots * hidden,
    )?;
    remote.rx_done.record(&ps).map_err(|e| e.to_string())?;
    // The peer matrix stage explicitly bound CUDART as well as the driver.
    // Restore both before the owner's raw merge and following shared-expert work.
    crate::dsv4_grouped::bind_matrix(owner)?;
    os.wait(&remote.rx_done).map_err(|e| e.to_string())?;
    unsafe {
        k::ck(
            "EP matrix original-slot return",
            k::memra_dsv4_ep_merge_slots(
                local.contribution.device_ptr_mut(&os).0 as *mut f32,
                remote.returned.device_ptr(&os).0 as *const f32,
                local.ids.device_ptr(&os).0 as *const i32,
                slots as i32,
                hidden as i32,
                bank.peer_first as i32,
                bank.count as i32,
                os.cu_stream().cast(),
            ),
        )?;
    }
    drain.complete = true;
    Ok(route_calls)
}

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

    #[test]
    fn route_stats_call_classification_keeps_unknown_counts_out_of_slot_totals() {
        let unknown = route_stats_call_delta(None, 1, 6);
        assert_eq!(unknown.calls, 1);
        assert_eq!(unknown.observed_calls, 0);
        assert_eq!(unknown.unobserved_calls, 1);
        assert_eq!(unknown.local_slots, 0);
        assert_eq!(unknown.peer_slots, 0);
        assert_eq!(unknown.busier_slots, 0);
        assert_eq!(unknown.local_hist, [0; 9]);
        assert_eq!(unknown.busier_hist, [0; 9]);

        let observed_one_row = route_stats_call_delta(Some((3, 5)), 1, 6);
        assert_eq!(observed_one_row.observed_calls, 1);
        assert_eq!(observed_one_row.unobserved_calls, 0);
        assert_eq!(observed_one_row.local_slots, 3);
        assert_eq!(observed_one_row.peer_slots, 5);
        assert_eq!(observed_one_row.busier_slots, 5);
        assert_eq!(observed_one_row.one_row_calls, 1);
        assert_eq!(observed_one_row.local_hist[3], 1);
        assert_eq!(observed_one_row.busier_hist[5], 1);

        let observed_multi_row = route_stats_call_delta(Some((12, 8)), 2, 6);
        assert_eq!(observed_multi_row.observed_calls, 1);
        assert_eq!(observed_multi_row.one_row_calls, 0);
        assert_eq!(observed_multi_row.local_hist, [0; 9]);

        let before = EpRouteStats {
            calls: 4,
            observed_calls: 2,
            unobserved_calls: 2,
            ..EpRouteStats::default()
        };
        let after = EpRouteStats {
            calls: 7,
            observed_calls: 3,
            unobserved_calls: 4,
            ..EpRouteStats::default()
        };
        let delta = after.delta(before);
        assert_eq!(delta.calls, 3);
        assert_eq!(delta.observed_calls, 1);
        assert_eq!(delta.unobserved_calls, 2);
        assert_eq!(delta.local_slots, 0);
        assert_eq!(delta.busier_slots, 0);
    }

    #[test]
    #[ignore = "requires an exclusively locked non-serving CUDA device"]
    fn cuda_ep_partition_matches_original_full_bank() {
        let gpu = Gpu::new(0).expect("gpu");
        unsafe { gpu.ctx.disable_event_tracking() };
        let stream = gpu.stream();
        let ne = 8usize;
        let topk = 6usize;
        for (hidden, inter) in [(128usize, 256usize), (4096, 2048)] {
            let wstride = hidden * inter / 2;
            let sstride = hidden * inter / 16;
            let weights: Vec<u8> = (0..ne * 3 * wstride)
                .map(|i| (i.wrapping_mul(73) + 31) as u8)
                .collect();
            let scales: Vec<u8> = (0..ne * 3 * sstride)
                .map(|i| 0x20 + (i % 64) as u8)
                .collect();
            let macro_scales: Vec<f32> = (0..ne * 3).map(|i| 2f32.powi(i as i32 % 7 - 3)).collect();
            let w = stream.clone_htod(&weights).unwrap();
            let sc = stream.clone_htod(&scales).unwrap();
            let s2 = stream.clone_htod(&macro_scales).unwrap();
            for rows in [1usize, 6, 32] {
                let slots = rows * topk;
                let ids: Vec<i32> = (0..slots)
                    .map(|i| ((i / topk * 7 + i % topk * 3) % ne) as i32)
                    .collect();
                let selected = stream.clone_htod(&ids).unwrap();
                for projection in 0..3i32 {
                    let (n, kdim, arows, a_group, per_slot) = if projection == 1 {
                        (hidden, inter, slots, 0, 1)
                    } else {
                        (inter, hidden, rows, topk as i32, 0)
                    };
                    let a: Vec<u8> = (0..arows * kdim)
                        .map(|i| {
                            let b = (i.wrapping_mul(11) + 37) as u8;
                            if b & 127 == 127 { 0 } else { b }
                        })
                        .collect();
                    let ascale: Vec<f32> = (0..arows * kdim / 128)
                        .map(|i| 2f32.powi(i as i32 % 9 - 5))
                        .collect();
                    let aq = stream.clone_htod(&a).unwrap();
                    let asc = stream.clone_htod(&ascale).unwrap();
                    for reduction in [0, 1] {
                        let sentinel = f32::from_bits(0x7fc54321);
                        let mut baseline =
                            stream.clone_htod(&vec![sentinel; slots * n + 17]).unwrap();
                        let mut left = stream.clone_htod(&vec![sentinel; slots * n + 17]).unwrap();
                        let mut right = stream.clone_htod(&vec![sentinel; slots * n + 17]).unwrap();
                        unsafe {
                            k::ck(
                                "full-bank anchor",
                                k::memra_dsv4_fp4_gemm_sel_g_arm(
                                    aq.device_ptr(&stream).0 as *const c_void,
                                    asc.device_ptr(&stream).0 as *const f32,
                                    w.device_ptr(&stream).0 as *const c_void,
                                    sc.device_ptr(&stream).0 as *const c_void,
                                    s2.device_ptr(&stream).0 as *const f32,
                                    selected.device_ptr(&stream).0 as *const i32,
                                    projection,
                                    per_slot,
                                    0,
                                    baseline.device_ptr_mut(&stream).0 as *mut f32,
                                    slots as i32,
                                    n as i32,
                                    kdim as i32,
                                    wstride as i64,
                                    sstride as i64,
                                    a_group,
                                    reduction,
                                    stream.cu_stream().cast(),
                                ),
                            )
                            .unwrap();
                            for (first, out) in [(0usize, &mut left), (ne / 2, &mut right)] {
                                k::ck(
                                    "partition",
                                    k::memra_dsv4_fp4_gemm_sel_ep(
                                        aq.device_ptr(&stream).0 as *const c_void,
                                        asc.device_ptr(&stream).0 as *const f32,
                                        (w.device_ptr(&stream).0 + (first * 3 * wstride) as u64)
                                            as *const c_void,
                                        (sc.device_ptr(&stream).0 + (first * 3 * sstride) as u64)
                                            as *const c_void,
                                        s2.device_ptr(&stream).0 as *const f32,
                                        selected.device_ptr(&stream).0 as *const i32,
                                        projection,
                                        per_slot,
                                        0,
                                        out.device_ptr_mut(&stream).0 as *mut f32,
                                        slots as i32,
                                        n as i32,
                                        kdim as i32,
                                        wstride as i64,
                                        sstride as i64,
                                        a_group,
                                        reduction,
                                        first as i32,
                                        (ne / 2) as i32,
                                        stream.cu_stream().cast(),
                                    ),
                                )
                                .unwrap();
                            }
                        }
                        let inactive = stream.clone_dtoh(&left).unwrap();
                        for slot in 0..slots {
                            if ids[slot] >= (ne / 2) as i32 {
                                assert!(
                                    inactive[slot * n..(slot + 1) * n]
                                        .iter()
                                        .all(|v| v.to_bits() == 0)
                                );
                            }
                        }
                        unsafe {
                            k::ck(
                                "ordered slot merge",
                                k::memra_dsv4_ep_merge_slots(
                                    left.device_ptr_mut(&stream).0 as *mut f32,
                                    right.device_ptr(&stream).0 as *const f32,
                                    selected.device_ptr(&stream).0 as *const i32,
                                    slots as i32,
                                    n as i32,
                                    (ne / 2) as i32,
                                    (ne / 2) as i32,
                                    stream.cu_stream().cast(),
                                ),
                            )
                            .unwrap();
                        }
                        let expected = stream.clone_dtoh(&baseline).unwrap();
                        let actual = stream.clone_dtoh(&left).unwrap();
                        assert!(
                            expected
                                .iter()
                                .zip(&actual)
                                .all(|(a, b)| a.to_bits() == b.to_bits()),
                            "partition mismatch h={hidden} i={inter} rows={rows} p={projection} reduce={reduction}"
                        );
                        assert!(
                            actual[slots * n..]
                                .iter()
                                .all(|v| v.to_bits() == sentinel.to_bits())
                        );
                    }
                }
                println!(
                    "PASS EP projection partitions hidden={hidden} intermediate={inter} rows={rows} both reductions, global scales, masks, original-slot merge and guards"
                );
            }
        }
    }
}