yasp-toolkit embedded

Build a TensorRT engine for your target GPU from a PyTorch model: turn the model into compile inputs, run the compile, verify the engine's numerics against torch, and check on tasks and the workers that ran them.

The router host defaults to https://compile.yasp.ai/embedded and can be overridden with $YASP_EMBEDDED_ROUTER_HOST or the advanced-auth --router-* flags shown in --help.

prepare

Generate the ONNX + NPZ inputs that embedded compile expects from a torch model.

Also captures the torch-eager outputs into a reference npz so embedded verify can compare compiled artifacts against ground truth without re-loading the torch model. Unless --bench-iters 0, the same npz additionally records per-iteration host-torch timings (on CUDA when available) that verify renders as the reference row of its benchmark table.

The model file must define a Model class (nn.Module subclass) and a get_fwd_args() -> tuple returning example positional inputs.

Usage:

$ yasp-toolkit embedded prepare [OPTIONS] {model_file}

Argument	Default	Description
model_file	required	Torch model .py defining Model and get_fwd_args.

Option	Default	Description
--onnx-out <path>	required	Destination path for the ONNX model.
--npz-out <path>	required	Destination path for the .npz sample inputs.
--reference-out <path>	—	Destination path for the torch-eager reference outputs npz. Default: .reference.npz alongside the ONNX file.
--bench-iters <int>	100	Timed host-torch iterations stored as benchmark_timings_ms in the reference npz, so embedded verify can show host latency next to each worker's. 0 disables.
--bench-warmup <int>	10	Warmup iterations before the timed host-torch benchmark.
--debug-taps <str>	—	Expose tap() markers in the model as extra reference outputs so verify compares intermediate tensors too. Repeat once per tap name (e.g. --debug-taps alpha --debug-taps beta), or --debug-taps all for every tap. Omit to expose none.
--debug-plugin-taps	—	Enables debug plugin taps.

See Debug taps for more information on taps.

compile

Build a TensorRT engine from a prepared model on your target hardware.

Phases: create the task, upload model.onnx + sample.npz, start, poll until the task reaches a terminal state, then optionally download the result.

The task id is printed as soon as it's obtained, so a failure mid-flow (upload, wait, download) is recoverable via the granular commands (embedded get <task-id>, download …).

Usage:

$ yasp-toolkit embedded compile [OPTIONS]

Option	Default	Description
-a, --arch-spec <str>	required	Worker arch_spec tag the task should route to (e.g. amd64-sm120-u2404-cu128-trt00).
--routing <str>	—	Pin the task to workers registered with this routing name (e.g. a dev/test pool registered with --tag routing=<name> / worker.tags.routing). Omit to use the shared stable pool.
-p, --precision <str>	fp32	Numerical precision (fp32, fp16, bf16, ...).
--onnx <file>	required	Local ONNX model to upload as model.onnx.
--npz <file>	required	Local NPZ sample input to upload as sample.npz.
-o, --output <file>	—	Where to write the resulting output.tgz (the demo log lives inside it as share/.../generated/demo_run.log — extract with transfer extract). Omit to skip the download step.
--dataset <str>	—	Name of a registered dataset (see dataset add/list) to attach to the task.
--refresh-manifest	—	Force a fresh dataset manifest even if a still-valid one is cached.
--calibration-cache <file>	—	Path to a pre-built INT8 calibration cache (TRT-XYZ-EntropyCalibration2 format, ~1 kB). Use with --precision int8 to skip live calibration on the worker. When omitted and the precision is int8, the worker calibrates against the dataset attached via --dataset (regular manifest path).
--custom-plugin <file>	—	Path to an extra TensorRT plugin shared library to register for this compile (repeatable).
--progress / --no-progress	progress	Render transfer progress bars.

The demo log lives inside output.tgz as share/.../generated/demo_run.log — pull it out with transfer extract.

calibrate

Produce an INT8 calibration cache from a dataset on your target hardware.

Phases: create the task, upload model.onnx + the dataset manifest, start, poll until the task reaches a terminal state, then optionally download the resulting calibration.cache.

Feed that cache to compile --calibration-cache (with --precision int8) to skip live calibration on subsequent builds.

The task id is printed as soon as it's obtained, so a failure mid-flow (upload, wait, download) is recoverable via the granular commands.

Usage:

$ yasp-toolkit embedded calibrate [OPTIONS]

