Custom Patterns#

The built-in TENSORRT_PATTERNS list is a good starting point, but real-world deployment often benefits from a custom pattern list that skips quantization on operators where INT8 hurts more than it helps.

This is the “mixed-precision” strategy: selectively place QDQ stubs based on the compute characteristics of each operator and its behavior on specific hardware.

Why customize?#

Two operators commonly benefit from being left in FP16:

Depthwise convolutions#

Depthwise convolutions (groups == in_channels) are memory-bound, not compute-bound. Quantizing them to INT8 forces TensorRT to insert data reformatting layers (INT8 → FP16 → INT8) around the convolution, and this reformatting overhead typically exceeds the compute savings from INT8.

This effect is especially pronounced on ConvNeXt, which uses depthwise 7×7 convolutions throughout.

Global average pooling#

AdaptiveAvgPool2d is an element-wise operation with no matrix multiplication. INT8 quantization adds QDQ stubs on both sides but provides negligible compute benefit while risking accuracy loss.

Writing a custom pattern#

A pattern is a subclass of Pattern with two class attributes:

  • tree — the node topology to match (a tuple of module types, predicates, and Wildcard entries)

  • graft — how to build the replacement (typically a FusedModule subclass whose constructor accepts the matched modules)

The base class provides default match() and replace() classmethods that handle tree matching and graph surgery automatically. Most patterns — including all built-in fusion patterns — need nothing else.

How tree and graft work#

The tree describes a linear chain of nodes in the FX graph. Each entry is one of:

  • A module type (e.g. nn.Conv2d) — matches via isinstance

  • A callable predicate (fx.Node) -> bool — for fine-grained checks

  • A Wildcard — an optional or variable-length slot ("?" = zero or one, "*" = zero or more, "+" = one or more)

The graft specifies what to put in place of the matched nodes. It takes one of two forms:

  • Module class (e.g. FusedConvBNAct) — the default replace() collects the matched modules from the tree (in order, with None for unmatched "?" wildcards) and passes them as positional arguments to the constructor. This is the common case for fusion patterns.

  • Tuple of ReplacementMaker callables — each callable receives a TreeMatch and returns a tuple of Replacement items (modules, existing nodes, or NodeInserter callables). Use this for conversions that produce multiple replacement nodes or need access to the TreeMatch to construct replacements. An empty tuple () is the special case that erases matched nodes with no replacement.

Built-in examples#

The simplest pattern is a single-node match:

class LinearPattern(Pattern):
    phase = Phase.FUSION
    tree: Tree = (nn.Linear,)
    graft: Graft = FusedLinear

A chain with an optional element:

_OPTIONAL_BN = Wildcard(nn.BatchNorm2d, quantifier="?")

class ConvBNActPattern(Pattern):
    phase = Phase.FUSION
    tree: Tree = (nn.Conv2d, _OPTIONAL_BN, ActivationLike)
    graft: Graft = FusedConvBNAct

FusedConvBNAct.__init__ receives (conv, bn_or_none, activation) — the three modules collected from the tree in order.

A callable predicate for more specific matching:

def _is_stem_conv(node: fx.Node) -> bool:
    module = get_module(node)
    return (
        isinstance(module, nn.Conv2d)
        and module.in_channels == 3
        and module.kernel_size == (7, 7)
    )

class StemConvBNActMaxPoolPattern(Pattern):
    phase = Phase.FUSION
    tree: Tree = (_is_stem_conv, _OPTIONAL_BN, ActivationLike, nn.MaxPool2d)
    graft: Graft = FusedConvBNActMaxPool

For branching topologies (e.g. residual add), use Fork:

import operator

from embedl_deploy._internal.core.tree.match import node_check

@node_check
def _is_add(node: fx.Node) -> bool:
    return node.op == "call_function" and node.target in (
        operator.add,
        operator.iadd,
    )

class ConvBNAddActPattern(Pattern):
    phase = Phase.FUSION
    tree: Tree = Fork(
        inputs=(
            (nn.Conv2d, nn.BatchNorm2d),  # conv branch
            (),                           # residual branch (any node)
        ),
        operator=_is_add,
        output=(ActivationLike,),
    )
    graft: Graft = FusedConvBNAddAct

An erasure pattern removes matched nodes with no replacement:

class RemoveAssertPattern(Pattern):
    phase = Phase.CONVERSION
    tree: Tree = (Wildcard(_is_assert_noise, quantifier="+"), _is_eq, _is_assert)
    graft: Graft = ()

For conversions that produce multiple replacement nodes, use a tuple of ReplacementMaker callables. Each receives a TreeMatch and returns a tuple of replacement items (nn.Module, fx.Node, or NodeInserter):

class DecomposeMultiheadAttentionPattern(Pattern):
    phase = Phase.CONVERSION
    tree: Tree = (_is_supported_mha,)
    graft: Graft = (_decompose_mha,)

Here _decompose_mha is a function (TreeMatch) -> tuple[Replacement, ...] that inspects the matched nn.MultiheadAttention and returns three replacement modules (MHAInProjection, ScaledDotProductAttention, nn.Linear).

Custom pattern: depthwise Conv without quantization#

Here is a complete custom pattern that matches depthwise convolutions and skips quantization by replacing them with a FusedModule that declares an empty inputs_to_quantize:

