Benchmark a Kernel with KCoral

Upload Python that runs your kernel, checks the output, and returns a timing report. For convenience, KCoral provides compile_tirx and benchmark in kcoral.builtins for your program to use directly. You can also use your own compilation and measurement code.

Prerequisites

Install the GPU worker environment and launch the server. The first example uses TIRx, TVM’s Python-embedded kernel language.

Compile, check and measure

KCORAL_URL=http://localhost:8000 python examples/benchmark_kernel.py

"""Compile a TIRx kernel, check its output, and measure GPU activity."""

import os

import numpy as np

from kcoral import Client, Program

SOURCE = r"""
from __future__ import annotations

import torch
from tvm.script import tirx as T
from kcoral.builtins import compile_tirx, benchmark

@T.jit
def add_one(A: T.Buffer((N,), "float32"), B: T.Buffer((N,), "float32"), *, N: T.constexpr):
    T.device_entry()
    i = T.cta_id([N])
    t = T.thread_id([1])
    B[i] = A[i] + 1.0


def evaluate(src):
    dst = torch.empty_like(src)
    compiled = compile_tirx(add_one, {"N": src.numel()})
    compiled(src, dst)
    torch.testing.assert_close(dst, src + 1.0, rtol=1e-2, atol=1e-3)
    return {"check": {"passed": True}, "timing": benchmark(compiled, src, dst)}
"""


def build_program() -> Program:
    program = Program()
    module = program.upload(id="module", kind="module", source=SOURCE)
    evaluate = program.get_function(id="evaluate", module=module, name="evaluate")
    values = np.arange(256, dtype=np.float32)
    src = program.upload(id="src", kind="tensor", value=values)
    report = program.run(id="report", fn=evaluate, args=[src])
    program.return_(key="report", value=report)
    return program


def main() -> None:
    with Client(os.environ.get("KCORAL_URL", "http://localhost:8000")) as client:
        result = client.execute(build_program(), timeout_seconds=120)
    if not result.completed:
        raise SystemExit(f"Benchmark failed: {result.error}")
    print(result.results["report"]["check"])
    print(result.results["report"]["timing"])


if __name__ == "__main__":
    main()

Download the example.

compile_tirx(kernel, bindings) specializes a TIRx kernel and compiles it for CUDA. bindings supplies constexpr values; a PrimFunc can be compiled without bindings. Compiled executables are cached by structural hash, up to 32 entries per process. The helper requires TVM in the worker environment.

benchmark(compiled, src, dst) measures GPU activity with CUPTI. Both functions use the GPU access of the Python call that invokes them. Importing them into a module also allows selecting them with get_function and calling them in separate instructions.

Where to compile

On a GPU server

Upload your kernel and Python compilation code. For TIRx, import compile_tirx from kcoral.builtins. CUDA C, CuTeDSL and Triton can use their own compiler APIs. The example below demonstrates all four languages. Client.health() reports installed versions.

Compilation can initialize or query CUDA and load GPU modules. Such calls run under the GPU lease, alongside any later kernel execution.

On a CPU server

A CPU server executes uploaded Python compilation code and returns library bytes for a subsequent request to a GPU server. Supply the target architecture from Client.target() on the GPU server. CPU workers also accept CUDA source uploads; selecting a name returns the source text and name for your compiler to consume.

Remote Compilation walks through building CUDA C on a CPU server, then uploading the resulting library to a GPU server for execution and measurement.

On the client

Build a library locally and upload it with kind="library". Build for the architecture reported by the server’s Client.target(). The library protocol describes export and linking requirements, and the example below demonstrates the build and submission.

Checking correctness

Uploaded code can use, for example, torch.testing.assert_close to compare output against a reference. Return any reports you want to inspect. An assertion failure stops subsequent instructions; results already selected by return_ remain available.

Measuring GPU activity

kcoral.builtins.benchmark measures each call from its first GPU activity to its last, including kernels, copies and memsets. Host work before and after those endpoints is excluded; gaps between GPU activities are included. This supports functions that launch multiple GPU operations.

