Skip to main content

moe_swiglu_seq_encode

Function moe_swiglu_seq_encode 

Source
pub fn moe_swiglu_seq_encode(
    encoder: &mut CommandEncoder,
    registry: &mut KernelRegistry,
    device: &DeviceRef,
    gate_up: &MlxBuffer,
    output: &MlxBuffer,
    intermediate: usize,
    top_k: usize,
    n_tokens: usize,
) -> Result<()>
Expand description

Multi-token SwiGLU for batched prefill.

Input: [n_tokens, top_k, 2*intermediate] Output: [n_tokens, top_k, intermediate]