Operator Fusions#

Fusions combine sequences of operators into single fused modules that map directly to hardware-accelerated kernels. After fusion, the model graph contains modules like FusedConvBNAct instead of separate Conv2d, BatchNorm2d, and ReLU layers.

Fused modules:

Are numerically equivalent to the original operator sequence (no weight folding at this stage — that’s handled by the hardware compiler).
Export cleanly to ONNX for downstream compilation.

Each FusedModule subclass declares inputs_to_quantize — the set of positional input indices where the quantization pipeline places QDQ stubs. See Quantization for the full table.

Convolution fusions#

ConvBNPattern#

Matches: Conv2d → [BatchNorm2d]

Produces: FusedConvBN

The most basic convolution fusion. The BatchNorm2d is optional — a bare Conv2d is also matched (useful for ensuring QDQ stub placement even on convolutions without batch normalization).

ConvBNActPattern#

Matches: Conv2d → [BatchNorm2d] → Activation

Produces: FusedConvBNAct

Fuses convolution, optional batch normalization, and an activation function. The pattern accepts any activation in ActivationLike: ReLU, ReLU6, LeakyReLU, ELU, GELU, SiLU, Mish, Hardswish, Hardsigmoid, PReLU, Sigmoid, Tanh.

StemConvBNActMaxPoolPattern#

Matches: Conv2d(3in, 7×7) → [BatchNorm2d] → Activation → MaxPool2d

Produces: FusedConvBNActMaxPool

Captures the classification network stem found in ResNet and similar architectures. The convolution is constrained to in_channels=3, kernel_size=(7,7) so only the actual stem is matched.

ConvBNAddActPattern#

Matches: Conv2d → BatchNorm2d → add(·, residual) → Activation

Produces: FusedConvBNAddAct

Captures the tail of ResNet-style bottleneck blocks where the convolution path merges with a skip connection before the final activation. This is a branching pattern — it matches a Fork topology with two inputs feeding into an operator.add node.

ConvBNActAddPattern#

Matches: Conv1d|2d → [BatchNorm1d|2d] → Activation → [Dropout] → add(·, residual)

Produces: FusedConvBNAct feeding FusedAdd

Captures the residual block whose activation precedes the skip add: a detector bottleneck built from Conv → BN → SiLU blocks, or a temporal convolution network’s residual tail. Like ConvBNAddActPattern, this is a branching pattern matching a Fork topology.

The block would otherwise be claimed by ActAddPattern, which pairs the activation with the add and so places a Q/DQ between the convolution and its own activation. TensorRT cannot fold an activation into a convolution kernel across that Q/DQ, so the activation and the add run as a separate floating-point kernel behind a pair of reformats. Keeping the activation with its convolution and quantizing both inputs of the add instead is the layout TensorRT’s own implicit quantization picks.

Every absorbed link must have exactly one user. A convolution, batch norm, activation or dropout output that another node still reads cannot move inside a fused module without being recomputed there, so those blocks stay with ConvBNActPattern and ActAddPattern.

ConvTransposeBNReLUPattern#

Matches: ConvTranspose1d|2d → [BatchNorm1d|2d] → [ReLU]

Produces: FusedConvTransposeBNReLU

Quantizes the activation input and transposed-convolution weight at either spatial rank. Both trailing modules are optional. Per-channel weight quantization uses axis 1 of the ONNX [input channels, output channels / groups, *spatial] layout. The normal upstream precision policy can still leave a matched chain in floating point.

This path supports INT8. FP8 and INT4 configurations fall back to floating point, as do grouped INT8 convolutions when either per-group channel count is not a multiple of four. Explicit output_size calls remain unfused; export recomposition can specialize their output padding first. Other activation functions remain outside the fused chain.

Linear fusions#

LinearActPattern#

Matches: Linear → Activation

Produces: FusedLinearAct

LinearPattern#

Matches: standalone Linear

Produces: FusedLinear

LayerNormPattern#

Matches: LayerNorm

Produces: FusedLayerNorm

ActAddPattern#

Matches: Fork(Activation, residual) → operator.add

Produces: FusedActAdd

Captures the inverted residual pattern where an activation precedes the skip-connection add (e.g. in transformer feed-forward blocks). Like ConvBNAddActPattern, this is a branching pattern matching a Fork topology. The activation must feed nothing but the add, or the dropout in front of it: one that also feeds another layer, or both operands of the add, is left unfused, since fusing it would only recompute it and quantize its pre-activation input instead. An activation fed by a single-user Conv → [BN] chain is claimed first by ConvBNActAddPattern, which keeps it with its convolution.

MulAddPattern#

Matches: Fork(operator.mul, residual) → operator.add

Produces: FusedMulAdd

Captures the layer-scale residual tail used by ConvNeXt-style blocks: a runtime tensor scaled by a constant (a parameter or a scalar literal), then added to a skip connection. Fully dynamic multiplies (e.g. SE-style gating) and adds whose other operand is a weight are not matched. The fused module declares no quantized inputs, keeping the layer-scale residual in floating point — INT8 Q/DQ on this operation measurably hurts ConvNeXt accuracy while the latency cost of floating point is negligible.