Pass an optional configuration dict after the callable’s arguments:

from kcoral.builtins import benchmark

timing = benchmark(kernel, src, dst, {"warmup_ms": 25, "repeat_ms": 100, "flush_l2": True})

Those are the defaults. The time budgets determine iteration counts from an initial estimate. Explicit warmup and repeat counts override their respective budgets. L2 flushing happens before each call and outside its measured span.

The report contains latency_ms_median, latency_ms_mean, latency_ms_min, latency_ms_max, warmup, repeat, flush_l2 and activities_stable. A false activities_stable means calls did not all launch the same GPU activities. The helper requires PyTorch and cupti-python in the worker environment.

Your Python can also invoke measurement tools such as NCU or Compute Sanitizer and return their reports with return_file or return_folder. Wait for GPU subprocesses to finish before returning or releasing the GPU lease.

Running host work without the GPU lease

get_function(..., cpu_only=True) declares that calls to the selected function use no GPU. Before its run, a GPU worker synchronizes outstanding work and releases the lease, allowing another request to use the GPU. Detected CUDA calls fail with gpu_access; this is a best-effort check.

The flag applies to the selected function’s run calls. The module’s top-level Python still executes under the lease when uploaded. To separate host compilation from GPU loading, return a path or bytes from a CPU-only build function, then load the compiled result in a function with GPU access.

Request timing

The execution response includes these timing fields, also exposed as attributes on the Python client’s ProgramResult. queue_ms measures waiting for a worker. After assignment, elapsed_ms includes GPU-lease waiting (lease_wait_ms), holding the lease (lease_held_ms), and other worker work. Holding a lease reserves the GPU but does not imply continuous GPU activity. Use the benchmark report for kernel latency.

Compilation examples

Compile on the GPU server

This client compiles, checks and measures TIRx, CUDA C, CuTeDSL and Triton kernels. Set KCORAL_URL to the GPU server address.

KCORAL_URL=http://localhost:8000 python examples/remote_compile_client.py

"""Compile, check and time TIRx, CuTeDSL, CUDA C and Triton kernels remotely.

The client uploads kernel and compiler code and needs no local CUDA toolchain.
Compilation retains the GPU lease because it
can call CUDA; the CPU reference explicitly releases it. Timing uses CUPTI GPU activity spans.
"""

from __future__ import annotations

import os

import numpy as np

from kcoral import Client, Program

N = 256

TIRX_KERNEL = r"""
from __future__ import annotations
from tvm.script import tirx as T


@T.jit
def main(
    A: T.Buffer((N,), "float32"),
    B: T.Buffer((N,), "float32"),
    *,
    N: T.constexpr,
):
    T.device_entry()
    i = T.cta_id([N])
    t = T.thread_id([1])
    B[i] = A[i] + 1.0
"""

CUTEDSL_KERNEL = r"""
import cutlass.cute as cute


@cute.kernel
def add_one_kernel(src: cute.Tensor, dst: cute.Tensor):
    tidx, _, _ = cute.arch.thread_idx()
    bidx, _, _ = cute.arch.block_idx()
    i = bidx * 256 + tidx
    if i < cute.size(src):
        dst[i] = src[i] + 1.0


@cute.jit
def add_one(src: cute.Tensor, dst: cute.Tensor):
    n = cute.size(src)
    add_one_kernel(src, dst).launch(grid=((n + 255) // 256, 1, 1), block=(256, 1, 1))
"""

# The selected function is exported through TVM FFI, so it takes TensorView parameters and
# returns void; the inline compiler adds the includes and export macro.
CUDA_KERNEL = r"""
__global__ void add_one_kernel(const float* x, float* y, int n) {
  int i = blockIdx.x * blockDim.x + threadIdx.x;
  if (i < n) y[i] = x[i] + 1.0f;
}

void add_one(tvm::ffi::TensorView x, tvm::ffi::TensorView y) {
  int n = static_cast<int>(x.numel());
  add_one_kernel<<<(n + 255) / 256, 256>>>(static_cast<const float*>(x.data_ptr()),
                                           static_cast<float*>(y.data_ptr()), n);
}
"""


