Inference Backends

TensorRT Inference Metadata

Use the TensorRTInferenceMetadata metadata type to compile a PyTorch nn.Module to TensorRT:

ai.py

from muna.beta import TensorRTInferenceMetadata
from torch import randn, Tensor
from torch.nn import Module

# Given a PyTorch model...
model: Module = ...
# With some example arguments...
example_args: list[Tensor] = [randn(1, 3, 224, 224)]

@compile(
    ...,
    metadata=[
        # Use TensorRT for model inference
        TensorRTInferenceMetadata(
            model=model,
            model_args=example_args,
            cuda_arch="sm_100",
            precision="int4"
        )
    ]
)
def predict() -> None:
    pass

The TensorRT inference backend is only available on Linux and Windows devices with compatible Nvidia GPUs.

We are working on adding support for consumer RTX GPUs with TensorRT for RTX.

Target CUDA Architectures

TensorRT engines must be compiled for specific target CUDA architectures. Below are CUDA architectures that our compiler supports:

CUDA Architecture	GPU Family
`sm_80`	Ampere (e.g. A100)
`sm_86`	Ampere
`sm_87`	Ampere
`sm_89`	Ada Lovelace (e.g. L40S)
`sm_90`	Hopper (e.g. H100)
`sm_100`	Blackwell (e.g. B200)

TensorRT Inference Precision

TensorRT allows for specifying the inference engine’s precision. Below are supported precision modes:

Precision	Notes
`fp32`	32-bit single precision inference.
`fp16`	16-bit half precision inference.
`int8`	8-bit quantized integer inference.

OnnxRuntime Inference Metadata

Use the OnnxRuntimeInferenceMetadata metadata type to compile a PyTorch nn.Module for inference with ONNXRuntime:

ai.py

from muna.beta import OnnxRuntimeInferenceMetadata
from torch import randn, Tensor
from torch.nn import Module

# Given a PyTorch model...
model: Module = ...
# With some example arguments...
example_args: list[Tensor] = [randn(1, 3, 224, 224)]

@compile(
    ...,
    metadata=[
        # Use ONNXRuntime for model inference
        OnnxRuntimeInferenceMetadata(
            model=model,
            model_args=example_args
        )
    ]
)
def predict() -> None:
    pass

OnnxRuntime Inference Session Metadata

Use the OnnxRuntimeInferenceSessionMetadata metadata type to compile an OnnxRuntime InferenceSession:

ai.py

from muna.beta import OnnxRuntimeInferenceSessionMetadata
from onnxruntime import InferenceSession

# Given an ONNXRuntime inference session...
model_path = "/path/to/model.onnx"
session = InferenceSession(model_path)

@compile(
    ...,
    metadata=[
        # Use ONNXRuntime for model inference
        OnnxRuntimeInferenceSessionMetadata(
            session=session,
            model_path=model_path
        )
    ]
)
def predict(...) -> None:
    pass

The ONNX model file must exist at the provided model_path within the compiler sandbox.

CoreML Inference Metadata

Use the CoreMLInferenceMetadata metadata type to compile a PyTorch nn.Module to CoreML:

ai.py

from muna.beta import CoreMLInferenceMetadata
from torch import randn, Tensor
from torch.nn import Module

# Given a PyTorch model...
model: Module = ...
# With some example arguments...
example_args: list[Tensor] = [randn(1, 3, 224, 224)]

@compile(
    ...,
    metadata=[
        # Use CoreML for model inference
        CoreMLInferenceMetadata(
            model=model,
            model_args=example_args
        )
    ]
)
def predict() -> None:
    pass

The CoreML inference backend is only available on iOS, macOS, and visionOS devices.

Llama.cpp Inference Metadata

Use the LlamaCppInferenceMetadata metadata type to compile a Llama instance:

llm.py

from muna.beta import LlamaCppInferenceMetadata
from llama_cpp import Llama

# Given an LLM
llm = Llama(...)

@compile(
    ...,
    metadata=[
        # Specify Llama.cpp inference metadata
        LlamaCppInferenceMetadata(
            model=llm,
            backends=["cuda"]
        )
    ]
)
def predict() -> None:
    pass

Llama.cpp Hardware Backends

Llama.cpp supports several hardware backends to accelerate model inference. Below are targets that are currently supported by Muna:

Backend	Notes
`cuda`	Nvidia CUDA backend. Linux only.

ExecuTorch Inference Metadata

Use the ExecuTorchInferenceMetadata metadata type to compile a PyTorch nn.Module for inference with ExecuTorch:

ai.py

from muna.beta import ExecuTorchInferenceMetadata
from torch import randn, Tensor
from torch.nn import Module

# Given a PyTorch model...
model: Module = ...
# With some example arguments...
example_args: list[Tensor] = [randn(1, 3, 224, 224)]

@compile(
    ...,
    metadata=[
        # Use ExecuTorch for model inference
        ExecuTorchInferenceMetadata(
            model=model,
            model_args=example_args,
            backend="xnnpack"
        )
    ]
)
def predict() -> None:
    pass

The ExecuTorch inference backend is only available on Android.

ExecuTorch Hardware Backends

ExecuTorch supports several hardware backends to accelerate model inference. Below are targets that are currently supported by Muna:

Backend	Notes
`xnnpack`	XNNPACK CPU backend. Always enabled.
`vulkan`	Vulkan GPU backend. Only supported on Android.