Attention fusions#

These patterns match the sub-modules produced by the DecomposeMultiheadAttentionPattern and DecomposeSwinAttentionPattern conversions.

MHAInProjectionPattern#

Matches: MHAInProjection

Produces: FusedMHAInProjection

ScaledDotProductAttentionPattern#

Matches: ScaledDotProductAttention

Produces: FusedScaledDotProductAttention

SwinAttentionPattern#

Matches: SwinAttention

Produces: FusedSwinAttention

Wraps the Swin Transformer window-attention module produced by DecomposeSwinAttentionPattern. The fused module carries shared spatial state (SwinSpatialState) for window partition/reverse geometry.

Pooling fusions#

AdaptiveAvgPoolPattern#

Matches: AdaptiveAvgPool2d

Produces: FusedAdaptiveAvgPool2d

Tip

In mixed-precision workflows the AdaptiveAvgPoolPattern is often omitted entirely from the pattern list so the pool is not fused and stays as a plain AdaptiveAvgPool2d in the graph. See Custom Patterns for details.

Fusion summary by architecture#

ResNet50#

Fused module	Count	Description
`FusedConvBNActMaxPool`	1	Stem: Conv(7×7) + BN + ReLU + MaxPool
`FusedConvBNAddAct`	16	Bottleneck residual blocks
`FusedConvBNAct`	16	Main-path Conv + BN + ReLU
`FusedConvBN`	17	Conv + BN without activation
`FusedAdaptiveAvgPool2d`	1	Global average pool
Total	51

Conversions applied: FlattenLinearToConv1x1Pattern converts the Flatten → Linear classifier into Conv2d(1×1) → Flatten.

ConvNeXt (Tiny/Base/Large)#

Fused module	Count	Description
`FusedConvBN`	23/41/41	Depthwise + downsampling convolutions
`FusedLayerNorm`	23/41/41	LayerNorm layers
`FusedLinearAct`	18/36/36	Linear + GELU chains
`FusedLinear`	18/36/36	Standalone linear layers
`FusedMulAdd`	18/36/36	Layer-scale residual tails (kept in FP)
`FusedAdaptiveAvgPool2d`	1/1/1	Global average pool

Counts increase with model depth (Base and Large share the block count and differ in channel width).

Conversions applied:

RemoveIdentityAdaptiveAvgPoolPattern removes identity pooling ops.
FlattenLinearToConv1x1Pattern converts the classifier head.

ConvNeXt uses depthwise separable convolutions extensively. The default pattern set quantizes all convolutions equally, but mixed-precision (see Custom Patterns) skips depthwise convolutions for better latency.

Vision Transformer (ViT-B/16)#

Fused module	Count	Description
`FusedConvBN`	1	Patch embedding Conv2d
`FusedMHAInProjection`	12	Q/K/V projections
`FusedScaledDotProductAttention`	12	Attention computation
`FusedLinear`	36+	Out-proj + MLP linear layers
`FusedLinearAct`	12	MLP hidden → GELU chains
`FusedLayerNorm`	25	Pre/post-norm layers

Conversions applied: DecomposeMultiheadAttentionPattern decomposes all 12 attention layers into explicit sub-modules.

Running fusions only#

from embedl_deploy import transform
from embedl_deploy.tensorrt import TENSORRT_FUSION_PATTERNS

result = transform(model, (example_input,), patterns=TENSORRT_FUSION_PATTERNS)

To inspect what was fused:

from collections import Counter

fused_counts = Counter(
    type(m).__name__
    for m in result.model.modules()
    if type(m).__name__.startswith('Fused')
)
for name, count in sorted(fused_counts.items()):
    print(f"  {name}: {count}")

Group normalization#

GroupNormPattern wraps nn.GroupNorm as FusedGroupNorm; exported aten.group_norm.default calls are first recomposed by AtenGroupNormPattern. GroupNorm keeps its input and affine parameters in floating point, matching LayerNorm’s boundary policy. A quantizable consumer can insert Q/DQ after normalization. The wrapper does not insert an output quantizer by itself or claim that normalization executes in INT8.

The boundary also stops upstream precision deferral: floating arithmetic before normalization does not automatically demote a convolution after it. Input rounding can strongly perturb low-variance groups, so input INT8 is not enabled. TensorRT may fuse output quantization into surrounding kernels; latency improvements depend on the downstream consumer and tensor shape.

Precision through indexed operations#

Precision-state propagation follows feature values through gather, scatter and grid sampling. Index and sampling-coordinate operands remain independent; their floating arithmetic does not force feature consumers to remain floating. Scatter retains the precision requirements of both its destination and source, including scalar value= inputs. Unknown operations conservatively retain all input dependencies.

This does not change FusedMulAdd’s floating-point policy: a MulAdd whose output contributes actual feature values can still demote downstream INT8 consumers. The rule also does not move Q/DQ through a sampling operation or quantize its coordinates.