import torch.nn as nn
from torch import fx

from embedl_deploy._internal.core.modules import FusedModule
from embedl_deploy._internal.core.patterns.main import Pattern, Phase
from embedl_deploy._internal.core.tree.types import Graft, Tree, Wildcard
from embedl_deploy._internal.core.tree.utils import get_module


class FusedDepthwiseConvBN(FusedModule):
    """Depthwise Conv2d → [BatchNorm2d] without quantization."""

    inputs_to_quantize: set[int] = set()

    def __init__(self, conv: nn.Conv2d, bn: nn.BatchNorm2d | None) -> None:
        super().__init__()
        self.conv = conv
        self.bn = bn

    def forward(self, x):
        x = self.conv(x)
        if self.bn is not None:
            x = self.bn(x)
        return x


def _is_depthwise_conv(node: fx.Node) -> bool:
    module = get_module(node)
    return (
        isinstance(module, nn.Conv2d)
        and module.groups > 1
        and module.groups == module.in_channels
    )


class DepthwiseConvBNPattern(Pattern):
    """Match depthwise Conv2d → [BatchNorm2d] without quantization."""

    phase = Phase.FUSION
    tree: Tree = (_is_depthwise_conv, Wildcard(nn.BatchNorm2d, quantifier="?"))
    graft: Graft = FusedDepthwiseConvBN

Key points:

  • The _is_depthwise_conv predicate narrows matching beyond a plain type check.

  • FusedDepthwiseConvBN declares inputs_to_quantize = set(), so no QuantStub instances are created and no stubs are placed around matched depthwise convolutions.

  • The default match() and replace() handle everything — the constructor receives (conv, bn_or_none) from the matched tree.

Overriding match() or replace()#

In rare cases the default methods are insufficient — for example, when pre-processing graph nodes before replacement or filtering matches with custom logic. You can override match() and/or replace() as classmethods. See DecomposeMultiheadAttentionPattern in the source for an example that calls super().replace() after unwrapping tuple outputs.

Building a custom pattern list#

A custom pattern list is assembled by combining built-in patterns with your custom ones. Pattern lists hold classes, not instances. Order matters — longer/more-specific patterns first:

# NOTE: These imports use internal APIs. Public exports may be added
# in a future release.
from embedl_deploy._internal.tensorrt.patterns.conversions import (
    DecomposeMultiheadAttentionPattern,
    FlattenLinearToConv1x1Pattern,
    RemoveIdentityAdaptiveAvgPoolPattern,
)
from embedl_deploy._internal.tensorrt.patterns.fusions import (
    AdaptiveAvgPoolPattern,
    ConvBNAddActPattern,
    ConvBNActPattern,
    ConvBNPattern,
    LayerNormPattern,
    LinearActPattern,
    LinearPattern,
    MHAInProjectionPattern,
    ScaledDotProductAttentionPattern,
    StemConvBNActMaxPoolPattern,
)

SMART_PATTERNS = [
    # -- Conversions (applied first, iteratively) --
    DecomposeMultiheadAttentionPattern,
    FlattenLinearToConv1x1Pattern,
    RemoveIdentityAdaptiveAvgPoolPattern,

    # -- Fusions (longest first) --
    StemConvBNActMaxPoolPattern,
    ConvBNAddActPattern,
    ConvBNActPattern,
    LinearActPattern,
    DepthwiseConvBNPattern,  # custom: no QDQ on depthwise
    ConvBNPattern,
    LinearPattern,
    LayerNormPattern,
    MHAInProjectionPattern,
    ScaledDotProductAttentionPattern,
    # NOTE: AdaptiveAvgPoolPattern intentionally omitted
    #       → no QDQ stubs around GlobalAveragePooling
]

Two deliberate omissions:

  1. No AdaptiveAvgPoolPattern — global average pooling stays in FP16.

  2. DepthwiseConvBNPattern — depthwise convolutions are fused but not quantized (FusedDepthwiseConvBN has empty inputs_to_quantize).

Using the custom pattern list#

import torch
from torch import nn

from embedl_deploy import transform
from embedl_deploy.quantize import (
    ModulesToSkip,
    Precision,
    QuantConfig,
    TensorQuantConfig,
    quantize,
)

# Define your model (e.g., from torchvision)
# my_model = torchvision.models.convnext_tiny(weights="DEFAULT")

# Transform with custom patterns
model = my_model.cpu().eval()
example_input = torch.randn(1, 3, 224, 224)
result = transform(model, (example_input,), patterns=SMART_PATTERNS)
fused_model = result.model

# Verify lossless fusion
with torch.no_grad():
    y_orig = model(example_input)
    y_fused = fused_model(example_input)
max_diff = (y_orig - y_fused).abs().max().item()
assert max_diff < 1e-4

# Quantize with custom skip rules
def forward_loop(model):
    for batch in calibration_batches[:32]:  # your calibration dataset
        model(batch)

quantized = quantize(
    fused_model,
    (example_input,),
    config=QuantConfig(
        activation=TensorQuantConfig(Precision.INT8, symmetric=True),
        weight=TensorQuantConfig(Precision.INT8, symmetric=True, per_channel=True),
        skip=ModulesToSkip(weight={nn.LayerNorm}),
    ),
    forward_loop=forward_loop,
    freeze_weights=True,
)