Write Triton,
Ship C++

One call, two artifacts, two engines.

Production runs in C++. Your Triton kernels do not.

Enterprise inference is served by pure C++ engines — ONNX Runtime and TensorRT — for latency, concurrency and whole-graph optimisation. But torch.onnx.export traces TorchScript or TorchDynamo IR that represents native PyTorch operators; an imperative Python call to a triton.JITFunction is opaque to it. The kernel is dropped or the export fails.

Doing it by hand means re-implementing the kernel in CUDA, writing the ONNX and TensorRT plugin classes, patching the PyTorch graph, and building and loading the result. A Triton kernel is not one operator — it presents five properties a C++ engine must recover.

What you would build by hand — and what KernelLens does instead.

To put one Triton kernel into either engine manually, you port it to CUDA, implement the ONNX classes (Ort::CustomOp, Ort::CustomOpApi, domain registration), implement the TensorRT classes (nvinfer1::IPluginV2DynamicExt, nvinfer1::IPluginCreator), write a CMake build, patch the FX graph with a C++ extension and @torch.library.custom_op, export, build the engine, and then drive each engine yourself — creating the session, binding raw CUDA pointers, creating the execution context, calling setTensorAddress and enqueueing.

The compile pipeline

kl.compile(model, sample_input, backends=[...]) produces a validated manifest, an ONNX graph and the plugin libraries. Each arrow below is a concrete artifact.

The run pipeline

At run time, PyTorch's GPU addresses are handed straight to the engine.

import torch, triton, triton.language as tl
import kernel_lens as kl

class RMSNorm(torch.nn.Module):
    def forward(self, x):
        B, M, N = x.shape
        out = torch.empty_like(x)
        rmsnorm_kernel[(B, M)](x, self.weight, out, x.stride(0), x.stride(1),
                               out.stride(0), out.stride(1),
                               N=N, eps=1e-6, BLOCK_SIZE=triton.next_power_of_2(N))
        return out

x = torch.randn(4, 512, 4096, device="cuda")
compiled = kl.compile(RMSNorm().cuda(), (x,), name="llama3_rmsnorm",
                      backends=["onnx", "tensorrt"])   # compile once

out_ort = compiled.run((x,), backend="onnx")            # run in ONNX Runtime
out_trt = compiled.run((x,), backend="tensorrt")        # run in TensorRT

Nine lines of Python, or 1,137 lines of C++.

The same thing — one Triton kernel running as a native operator in ONNX Runtime and TensorRT — from either side. On the left is everything the user writes. On the right is what KernelLens generates, and what you would otherwise write and maintain by hand.

import kernel_lens as kl

# the Triton kernel + module you already wrote
compiled = kl.compile(
    RMSNorm(4096).cuda(), (x,),
    name="llama3_rmsnorm",
    backends=["onnx", "tensorrt"],
)

out_ort = compiled.run((x,), backend="onnx")
out_trt = compiled.run((x,), backend="tensorrt")

// 1. port the kernel to CUDA  (rmsnorm.cu)
// 2. OrtCustomOp for ONNX Runtime
struct RmsnormOp : Ort::CustomOpBase<...> {
  const char* GetName() const { return "rmsnorm"; }
  ONNXTensorElementDataType GetInputType(size_t i) const;
  // ... + a Compute() that rebuilds the PTX argument array
};
// 3. TensorRT plugin
class RmsnormPlugin : public nvinfer1::IPluginV2DynamicExt {
  DimsExprs getOutputDimensions(...) noexcept override;
  bool supportsFormatCombination(...) noexcept override;
  int32_t enqueue(...) noexcept override;      // launch
  size_t getSerializationSize() const noexcept override;
  void serialize(void*) const noexcept override;
  IPluginV2DynamicExt* clone() const noexcept override;
};
class RmsnormPluginCreator : public nvinfer1::IPluginCreator { /* ... */ };
// 4. CMakeLists.txt + nvcc + g++
// 5. a PyTorch C++ extension that patches the FX graph
// 6. drive the engines: Ort::Session + RunWithBindings,
//    IExecutionContext + setTensorAddress + enqueue

Fold the standard layers. Keep the kernels.

The headline: a block whose attention is a custom squared-ReLU flash kernel, whose normalisations and gated activation are custom kernels, and whose projections and residuals are standard operators. Measured on an RTX 5060 Ti (Blackwell, sm_120a).

Folding + kernels vs the alternatives

Same block, three ways to run it. A. Folded only expresses the attention with standard operators and lets the compiler fold everything (what you get if you cannot ship the kernel). B. Kernels only runs the custom kernels with the standard layers unfused (the usual way to use a Triton kernel) — in its torch.compile form the kernel is captured, not reimplemented: since PyTorch 2.1 TorchDynamo records a @triton.jit launch as a triton_kernel_wrapper_mutation node with no graph break, and Inductor emits a call to your kernel while fusing only the surrounding operators. C. Both is KernelLens.

