graph_rewrites#

Shared ONNX graph rewrites for quantization.

Functions

cast_custom_ops

Adds cast_to_fp16 nodes to the inputs and cast_to_fp32 to the outputs of a layer in the requested indices.

convert_fp16_io

Convert graph I/O to FP16.

insert_fp8_mha_casts

Insert three cast ops.

insert_matmul_casts

Insert three cast nodes for MatMul's two inputs and output.

remove_output_initializers

Remove initializers that are also listed as graph outputs.

remove_redundant_cast_nodes

Remove redundant Cast nodes from the ONNX graph to optimize model performance.
cast_custom_ops(onnx_model, ops_to_cast)#

Adds cast_to_fp16 nodes to the inputs and cast_to_fp32 to the outputs of a layer in the requested indices.

Parameters:

  • onnx_model (ModelProto)

  • ops_to_cast (dict)

Return type:

ModelProto

convert_fp16_io(graph)#

Convert graph I/O to FP16.

insert_fp8_mha_casts(onnx_model)#

Insert three cast ops.

The first cast will be added before the input0 of MatMul to cast fp16 to fp32. The second cast will be added before the input1 of MatMul to cast fp16 to fp32. The third cast will be added after the output of MatMul to cast fp32 back to fp16. The insertion of Cast ops in the FP8 MHA part actually forbids the MHAs to run with FP16 accumulation because the compiler only has FP32 accumulation kernels for FP8 MHAs.

insert_matmul_casts(graph, matmul_node)#

Insert three cast nodes for MatMul’s two inputs and output.

remove_output_initializers(graph, graph_initializers)#

Remove initializers that are also listed as graph outputs.

Having initializers (constant tensors) that are also marked as outputs can lead to ONNX Runtime or conversion tool errors, particularly related to ambiguous ‘dtype’ or shape inference. This step ensures compatibility by detaching such initializers from the graph’s outputs.

Parameters:

  • graph (Graph)

  • graph_initializers (list)

remove_redundant_cast_nodes(graph)#

Remove redundant Cast nodes from the ONNX graph to optimize model performance.

This function identifies and removes two types of redundant Cast nodes:

  1. Cast nodes where input and output types are identical - Before: t1 (dtype=fp16) -> cast (to=fp16) -> t2 -> Op - After: t1 (dtype=fp16) -> Op

  2. Cast nodes that can be fused with initializers - Before: (initializer) t1 (dtype=fp32) -> cast (to=fp16) -> t2 -> Op - After: (initializer) t1 (dtype=fp16) -> Op

The function preserves Cast nodes that: - Have outputs that are graph outputs - Are necessary for type conversion - Have dynamic inputs (not initializers)

Parameters:

graph (GraphProto) – ONNX graph to optimize. The graph will be modified in-place.

Return type:

None

Note

  • This optimization is particularly useful for models with many Cast operations

  • The function modifies the graph in-place

  • All tensor consumers are updated to maintain graph connectivity

  • Initializer data types are converted when possible to eliminate Cast nodes