1use onnx_runtime_ir::{
2 Attribute, DataType, Dim, Graph, Node, NodeId, TensorData, ValueId, WeightRef, static_shape,
3};
4use onnx_runtime_optimizer::{
5 OptimizationPass, OptimizerError, PassContext, Result as OptimizerResult,
6};
7
8const PROJECTION_FUSION_ENV: &str = "ONNX_GENAI_PROJECTION_FUSION";
9const MICROSOFT_DOMAIN: &str = "com.microsoft";
10
11pub struct ProjectionFusion {
16 enabled: bool,
17}
18
19impl ProjectionFusion {
20 pub fn new() -> Self {
21 Self {
22 enabled: std::env::var_os(PROJECTION_FUSION_ENV).is_some_and(|value| value == "1"),
23 }
24 }
25
26 pub fn enabled(&self) -> bool {
27 self.enabled
28 }
29}
30
31impl Default for ProjectionFusion {
32 fn default() -> Self {
33 Self::new()
34 }
35}
36
37pub fn cpu_optimization_passes() -> Vec<Box<dyn OptimizationPass>> {
39 let projection_fusion = ProjectionFusion::new();
40 if projection_fusion.enabled() {
41 vec![Box::new(projection_fusion)]
42 } else {
43 Vec::new()
44 }
45}
46
47impl OptimizationPass for ProjectionFusion {
48 fn name(&self) -> &str {
49 "CpuProjectionFusion"
50 }
51
52 fn run(&self, graph: &mut Graph, ctx: &PassContext) -> OptimizerResult<()> {
53 if !self.enabled {
54 return Ok(());
55 }
56
57 let silu_nodes: Vec<NodeId> = graph
58 .nodes
59 .iter()
60 .filter_map(|(id, node)| {
61 (node.op_type == "Silu"
62 && node.domain == MICROSOFT_DOMAIN
63 && node.inputs.len() == 1
64 && node.outputs.len() == 1)
65 .then_some(id)
66 })
67 .collect();
68
69 let mut changed = false;
70 for silu_id in silu_nodes {
71 let Some(candidate) = GateUpCandidate::match_from_silu(graph, silu_id, ctx) else {
72 continue;
73 };
74 candidate.apply(graph);
75 changed = true;
76 }
77
78 if changed {
79 graph.validate().map_err(OptimizerError::from)?;
80 }
81 Ok(())
82 }
83}
84
85struct Projection {
86 node_id: NodeId,
87 activation: ValueId,
88 output: ValueId,
89 weight: ValueId,
90 scales: ValueId,
91 k: usize,
92 n: usize,
93 bits: usize,
94 block_size: usize,
95 accuracy_level: i64,
96 weight_bytes: Vec<u8>,
97 scale_bytes: Vec<u8>,
98}
99
100impl Projection {
101 fn parse(graph: &Graph, node_id: NodeId, ctx: &PassContext) -> Option<Self> {
102 let node = graph.try_node(node_id)?;
103 if node.op_type != "MatMulNBits"
104 || node.domain != MICROSOFT_DOMAIN
105 || node.outputs.len() != 1
106 || !(3..=6).contains(&node.inputs.len())
107 || node.inputs.iter().skip(3).any(Option::is_some)
108 {
109 return None;
110 }
111 const ALLOWED_ATTRIBUTES: &[&str] = &[
112 "K",
113 "N",
114 "bits",
115 "block_size",
116 "accuracy_level",
117 "weight_prepacked",
118 ];
119 if node
120 .attributes
121 .keys()
122 .any(|attribute| !ALLOWED_ATTRIBUTES.contains(&attribute.as_str()))
123 {
124 return None;
125 }
126
127 let activation = node.inputs[0]?;
128 let weight = node.inputs[1]?;
129 let scales = node.inputs[2]?;
130 let output = node.outputs[0];
131 let k = positive_attr(node, "K")?;
132 let n = positive_attr(node, "N")?;
133 let bits = optional_nonnegative_attr(node, "bits", 4)?;
134 let block_size = positive_attr(node, "block_size")?;
135 let accuracy_level = node
136 .attr("accuracy_level")
137 .and_then(Attribute::as_int)
138 .unwrap_or(0);
139 let weight_prepacked = optional_nonnegative_attr(node, "weight_prepacked", 0)?;
140 if bits != 4
141 || weight_prepacked != 0
142 || block_size < 16
143 || !block_size.is_power_of_two()
144 || graph.value(activation).dtype != DataType::Float32
145 || graph.value(activation).shape.last() != Some(&Dim::Static(k))
146 || graph.value(output).dtype != DataType::Float32
147 {
148 return None;
149 }
150
151 let output_shape = &graph.value(output).shape;
152 if output_shape.last() != Some(&Dim::Static(n)) {
153 return None;
154 }
155
156 let k_blocks = k.div_ceil(block_size);
157 let packed_block_bytes = block_size.checked_mul(bits)?.checked_div(8)?;
158 let expected_weight_bytes = n.checked_mul(k_blocks)?.checked_mul(packed_block_bytes)?;
159 let expected_scale_values = n.checked_mul(k_blocks)?;
160 let expected_scale_bytes =
161 expected_scale_values.checked_mul(DataType::Float32.byte_size())?;
162
163 let weight_ref = graph.initializers.get(&weight)?;
164 if weight_ref.dtype() != DataType::Uint8
165 || weight_ref.dims() != [n, k_blocks, packed_block_bytes]
166 {
167 return None;
168 }
169 let scale_ref = graph.initializers.get(&scales)?;
170 if scale_ref.dtype() != DataType::Float32
171 || !(scale_ref.dims() == [n, k_blocks] || scale_ref.dims() == [expected_scale_values])
172 {
173 return None;
174 }
175
176 let weight_bytes = ctx.initializer_bytes(weight_ref)?;
177 let scale_bytes = ctx.initializer_bytes(scale_ref)?;
178 if weight_bytes.len() != expected_weight_bytes || scale_bytes.len() != expected_scale_bytes
179 {
180 return None;
181 }
182
183 Some(Self {
184 node_id,
185 activation,
186 output,
187 weight,
188 scales,
189 k,
190 n,
191 bits,
192 block_size,
193 accuracy_level,
194 weight_bytes: weight_bytes.to_vec(),
195 scale_bytes: scale_bytes.to_vec(),
196 })
197 }
198
199 fn compatible_with(&self, other: &Self, graph: &Graph) -> bool {
200 self.activation == other.activation
201 && self.k == other.k
202 && self.bits == other.bits
203 && self.block_size == other.block_size
204 && self.accuracy_level == other.accuracy_level
205 && graph.value(self.output).shape[..graph.value(self.output).shape.len() - 1]
206 == graph.value(other.output).shape[..graph.value(other.output).shape.len() - 1]
207 }
208}
209
210struct GateUpCandidate {
211 gate: Projection,
212 up: Projection,
213 total_n: usize,
214 fused_weight: Vec<u8>,
215 fused_scales: Vec<u8>,
216}
217
218impl GateUpCandidate {
219 fn match_from_silu(graph: &Graph, silu_id: NodeId, ctx: &PassContext) -> Option<Self> {
220 let silu = graph.try_node(silu_id)?;
221 let gate_output = silu.inputs[0]?;
222 let silu_output = silu.outputs[0];
223 if graph.outputs.contains(&gate_output)
224 || graph.outputs.contains(&silu_output)
225 || graph.consumers(gate_output) != [silu_id]
226 {
227 return None;
228 }
229
230 let silu_consumers = graph.consumers(silu_output);
231 if silu_consumers.len() != 1 {
232 return None;
233 }
234 let mul_id = silu_consumers[0];
235 let mul = graph.try_node(mul_id)?;
236 if mul.op_type != "Mul"
237 || !matches!(mul.domain.as_str(), "" | "ai.onnx")
238 || mul.inputs.len() != 2
239 || mul.outputs.len() != 1
240 {
241 return None;
242 }
243 let up_output = if mul.inputs[0] == Some(silu_output) {
244 mul.inputs[1]?
245 } else if mul.inputs[1] == Some(silu_output) {
246 mul.inputs[0]?
247 } else {
248 return None;
249 };
250 if up_output == gate_output
251 || graph.outputs.contains(&up_output)
252 || graph.consumers(up_output) != [mul_id]
253 {
254 return None;
255 }
256
257 let gate_id = graph.value(gate_output).producer?;
258 let up_id = graph.value(up_output).producer?;
259 if gate_id == up_id {
260 return None;
261 }
262 let mut gate = Projection::parse(graph, gate_id, ctx)?;
263 let mut up = Projection::parse(graph, up_id, ctx)?;
264 if gate.output != gate_output || up.output != up_output || !gate.compatible_with(&up, graph)
265 {
266 return None;
267 }
268 let total_n = gate.n.checked_add(up.n)?;
269 i64::try_from(total_n).ok()?;
270 let fused_weight_len =
271 checked_combined_initializer_len(gate.weight_bytes.len(), up.weight_bytes.len())?;
272 let fused_scale_len =
273 checked_combined_initializer_len(gate.scale_bytes.len(), up.scale_bytes.len())?;
274 let fused_weight = concatenate_initializer_bytes(
275 std::mem::take(&mut gate.weight_bytes),
276 std::mem::take(&mut up.weight_bytes),
277 fused_weight_len,
278 )?;
279 let fused_scales = concatenate_initializer_bytes(
280 std::mem::take(&mut gate.scale_bytes),
281 std::mem::take(&mut up.scale_bytes),
282 fused_scale_len,
283 )?;
284
285 Some(Self {
286 gate,
287 up,
288 total_n,
289 fused_weight,
290 fused_scales,
291 })
292 }
293
294 fn apply(self, graph: &mut Graph) {
295 let first_id = self.gate.node_id.0;
296 let k_blocks = self.gate.k.div_ceil(self.gate.block_size);
297 let packed_block_bytes = self.gate.block_size * self.gate.bits / 8;
298 let total_n = self.total_n;
299
300 let fused_weight = self.fused_weight;
301 let fused_scales = self.fused_scales;
302
303 let weight_name = format!("__nxrt_fused_projection_{first_id}_weight");
304 let fused_weight_value = graph.create_named_value(
305 weight_name,
306 DataType::Uint8,
307 static_shape([total_n, k_blocks, packed_block_bytes]),
308 );
309 graph.set_initializer(
310 fused_weight_value,
311 WeightRef::Inline(TensorData::from_raw(
312 DataType::Uint8,
313 vec![total_n, k_blocks, packed_block_bytes],
314 fused_weight,
315 )),
316 );
317
318 let scales_name = format!("__nxrt_fused_projection_{first_id}_scales");
319 let fused_scale_value = graph.create_named_value(
320 scales_name,
321 DataType::Float32,
322 static_shape([total_n, k_blocks]),
323 );
324 graph.set_initializer(
325 fused_scale_value,
326 WeightRef::Inline(TensorData::from_raw(
327 DataType::Float32,
328 vec![total_n, k_blocks],
329 fused_scales,
330 )),
331 );
332
333 let split_name = format!("__nxrt_fused_projection_{first_id}_split");
334 let split_value = graph.create_named_value(split_name, DataType::Int64, static_shape([2]));
335 let mut split_bytes = Vec::with_capacity(2 * std::mem::size_of::<i64>());
336 split_bytes.extend_from_slice(&(self.gate.n as i64).to_le_bytes());
337 split_bytes.extend_from_slice(&(self.up.n as i64).to_le_bytes());
338 graph.set_initializer(
339 split_value,
340 WeightRef::Inline(TensorData::from_raw(DataType::Int64, vec![2], split_bytes)),
341 );
342
343 let mut fused_shape = graph.value(self.gate.output).shape.clone();
344 *fused_shape
345 .last_mut()
346 .expect("MatMulNBits output has a last dimension") = Dim::Static(total_n);
347 let fused_output = graph.create_named_value(
348 format!("__nxrt_fused_projection_{first_id}_output"),
349 DataType::Float32,
350 fused_shape,
351 );
352
353 let mut fused_node = graph.node(self.gate.node_id).clone();
354 fused_node.name = format!("fused_projection_gate_up_{first_id}");
355 fused_node.inputs = vec![
356 Some(self.gate.activation),
357 Some(fused_weight_value),
358 Some(fused_scale_value),
359 ];
360 fused_node.outputs = vec![fused_output];
361 fused_node
362 .attributes
363 .insert("N".to_string(), Attribute::Int(total_n as i64));
364
365 graph.remove_node(self.gate.node_id);
366 graph.remove_node(self.up.node_id);
367 graph.insert_node(fused_node);
368
369 let mut split = Node::new(
370 NodeId(0),
371 "Split",
372 vec![Some(fused_output), Some(split_value)],
373 vec![self.gate.output, self.up.output],
374 );
375 split.name = format!("fused_projection_split_{first_id}");
376 split
377 .attributes
378 .insert("axis".to_string(), Attribute::Int(-1));
379 graph.insert_node(split);
380
381 remove_orphan_initializer(graph, self.gate.weight);
382 remove_orphan_initializer(graph, self.gate.scales);
383 remove_orphan_initializer(graph, self.up.weight);
384 remove_orphan_initializer(graph, self.up.scales);
385 }
386}
387
388fn positive_attr(node: &Node, name: &str) -> Option<usize> {
389 usize::try_from(node.attr(name)?.as_int()?)
390 .ok()
391 .filter(|&v| v > 0)
392}
393
394fn optional_nonnegative_attr(node: &Node, name: &str, default: usize) -> Option<usize> {
395 match node.attr(name) {
396 Some(value) => usize::try_from(value.as_int()?).ok(),
397 None => Some(default),
398 }
399}
400
401fn checked_combined_initializer_len(first: usize, second: usize) -> Option<usize> {
402 first
403 .checked_add(second)
404 .filter(|&combined| combined <= isize::MAX as usize)
405}
406
407fn concatenate_initializer_bytes(
408 mut first: Vec<u8>,
409 second: Vec<u8>,
410 combined_len: usize,
411) -> Option<Vec<u8>> {
412 first.try_reserve_exact(second.len()).ok()?;
413 first.extend_from_slice(&second);
414 debug_assert_eq!(first.len(), combined_len);
415 Some(first)
416}
417
418fn remove_orphan_initializer(graph: &mut Graph, value: ValueId) {
419 if !graph.has_uses(value)
420 && !graph.inputs.contains(&value)
421 && !graph.outputs.contains(&value)
422 {
423 graph.initializers.remove(&value);
424 }
425}
426
427#[cfg(test)]
428mod tests {
429 use super::checked_combined_initializer_len;
430
431 #[test]
432 fn fused_initializer_capacity_rejects_overflow_and_isize_excess() {
433 assert_eq!(checked_combined_initializer_len(usize::MAX, 1), None);
434 assert_eq!(
435 checked_combined_initializer_len(isize::MAX as usize, 1),
436 None
437 );
438 assert_eq!(
439 checked_combined_initializer_len(0, isize::MAX as usize),
440 Some(isize::MAX as usize)
441 );
442 }
443}