Two-layer block, D=512, H=8, S=4096, batch 1. Relative is versus the folded-only torch.compile row. The kernel wins where folding would materialise; folding wins on the element-wise chain around the projections. KernelLens is the only configuration that has both.

Single-operator latency

The C++ plugin runs the same kernel behind an operator boundary, within 6–13% of native Triton on the three larger operators. The value is deployment, not raw speed on a single op.

End-to-end 4-layer LLaMA decoder

The TensorRT plugin nearly halves time-to-first-token (4.14 ms vs 7.85/7.92 ms) and attains the lowest peak VRAM, with total generation on par with the baselines.

Numerical parity

Parity is against native Triton execution of the same kernel.

Coverage: module patterns

Parity on specific operators is not enough — the compiler has to handle the many ways a module can call a kernel. A conformance suite (tests/test_conformance.py) exercises fourteen patterns against both backends; all pass at 0.00e+00.

Real pretrained models

The result splits by precision. In fp32, torch.compile barely beats eager (1.02–1.12×) because it leaves the operators around the kernel unfused, while KernelLens fuses them: 3.2–5.5× faster, growing with model size. In fp16, torch.compile optimises well and KernelLens only matches it — TensorRT edges it by 1.06–1.22×, and ONNX Runtime is slower (0.67–0.84×). Keeping kernels pays when the baseline leaves the graph unfused; when it already fuses, we break even. Accuracy tracks precision rather than speed: fp32 holds the same greedy token on ≥98.8% of positions, fp16 under TensorRT on 93.4–97.9%, and OPT-350M in fp16 under ORT is bit-exact.

One engine, many shapes

kl.compile(..., dynamic_shapes=True) exports the leading dimension as a symbolic axis, so a single engine serves any size in a declared range — instead of one engine per shape.

TensorRT is not always faster. Here is when it is.

torch.compile is a strong baseline — it fuses, and since PyTorch 2.1 it captures your Triton kernels rather than breaking on them (verified: zero graph breaks). So the honest answer to "why an engine at all?" is not "raw speed". Here is the head-to-head, on the same RTX 5060 Ti, with the custom kernels captured on both sides.

TensorRT wins on GEMM-heavy transformer prefill, where its kernels and folding beat Inductor's generated GEMMs. On convolution-heavy graphs it is still behind, even after two fixes: the plugin now declares channel-last formats for pointwise kernels (so the engine no longer inserts reformats around it), and it now enforces the dtype the kernel was compiled for (previously the engine silently ran fp16 models in fp32). Together these took the fp16 conv case from 10.86 to 8.20 ms with exact parity; the remaining gap is the engine's convolution performance, not the plugin.

Three lines to a C++ engine.

pip install kernel-lens          # core compiler
pip install kernel-lens[ort]     # + ONNX Runtime backend
pip install kernel-lens[trt]     # + TensorRT backend
pip install kernel-lens[all]     # both

git clone https://github.com/Just1truc/kernel_lens.git
cd kernel_lens
pip install -e .[all]

Requirements: Linux x86-64, Python ≥ 3.10, CUDA ≥ 12.0, PyTorch ≥ 2.0, Triton ≥ 2.1, and nvcc + g++ on PATH. TensorRT 8.6 through 11 is supported — see which TensorRT versions. Environment: KERNEL_LENS_CACHE_DIR (put it on a persistent volume), KERNEL_LENS_PLUGIN_API (auto, v2 or v3), KERNEL_LENS_HEADER_CACHE, KERNEL_LENS_CCBIN (host compiler for nvcc; CUDA toolkits before 12.9 reject GCC 13’s exception specifiers, so set this to g++-12 there), KERNEL_LENS_DEBUG, TENSORRT_INCLUDE_DIR.

Where the work stands.

• Tracing, AST taxonomy, manifest assemblydone
• ONNX injection, CPU scalars, in-place aliasingdone
• OrtCustomOp + IPluginV2DynamicExt synthesisdone
• Zero-copy binding, cold-start cachingdone
• 16-byte alignment guard, PTX ISA clampingdone
• FP8 / INT4 quantized GEMM (exact)done
• Folding + kernels benchmarkdone
• Second architecture (A100 / sm_80): 14/14 conformance, bit-exactdone
• Real pretrained models: OPT-125M → 6.7B, both engines, both precisionsdone
• TensorRT-LLM backend (non-attention kernels)next
• Auto-tune dispatch, CUDA Graphs, NCCL, Metal/WebGPUnext

Fair objections.

Compiling PyTorch Triton Kernels into C++ Inference Plugins.

The technical report derives the parameter taxonomy, the manual pipeline, and the KernelLens compile and run pipelines, and reports the folding, parity, latency, end-to-end and quantized results. A TMLR-formatted variant is generated from the same source.

Read the paper (PDF) TMLR version Source & reproduction scripts

@article{duc2026kernellens,
  title       = {KernelLens: Compiling PyTorch Triton Kernels into C++ Inference Plugins},
  author      = {Duc, Justin},
  institution = {Galiad Research},
  year        = {2026},
  version     = {1.2.1},
  doi         = {10.5281/zenodo.23194523}
}