hk-dsl — HipKittens kernels in Python, over an MLIR dialect

Write Radeon kernels in Python, get HipKittens C++, get a .so, call it from torch. The point is not convenience: it is that three classes of silently wrong kernel on RDNA3 stop being expressible.

import hk

@hk.kernel(grid=lambda p: (hk.cdiv(p.o.cols, 64), hk.cdiv(p.o.rows, 16), p.o.batch))
def add(a: hk.Tensor[hk.bf16], b: hk.Tensor[hk.bf16], o: hk.Tensor[hk.bf16]):
    idx = hk.tile_coord(hk.arch.block_idx.z, 0, hk.arch.block_idx.y,
                        hk.arch.block_idx.x)
    t = hk.RT(hk.bf16, 16, 64)
    o[idx] = a[idx].load(t) + b[idx].load(t)

add(a, b, out)          # traces, verifies, translates, compiles (once), launches

The three traps

Each of these cost days to find by hand, and none of them produces a diagnostic:

• s_waitcnt carries no register dependence, so the compiler will hoist a WMMA above the wait that was supposed to guard it. The translation binds the fragment with a volatile asm after every wait (IR).
• A VGPR spill is silent corruption, not slowness, because the library hand-places s_waitcnt and scratch traffic reorders against it. k.build() refuses to return a spilling kernel (runtime).
• Only row-layout 16-bit operands reach ds_read_b128; column layout degrades to eight times the instructions. The passes flag it while tracing, not at benchmark time (IR).

The whole argument is in why this exists.

Install

From a wheel. One per Python interpreter, because it carries the MLIR bindings for one CPython; it includes the HipKittens headers, so at run time only hipcc is needed. CPython 3.10 and 3.12 wheels are attached to each GitHub release, or build one with scripts/build_wheel.sh (building):

pip install hk_dsl-0.3.1-cp310-cp310-manylinux_2_35_x86_64.whl

From source. The backend is the third_party/HipKittens submodule, so clone with --recursive:

git clone --recursive https://github.com/vivienfanghuagood/hk-dsl.git && cd hk-dsl
scripts/build_mlir.sh       # MLIR 23.1.3 only, X86 target, Python bindings (~20 min on 64 cores)
scripts/build_hk_mlir.sh    # the dialect, passes, translation and hk._mlir (~1 min)

torch is not a dependency: tracing and compiling do not need it, launching does. Details, memory limits and the variables the scripts take: getting started.

First kernel

With add above, each step needs more than the one before it:

print(add.trace().dump())     # the hk MLIR module, after the passes -- needs nothing
print(add.source())           # the HipKittens C++ it translates to -- needs nothing
print(add.build().report())   # hipcc + the resource gate -- no GPU
add(a, b, out)                # launch -- a Radeon and torch

build() needing no GPU is deliberate: the failure most likely to bite — generated code that compiles but spills — is caught on a laptop. The full walk-through, with tensors and a check against torch, is getting started.

Results

On one W7900D (gfx1100). Each number is from the development record, linked; performance has the full tables, the methods, and the losses.

Repository layout

hk/dsl/         the language: @kernel / @func, values and operators, the AST
                rewrite, control flow, hk.arch / lds / math / group / collective
hk/types.py     dtypes and the tile / vector / tensor types a kernel is written in
hk/compiler/    the MLIR context, the pass pipeline, the translation entry point
hk/_mlir/       the built bindings (hk dialect + MLIR core), not checked in
hk/target/      gfx1100 / gfx1201 facts, as data
hk/codegen/     the module boundary around a translated kernel (scaffold.py)
hk/runtime/     arch detection, hipcc + content-hash cache, the resource gate,
                and the torch.ops registration (torch_ext.py)
hk/autotune/    schedule search: build the space, keep what passes the gate, time
hk/tuned/       the schedules that search picked, and the measured routing
                tables, as shipped data
hk/ops/         kernels written in the DSL, plus the SDPA drop-in (sdpa.py)
hk/integration/ patching a running vLLM or SGLang; engine.py is what every
                patch shares (op, routes, counters, CUDA-graph contract)
mlir/           the hk dialect (ODS), its passes, the C++ translation,
                hk-opt / hk-translate, lit tests
include/        hk's own headers (pyutils/hk_bind.cuh, the fast pybind entry)
third_party/HipKittens/
                the backend: kittens.cuh and include/rdna3, include/rdna4;
                also the C++ kernels hk is measured against (kernels/rdna3)
scripts/        building MLIR, the dialect and the wheels
tools/hk-bench/ the benchmarks and engine harnesses behind every number here
examples/       complete kernels, each checked against torch and timed
docs/           the documentation, and the development record (docs/history)

The hk/ paths are under python/; the tests are in tests/hk/.

Documentation

• Documentation index, and why this exists
• Getting started — install, a first kernel, running the tests and examples
• The language — parameters, types, operators, memory, control flow, limits
• The IR — the hk dialect, its passes, the C++ translation, hk-opt / hk-translate
• The runtime — the cache, the resource gates, warming, autotune, torch.ops
• Integration — vLLM and SGLang: the patches, their switches, routing, CUDA graphs
• Performance — every result, with its hardware, baseline and source
• Building — MLIR, the dialect, the wheels
• Testing — the four tiers and the gates
• Development record — how it got here, phase by phase, with the negative results

Examples

Each runs on the GPU (checks against torch, then times), and with --dump ir|cpp or --build on a machine without one. See examples/README.md.

License

MIT; see LICENSE. hk-dsl began as part of a fork of HipKittens by HazyResearch, and it generates C++ against the HipKittens headers, which remain HipKittens', under its MIT License. See NOTICE.