Struct MoeV41 

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

Fields

last: Cell<MoeV41Timing>

The last forward’s timing.

cfg: MoeV41Cfg

Implementations

impl MoeV41

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

The device bias for a layer’s router (MoeV41LayerWeights::gate_bias_dev).

pub fn header_words(&self) -> usize

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)

The arena facts the selection kernel builds pointers from.

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

The device selection for a replayed step: the slot table brought up to date, the kernel launched with its header at layer w.layer’s row and NO read-back (read_headers after the segment). Device work only past the table update.

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

Every layer’s header, one read-back (drains stream).

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

(miss flag, picks, weights) of layer layer out of read_headers.

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>)>

One layer’s MoE for m tokens: routed experts from the cache, plus the shared expert. Returns the bf16 [m, dim] output and the routing.

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>)>

Router logits on the GPU (f32-accumulated), the selection on the CPU.

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

The router GEMM/GEMV into self.logits: device work only (the half a CUDA graph can hold).

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

The selection: drain the stream, download the logits route_launch wrote, pick the top-k on the CPU as the reference does.

impl MoeV41

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

The next layer’s router GEMV on this layer’s single-token input, into pred_logits. Device work only; issue it right after route_launch so one drain serves both read-backs.

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

The k best experts of the predicted layer by score + bias, as cache keys, from the logits predict_launch wrote (one read-back).

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

Diagnostics: the router of this layer on x (one token), the k best experts by score + bias in the reference’s order. Drains the stream; clobbers self.logits.

impl MoeV41

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

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>

The HOST span of a single-token step between the router GEMV and the expert compute: the next layer’s router on x (prediction, if next), the selection (drains the stream), the cache fetch with the predicted experts started in the background, 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>

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<()>

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 (its own h scratch, as the eager step’s side stream uses): 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>

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 same accumulate and finish. The eager step’s order and bits.