TRITON_KERNEL = r"""
import triton
import triton.language as tl


@triton.jit
def add_one(x_ptr, y_ptr, n, BLOCK: tl.constexpr):
    offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
    mask = offs < n
    tl.store(y_ptr + offs, tl.load(x_ptr + offs, mask=mask) + 1.0, mask=mask)
"""


# The reference needs no GPU, so it is declared cpu_only when selected: the worker
# hands the GPU over for the call and fails it should it touch CUDA after all.
# assert_close takes its CPU result as is.
CPU_REFERENCE = r"""
import torch


def expected(n):
    return torch.arange(n, dtype=torch.float32) + 1.0
"""


# These functions and imports are uploaded to the server.
OPERATIONS = r"""
from kcoral.builtins import benchmark, compile_tirx


def empty(spec):
    import torch

    return torch.empty(spec["shape"], dtype=getattr(torch, spec["dtype"]), device="cuda")


def assert_close(actual, expected):
    import torch

    torch.testing.assert_close(actual.cpu(), expected.cpu(), rtol=1e-2, atol=1e-3)
    return {"ok": True}


def compile_cuda_binary(source, cfg):
    import os
    from pathlib import Path

    import tvm_ffi.cpp

    arch = cfg["arch"].removeprefix("sm_")
    suffix = "a" if arch.endswith("a") else ""
    digits = arch.removesuffix("a")
    key = "TVM_FFI_CUDA_ARCH_LIST"
    previous = os.environ.get(key)
    os.environ[key] = f"{int(digits[:-1])}.{digits[-1]}{suffix}"
    try:
        path = tvm_ffi.cpp.build_inline(
            name=f"example_{source.name}",
            cuda_sources=source.source,
            functions=source.name,
            backend="cuda",
            extra_cuda_cflags=cfg.get("extra_cuda_cflags"),
        )
        return Path(path).read_bytes()
    finally:
        if previous is None:
            os.environ.pop(key, None)
        else:
            os.environ[key] = previous


def compile_cuda(source):
    import tempfile
    from pathlib import Path

    import torch
    import tvm_ffi

    major, minor = torch.cuda.get_device_capability()
    arch = f"sm_{major}{minor}" + ("a" if major >= 9 else "")
    data = compile_cuda_binary(source, {"arch": arch})
    with tempfile.TemporaryDirectory() as directory:
        path = Path(directory) / "kernel.so"
        path.write_bytes(data)
        module = tvm_ffi.load_module(str(path))
    function = module.get_function(source.name)

    # Keep the defining module alive with its callable.
    def invoke(*args):
        _ = module
        return function(*args)

    return invoke


def compile_cutedsl(kernel, *tensors):
    import cutlass.cute as cute
    from cutlass.cute.runtime import from_dlpack

    return cute.compile(kernel, *(from_dlpack(tensor) for tensor in tensors))


def compile_triton(kernel, *args):
    *operands, cfg = args
    options = dict(cfg)
    grid = tuple(options.pop("grid"))
    kernel.warmup(*operands, grid=grid, **options)
    return lambda *values: kernel[grid](*values, **options)
"""


def check_against_cpu_reference(program: Program, dst, assert_close) -> None:
    module = program.upload(id="reference_module", kind="module", source=CPU_REFERENCE)
    reference = program.get_function(id="reference", module=module, name="expected", cpu_only=True)
    expected = program.run(id="expected", fn=reference, args=[N])
    program.run(id="check", fn=assert_close, args=[dst, expected])


