Optimize your model

TensorRT builds a fast engine out of the kernels it already has. When a region of your model has no good native kernel — a fused attention block, a depthwise convolution chain, an op TensorRT cannot convert at all — that region sets the floor on your latency, and no amount of recompiling moves it.

This flow removes that floor by joining the two products end to end: Kronos measures the engine on the target and carves the slow regions out of the ONNX as custom-op nodes, Gaia generates a CUDA kernel for each one, and Kronos cross-compiles those kernels into TensorRT plugins and links them into the next build.

Nothing is installed on the target machine. The plugins travel with the compile as .so files.

graph LR
    A[1. Baseline<br/>prepare + compile] --> B[2. propose<br/>where to plug]
    B --> C[3. prepare --plugin<br/>carve the regions]
    C --> D[4. kernelgen compile<br/>generate kernels]
    D --> E[5. build-plugin<br/>cross-compile .so]
    E --> F[6. compile --custom-plugin<br/>rebuild the engine]
    F --> G[7. verify<br/>measure the win]

Steps 1–3 and 5–7 are yasp-toolkit embedded (Kronos); step 4 is yasp-toolkit kernelgen (Gaia). Every example below is also spelled with the short alias ytk, which is the same CLI.

Prerequisites

A Kronos worker on the target hardware — see Kronos workers. It builds the engines and profiles them, which is where the measurements come from.
A Gaia eval worker on a GPU of the same architecture as the target — see Gaia workers. Generated kernels are scored on real hardware, so this GPU decides whether a kernel is good.
Your API key (get one).

Plugin building needs nothing from you: build-plugin cross-compiles in yasp's cloud for the arch spec you name.

The worked example

The walkthrough uses EfficientViT-B1, a computer-vision backbone, at batch 4 and 1088×1920, compiled fp16 for a DRIVE AGX Orin:

export ARCH=arm64-sm87-dos6010-cu114-trt86

It is an ordinary model.py — a Model class and a get_fwd_args(), exactly as in Compile your first model. Plugin selection never touches the model file; it all happens on the command line.

Step 1 — Baseline the engine

Prepare and compile as usual. The baseline is not optional: every later decision is made from measured time on the target, not from reading the graph.

ytk embedded prepare model.py \
  --onnx-out      base/model.onnx \
  --npz-out       base/model.npz \
  --reference-out base/model.reference.npz

ytk embedded compile -a "$ARCH" -p fp16 \
  --onnx base/model.onnx --npz base/model.npz -o base/base.tgz

Alongside the usual outputs, prepare writes base/fx_nodes.json — the traced FX graph's nodes, each carrying the torch module that owns it. The frontend exports a flat ONNX graph with no module scopes in the layer names, so this file is what later lets a TensorRT layer name be mapped back to the nn.Module it came from. Keep it next to the ONNX it was produced with.

The compile bundle already contains a per-layer profile captured on the target (Profile the model), so verify gives you the baseline immediately:

ytk embedded verify base/model.reference.npz base/base.tgz

│ base              │    10 │ 70.283 │ 70.3… │ 70.343 │ 70.4… │ 70.466 │ 70.4… │

70.3 ms. The per-layer table is in the same output, but TensorRT layer names (/conv_7/Conv + PWN(...)) are rarely enough on their own to tell you which part of your model to attack. That is the next command's job.

Step 2 — Ask where to plug

propose reads the profile straight out of the compile bundle, attributes every millisecond back to the torch modules it came from, and prints the regions worth replacing:

ytk embedded propose \
  --profile   base/base.tgz \
  --onnx      base/model.onnx \
  --fx-nodes  base/fx_nodes.json \
  --model     model.py \
  --emit-prepare run_prepare.sh