Option	Default	Description
-a, --arch-spec <str>	required	Worker arch_spec tag the task should route to (e.g. amd64-sm120-u2404-cu128-trt00).
--routing <str>	—	Pin the task to workers registered with this routing name (e.g. a dev/test pool registered with --tag routing=<name> / worker.tags.routing). Omit to use the shared stable pool.
-p, --precision <str>	fp32	Numerical precision (fp32, fp16, bf16, ...).
--onnx <file>	required	Local ONNX model to upload as model.onnx.
--dataset <str>	required	Name of a registered dataset (see dataset add/list) the worker calibrates against.
--input-name <str>	—	Name of the model input the dataset feeds; disambiguates models with multiple inputs.
-o, --output <file>	—	Where to write the resulting binary calibration.cache. Omit to skip the download step.
--refresh-manifest	—	Force a fresh dataset manifest even if a still-valid one is cached.
--custom-plugin <file>	—	Path to an extra TensorRT plugin shared library to register for this calibration (repeatable).
--progress / --no-progress	progress	Render transfer progress bars.

build-plugin

Cross-compile a TensorRT plugin from source for your target hardware.

Phases: create the build, upload the source file, start, poll until the build reaches a terminal state, then download the resulting shared object.

The build id is printed as soon as it's obtained, so a failure mid-flow (upload, wait, download) is recoverable via download.

Usage:

$ yasp-toolkit embedded build-plugin [OPTIONS]

Option	Default	Description
-a, --arch-spec <str>	required	Build spec to cross-compile for; routes to plugin builders advertising this tag (e.g. jp61-cu126-trt103-py311).
-s, --source <file>	required	Plugin source file to build.
--input-format <str>	cuda	Source language of --source. The build backend validates the values it supports.
-o, --output <file>	required	Where to write the compiled shared object (plugin.so).
--progress / --no-progress	progress	Render transfer progress bars.

get

Fetch a task's (compile, calibration or plugin build) current state.

With --wait, poll until terminal.

Usage:

$ yasp-toolkit embedded get [OPTIONS] {task_id}

Argument	Default	Description
task_id	required	Task ID (a compile, calibration or plugin-build id).

Option	Default	Description
--wait	—	Poll every 15s until the task is completed or failed.

key

Print a scoped key a Kronos worker can register with.

Use the JWT wherever a Kronos worker asks for YASP_API_KEY:

export WORKER_KEY=$(yasp-toolkit embedded key)
kubectl create secret generic yasp-api-token --from-literal=token="$WORKER_KEY"

Usage:

$ yasp-toolkit embedded key [OPTIONS]

Option	Default	Description
--audience <str>	embedded-api	Token audience claim.
--scope <str>	—	Token scope (repeatable). Defaults to embedded-register only.

workers

List embedded workers (optionally filtered by tag).

Usage:

$ yasp-toolkit embedded workers [OPTIONS]

Option	Default	Description
-t, --tag <str>	—	KEY=VALUE filter on worker tags. Repeatable.
--detailed	—	Include worker ID and user columns.

verify

Compare each compiled .tgz's outputs against the torch reference.

Reads the reference npz produced by embedded prepare, then extracts share/.../sample_inputs.npz and share/.../generated/sample_outputs.npz from each tgz. Sanity-checks that the worker saw the same bytes across every tgz (otherwise comparing outputs is meaningless), then scores each output tensor positionally (by _<i> suffix) — MAPE / MSE / max-diff / cosine, and top-K agreement when the output shape looks classification-y.

When a tgz's outputs npz carries the worker's benchmark_timings_ms telemetry array, a per-tgz benchmark table (min/mean/p50/p95/p99/max ms) is appended after the accuracy table. A reference npz produced by embedded prepare with --bench-iters > 0 carries host-torch timings under the same key; those render as the table's first row, so each target's latency reads against the host baseline.

Usage:

$ yasp-toolkit embedded verify [OPTIONS] {reference} {tgzs}...

Argument	Default	Description
reference	required	Reference outputs npz produced by embedded prepare.
tgzs...	required	One or more compiled output .tgz files to score against the reference.

Option	Default	Description
--max-mape <float>	—	Exit 1 if any output tensor's MAPE exceeds this percentage.
--remove-outliers	—	Drop MAPE outliers via 1.5x IQR before averaging (mirrors the inference notebook).
--topk <int>	—	Report top-K classification agreement vs reference. Auto-on for 2D outputs with last-dim > 100.
--json	—	Emit machine-readable JSON instead of a Rich table.
--verbose	—	Output all missing outputs and profiling layers.

verify compares every array in the reference against the corresponding array in each build's outputs — the model outputs plus any debug taps exposed at prepare time — skipping only the benchmark telemetry. Each becomes a row scored with MAPE / MSE / max-diff / cosine (and top-K when a shape looks classification-y).

After the accuracy table, verify also prints a benchmark table (min/mean/p50/p95/p99/max ms). Each .tgz contributes a row from the worker's benchmark_timings_ms telemetry, and — when the reference npz was produced by prepare with --bench-iters > 0 (the default) — a torch host row from the host-torch timings, so you can read the compiled-vs-eager speedup directly. This replaces pulling demo_run.log by hand for latency.