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]