def tirx_program() -> Program:
    program = Program()
    operations = program.upload(id="operations", kind="module", source=OPERATIONS)
    empty = program.get_function(id="empty", module=operations, name="empty")
    assert_close = program.get_function(id="assert_close", module=operations, name="assert_close")
    benchmark = program.get_function(id="benchmark", module=operations, name="benchmark")
    compile_tirx = program.get_function(id="compile_tirx", module=operations, name="compile_tirx")
    module = program.upload(id="kernel_module", kind="module", source=TIRX_KERNEL)
    kernel = program.get_function(id="kernel", module=module, name="main")
    src = program.upload(id="src", kind="tensor", value=np.arange(N, dtype=np.float32))
    dst = program.run(id="dst", fn=empty, args=[{"shape": [N], "dtype": "float32"}])

    # `bindings` supplies the T.constexpr values the @T.jit kernel specializes on.
    compiled = program.run(id="compiled", fn=compile_tirx, args=[kernel, {"N": N}])
    program.run(id="invoke", fn=compiled, args=[src, dst])
    check_against_cpu_reference(program, dst, assert_close)
    timing = program.run(
        id="timing",
        fn=benchmark,
        args=[compiled, src, dst, {"warmup_ms": 25, "repeat_ms": 100}],
    )
    program.return_(key="timing", value=timing)
    program.return_(key="dst", value=dst)
    return program


def cutedsl_program() -> Program:
    program = Program()
    operations = program.upload(id="operations", kind="module", source=OPERATIONS)
    empty = program.get_function(id="empty", module=operations, name="empty")
    assert_close = program.get_function(id="assert_close", module=operations, name="assert_close")
    benchmark = program.get_function(id="benchmark", module=operations, name="benchmark")
    compile_cutedsl = program.get_function(
        id="compile_cutedsl", module=operations, name="compile_cutedsl"
    )
    module = program.upload(id="kernel_module", kind="module", source=CUTEDSL_KERNEL)
    kernel = program.get_function(id="kernel", module=module, name="add_one")
    src = program.upload(id="src", kind="tensor", value=np.arange(N, dtype=np.float32))
    dst = program.run(id="dst", fn=empty, args=[{"shape": [N], "dtype": "float32"}])

    # CuTeDSL specializes on the tensors, so compiling takes them too; what comes
    # back is called with the same plain ones.
    compiled = program.run(id="compiled", fn=compile_cutedsl, args=[kernel, src, dst])
    program.run(id="invoke", fn=compiled, args=[src, dst])
    check_against_cpu_reference(program, dst, assert_close)
    timing = program.run(
        id="timing",
        fn=benchmark,
        args=[compiled, src, dst, {"warmup_ms": 25, "repeat_ms": 100}],
    )
    program.return_(key="timing", value=timing)
    program.return_(key="dst", value=dst)
    return program


def cuda_program() -> Program:
    program = Program()
    operations = program.upload(id="operations", kind="module", source=OPERATIONS)
    empty = program.get_function(id="empty", module=operations, name="empty")
    assert_close = program.get_function(id="assert_close", module=operations, name="assert_close")
    benchmark = program.get_function(id="benchmark", module=operations, name="benchmark")
    compile_cuda = program.get_function(id="compile_cuda", module=operations, name="compile_cuda")
    # `language` makes this a CUDA C source module; selecting a function from it
    # creates the source descriptor consumed by compile_cuda in OPERATIONS.
    module = program.upload(id="kernel_module", kind="module", source=CUDA_KERNEL, language="cuda")
    kernel = program.get_function(id="kernel", module=module, name="add_one")
    src = program.upload(id="src", kind="tensor", value=np.arange(N, dtype=np.float32))
    dst = program.run(id="dst", fn=empty, args=[{"shape": [N], "dtype": "float32"}])

    # Built for the worker GPU's arch, and cached on disk by source and flags, so
    # recompiling the same source is much cheaper.
    compiled = program.run(id="compiled", fn=compile_cuda, args=[kernel])
    program.run(id="invoke", fn=compiled, args=[src, dst])
    check_against_cpu_reference(program, dst, assert_close)
    timing = program.run(
        id="timing",
        fn=benchmark,
        args=[compiled, src, dst, {"warmup_ms": 25, "repeat_ms": 100}],
    )
    program.return_(key="timing", value=timing)
    program.return_(key="dst", value=dst)
    return program


