Internals

How the Gensyn Compiler and RepOp kernels achieve bitwise reproducibility under the hood.

Gensyn Compiler

The Gensyn Compiler converts ONNX-serialized ML models into PyTorch modules, optionally with reproducible RepOp kernels replacing standard operations.

It is an MLIR-based, multi-stage compiler with a Python execution layer.

How It Works

The compiler uses MLIR dialects to reason about the incoming model:

1. A dialect that determines which operations need to be lowered to RepOp kernels (rather than standard PyTorch kernels).

2. A dialect that generates the final PyTorch module from a given set of operations.

Python API

from gensyn_mjolnir import convert, CompileOptions

# Convert an ONNX model to a PyTorch module with reproducible kernels
module = convert(
    "model/model.onnx",
    options=CompileOptions(requires_reproducibility=True)
)

Convert()

convert() is the core function of the Gensyn Compiler. It takes an ONNX model and converts it into a PyTorch module that can be used for inference.

When requires_reproducibility is enabled (which it is by default), the compiler replaces standard PyTorch operations with RepOp kernels that guarantee bitwise-identical results across hardware.

It has two parameters:

• onnx_model_or_path: A str, Path or ModelProto type. This is either a file path to an ONNX model on a disk or an in-memory ModelProto object.

• options (CompileOptions) is the configuration for the compilation process, which you can read more about below. It defaults to reproducible mode.

def convert(
    onnx_model_or_path: str | Path | ModelProto,
    options: CompileOptions = CompileOptions(),
) -> torch.nn.Module

CompileOptions

CompileOptions controls how the compiler processes the model. In most cases the defaults are what you want, which corresponds to reproducible mode with symlinked tensors and a temporary artifacts directory.

The fields are:

• artifacts_dir: always str or None. This is the directory where the compiler writes intermediate artifacts. If not set, a temporary directory is used (preserved when the MJOLNIR_DEBUG environment variable is set, which is useful for inspecting compiler output during debugging).

• colocate_tensors: A bool that is False by default. When set to True it copies external tensor files into the artifacts directory. When False, it creates symbolic links instead.

Symlinking is faster and saves disk space, but copying may be needed if you plan to move the artifacts directory to another location.

• requires_reproducibility: Also a bool but set to True by default. When True, the compiler replaces standard PyTorch operations with RepOp kernels for cross-hardware reproducibility. When False, it uses standard PyTorch kernels which are faster but not reproducible across different hardware.

@dataclass(frozen=True, slots=True, kw_only=True)
class CompileOptions:
    artifacts_dir: str | None = None
    # Directory for compiler artifacts. Uses a temp directory if not specified
    # (preserved when MJOLNIR_DEBUG env var is set).

    colocate_tensors: bool = False
    # When True, copies external tensor files into the artifacts directory.
    # When False, creates symbolic links instead.

    requires_reproducibility: bool = True
    # When True, compiles with RepOp kernels for cross-hardware reproducibility.
    # When False, uses standard PyTorch kernels.

Inference Sessions

REE v0.3.0 introduces InferenceSession, an SDK abstraction for managing the lifecycle of an inference workflow.

An inference session is responsible for setup, execution, and teardown. This lets applications reuse session state across multiple inference calls within the same container instance instead of repeating setup steps for every generation.

This is useful for:

• Running multiple prompts against the same prepared environment.

• Building multi-turn inference workflows.

• Implementing tool-call loops where the model may request an external tool before producing a final answer.

• Avoiding repeated compiler initialization where session reuse is supported.

RepOps

RepOps (Reproducible Operators) are purpose-built GPU kernels that guarantee bitwise-identical outputs regardless of hardware architecture. They cover the full set of operators needed for neural network inference and training.

For a standalone demo of RepOp kernels, see the RepOps Demo repository.

How RepOps Achieve Cross-Hardware Reproducibility

• Fixed reduction ordering: Every kernel accumulates values in a single canonical order. The reduction tile size is fixed across all GPU architectures. All accumulation is in FP32 using fused multiply-add instructions.

• Correctly rounded transcendentals: Custom implementations of exp, sin, tanh, etc. that produce identical results on every CUDA-capable GPU.

• Extended-precision arithmetic: Operations like the error function (used in GELU) use extended-precision fixed-point arithmetic for cross-hardware consistency.

• Architecture-adaptive output tiling: Kernels adapt output tile dimensions to different GPU architectures (using available shared memory), but never change the reduction dimension, so reproducibility is preserved.

Pipeline Parallelism

Large models often exceed the memory of a single GPU. Pipeline parallelism splits the model's layers into sequential partitions (also called stages), each placed on a different GPU.

Pipeline Parallelism is only possible on multi-GPU devices. Using the flag to enable this mechanism on a single-GPU host will result in a failed run.

A forward pass walks the input through partition 1, then passes its activations to partition 2, and so on, with each GPU holding only its own slice of the weights. This trades a single large memory footprint for several smaller ones, making it possible to run models that would otherwise be impossible to load.

Pipeline parallelism is orthogonal to the operation set. You can combine --n-partitions with default, deterministic, or reproducible mode. Only reproducible mode guarantees bitwise identity across different hardware.

REE exposes this through the --n-partitions flag, which controls how many partitions the model is divided into. On a host with enough aggregate GPU memory, this lets REE run models up to 72B parameters while preserving the reproducibility guarantees provided by RepOps.

Partition boundaries are deterministic for a given model and partition count, so splitting a model does not introduce new sources of numerical drift: the same --n-partitions value on different supported hardware produces bitwise-identical output, and cross-partition-count runs of the same model also match when using --operation-set reproducible.

Container Details

Property	Value
REE Image	gensynai/ree:v0.4.0
REE Version	0.4.0
SDK Version	gensyn-sdk 0.1.0
Compiler Version	gensyn-compiler 0.1.0
Base OS	Ubuntu 24.04.1 LTS
Python	3.11.15
PyTorch	2.10.0
Transformers	4.51.0
ONNX	1.21.0
Entrypoint	/runtime/bin/gensyn-sdk
User	gensyn (non-root)
Working Dir	/home/gensyn

Inference Sessions

InferenceSession is a stateful SDK wrapper that loads a prepared task directory once and can run multiple completions against it.

At initialization, the session reads prepare metadata, loads the tokenizer and generation config, configures the selected operation set, selects the device, and builds the inference backend. Calls to complete() reuse that state.

InferenceSession.complete() supports either a plain prompt or chat-style messages. Tool definitions can be supplied only with messages, and are forwarded into the tokenizer's chat template when supported.

Interactive Mode

To explore REE's components directly (SDK, Compiler), start the container in interactive mode:

docker run -it --entrypoint bash -v ~/.cache:/home/gensyn/.cache gensynai/ree:v0.4.0

From inside the container, you can run gensyn-sdk commands directly and inspect intermediate artifacts.