Skip to content

Debug taps

When a compiled engine's accuracy drifts from the torch model, verify tells you that the final outputs diverged — not where. Debug taps expose selected intermediate tensors so verify compares them alongside the real outputs, narrowing a single failing output down to the region of the model where torch and the engine parted ways. The trail is as fine-grained as you choose to make it: tap the points you already suspect, or take every yasp plugin boundary for free — each one becomes its own scored row in the accuracy table, wherever torch has a value to compare against.

There are two kinds, and they combine freely:

  • Manual taps — tensors you mark yourself with tap(x, "name") and expose with --debug-taps. Reach for these when you have an idea of where to look.
  • Plugin taps — every input and output of every yasp TensorRT plugin, tapped automatically by --debug-plugin-taps with no model changes at all. Reach for these when the drift is around a custom op and you don't yet know which tensor to probe.

Both are off by default and zero cost when off — an untapped compile is bit-identical to one whose taps you never turned on. You opt in at prepare time, then compile and verify exactly as usual.

graph LR
    A[1. Mark tensors<br/>with tap<br/>manual taps only] --> B[2. prepare<br/>--debug-taps<br/>--debug-plugin-taps] --> C[3. compile] --> D[4. verify<br/>taps as extra rows]

Step 1 — Mark tensors in your model

Manual taps only — skip straight to step 2 if you just want the plugin boundaries, which need no markers.

Import tap and wrap any intermediate tensor with tap(x, "name"). At runtime tap just returns its input unchanged, so it never alters your model's numerics — it only plants a named anchor the exporter can promote to an output later.

import torch
from torch import nn
from yasp.inference.frontend import tap


class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.backbone = nn.Sequential(nn.Conv2d(3, 16, 3, padding=1), nn.ReLU())
        self.head = nn.Linear(16, 10)

    def forward(self, x):
        x = self.backbone(x)
        x = tap(x, "backbone_out")  # anchor the feature map before the head
        pooled = x.mean(dim=(2, 3))
        return self.head(pooled)


def get_fwd_args():
    return (torch.randn(1, 3, 32, 32),)

Inside an nn.Sequential — where there's no forward to edit — wrap tap in a tiny module and drop it between layers:

class Tap(nn.Module):
    def __init__(self, name: str):
        super().__init__()
        self.name = name

    def forward(self, x):
        return tap(x, self.name)


self.net = nn.Sequential(
    nn.Conv2d(3, 16, 3, padding=1),
    nn.ReLU(),
    Tap("test_1"),  # tap the activation here
    nn.AdaptiveAvgPool2d(1),
    nn.Flatten(),
    Tap("test_2"),  # …and again after flatten
    nn.Linear(16, 10),
)

Tap names must be unique within a model — each one becomes a distinct comparison in verify.

Step 2 — Prepare with --debug-taps

Pass --debug-taps to prepare to expose your taps. It takes a single tap name and is repeatable — pass it once per tap you want:

yasp-toolkit embedded prepare model.py \
  --onnx-out      build/model.onnx \
  --npz-out       build/model.npz \
  --reference-out build/model.reference.npz \
  --debug-taps    all
--debug-taps Meaning
omitted (default) Taps off. Every tap() marker is stripped; the ONNX is identical to a tap-free model.
--debug-taps all Expose every tap() in the model. all wins wherever it appears, even alongside names.
--debug-taps test_1 --debug-taps test_2 Expose only the named subset — the rest are stripped.

prepare records each exposed tap's torch value into the reference npz next to the model's normal outputs, and stamps the ONNX so the compiled engine emits the same tensors. Because the tap manifest rides inside the ONNX, compile and verify need no extra flags — just run them as in Compile your first model.

Plugin taps with --debug-plugin-taps

--debug-plugin-taps is a plain on/off flag that needs no tap() markers anywhere. It tells the exporter to tap every boundary tensor of every yasp TensorRT plugin — all inputs and all outputs of each custom op — and promote them to engine outputs:

yasp-toolkit embedded prepare model.py \
  --onnx-out      build/model.onnx \
  --npz-out       build/model.npz \
  --reference-out build/model.reference.npz \
  --debug-plugin-taps

It is independent of --debug-taps — pass either, both, or neither. Each tapped boundary is named after the plugin node and the boundary's role (…Argsort1D__out0), which is the name it carries in the verify table.

The difference from a manual tap is what happens when there's no torch value to compare against. A plugin boundary often has none: it can sit inside a region torch computes as one fused kernel, so that intermediate never exists eagerly. prepare captures what it can — from the eager run, or by recomputing the boundary with onnxruntime — and leaves the rest uncompared rather than failing.

--debug-taps --debug-plugin-taps
Needs tap() markers in the model yes no
Picks the tensors you do, by name every plugin boundary, automatically
A probe with no reference value fails prepare reported, the flow continues
Engine-side values in the build always always, compared or not

So an uncompared plugin tap still hands you the engine's own value for that tensor, in the build's sample_outputs.npz, to read by hand.

A requested manual tap must survive

Taps are off unless you ask for them, so when you do, every named tap must reach a comparable output or prepare fails with a message naming the probe. That happens when a rewrite pass optimizes the tapped tensor away, when the name matches no tap() in the traced graph (usually a typo, or a tap on a branch the trace didn't take), or when the value can't be captured — e.g. a bfloat16 tensor numpy can't represent. The fix is to relocate the tap to a tensor that survives (a module boundary is a safe bet), correct the name, or cast to float32 before tap().

Plugin taps are exempt — you didn't place them, so one that can't be compared is reported, not fatal. The one exception is a plugin tap that did have a reference and then lost its engine binding during the build; that fails like a manual probe.

Step 3 — Verify

Run verify unchanged:

yasp-toolkit embedded verify \
  build/model.reference.npz \
  "build/model.$ARCH.fp16.tgz"

verify compares every array in the reference against the build's outputs — model outputs and taps alike — so each tap becomes its own row in the accuracy table, scored with the same MAPE / MSE / max-diff / cosine metrics as a real output. Read the rows in model order: the first tap whose numbers fall off is the layer where the engine started to diverge. A tap that diverges is reported, never fatal — that divergence is exactly the signal a tap exists to surface (only --max-mape, which you set yourself, can make verify exit non-zero).

Plugin taps score the same way under their plugin-derived names, for every boundary that got a reference. A boundary that got none has nothing in the reference to pair with, so verify says so and moves on:

[model.sm90.fp16]: 7 outputs are missing in the reference, this can occur when plugin taps with no tensor outputs are enabled

That warning is expected with --debug-plugin-taps, not an error — pass --verbose to list each unpaired name as its own row in the table. Their engine-side values are still in the build's sample_outputs.npz (share/yasp_binary/examples/generated/, extract it with transfer extract), and the worker's own per-plugin report — which names the exact reason each boundary went uncompared — is written to the compile logs.

How it works

  • prepare --debug-taps captures each exposed manual tap's torch-eager value during the reference run and appends it to the reference npz, keyed by tap name — right alongside the model outputs and the host-torch benchmark_timings_ms telemetry.
  • prepare --debug-plugin-taps instead finds the tensors itself, tapping every input and output of every yasp plugin in the exported graph. Boundaries it can't reference are still tapped — just not written to the reference npz.
  • Either way the run stamps the tap manifest into the ONNX, so when the worker builds the engine it promotes those taps to engine outputs. They land in the build's sample_outputs.npz next to the real outputs.
  • verify loads both npz files and pairs their arrays by name. Taps therefore need no special handling — they're just more arrays to score. A name present on one side only is what produces the "missing in the reference" / "missing in the sample output" warnings.