mlx_c_rust/compile/
norm_fusion.rs

1//! Epsilon-specialized compiled composite fusing the gated-delta head
2//! normalizations.
3//!
4//! The query/key head RMS normalizations of the gated-delta prework ran as
5//! two launches per layer; one dual-output replay produces both. Every
6//! primitive and its order match the composed path — `mlx_fast_rms_norm` is
7//! the identical Metal kernel over its own input — and both the elementwise
8//! fusion parity contract and the upstream greedy trajectory replay prove
9//! the bits unchanged (issue #1366: decode is bound by launches per token).
10//!
11//! A second fusion was tried here and rejected by the trajectory replay:
12//! chaining the output RMS normalization into the precise SwiGLU epilogue
13//! passed synthetic bit parity at small widths yet diverged the real
14//! upstream trajectory, because MLX's compiler reshapes the reduction
15//! schedule when the elementwise epilogue is fused at the model's real
16//! widths. Reduction-into-epilogue fusions stay composed until a kernel-
17//! level contract can pin the summation order.
18//!
19//! Graph builders are plain function pointers that cannot capture the RMS
20//! epsilon, and `mlx_fast_rms_norm` bakes its epsilon into the compiled
21//! graph at trace time, so the epsilon rides in a process-wide bit slot for
22//! the duration of each compilation. Model loads are sequential in a worker
23//! process — one resident model at a time — which is the contract that keeps
24//! the slot race-free; the compiled graphs keep their baked value forever
25//! and later slot writes cannot affect them.
26
27use std::os::raw::c_int;
28use std::sync::atomic::{AtomicU32, Ordering};
29
30use crate::MlxCError;
31use crate::compile::graph::{
32    MlxCompiledGraph, array_from_vector, graph_output_array, set_graph_output_vector,
33};
34use crate::{MlxArray, MlxBindingsContext, MlxStream, raw};
35
36static BAKED_RMS_NORM_EPSILON_BITS: AtomicU32 = AtomicU32::new(0);
37
38fn baked_rms_norm_epsilon() -> f32 {
39    f32::from_bits(BAKED_RMS_NORM_EPSILON_BITS.load(Ordering::Relaxed))
40}
41
42const APPLY_DUAL_HEAD_NORM_OPERATION: &str = "apply the compiled MLX dual head norm";
43const COMPILE_DUAL_HEAD_NORM_OPERATION: &str = "compile the shapeless MLX dual head norm graph";
44
45/// Owns the compiled dual-output query/key head-normalization graph.
46///
47/// Output order is fixed by the builder: the normalized queries first, the
48/// normalized keys second. Callers consume them by that order.
49#[derive(Debug)]
50pub struct MlxCompiledDualHeadNorm {
51    compiled_graph: MlxCompiledGraph,
52}
53
54impl MlxCompiledDualHeadNorm {
55    /// Compiles the shapeless dual head-norm graph with the model's RMS
56    /// epsilon baked in.
57    pub fn new(rms_norm_epsilon: f32) -> Result<Self, MlxCError> {
58        BAKED_RMS_NORM_EPSILON_BITS.store(rms_norm_epsilon.to_bits(), Ordering::Relaxed);
59        Ok(Self {
60            compiled_graph: MlxCompiledGraph::new(
61                build_dual_head_norm_graph,
62                COMPILE_DUAL_HEAD_NORM_OPERATION,
63                true,
64            )?,
65        })
66    }
67}
68
69impl MlxBindingsContext {
70    /// Applies the dual head-norm graph, returning the normalized queries
71    /// and keys one replay produced.
72    pub fn apply_compiled_dual_head_norm(
73        &self,
74        dual_head_norm: &MlxCompiledDualHeadNorm,
75        queries: &MlxArray,
76        keys: &MlxArray,
77        query_scale_weight: &MlxArray,
78        key_scale_weight: &MlxArray,
79    ) -> Result<(MlxArray, MlxArray), MlxCError> {
80        let mut graph_outputs = dual_head_norm.compiled_graph.apply_multi(
81            &[queries, keys, query_scale_weight, key_scale_weight],
82            APPLY_DUAL_HEAD_NORM_OPERATION,
83        )?;
84        if graph_outputs.len() != 2 {
85            return Err(MlxCError {
86                operation: APPLY_DUAL_HEAD_NORM_OPERATION,
87                description: format!(
88                    "compiled MLX dual head norm returned {} outputs, expected 2",
89                    graph_outputs.len()
90                ),
91            });
92        }
93        let normalized_keys = graph_outputs
94            .pop()
95            .expect("the length check proved two outputs");
96        let normalized_queries = graph_outputs
97            .pop()
98            .expect("the length check proved two outputs");
99        Ok((normalized_queries, normalized_keys))
100    }
101}
102
103unsafe extern "C" fn build_dual_head_norm_graph(
104    output_vector: *mut raw::mlx_vector_array,
105    input_vector: raw::mlx_vector_array,
106) -> c_int {
107    if output_vector.is_null() || unsafe { raw::mlx_vector_array_size(input_vector) } != 4 {
108        return 1;
109    }
110    let queries = match array_from_vector(input_vector, 0) {
111        Ok(queries) => queries,
112        Err(get_status) => return get_status,
113    };
114    let keys = match array_from_vector(input_vector, 1) {
115        Ok(keys) => keys,
116        Err(get_status) => return get_status,
117    };
118    let query_scale_weight = match array_from_vector(input_vector, 2) {
119        Ok(query_scale_weight) => query_scale_weight,
120        Err(get_status) => return get_status,
121    };
122    let key_scale_weight = match array_from_vector(input_vector, 3) {
123        Ok(key_scale_weight) => key_scale_weight,
124        Err(get_status) => return get_status,
125    };
126    let gpu_stream = match MlxStream::default_gpu() {
127        Ok(gpu_stream) => gpu_stream,
128        Err(_) => return 1,
129    };
130    let normalized_queries = match graph_output_array(|output_array| {
131        // SAFETY: Inputs and stream are live, and the output is uniquely writable.
132        unsafe {
133            raw::mlx_fast_rms_norm(
134                output_array,
135                queries.raw(),
136                query_scale_weight.raw(),
137                baked_rms_norm_epsilon(),
138                gpu_stream.raw(),
139            )
140        }
141    }) {
142        Ok(normalized_queries) => normalized_queries,
143        Err(build_status) => return build_status,
144    };
145    let normalized_keys = match graph_output_array(|output_array| {
146        // SAFETY: Inputs and stream are live, and the output is uniquely writable.
147        unsafe {
148            raw::mlx_fast_rms_norm(
149                output_array,
150                keys.raw(),
151                key_scale_weight.raw(),
152                baked_rms_norm_epsilon(),
153                gpu_stream.raw(),
154            )
155        }
156    }) {
157        Ok(normalized_keys) => normalized_keys,
158        Err(build_status) => return build_status,
159    };
160    // Output order is the caller's ABI: normalized queries first, keys
161    // second. Multi-output tracing only supports publishing the complete
162    // result vector once at the end of the builder.
163    // SAFETY: The output vector is unique and live for this callback.
164    unsafe { set_graph_output_vector(output_vector, &[&normalized_queries, &normalized_keys]) }
165}
mlx_c_rust/compile/norm_fusion.rs

mlx_c_rust/compile/
norm_fusion.rs