measured: base.tgz — every ms below is TensorRT layer time attributed to a
region (today's cost = the most a plugin can recover). model total 70.3 ms.

WHAT TO PLUG                                 SITES   TRT ms   plug as (names are yours)
LiteMLA region (C=384)                           3      7.5   yasp::LiteMLAC384=LiteMLA:around@where=*m.stages.2.blocks.*.context_module.main
LiteMLA region (C=768)                           4      5.1   yasp::LiteMLAC768=LiteMLA:around@where=*m.stages.3.blocks.*.context_module.main
depthwise 3x3 s1  (C=16..1024, 4 families)  (dynamic remainder)     9      5.9   yasp::DwConv3x3s1Dyn=ConvNormAct@where=*.0.blocks.1.main.depth_conv,*.local_module.main.depth_conv,*.res0.main.depth_conv
depthwise 5x5 s1  (C=384)                        3      3.7   yasp::DwConv5x5s1C384=Conv2d@where=*m.stages.2.blocks.*.context_module.main.aggreg.0.0
depthwise 5x5 s1  (C=768)                        4      3.0   yasp::DwConv5x5s1C768=Conv2d@where=*m.stages.3.blocks.*.context_module.main.aggreg.0.0
depthwise 3x3 s1  (C=256)                        2      2.1   yasp::DwConv3x3s1C256=ConvNormAct@where=*.1.blocks.1.main.depth_conv,*.2.main.depth_conv

NOT WORTH PLUGGING
1x1 / dense GEMMs         51 sites     19.9 ms   TRT at roofline — generated kernels lost every attempt
families under 1 ms        3 sites      1.6 ms   a kernelgen run cannot repay this
unattributed glue                —     19.0 ms   reformats between regions; shrinks as regions get plugged
depthwise 3x3 s2x2  (C=64)   1 sites      1.8 ms   strided: bundle grammar cannot express strided output extents

Six candidates covering ~27 ms of the 70 ms — roughly 40% of the model's time sits in regions a custom kernel could take over. Every TRT ms figure is what that region costs today, which is the most a perfect plugin could recover.

Read the second table just as carefully. NOT WORTH PLUGGING is where the command earns its keep: the 19.9 ms of 1×1 GEMMs is the single biggest block in the model and is exactly the thing you should not write a kernel for — TensorRT already runs it at the memory roofline.

A few things worth knowing about how the carve is chosen:

Regions are split per shape family, because a kernel specialised on constant channel counts measurably beats a channel-general one on fixed-shape engines. Everything too small for its own kernel folds into one dynamic remainder plugin per op (DwConv3x3s1Dyn above, 9 sites).
Pass --merge-shape-families for one dynamic candidate per op instead — only right if your deployment genuinely pushes several shapes through one engine.
Pass --plugin <selector> (repeatable) to attribute the measured cost onto a carve you already ship, and report exactly those plugins.
--verbose adds the torch-class time table and a per-proposal detail block; --top/--depth control how much is shown.

--emit-prepare run_prepare.sh writes a runnable script carrying exactly those selectors, with the plugin names parameterised at the top so you can rename them to whatever your .so files will register.

See embedded propose for every flag.

The names are yours

TensorRT resolves a plugin by (namespace, name, version). The name on the left of = is the ONNX node prepare inserts and the name the built plugin must report. propose fills in a derived default; whatever you settle on, you pass it again as build-plugin --plugin-name in step 5. Everything to the right of = is derived from the measurement and should be left alone.

Step 3 — Carve the regions out

Run prepare again with one --plugin per candidate. This is the only change to the export:

ytk embedded prepare model.py \
  --onnx-out      gaia/model.onnx \
  --npz-out       gaia/model.npz \
  --reference-out gaia/model.reference.npz \
  --plugin 'yasp::LiteMLAC384=LiteMLA:around@where=*m.stages.2.blocks.*.context_module.main' \
  --plugin 'yasp::LiteMLAC768=LiteMLA:around@where=*m.stages.3.blocks.*.context_module.main' \
  --plugin 'yasp::DwConv3x3s1Dyn=ConvNormAct@where=*.0.blocks.1.main.depth_conv,*.local_module.main.depth_conv,*.res0.main.depth_conv' \
  --plugin 'yasp::DwConv5x5s1C384=Conv2d@where=*m.stages.2.blocks.*.context_module.main.aggreg.0.0' \
  --plugin 'yasp::DwConv5x5s1C768=Conv2d@where=*m.stages.3.blocks.*.context_module.main.aggreg.0.0' \
  --plugin 'yasp::DwConv3x3s1C256=ConvNormAct@where=*.1.blocks.1.main.depth_conv,*.2.main.depth_conv'

(Or just ./run_prepare.sh, which propose generated with these arguments already filled in.)

Each matched region is replaced in the ONNX by a single custom-op node, and the export gains two things:

Artifact	What it is
`plugins.json`	The machine-readable carve — one entry per plugin with its placeholder op, site count, selector, and the path to its Gaia input. `propose` reads it back to recognise plugins in a later profile.
`plugin_candidates/<Name>.pt2`	The kernel-generation reference: the carved region exported as a `torch.export` program with weights bundled. This is what Gaia optimizes against.

Also written per candidate: <Name>.io.npz and <Name>.inputs.pt (the captured I/O, so the region's numerics can be checked), <Name>.reference.py (the torch source of the region), and candidates.json.

Re-running prepare refreshes only its own artifacts in that directory — generated kernels, reports and built .so files are left alone.

prepare also prints the carve as a table when it finishes — which regions matched, how many sites each covers, and which ones carry a kernel-generation input — and candidates.json records the same thing per candidate: the per-site input/output shapes, the region's op inventory, and which axes are constant across sites.

See embedded prepare for every flag.

Selector grammar

A selector reads left to right as "insert a plugin named PluginName in namespace ns, at the region called OpsName":

<namespace>::<PluginName>=<OpsName>[:around][@modifier=...]

Part	Meaning
`<namespace>`	The TensorRT plugin namespace the `.so`'s creator registers under. `yasp` is the builder default; `abc::X` targets a creator under `abc`; `::X` targets the empty namespace (a hand-built `.so` that never called `setPluginNamespace`).
`<PluginName>`	The op emitted into the ONNX, and the name the plugin must report. Decoupled from `OpsName`, so you can name it for the target.
`<OpsName>`	The region — the torch class shown in propose's `BLOCK` column. Bare, it replaces every instance of that block wholesale.
`:around`	Target the heavy op(s) inside the block and grow the region to absorb the plumbing around them (transposes, slices, casts, pads, pointwise, and the joining concat). The emitted `.pt2` then carries those reformats, so the generated kernel can replace them too.

Modifiers, appended in any order:

Modifier	Effect
`@where=<globs>`	Scope a class match to module-path globs, comma-separated (OR); a `!`-prefixed glob excludes. Essential when one class is the right operator but the wrong scope — every depthwise conv here is the same `ConvNormAct`, but the stride-2 one needs a different kernel.
`@lift=<how>`	Pass the swapped module's weights to the plugin as real inputs. Automatic for a swapped module containing a conv (`fold_conv_bn`); `@lift=none` opts out. A region with more than one conv needs `@lift=fold_conv_bn_seq`, which emits one folded `(weight, bias)` pair per conv in call order — `fold_conv_bn` refuses such a region, since one folded pair cannot describe two convolutions.
`@nojoin`	Stop an `:around` region from swallowing the concat that joins its inputs. Absorbing it is the default and is worth ~0.6 ms/site here.

Escape hatches for regions no class name describes: call:<dotted.path> (intercept an op during export), A...B (a dataflow span between two fx_nodes.json nodes), a comma-separated node list, or mod:<suffix> (a module qualified-name suffix).

Step 4 — Generate the kernels (Gaia)

One kernelgen compile per candidate:

ytk kernelgen compile \
  -r gaia/plugin_candidates/DwConv3x3s1C256.pt2 \
  -g Orin \
  --config deployment_target=tensorrt \
  -o gaia/sources/DwConv3x3s1C256.module_source.py

deployment_target=tensorrt is required

Every kernelgen compile in this flow must pass --config deployment_target=tensorrt. It tells the generator the kernel is destined for a TensorRT plugin rather than a standalone torch extension, which is what makes the emitted bundle something build-plugin can turn into a .so. Omit it and the run produces a kernel you cannot plug in — a wasted job, since generation is not cheap.

-g names the eval worker's GPU, and it should match the target's architecture — the DRIVE AGX Orin here is sm87, so kernels are scored on an Orin. Gaia optimizes against the torch reference in the .pt2 and keeps the candidate that measured fastest on that hardware.

Torch speedup is not TensorRT speedup

Gaia scores candidates against the torch reference, and torch is typically several times slower than TensorRT for the same region. A large speedup over torch does not by itself mean the kernel beats the TensorRT layer it replaces. The real bar is the measured per-site TensorRT time of the region — propose --verbose prints it for every already-plugged region — so read report.json against that number, not against the torch speedup.

Two properties of the generated kernel are hard requirements rather than preferences, because getting either wrong produces a kernel TensorRT cannot insert at all — check them on the result before spending a build-plugin on it:

Input arity — the node's inputs, weights included. A region whose weights were lifted carries them as real inputs, not baked constants. A mismatch surfaces as TensorRT Error 9, "could not find any supported formats consistent with input/output data types" — a format error for what is really an arity problem.
Layout and precision — e.g. kHWC8 (NHWC, channels padded to a multiple of 8) and fp16 in/out with fp32 accumulation. An NCHW kernel needs a reformat on every edge that costs more than the kernel saves, and an fp32-only kernel cannot go into an fp16 engine.

The output is a single Python file bundling everything a plugin needs: the CUDA header, the CUDA source with the generated kernel, and a plugin_meta.yaml block declaring the TensorRT contract (input/output dtypes and formats, output shape expressions, extents, launcher symbol, workspace function, and constraints).

See kernelgen compile for routing and --config options.

Step 5 — Build the plugins

build-plugin cross-compiles the generated source for your target architecture in yasp's cloud — no CUDA toolchain on your machine, and nothing installed on the target:

ytk embedded build-plugin -a "$ARCH" \
  --plugin-name DwConv3x3s1C256 \
  -s gaia/sources/DwConv3x3s1C256.module_source.py \
  -o gaia/plugins/DwConv3x3s1C256.so

It prints the --custom-plugin line to paste into the next step when it finishes.

Always pass --plugin-name

TensorRT resolves a plugin by (namespace, name, version), so the .so has to report exactly the <PluginName> your prepare --plugin selector put into the ONNX — DwConv3x3s1C256 here. Set it yourself on every build-plugin call rather than relying on the name the source arrives with: a plugin whose name does not match its node is one TensorRT cannot resolve, and you find out at step 6 after paying for a compile.

--plugin-name rewrites plugin_name: in a copy of the source before uploading, so the generated file — which cost a job to produce — is never modified in place. Use it instead of hand-editing generated source.

See embedded build-plugin.

Step 6 — Compile with the plugins

Pass every .so to the compile. They are registered with TensorRT before the build starts, so the builder resolves the custom-op nodes prepare inserted:

ytk embedded compile -a "$ARCH" -p fp16 \
  --onnx gaia/model.onnx --npz gaia/model.npz -o gaia/gaia.tgz \
  --custom-plugin gaia/plugins/DwConv3x3s1C256.so \
  --custom-plugin gaia/plugins/DwConv3x3s1Dyn.so \
  --custom-plugin gaia/plugins/DwConv5x5s1C384.so \
  --custom-plugin gaia/plugins/DwConv5x5s1C768.so \
  --custom-plugin gaia/plugins/LiteMLAC384.so \
  --custom-plugin gaia/plugins/LiteMLAC768.so

The plugins ship with the compile request. The target machine needs no plugin installation, no toolchain, and no change to how it runs the engine.

See embedded compile.

Step 7 — Measure

verify takes several bundles and lines their benchmark rows up, so the before/after is one command:

ytk embedded verify base/model.reference.npz base/base.tgz gaia/gaia.tgz

┃ Implementation    ┃ Iter… ┃   Min  ┃  Mean ┃    p50 ┃   p95 ┃    p99 ┃   Max ┃
│ torch host (cuda) │   100 │ 93.618 │ 109.… │ 107.8… │ 136.… │ 147.9… │ 150.… │
│ base              │    10 │ 70.283 │ 70.3… │ 70.343 │ 70.4… │ 70.466 │ 70.4… │
│ gaia              │    10 │ 42.931 │ 42.9… │ 42.951 │ 42.9… │ 42.976 │ 42.9… │

70.3 ms → 42.9 ms, a 1.64× end-to-end speedup from six generated kernels, with no change to model.py and nothing installed on the device. The torch host row is your reference machine's GPU, not the target, so ignore it when the two are different hardware.

verify also scores numerics against the torch reference for every output; a plugin changes how a region computes, so read that table too — see embedded verify.

Then ask again

propose recognises plugins that are already in the profiled engine — it reads plugins.json from beside the profile (or from --plugins), so keep it with the prepare output it came from:

ytk embedded propose --profile gaia/gaia.tgz --onnx gaia/model.onnx \
  --fx-nodes gaia/fx_nodes.json --model model.py

measured: gaia.tgz — every ms below is TensorRT layer time attributed to a
region (today's cost = the most a plugin can recover). model total 43.2 ms.

WHAT TO PLUG                                 SITES   TRT ms   plug as (names are yours)
DwConv3x3s1C256 (already plugged)                2      1.6   re-optimize: 0.787 ms/site to beat
DwConv3x3s1Dyn (already plugged)                 9      6.1   re-optimize: 0.674 ms/site to beat
DwConv5x5s1C384 (already plugged)                3      1.4   re-optimize: 0.460 ms/site to beat
DwConv5x5s1C768 (already plugged)                4      1.0   re-optimize: 0.241 ms/site to beat
LiteMLAC384 (already plugged)                    3      1.9   re-optimize: 0.622 ms/site to beat
LiteMLAC768 (already plugged)                    4      1.3   re-optimize: 0.333 ms/site to beat

NOT WORTH PLUGGING
1x1 / dense GEMMs         51 sites     20.0 ms   TRT at roofline — generated kernels lost every attempt
families under 1 ms        3 sites      1.6 ms   a kernelgen run cannot repay this
unattributed glue                —      5.9 ms   reformats between regions; shrinks as regions get plugged
depthwise 3x3 s2x2  (C=64)   1 sites      1.8 ms   strided: bundle grammar cannot express strided output extents

Two useful readings. The plugged regions now report a per-site time to beat — with --verbose, each one names the .pt2 to re-run against, which turns "write a kernel" into the much better-posed "beat this kernel". And the unattributed glue bucket fell from 19.0 ms to 5.9 ms: absorbing reformats into the plugin regions removed work that was never attributable to any layer in the first place.

Where to next?

Profile the model — the per-layer profile propose consumes, and how to collect one yourself on the device.
Generate your first kernel — the Gaia flow on its own, for kernels that are not part of an engine.
Debug taps — when a plugged engine drifts from torch, expose every yasp plugin boundary as a comparison point and let verify pinpoint the region.
embedded reference — every flag for prepare, propose, compile, build-plugin and verify.
kernelgen reference — export, compile, and get.