def triton_program() -> Program:
    program = Program()
    operations = program.upload(id="operations", kind="module", source=OPERATIONS)
    empty = program.get_function(id="empty", module=operations, name="empty")
    assert_close = program.get_function(id="assert_close", module=operations, name="assert_close")
    benchmark = program.get_function(id="benchmark", module=operations, name="benchmark")
    compile_triton = program.get_function(
        id="compile_triton", module=operations, name="compile_triton"
    )
    module = program.upload(id="kernel_module", kind="module", source=TRITON_KERNEL)
    kernel = program.get_function(id="kernel", module=module, name="add_one")
    src = program.upload(id="src", kind="tensor", value=np.arange(N, dtype=np.float32))
    dst = program.run(id="dst", fn=empty, args=[{"shape": [N], "dtype": "float32"}])

    # A Triton kernel computes its grid at launch, so the grid travels as data
    # rather than as a launcher the client writes. Every other `cfg` key is a
    # launch keyword — num_warps, num_stages, a constexpr by name.
    compiled = program.run(
        id="compiled",
        fn=compile_triton,
        args=[kernel, src, dst, N, 256, {"grid": [1], "num_warps": 4}],
    )
    program.run(id="invoke", fn=compiled, args=[src, dst, N, 256])
    check_against_cpu_reference(program, dst, assert_close)
    timing = program.run(
        id="timing",
        fn=benchmark,
        args=[compiled, src, dst, N, 256, {"warmup_ms": 25, "repeat_ms": 100}],
    )
    program.return_(key="timing", value=timing)
    program.return_(key="dst", value=dst)
    return program


def main() -> None:
    expected = np.arange(N, dtype=np.float32) + 1.0
    with Client(os.environ.get("KCORAL_URL", "http://localhost:8000")) as client:
        programs = (
            ("TIRx", tirx_program()),
            ("CuTeDSL", cutedsl_program()),
            ("CUDA C", cuda_program()),
            ("Triton", triton_program()),
        )
        for language, program in programs:
            result = client.execute(program, timeout_seconds=120)
            if result.status != "COMPLETED":
                print(f"{language}: {result.status} — {result.error}")
                continue
            # Only lease_held_ms occupied the GPU; the CPU reference ran off it.
            print(
                f"{language}: {result.elapsed_ms:.0f} ms total, "
                f"{result.lease_held_ms:.0f} ms on the GPU, "
                f"{result.elapsed_ms - result.lease_held_ms:.0f} ms off it"
            )
            print(f"  kernel {result.results['timing']['latency_ms_median'] * 1e3:.1f} us")
            np.testing.assert_allclose(result.results["dst"], expected)


if __name__ == "__main__":
    main()

Download the client.

Build and upload a library

This client builds a CUDA C library locally for the server’s target architecture, then uploads, checks and measures it. The client machine needs the CUDA toolkit, a host C++ compiler and TVM FFI from the compiler environment.

KCORAL_URL=http://localhost:8000 python examples/library_upload_client.py

"""Build a kernel locally, upload the library, check its output, and time it.

The client owns the compiler flags and builds for the architecture reported by
GET /health. Allocation, comparison and CUPTI measurement run on the server.
The client needs a CUDA toolchain.
"""

from __future__ import annotations

import os
import pathlib
import tempfile

from kcoral import Client, Program

N = 1 << 20

# The export macro is what makes the object loadable: it emits the
# `__tvm_ffi_add_one` symbol the server looks up by function name.
SOURCE = r"""
#include <tvm/ffi/container/tensor.h>

__global__ void add_one_kernel(const float* x, float* y, int n) {
  int i = blockIdx.x * blockDim.x + threadIdx.x;
  if (i < n) y[i] = x[i] + 1.0f;
}

void add_one(tvm::ffi::TensorView x, tvm::ffi::TensorView y) {
  int n = static_cast<int>(x.numel());
  add_one_kernel<<<(n + 255) / 256, 256>>>(static_cast<const float*>(x.data_ptr()),
                                           static_cast<float*>(y.data_ptr()), n);
}

TVM_FFI_DLL_EXPORT_TYPED_FUNC(add_one, add_one);
"""

