Struct MoeV41 

pub struct MoeV41 {
    pub last: Cell<MoeV41Timing>,
    pub cfg: MoeV41Cfg,
    /* private fields */
}

Fields

last: Cell<MoeV41Timing>

2026-09-25: The last forward’s timing.

cfg: MoeV41Cfg

Implementations

impl MoeV41

pub fn upload_bias(gpu: &dyn GpuBackend, bias: &[f32]) -> Result<DevicePtr>

2026-09-25: The device bias for a layer’s router (MoeV41LayerWeights::gate_bias_dev).

pub fn header_words(&self) -> usize

2026-09-25: Words in one layer’s header: the miss flag, the picks, the weight bits, the plan slots.

pub fn arena_of(lru: &ExpertLru) -> (u64, u64, u64, u64, u64)

2026-09-25: The arena facts the selection kernel builds pointers from: (base, slot bytes, gate / up / down offsets).

pub fn route_select_deferred( &self, gpu: &dyn GpuBackend, w: &MoeV41LayerWeights, arena: (u64, u64, u64, u64, u64), stream: u64, ) -> Result<()>

2026-09-25: The device selection for a replayed step: the kernel launched with its header in layer w.layer’s row, and no read-back (the caller runs read_headers after the segment). The slot table must already be up to date. Device work only; an error when device_route_ok refuses the layer.

pub fn read_headers(&self, gpu: &dyn GpuBackend, stream: u64) -> Result<Vec<u8>>

2026-09-25: Every layer’s header in one read-back ordered on stream.

pub fn parse_header( &self, hdr: &[u8], layer: u32, ) -> Result<(bool, Vec<usize>, Vec<f32>)>

2026-09-25: (miss flag, picks, weights) of layer layer out of read_headers; an error when a pick is outside 0..n_routed.

impl MoeV41

pub fn forward<S: ExpertSource + ?Sized>( &self, gpu: &dyn GpuBackend, w: &MoeV41LayerWeights, lru: &mut ExpertLru, src: &S, x: DevicePtr, m: usize, reader_threads: usize, next: Option<&RouterWeights>, stream: u64, ) -> Result<(DevicePtr, Vec<f32>, Vec<usize>)>

2026-09-25: One layer’s MoE for m tokens: routed experts from the cache, plus the shared expert. Returns the bf16 [m, dim] output (the out buffer) and the routing weights and indices, [m, topk] each.

impl MoeV41

pub fn new(gpu: &dyn GpuBackend, cfg: MoeV41Cfg) -> Result<Self>

pub fn out_ptr(&self) -> DevicePtr

pub fn free(self, gpu: &dyn GpuBackend) -> Result<()>

impl MoeV41

pub fn route( &self, gpu: &dyn GpuBackend, w: &MoeV41LayerWeights, x: DevicePtr, m: usize, stream: u64, ) -> Result<(Vec<f32>, Vec<usize>)>

2026-09-25: Router logits on the GPU (f32 out), the selection on the CPU. Drains stream.

pub fn route_launch( &self, gpu: &dyn GpuBackend, w: &MoeV41LayerWeights, x: DevicePtr, m: usize, stream: u64, ) -> Result<()>

2026-09-25: The router GEMM/GEMV into self.logits: device work only, so a CUDA graph can hold it (segment A of the step graph ends with it). Errors when m is outside 1..=max_tokens.

pub fn route_select( &self, gpu: &dyn GpuBackend, w: &MoeV41LayerWeights, m: usize, stream: u64, ) -> Result<(Vec<f32>, Vec<usize>)>

2026-09-25: The selection: drain the stream, download the logits route_launch wrote, pick the top-k on the CPU with route_from_logits. Errors when the host bias length is not n_routed.

impl MoeV41

pub fn predict_launch( &self, gpu: &dyn GpuBackend, next: &RouterWeights, x: DevicePtr, stream: u64, ) -> Result<()>

2026-09-25: The next layer’s router GEMV on this layer’s single-token input, into pred_logits. Device work only; predict_select reads it back.

pub fn predict_select( &self, gpu: &dyn GpuBackend, next: &RouterWeights, k: usize, stream: u64, ) -> Result<Vec<(u32, u32)>>

2026-09-25: The k best experts of the predicted layer by score + bias, as (layer, expert) cache keys, from the logits predict_launch wrote. Drains stream.

pub fn route_predict( &self, gpu: &dyn GpuBackend, w: &MoeV41LayerWeights, x: DevicePtr, k: usize, stream: u64, ) -> Result<Vec<usize>>

2026-09-25: The router of this layer on x (one token), the k best experts by score + bias in route_from_logits order. Drains the stream; overwrites self.logits.

impl MoeV41

pub fn ffn_input_q8_m1( &self, gpu: &dyn GpuBackend, x: DevicePtr, stream: u64, ) -> Result<()>

2026-09-25: The token’s q8_1 rows into a_q8, once for the routed experts and the shared expert’s two input projections (x_pre at the readers).

pub fn stage_m1<S: ExpertSource + ?Sized>( &self, gpu: &dyn GpuBackend, w: &MoeV41LayerWeights, lru: &mut ExpertLru, src: &S, reader_threads: usize, next: Option<&RouterWeights>, x: DevicePtr, stream: u64, ) -> Result<MoeV41Stage>

2026-09-25: The host span of a single-token step between the router GEMV and the expert compute: the next layer’s router on x (when next is given, prefetch_k() > 0 and the cache has a reader pool), the selection (drains the stream), the cache fetch with the predicted experts, and the rows / weights / pointer-table uploads into fixed buffers. route_launch must have been issued on stream first.

pub fn compute_m1( &self, gpu: &dyn GpuBackend, w: &MoeV41LayerWeights, x: DevicePtr, ne: usize, stream: u64, ) -> Result<DevicePtr>

2026-09-25: The single-token expert compute after stage_m1: the accumulator cleared, the routed experts from the pointer table, the shared expert, the finish into out. Device work only, no per-token scalar arguments (ne is the top-k: one token’s picks are distinct experts), so a CUDA graph of it replays for any token once stage_m1 has refilled rows_dev / weight_dev / ptrs_dev.

pub fn shared_expert_on( &self, gpu: &dyn GpuBackend, w: &MoeV41LayerWeights, x: DevicePtr, side: u64, ) -> Result<()>

2026-09-25: The single-token shared expert up to sd on side, reading the token’s q8_1 rows that Self::ffn_input_q8_m1 wrote on the forking stream, with its own q8_1 scratch sh_q8 (as forward’s side-stream arm): the fork half of Self::compute_m1_joined. Device work only.

pub fn compute_m1_joined( &self, gpu: &dyn GpuBackend, x: DevicePtr, ne: usize, join: u64, stream: u64, ) -> Result<DevicePtr>

2026-09-25: Self::compute_m1 with the shared expert forked onto a side stream by Self::shared_expert_on: the routed experts, then the wait on join (recorded on the side after sd), then the accumulate and finish, the same launch order as forward’s side-stream arm. The token’s q8_1 rows must already be in a_q8 (Self::ffn_input_q8_m1).