Skip to content

Generate your first kernel

With your reference machine set up, generating a kernel is four steps: start a runner on the target machine, then export, generate, and inspect from the reference machine.

graph LR
    A[1. Start a runner] --> B[2. Export] --> C[3. Generate] --> D[4. Inspect]

Run them yourself below, or skip to the agent and describe the job in plain English.

We'll use this W4A16 weight-only quantized GEMM throughout — a naive reference (it unpacks the 4-bit weights to a full matrix and runs torch.matmul) whose whole point is to be fused, so it's the perfect first model. A model file needs a Model class plus two factories — get_inputs() returning the example forward inputs and get_init_inputs() returning the constructor args (Model(M, N, K) here). Save it as source.py:

source.py — W4A16 weight-only quantized GEMM 
# source.py
from __future__ import annotations

import torch
import torch.nn as nn


GROUP_SIZE = 128


def _pack_int4(w_q: torch.Tensor) -> torch.Tensor:
    """Pack (K, N) uint8 in [0,15] into (K//2, N) uint8 (even rows low nibble, odd rows high)."""
    K, N = w_q.shape
    assert K % 2 == 0
    lo = w_q[0::2].to(torch.uint8) & 0xF
    hi = w_q[1::2].to(torch.uint8) & 0xF
    return (lo | (hi << 4)).contiguous()


def _unpack_int4(w_packed: torch.Tensor, K: int) -> torch.Tensor:
    """Unpack (K//2, N) uint8 -> (K, N) uint8 in [0,15]."""
    Kh, N = w_packed.shape
    assert Kh * 2 == K
    out = torch.empty((K, N), dtype=torch.uint8, device=w_packed.device)
    out[0::2] = w_packed & 0xF
    out[1::2] = (w_packed >> 4) & 0xF
    return out


class Model(nn.Module):
    """W4A16 GEMM: y = x @ dequant(w_q, scales, zeros)."""

    def __init__(self, M: int, N: int, K: int, group_size: int = GROUP_SIZE):
        super().__init__()
        assert K % group_size == 0, "K must be divisible by group_size"
        assert K % 2 == 0, "K must be even (int4 packing)"
        self.M, self.N, self.K = M, N, K
        self.group_size = group_size
        n_groups = K // group_size

        torch.manual_seed(0xC0DE ^ (M * 1315423911 + N * 2654435761 + K))
        w_full = torch.randn(K, N, dtype=torch.float32) * 0.02

        w_g = w_full.view(n_groups, group_size, N)
        w_min = w_g.min(dim=1, keepdim=True).values
        w_max = w_g.max(dim=1, keepdim=True).values
        scales = (w_max - w_min).clamp_min(1e-8) / 15.0
        zeros = (-w_min / scales).round().clamp(0, 15)
        w_q = ((w_g / scales) + zeros).round().clamp(0, 15).to(torch.uint8).view(K, N)

        self.register_buffer("w_q", _pack_int4(w_q))  # uint8, (K//2, N)
        self.register_buffer("scales", scales.squeeze(1).to(torch.bfloat16))  # bf16, (n_groups, N)
        self.register_buffer("zeros", zeros.squeeze(1).to(torch.bfloat16))  # bf16, (n_groups, N)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        K = self.K
        w_unpacked = _unpack_int4(self.w_q, K).to(torch.bfloat16)
        scales = self.scales.repeat_interleave(self.group_size, dim=0)
        zeros = self.zeros.repeat_interleave(self.group_size, dim=0)
        w_bf = (w_unpacked - zeros) * scales
        return x.to(torch.bfloat16) @ w_bf


M = 1
N = 12288
K = 4096


def get_inputs():
    return [torch.randn(M, K, dtype=torch.bfloat16)]


def get_init_inputs():
    return [M, N, K]

Step 1 — Start a runner

A kernel is scored on the target GPU, not in the cloud — so you first put an eval worker there: a yasp-eval container that detects the hardware, registers with the Eval router, and then benchmarks every candidate kernel KernelGen routes to it. You deploy it once; it stays up and serves all your generations.

Deploy one on the target — see Gaia workers (single host with Docker, or Helm for a Kubernetes cluster). Then confirm it registered and note its GPU Name — the value you pass to --gpu-name in step 3:

yasp-toolkit eval workers

Step 2 — Export the reference

On the reference machine, turn source.py into the reference.pt2 the generator consumes — a torch.export program with its weights bundled:

yasp-toolkit kernelgen export \
  --input  source.py \
  --output build/source.pt2

export traces the model with torch.export in inference mode, decomposes it to Core ATen IR, and saves it with the weights bundled. It creates build/ itself. See kernelgen export.

Model won't export?

If export can't trace your model, send the exact error (and your source.py if you can share it) to support@yasp.ai.

Step 3 — Generate

One command creates the task, uploads source.pt2, runs the generation, and downloads the result. Use your worker's GPU Name from step 1:

export GPU="<GPU Name from `yasp-toolkit eval workers`>"

yasp-toolkit kernelgen compile \
  --reference build/source.pt2 \
  --gpu-name  "$GPU" \
  --output    build/module_source.py \
  --report    build/report.json

The task id prints immediately, then a progress bar streams the upload and a spinner tracks the generation. On success build/module_source.py holds the optimized kernel and build/report.json the compilation report. See kernelgen compile for --config tuning and routing options.

Generation failed?

Check the job's state with yasp-toolkit kernelgen get <TASK_ID> --wait (the task id is printed when the task starts) — it polls until the job reaches a terminal state and exits non-zero if it failed. Send the id and the exact error to support@yasp.ai.

Step 4 — Inspect

Read the generated kernel and its report:

cat build/module_source.py     # the optimized, fused kernel source
cat build/report.json          # timing, accuracy, and what the generator changed

report.json records how the generated kernel scored on the worker's GPU against the reference; module_source.py holds the generated source — the W4A16 unpack and the GEMM fused instead of materializing the full bf16 weight matrix.

Prefer the agent?

The same four steps, by conversation. Launch the shell and describe the job:

yasp

A startup check verifies your environment first; then the bundled skills handle the rest — e.g. "Deploy an eval worker to gpu-lab-03, then generate a fused kernel from source.py and show me the speedup." It deploys the worker, exports the reference, runs the end-to-end generation, and reports how the kernel scored.

Where to next?

Your own models follow the exact same four steps — point export at a different source.py and reuse the running worker for every generation. From here:

  • Gaia workers overview — deploy and manage the eval workers that back generation.
  • kernelgen reference — export, compile, and get in one place.
  • eval reference — score candidate kernels against a reference with eval compare, and inspect evaluators.
  • CLI reference — every yasp and yasp-toolkit command, flag by flag. Session management (yasp --list, --resume) lives on the yasp page.