REFERENCE = "def main(a):\n    return a + 1.0\n"


# These functions and imports are uploaded to the server.
OPERATIONS = r"""
from kcoral.builtins import benchmark


def empty(spec):
    import torch

    return torch.empty(spec["shape"], dtype=getattr(torch, spec["dtype"]), device="cuda")


def randn(spec):
    import torch

    generator = torch.Generator(device="cuda").manual_seed(spec.get("seed", 0))
    return torch.randn(
        spec["shape"], dtype=getattr(torch, spec["dtype"]), device="cuda", generator=generator
    )


def assert_close(actual, expected):
    import torch

    torch.testing.assert_close(actual.cpu(), expected.cpu(), rtol=1e-2, atol=1e-3)
    return {"ok": True}
"""


def build_library(arch: str, directory: str) -> bytes:
    """Compile SOURCE for `arch` — the server's, not this machine's."""
    import tvm_ffi.cpp

    source_path = os.path.join(directory, "add_one.cu")
    pathlib.Path(source_path).write_text(SOURCE)
    # Whatever the local toolchain accepts belongs here; this freedom is the
    # reason to upload a library rather than let the server build one.
    library = tvm_ffi.cpp.build(
        "add_one",
        cuda_files=source_path,
        extra_cuda_cflags=[
            f"-gencode=arch=compute_{arch.removeprefix('sm_')},code={arch}",
            "-O3",
        ],
        build_directory=directory,
        output=os.path.join(directory, "add_one.so"),
    )
    return pathlib.Path(library).read_bytes()


def build_program(library: bytes) -> Program:
    program = Program()
    operations = program.upload(id="operations", kind="module", source=OPERATIONS)
    empty = program.get_function(id="empty", module=operations, name="empty")
    randn = program.get_function(id="randn", module=operations, name="randn")
    assert_close = program.get_function(id="assert_close", module=operations, name="assert_close")
    benchmark = program.get_function(id="benchmark", module=operations, name="benchmark")
    # No compile instruction follows: get_function binds the precompiled callable.
    module = program.upload(id="kernel_module", kind="library", value=library)
    kernel = program.get_function(id="kernel", module=module, name="add_one")
    reference_module = program.upload(id="reference_module", kind="module", source=REFERENCE)
    reference = program.get_function(id="reference", module=reference_module, name="main")

    src = program.run(id="src", fn=randn, args=[{"shape": [N], "dtype": "float32", "seed": 0}])
    dst = program.run(id="dst", fn=empty, args=[{"shape": [N], "dtype": "float32"}])
    program.run(id="invoke", fn=kernel, args=[src, dst])

    # Compared on the server against a plain-Python reference, so the output
    # tensor never travels; assert_close stops the program before timing a
    # kernel that is wrong.
    expected = program.run(id="expected", fn=reference, args=[src])
    check = program.run(id="check", fn=assert_close, args=[dst, expected])
    timing = program.run(id="timing", fn=benchmark, args=[kernel, src, dst])
    program.return_(key="check", value=check)
    program.return_(key="timing", value=timing)
    return program


def main() -> None:
    with Client(os.environ.get("KCORAL_URL", "http://localhost:8000")) as client:
        arch = client.target()["arch"]
        with tempfile.TemporaryDirectory() as directory:
            library = build_library(arch, directory)
        print(f"built {len(library) / 1024:.0f} KiB for {arch}")

        result = client.execute(build_program(library), timeout_seconds=120)
        if result.status != "COMPLETED":
            raise SystemExit(f"{result.status} — {result.error}")
        print(
            f"{result.elapsed_ms:.0f} ms total, "
            f"{result.lease_held_ms:.0f} ms on the GPU, "
            f"kernel {result.results['timing']['latency_ms_median'] * 1e3:.1f} us"
        )


if __name__ == "__main__":
    main()

Download the client.

For a separate compilation server, follow Remote Compilation.