yasp-toolkit kernelgen

Generate optimized GPU kernels from a PyTorch reference model: export the reference, run the generation flow, and check on jobs in progress.

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

export

The model file must define a model class (by default Model, an nn.Module subclass) and a forward-inputs factory (by default get_inputs or get_fwd_args) returning the example positional inputs. It may also define a constructor-args factory (get_init_inputs) returning the model's positional constructor args, and a dynamic-shapes factory (get_dynamic_shapes) returning a torch.export spec. Each lookup tries its candidate names in order, so multiple conventions are accepted.

Export a torch model .py to a .pt2 reference ready for kernelgen compile.

The model file must define a model class (by default Model, an nn.Module subclass) and a forward-inputs factory (by default get_inputs or get_fwd_args) returning the example positional inputs. It may also define a constructor-args factory (get_init_inputs) returning Model(...)'s positional args, and a dynamic-shapes factory (get_dynamic_shapes) returning a torch.export spec. Each lookup tries its candidate names in order, so multiple conventions are accepted.

The model is traced with torch.export in inference mode, decomposed supported yasp IR as the reference.pt2 thatkernelgen compile --reference expects.

Usage:

$ yasp-toolkit kernelgen export [OPTIONS]

Option	Default	Description
-i, --input <file>	required	Torch model .py defining a Model class and its input-factory functions.
-o, --output <file>	—	Where to write the .pt2 reference. Default: <input stem>.pt2 alongside the model.
--model-cls <str>	Model	Candidate names for the model class to instantiate. The first one defined in the module is used.
--inputs-fn <str>	get_inputs, get_fwd_args	Candidate names for the forward-inputs factory (returns the positional inputs). The first one defined in the module is used.
--init-fn <str>	get_init_inputs	Candidate names for the constructor-args factory (returns Model(...)'s positional args). The first one defined is used; when none is defined, Model() is built with no arguments.
--dynamic-shapes-fn <str>	get_dynamic_shapes	Candidate names for the dynamic-shapes factory forwarded to torch.export. The first one defined is used; when none is defined, shapes are fully static.

compile

Get an optimized kernel for a reference model, scored on real hardware.

Phases: create the job, upload reference.pt2, start, poll until the job reaches a terminal state, then optionally download the optimized kernel source to -o and/or the compilation report.json to --report.

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

Usage:

$ yasp-toolkit kernelgen compile [OPTIONS]

Option	Default	Description
-r, --reference <file>	required	Reference module (reference.pt2) to compile.
-g, --gpu-name <str>	required	Hardware target the job routes to (e.g. NVIDIA RTX PRO 6000 Blackwell).
-c, --config <str>	—	KEY=VALUE compile setting forwarded to the worker as --config key=value. Repeatable. A value of the form @PATH is loaded from that file (use @@ for a literal leading @) -- handy for long free-form settings like guidance.
--eval-routing <str>	—	Eval-router routing key targeting a specific evaluator for the in-compile evaluations.
-f, --format <str>	python	Output kernel format. Free-form; the compile backend validates supported values.
-o, --output <file>	—	Where to write the optimized kernel source (module_source.py). Omit to skip its download.
-R, --report <file>	—	Where to write the compilation report.json. Omit to skip its download.
--progress / --no-progress	progress	Render transfer progress bars.
--vllm	—	Append the vLLM put-back guidance to --config guidance: cuda-graph / torch.compile safety prompts + sample kernels + the role's C implementation (--c-implementation, if given). Without --vllm, guidance is your --config guidance only.
--c-implementation <file>	—	The role's C implementation to embed in the --vllm guidance (e.g. its sidecar.py). Omit for pure-aten roles.

get

Fetch a kernel-generation job's current state. With --wait, poll until terminal.

Usage:

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

Argument	Default	Description
task_id	required	Kernel-generation task ID.

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