LiteRT Inference Metadata

Use the LiteRTInferenceMetadata metadata type to compile a PyTorch nn.Module for inference with LiteRT:

ai.py

from muna.beta import LiteRTInferenceMetadata
from torch import randn, Tensor
from torch.nn import Module

# Given a PyTorch model...
model: Module = ...
# With some example arguments...
example_args: list[Tensor] = [randn(1, 3, 224, 224)]

@compile(
    ...,
    metadata=[
        # Use LiteRT for model inference
        LiteRTInferenceMetadata(
            model=model,
            model_args=example_args
        )
    ]
)
def predict() -> None:
    pass

TensorFlow Lite Interpreter Metadata

Use the TFLiteInterpreterMetadata metadata type to compile a TensorFlow Lite Interpreter:

ai.py

from muna.beta import TFLiteInterpreterMetadata
from tensorflow import lite

# Given a TFLite interpreter...
model_path = "/path/to/model.tflite"
interpreter = lite.Interpreter(model_path)

@compile(
    ...,
    metadata=[
        # Use TensorFlow Lite for model inference
        TFLiteInterpreterMetadata(
            interpreter=interpreter,
            model_path=model_path
        )
    ]
)
def predict(...) -> None:
    pass

The TensorFlow Lite model file must exist at the provided model_path within the compiler sandbox.

QNN Inference Metadata

Use the QnnInferenceMetadata metadata type to compile a PyTorch nn.Module to a Qualcomm QNN context binary:

ai.py

from muna.beta import QnnInferenceMetadata
from torch import randn, Tensor
from torch.nn import Module

# Given a PyTorch model...
model: Module = ...
# With some example arguments...
example_args: list[Tensor] = [randn(1, 3, 224, 224)]

@compile(
    ...,
    metadata=[
        # Use QNN for model inference
        QnnInferenceMetadata(
            model=model,
            model_args=example_args,
            backend="gpu",
            quantization=None
        )
    ]
)
def predict() -> None:
    pass

The QNN inference backend is only available on Android and Windows devices with Qualcomm processors.

QNN Hardware Backends

QNN requires that a hardware device backend is specified ahead of time. Below are supported backends:

Backend	Notes
`cpu`	Reference `aarch64` CPU backend.
`gpu`	Adreno GPU backend, accelerated by OpenCL.
`htp`	Hexagon NPU backend.

Learn more about QNN hardware backends.

QNN Model Quantization

When using the htp backend, you must specify a model quantization mode as the Hexagon NPU only supports running integer-quantized models. Below are supported quantization modes:

Quantization	Notes
`w8a8`	Weights and activations are quantized to `uint8`.
`w8a16`	Weights are quantized to `uint8` while activations are quantized to `uint16`.
`w4a8`	Weights are quantized to `uint4` while activations are quantized to `uint8`.
`w4a16`	Weights are quantized to `uint4` while activations are quantized to `uint16`.

OpenVINO Inference Metadata

Use the OpenVINOInferenceMetadata metadata type to compile a PyTorch nn.Module to OpenVINO IR:

ai.py

from muna.beta import OpenVINOInferenceMetadata
from torch import randn, Tensor
from torch.nn import Module

# Given a PyTorch model...
model: Module = ...
# With some example arguments...
example_args: list[Tensor] = [randn(1, 3, 224, 224)]

@compile(
    ...,
    metadata=[
        # Use OpenVINO for model inference
        OpenVINOInferenceMetadata(
            model=model,
            model_args=example_args
        )
    ]
)
def predict() -> None:
    pass

At runtime, the OpenVINO IR will be used for inference with the OpenVINO toolkit.

The OpenVINO inference backend is only available on Linux and Windows x86_64 devices with Intel processors.

IREE Inference Metadata

Use the muna.beta.IREEInferenceMetadata metadata type to compile a PyTorch nn.Module for inference with IREE:

ai.py

from muna.beta import IREEInferenceMetadata
from torch import randn, Tensor
from torch.nn import Module

# Given a PyTorch model...
model: Module = ...
# With some example arguments...
example_args: list[Tensor] = [randn(1, 3, 224, 224)]

@compile(
    ...,
    metadata=[
        # Use IREE for model inference
        IREEInferenceMetadata(
            model=model,
            model_args=example_args,
            backend="vulkan"
        )
    ]
)
def predict() -> None:
    pass

The IREE inference backend is only available on Android devices.

IREE HAL Target Backends

IREE supports several HAL target backends that the model can be compiled against. Below are targets that are currently supported by Muna:

Target	Notes
`vulkan`	Vulkan GPU backend. Only supported on Android.

MIGraphX Inference Metadata

Coming soon 🤫.

Compiling Models

Insiders

Inference Backends

Target CUDA Architectures

TensorRT Inference Precision

Llama.cpp Hardware Backends

ExecuTorch Hardware Backends

QNN Hardware Backends

QNN Model Quantization

IREE HAL Target Backends

​Target CUDA Architectures

​TensorRT Inference Precision

​Llama.cpp Hardware Backends

​ExecuTorch Hardware Backends

​QNN Hardware Backends

​QNN Model Quantization

​IREE HAL Target Backends

Target CUDA Architectures

TensorRT Inference Precision

Llama.cpp Hardware Backends

ExecuTorch Hardware Backends

QNN Hardware Backends

QNN Model Quantization

IREE HAL Target Backends