mlx_c_rust/compile/
swiglu.rs

1use std::os::raw::c_int;
2
3use crate::MlxBindingsContext;
4use crate::{MlxArray, MlxStream, raw};
5use crate::{
6    MlxCError,
7    compile::graph::{MlxCompiledGraph, array_from_vector, graph_output_array, set_graph_output},
8};
9
10const APPLY_COMPILED_SWIGLU_OPERATION: &str = "apply a compiled MLX SwiGLU graph";
11const COMPILE_SWIGLU_OPERATION: &str = "compile a shapeless MLX SwiGLU graph";
12
13/// Owned shapeless MLX compilation of `silu(gate) * input`.
14#[derive(Debug)]
15pub struct MlxCompiledSwiGlu {
16    compiled_graph: MlxCompiledGraph,
17}
18
19impl MlxCompiledSwiGlu {
20    /// Compiles one shape-polymorphic SwiGLU graph for reuse across model layers.
21    pub fn new() -> Result<Self, MlxCError> {
22        Ok(Self {
23            compiled_graph: MlxCompiledGraph::new(
24                build_swiglu_graph,
25                COMPILE_SWIGLU_OPERATION,
26                true,
27            )?,
28        })
29    }
30}
31
32impl MlxBindingsContext {
33    /// Applies a retained shapeless SwiGLU compilation to broadcast-compatible inputs.
34    pub fn apply_compiled_swiglu(
35        &self,
36        compiled_swiglu: &MlxCompiledSwiGlu,
37        gate: &MlxArray,
38        input: &MlxArray,
39    ) -> Result<MlxArray, MlxCError> {
40        compiled_swiglu
41            .compiled_graph
42            .apply(&[gate, input], APPLY_COMPILED_SWIGLU_OPERATION)
43    }
44}
45
46unsafe extern "C" fn build_swiglu_graph(
47    output_vector: *mut raw::mlx_vector_array,
48    input_vector: raw::mlx_vector_array,
49) -> c_int {
50    if output_vector.is_null() {
51        return 1;
52    }
53    if unsafe { raw::mlx_vector_array_size(input_vector) } != 2 {
54        return 1;
55    }
56    let gate = match array_from_vector(input_vector, 0) {
57        Ok(gate) => gate,
58        Err(status) => return status,
59    };
60    let input = match array_from_vector(input_vector, 1) {
61        Ok(input) => input,
62        Err(status) => return status,
63    };
64    let gpu_stream = match MlxStream::default_gpu() {
65        Ok(gpu_stream) => gpu_stream,
66        Err(_) => return 1,
67    };
68    let sigmoid_gate = match graph_output_array(|output| {
69        // SAFETY: Inputs and stream are live, and `output` is uniquely writable.
70        unsafe { raw::mlx_sigmoid(output, gate.raw(), gpu_stream.raw()) }
71    }) {
72        Ok(sigmoid_gate) => sigmoid_gate,
73        Err(status) => return status,
74    };
75    let activated_gate = match graph_output_array(|output| {
76        // SAFETY: Inputs and stream are live, and `output` is uniquely writable.
77        unsafe { raw::mlx_multiply(output, gate.raw(), sigmoid_gate.raw(), gpu_stream.raw()) }
78    }) {
79        Ok(activated_gate) => activated_gate,
80        Err(status) => return status,
81    };
82    let swiglu_output = match graph_output_array(|output| {
83        // SAFETY: Inputs and stream are live, and `output` is uniquely writable.
84        unsafe { raw::mlx_multiply(output, activated_gate.raw(), input.raw(), gpu_stream.raw()) }
85    }) {
86        Ok(swiglu_output) => swiglu_output,
87        Err(status) => return status,
88    };
89    // SAFETY: The output vector is unique and live for this callback.
90    unsafe { set_graph_output(output_vector, &swiglu_output) }
91}
mlx_c_rust/compile/swiglu.rs

mlx_c_rust/compile/
swiglu.rs
