[dependencies]
gpukit = "^0.11.13"

gpukit — the GPU-compute toolkit for NURL

One dependency for running compute on the GPU without the marshalling boilerplate. The gpu package is the low-level interface — open a device, JIT-compile a CUDA-C kernel (NVRTC, or the host-C++ CPU backend), allocate device buffers, upload/download, build the argument vector, compute a grid, launch, sync. Every GPU-using package re-writes the same ~35 lines of that dance for each kernel. gpukit is that glue, as a facade.

Before / after

Hand-marshalled (the pattern in onnx / anomaly today):

<!-- doctest-prelude $ deps/gpu/src/gpu.nu -->

: GpuBuffer b_x   ( gpu_alloc g * n 8 )
: GpuBuffer b_out ( gpu_alloc g * n 8 )
( gpu_upload b_x # *u ( vec_data [f] xs ) )
: (Vec i) args ( vec_new [i] )
( vec_push [i] args (gpu_arg_buffer b_x) )
( vec_push [i] args (gpu_arg_i64 n) )
( vec_push [i] args (gpu_arg_buffer b_out) )
( gpu_launch k (gpu_grid n 256) 256 args )
( gpu_sync g )
( gpu_download # *u ( vec_data [f] out ) b_out )

With gpukit:

<!-- doctest-file-prelude $ deps/gpukit/src/kernels.nu -->

: (Vec GkArg) call ( vec_new [GkArg] )
( vec_push [GkArg] call ( gk_in_f  xs ) )
( vec_push [GkArg] call ( gk_i64   n  ) )
( vec_push [GkArg] call ( gk_out_f out ) )
( gk_run kit src `my_kernel` ( gk_grid n 256 ) 256 call )

gk_run compiles-and-caches the kernel by name, allocates a device buffer per buffer binding, uploads inputs, marshals the args in binding order, launches, syncs and downloads outputs; the device buffers go with the call.

Ready-made kernels

For the common cases you don't write a kernel at all:

: GpuKit kit ( gk_open 0 )
( gk_add_f kit out a b )                 // out = a + b   (also sub/mul/div)
( gk_map_f kit `relu` out x `x>0.0?x:0.0` )   // out[i] = f(x)
( gk_matmul_f kit c a b m k n )          // C[m×n] = A[m×k]·B[k×n]
: ?f s ( gk_reduce_sum_f kit x )         // Σ x
: ?f d ( gk_dot_f kit a b )              // a·b
// nothing to close: the kit goes with its last owner

API

Lifecycle:

Call
( gk_open ordinal ) → GpuKit	open a device (CUDA, else CPU backend)
( gk_open_best ) → GpuKit	the device a caller who does not care should get: $NURL_GPU_DEVICE, else highest compute capability, memory breaking ties
( gk_ok kit ) → b, ( gk_backend kit ) → s, ( gk_device_name kit ) → s
( gk_bind_thread kit ) → b	make the device current on THIS thread (CUDA contexts are thread-local)
( gk_close kit )	early release (optional): the kit — cached kernels, pool, device — goes with its last owner
( gk_gpu kit ) → Gpu	the kit's device, for the gpu package's own calls (graphs, timers)
( gk_grid n block ) → i	grid size for n threads

Bindings (f is a C double, i a C long long):

Call
( gk_in_f v ) / ( gk_in_i v )	input double[] / long long[]
( gk_out_f v ) / ( gk_out_i v )	output (caller pre-sizes)
( gk_buf_in host bytes ) / ( gk_buf_out host bytes )	raw buffer (e.g. packed f32)
( gk_i64 v ) / ( gk_i32 v ) / ( gk_f32 v )	scalar

Workhorse:

Call
( gk_run kit src name grid block call ) → b	compile-cached, marshal, launch, sync, download

The call is a Vec GkArg listing the kernel's arguments in declaration order (buffers and scalars interleaved exactly as the signature expects).

Ready-made f64 kernels: gk_add_f / gk_sub_f / gk_mul_f / gk_div_f, gk_map_f, gk_matmul_f, gk_reduce_sum_f, gk_dot_f.

Performance

Three things matter more than any individual kernel, and all three are handled here rather than in every caller:

Memory is pooled. cuMemAlloc/cuMemFree synchronise and cost ~315 us per 12 MB pair on a 4090. A GkBuf's last owner retires its block for reuse by the next gk_dbuf_new of the same size instead of returning it to the driver — a forward pass that allocates its scratch per layer was otherwise spending a third of its time in the allocator with the device idle. gk_pool F turns it off; gk_pool_release kit hands idle blocks back, which is also what an allocation failure does before reporting out-of-memory. The idle side is bounded: $NURL_GK_POOL_MAX bytes, else a quarter of device memory, else whatever gk_pool_budget kit bytes says (0 = unlimited). Past it the least recently idled blocks go back to the driver, so a caller whose shapes follow its input cannot ratchet the card full; gk_pool_count and gk_pool_idle_bytes say what is held.

You can see where the time went. gk_prof kit T brackets every launch with a CUDA event pair and accumulates device time per kernel; gk_prof_report kit prints them slowest-first with call counts and shares. Events measure the GPU — a host clock around an asynchronous launch measures the launch.

gpukit profile: 1113 ms of device time over 26 kernels
  30%  342121 us  1728 calls  197987 ns/call  gk32_gemm_smem
  21%  240305 us   432 calls  556263 ns/call  gk32_attn64_64
  18%  207653 us   180 calls 1153633 ns/call  gk32_conv2d

The kernels are shaped for the machine. Two different machines, in fact: matmul/bmm/gemm carry a register-tiled body for the CPU backend (one thread owns an 8x32 block of C — see 0.5.0) and a shared-memory tiled one for CUDA, because 256 accumulators per thread is right on a core and far past a CUDA thread's register budget. The CUDA numbers, measured on an RTX 4090 in f32, all bit-identical to the per-element kernels they replaced (see Numerics) and all of them checked kernels:

	before	after
gkd_gemm 783x1024x1024	3.1 TFLOP/s	24.9
the same with transb=1	0.6	24.9
gkd_bmm 16 x 783x64x783	3.9	20.6
gkd_layernorm [783, 1024]	300 us	53
gkd_conv2d 3x3 256→256	2.5 ms	1.15
a DPT depth head's 30 convolutions	39 ms	22
gkd_convtranspose2d 4x4 stride 4	14.2 ms	0.35
gkd_perm [783,3,16,64] → [3,16,783,64]	87 us	15

gkd_attention is the fused scaled-dot-product attention: softmax(scale · Q·Kᵀ)·V over [heads, n, hd] operands, without ever materialising the [heads, n, nkv] score matrix. Against the composed bmm/scale/softmax/bmm it is 3.8x at nkv = n and 6.2x at nkv = 4n — and at a 72-frame streaming window the score matrix it does not allocate is 2.8 GB. gkd_attention_ok kit hd reports in advance whether it will run, so a caller can size its workspace before the first call. gkd_attention_masked is the same kernel with a per-key additive bias (heads/hpb rows of nkv floats, 0 keeps a key, -1e30 drops it) — what a padded batch needs to leave its real rows unchanged. It is a checked kernel like the rest: its tiles are shared arrays read through windows.

Numerics

gpukit adds no numerics of its own — it only marshals — so a kernel runs bit-for-bit the same through gk_run as through hand-written gpu_* calls, and the gpu package's CUDA / CPU-backend / pure equivalence is preserved. The elementwise ops and gk_matmul_f accumulate in the same order a sequential host loop would, so they are bit-identical to a naive host implementation. That holds for the tiled kernels too: blocking a matmul over shared memory changes which thread visits an element and when, not the order the products are summed in, and an out-of-range staging slot holds an exact zero. The one deliberate exception is gkd_attention, whose online softmax rescales a running sum over key tiles rather than running left to right — the same value to f32 rounding, and the same algorithm torch's SDPA uses, but a golden taken with gkd_bmm + gkd_softmax_ax will not reproduce bit for bit. gk_reduce_sum_f / gk_dot_f combine per-thread partial sums (a parallel order), matching a host sum to rounding but not guaranteed bit-identical to a strictly sequential accumulation.

Memory

Nothing is released by hand. A GpuKit is a handle: every copy of it (a struct field, a Vec element, a capture) is the same kit, and its last owner releases the cached kernels, the memory pool and the device. A GkBuf holds its block and the kit it came from: its last owner — the buffer, a copy, a view (gk_buf_view) — hands the block back to the kit's pool, and the device stays open until every buffer from it is gone, whatever order the owners go in. gk_dbuf_free and gk_close remain as optional early releases of one owner.

Device-resident buffers (src/dev.nu)

gk_run marshals host↔device per call; chained pipelines want data to STAY on the device. GkBuf is an element-typed device allocation (GK_F32 | GK_F64 | GK_I64) with gk_dbuf_new/_upload/_download (. b dptr, . b n, . b dtype are plain fields), views that hold their block (gk_buf_view b off n, gk_buf_view_as b byte_off n dtype), an empty buffer (gk_buf_none dtype) and an unowned wrapper over foreign device memory (gk_buf_wrap dptr n dtype), and gk_run_dev launches a cached kernel over raw device args with zero copies. Ready-made dtype-generic kernels: gkd_add/sub/mul/div (1-element scalar broadcast via the same kernel), gkd_relu/sigmoid/ exp/tanh/sqrt/log, gkd_matmul, gkd_softmax_rows, gkd_sum. gk_dbuf_upload_raw takes bytes already in the buffer's element type — a memory-mapped f32 checkpoint straight to an f32 buffer, no conversion and no staging copy. GK_F32 computes in true float32 (accumulation included); tests/devcheck.nu verifies 12/12 vs numpy per dtype on CUDA and the CPU backend.

Tests

./tests/gpukit_test.sh builds tests/demo.nu and runs it on the default backend and, when a host C++ compiler is present, on the CPU backend — checking every bit-identical op with == against a host computation. On an RTX 4090 and on the CPU backend: 14 / 14.

Checked kernels

gk_run compiles whatever kernel source it is given, so only unsafe code may call it. gk_run_checked (src/checked.nu) takes a kernel body written in a subset that cannot address memory it was not handed. Buffers are reached only through LD(name, i) and ST(name, i, v), which check every index against that buffer's length. The body sees its scalars by name, plus gid and nthreads. It may declare scalar locals, branch, loop and call math functions and the __d*_rn rounding intrinsics. Anything that makes or follows an address is refused, with the reason in gk_checked_error, before anything compiles. That covers pointers, dereferences, &, [], member access, the preprocessor, asm, goto, reserved names and other calls. So a safe package can launch its own kernels:

: ~ GkCall c ( gk_call_new )
( gk_call_in_f c `x` xs )
( gk_call_out_f c `y` ( vec_zeroed [f] n ) )
( gk_call_i64 c `n` n )
: b ok ( gk_run_checked kit c `scale2` `ST(y, gid, 2.0 * LD(x, gid));` n 256 )
: ( Vec f ) ys ( gk_call_take_f c `y` )

A load past the end reads the buffer's last element, and a store past it is dropped. An output starts as the Vec it was bound with. arima's device evaluator is written this way and runs as fast as its unchecked predecessor did.

On device buffers

gk_run_checked_dev (and gk_record_dev / gk_replay, which bind and check once and then only launch) take device-resident bindings: gk_call_dev_in / gk_call_dev for f32/f64/i64 GkBufs, gk_call_half_in for half or bfloat16 GkHBufs, and gk_call_bytes_in / gk_call_bytes for GkBBuf byte buffers. On a byte buffer LD(q, i) is the byte, LD16 / LD32 the aligned 16- or 32-bit word holding byte i, and F16(bits) / BF16(bits) turn 16 bits into a float; ST stores the low byte. A quantised weight format is read this way, without a pointer. LD4(x, i, a, b, c, d) loads the aligned group of four f32 holding element i in one instruction (the buffer must be 16-byte aligned). Cooperative kernels have SYNC(), SHARED_* arrays (up to 48 KB a block) and SHFL_DOWN / SHFL_XOR, under a rule the checker proves: every if / for / while around one decides the same for the whole block.

Windows make a loop's checks free. WIN(w, buf, off, len); (or WINS(w, buf, off, len, stride);) checks a range of a device or byte buffer once; LD(w, i) / ST(w, i, v) then index it from 0 and LEN(w) is its length. A window that does not fit its buffer is empty. A loop bounded by LEN(w) lets the compiler drop every per-load check:

WIN(xr, x, row * n, n);
float ss = 0.0f;
for (int j = 0; j < LEN(xr); j++) { ss += LD(xr, j) * LD(xr, j); }

The check folds for a signed counter compared to LEN, so count with an int (NVRTC also unrolls an int loop four times further than a long long one). Inside if (LEN(w) == 32) { … } a constant bound works too. WIN2(w, buf, off, rows, stride, cols); is the 2-D form (LD2(w, r, c), ROWS(w), COLS(w)): a tile of a matrix, checked once. WIN3(w, buf, off, planes, pstride, rows, stride, cols); is the 3-D one (LD3(w, p, r, c), PLANES(w)): an NCHW feature map, say. Each index is checked against its own extent, so each check folds against whatever bounds that index alone — a loop counter, a continue past the padding — where one flattened index would need the compiler to reason about a sum. A window may look into a shared array too: WIN2(t, tile, 0, 64, 65, 64); then LD2(t, row, d) reads a padded tile with the column checked against the loop that counts it and the row against the tile, so a thread's row masked to the tile ((tid / 16) & 15) folds as well. nurllama's RMSNorm, attention and quantised matvecs are written this way, on both backends: within 4 % of their raw originals, or faster.

A kernel compiled per shape can go further: gk_call_exact c builds it for its buffers' lengths (each a constant in the source and part of the cache key; other lengths build another kernel), so with the dimensions constant too every window is checked at compile time.