Struct MoeBucketPlan 

pub struct MoeBucketPlan {
    pub expert_offsets: Vec<usize>,
    pub packed_token_idx: Vec<u32>,
    pub pairs_by_token: Vec<i32>,
    pub pair_weights: Vec<f32>,
    /* private fields */
}

Bucket plan: per-expert lists of which (token, k_slot) pairs route through that expert. Built host-side from the router output and used by moe_forward_bucketed to issue ONE m=tokens_per_expert Marlin GEMM per active expert instead of batch * top_k m=1 GEMMs.

Fields

expert_offsets: Vec<usize>

expert_offsets[e+1] - expert_offsets[e] = tokens routed to expert e. Length: num_experts + 1. expert_offsets[num_experts] = total_pairs (always batch * top_k).

packed_token_idx: Vec<u32>

[total_pairs] flat: which input token each packed-row gathers from. Index into x[batch, hidden].

pairs_by_token: Vec<i32>

[batch, top_k] row-major: for each (b, k_slot), which row of the packed buffers carries that pair’s contribution. Used by B::moe_combine to scatter weighted sums back to out[b].

pair_weights: Vec<f32>

[batch, top_k] row-major: combine weight for the (b, k_slot) pair, copied verbatim from the router output. Used by B::moe_combine.

Implementations

impl MoeBucketPlan

pub fn empty() -> Self

Empty plan with no allocation. Use Self::rebuild_into before reuse — this is the cheap constructor for putting the plan in a scratch struct.

pub fn build( route: &RouterOutput, batch: usize, num_experts: usize, top_k: usize, ) -> Self

Allocate a fresh plan. Convenience wrapper over Self::rebuild_into for tests and code paths that don’t care about reuse.

pub fn rebuild_into( &mut self, route: &RouterOutput, batch: usize, num_experts: usize, top_k: usize, )

Allocation-free rebuild. Reuses the existing expert_offsets, packed_token_idx, pairs_by_token, pair_weights buffers via clear() + resize(). Uses the trailing tail of expert_offsets as the host-side cursor scratch (saves the per-call cursors.clone()).