gpukit.nu

gpukit/gpukit.nu — the ergonomic GPU-compute facade over the gpu package.

The gpu package is the low-level interface: open a device, compile a CUDA-C kernel (JIT via NVRTC, or the host-C++ CPU backend), allocate device buffers, upload/download, build the argument vector, compute a grid, launch, sync, free. Every GPU-using package (onnx, objdet, anomaly, yoloe) hand- writes the same ~35 lines of that marshalling for each kernel it runs:

: GpuBuffer b_x ( gpu_alloc g ( n 8) ) // one per array … ( gpu_upload b_x (# u (vec_data [f] xs)) ) // upload each input … : (Vec i) args (vec_new [i]) // build the arg vector … ( vec_push [i] args (gpu_arg_buffer b_x) ) … ( gpu_launch k (gpu_grid n 256) 256 args ) ( gpu_sync g ) ( gpu_download (# *u (vec_data [f] out)) b_out )

gpukit collapses that to a list of typed bindings and one gk_run:

: GpuKit kit ( gk_open 0 ) : (Vec GkArg) call ( vec_new [GkArg] ) ( vec_push [GkArg] call ( gk_in_f xs ) ) // input double[] ( vec_push [GkArg] call ( gk_i64 n ) ) // long long scalar ( vec_push [GkArg] call ( gk_out_f out) ) // output double[] ( 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, builds the arg vector in binding order, launches, syncs, downloads outputs, and lets every device buffer go. Kernel sources are cached on the kit, so a hot path compiles each kernel once.

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.

Memory: a GpuKit is a handle on the kit's state (an rcbox): 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 device-memory pool and the device. Every GkBuf the kit hands out (dev.nu) holds the kit too, so the device outlives its memory whichever order the owners go in. Nothing is released by hand; gk_close is an optional early release of one owner.

API

: ~ i g_pool_n 0 // rows, every kit's pools together

The device-memory pool's process-wide counters and switches (see the pool section below).

: ~ b g_pool_on T

: ~ i g_pool_idle 0 // bytes currently held idle

: ~ i g_pool_max -1 // budget in bytes; <0 = not set yet, 0 = unlimited

: ~ i g_pool_clock 0 // LRU stamp source

: GkArg

: GkArg {
    i kind
    * u host  // host data pointer for a buffer binding
    i bytes  // buffer byte size
    i argval  // pre-encoded scalar arg (gpu_arg_*) for a scalar binding
}

A single argument to a kernel launch. kind 0 = input buffer 1 = output buffer 2 = scalar

: GkKernelEntry

: GkKernelEntry {
    String name
    GpuKernel kernel
    i calls  // launches, counted only while gk_prof_on
    i ns  // device time in those launches (CUDA events)
}

: GpuKitImpl

: GpuKitImpl {
    Gpu gpu
    b ok
    ( Vec GkKernelEntry ) cache
    ( Vec GpuBuffer ) pbuf
    ( Vec i ) pdptr
    ( Vec i ) pbytes
    ( Vec i ) pinuse
    ( Vec i ) pstamp
    GpuTimer ev0  // the profiler's event pair (gk_prof)
    GpuTimer ev1
    b census  // a census kit (gk_open_census): records, never launches
    ( Vec String ) csrc  // census only: each cache slot's kernel source
}

The kit's state. The pool is five parallel columns, one row per device block the kit allocated (in use by a GkBuf, or idle): the blocks themselves (owned — dropping a row frees its memory) and the plain columns a scan reads.

% Drop GpuKitImpl

% Drop GpuKitImpl {
    @ drop GpuKitImpl x → v {
        = g_pool_n - g_pool_n ( vec_len [i] . x pdptr )
        : ~ i k 0
        ~ < k ( vec_len [i] . x pdptr ) {
            ? == ( __gk_col . x pinuse k ) 0 { = g_pool_idle - g_pool_idle ( __gk_col . x pbytes k ) } {}
            = k + k 1
        }
    }
}

The process-wide pool counters (gk_pool_count / gk_pool_idle_bytes) sum every kit's pool; a kit that goes takes its rows out of them. Its fields — the pool's blocks, the cached kernels, the timers, the device — are released by the compiler after this (drop glue).

: GpuKit { s ctl }

A GpuKit is a handle on its state in an rcbox (stdlib/core/rcbox.nu): every copy is the same kit, and the last owner releases it.

@ GpuKit_share GpuKit h → GpuKit

@ GpuKit_drop sink GpuKit h → v

@ _GpuKit_ptr GpuKit h → *GpuKitImpl

@ gk_open i ordinal → GpuKit

Open device ordinal (CUDA when present, else the gpu package's CPU backend). The kit is a long-lived handle (its kernel cache and pool persist across calls). Check gk_ok; a failed open is still a kit.

@ gk_open_best → GpuKit

Open the device a caller who does not care should get: $NURL_GPU_DEVICE when set, else the highest compute capability with memory breaking ties. gk_open 0 binds ordinal 0, which on a box with an old card beside a new one is whichever the driver enumerates first — a 4 GB GTX 970 in front of a 24 GB RTX 4090, say, where a model that fits on one fails to allocate on the other. Prefer this everywhere the ordinal is not a user choice.

@ gk_ok GpuKit kit__h → b

@ gk_gpu GpuKit kit__h → Gpu

The kit's device, for the gpu package's own calls (graphs, timers): another owner of it, so it stays valid however long it is kept.

@ gk_backend GpuKit kit__h → s

"cuda" or "cpu". A census kit answers "cpu": it takes the static backend's branches (see gk_open_census).

@ gk_jit GpuKit kit__h → b

T when this kit compiles kernels from source at run time — CUDA (NVRTC) or the CPU backend (host C++) — so a shape-specialised kernel, whose name and body are generated per call, costs one compile and nothing else. F on the static and WebGPU backends, whose kernel sets are fixed when the program is BUILT and looked up by entry name: a name carrying a run-time shape can never be in such a set, so a wrapper that specialises must ask this first and fall back to its generic kernel. F on a census kit, which records what a static build must link.

@ gk_open_census → GpuKit

── Kernel census ─────────────────────────────────────────────────────

A census kit opens no device and launches nothing. Every gk_run / gk_run_dev — so every gkd_* wrapper — validates its arguments exactly as it would on a device, then RECORDS the kernel it would compile (entry name and the exact source, once per name) and reports success. It answers gk_backend "cpu" and gk_jit F, which are the branches the gpu package's STATIC backend takes — so driving a program's kernel calls through a census kit yields precisely the kernel set a static build (a precompiled kernels_static.c, native or wasm32) has to link, from the same source builders that run on a device: nothing mirrored by hand, nothing to rot. Buffers for the calls can be gk_buf_wrap views over any nonzero address; nothing is dereferenced.

@ gk_is_census GpuKit kit__h → b

@ gk_census_name GpuKit kit__h i k → s

What a census kit recorded: gk_kernel_count kernels, slot k's entry name and source (borrowed from the kit; "" out of range).

@ gk_census_src GpuKit kit__h i k → s

@ gk_device_name GpuKit kit__h → s

@ gk_bind_thread GpuKit kit__h → b

Make the kit's device current on the CALLING thread — see gpu_bind_thread. A single-threaded program never needs it; anything that hands device work to a pool or a fiber runtime does.

@ gk_mem_free GpuKit kit__h → i

Free / total device memory in bytes; 0 means the backend cannot say. CUDA asks the driver; the CPU backends report host RAM.

@ gk_mem_total GpuKit kit__h → i

@ gk_pool b on → v

@ gk_pool_enabled → b

@ gk_pool_count → i

Blocks the pools hold, in use or idle — for tests and for anyone wondering where the VRAM went.

@ gk_pool_idle_bytes → i

Bytes the pools are holding that nothing is using.

@ gk_pool_budget GpuKit kit__h i bytes → v

Cap the idle side at bytes (0 = unlimited). Trims immediately.

@ gk_pool_budget_bytes → i

@ _gk_pool_default * GpuKitImpl kit → v

The budget a caller who never sets one gets: $NURL_GK_POOL_MAX (bytes) when set, else a quarter of the device's memory — enough that a steady shape mix never evicts, small enough that a drifting one cannot eat the card. Resolved once, on the first allocation.

@ _gk_pool_take * GpuKitImpl kit i bytes → i

An idle block of exactly bytes, marked in-use — its device pointer, or 0.

@ _gk_pool_add * GpuKitImpl kit GpuBuffer gb → v

A fresh driver allocation joins the table, in use.

@ _gk_pool_give * GpuKitImpl kit i dptr → v

The last owner of the block at dptr let go of it: idle in the pool, or — with the pool off — straight back to the driver.

@ gk_pool_release GpuKit kit__h → v

Hand every idle block back to the driver and drop it from the table. In-use blocks stay — this is a trim, not a reset.

@ _gk_pool_release * GpuKitImpl kit → v

@ gk_close sink GpuKit kit → v

Let go of kit now rather than at the end of its owner's scope: the device closes with the kit's last owner — which includes every GkBuf still alive from it.

@ gk_grid i n i block → i

Grid size for n threads at the given block size (re-export of gpu_grid).

@ gk_in_f ( Vec f ) v → GkArg

@ gk_in_i ( Vec i ) v → GkArg

@ gk_out_f ( Vec f ) v → GkArg

Output buffers must be pre-sized by the caller; results are copied back in.

@ gk_out_i ( Vec i ) v → GkArg

@ gk_buf_in * u host i bytes → GkArg

Raw buffers, for element layouts other than f64/i64 (e.g. a packed float32 host buffer): pass the host pointer and byte size directly.

@ gk_buf_out * u host i bytes → GkArg

@ gk_i64 i v → GkArg

@ gk_i32 i v → GkArg

@ gk_f32 f v → GkArg

@ _gk_kernel_slot * GpuKitImpl kit s src s name → i

Compile src (entry name) once per kit; subsequent calls with the same name reuse the cached kernel. A failed compile returns a not-ok kernel and is not cached (so a fixed source can be retried). The cache SLOT for name, compiling on a miss; -1 when the compile failed. A launch borrows the slot's kernel in place (gkslot_launch); the slot is also what the profiler accumulates into.

@ _gk_slot_launch * GpuKitImpl kit i slot i grid i block ( Vec i ) args → i

Launch the kernel in cache slot slot, borrowed in place (no copy of the handle per launch). Nonzero = the gpu_launch error, or 1 for an empty slot.

@ gk_kernel_count GpuKit kit__h → i

Kernels in the kit's cache (each compiled once, by name).

@ gk_compile GpuKit kit__h s src s name → b

Warm the cache: compile src (entry name) into the kit now and report whether it succeeded, so a long-lived caller can detect a bad kernel at setup instead of on the first launch.

@ gk_run GpuKit kit__h s src s name i grid i block ( Vec GkArg ) call → b

Compile-cached, marshal, launch, sync, download. call lists the kernel's arguments in declaration order (buffers and scalars interleaved exactly as the kernel signature expects). Returns F on any device error.

---

dev.nu

gpukit/dev.nu — device-RESIDENT buffers + dtype-aware kernels.

gk_run marshals host↔device around every call, which is right for one-shot compute but wrong for chained pipelines (an ndarray expression, a NN forward pass): the data should stay on the device between ops. This layer adds exactly that seam:

GkBuf an element-typed device allocation (GK_F32 | GK_F64 | GK_I64) gk_dbuf_new / free / upload / download (f64 host view, converting) gkdbuf_upload_i / downloadi (exact i64 host view, GK_I64 only) gk_run_dev cached-compile + launch + sync over RAW device args gkd_* ready-made dtype-generic kernels over GkBuf: elementwise (scalar + full N-D stride broadcast), unary maps, matmul, batched matmul, gather/scatter, row softmax, sum reduction

Numerics: GK_F64 kernels compute in double exactly like kernels.nu. GK_F32 buffers hold real float32 and kernels compute IN float32 (accumulation included) — true float32 semantics, matching numpy float32 / onnxruntime, NOT the host-tensor trick of f64-compute + grid-rounding. Uploads convert f64 → f32 via a staging buffer. GK_I64 holds long long (index tensors: gather/scatter/argmax); the f64 host view converts by C truncation, the _i entry points are exact.

API

: i GK_F64 0

: i GK_F32 1

: i GK_I64 2

: GkBuf

: GkBuf {
    i dptr
    i n
    i dtype
    GkMem mem
}

An element-typed device allocation: dptr, n and dtype are plain fields (read them freely), mem owns the block. dptr 0 = failed or empty. Every copy of a GkBuf — a struct field, a Vec element, a view into it — holds the block, and its LAST owner hands it back to the kit's pool (gpukit.nu): nothing is freed by hand. gk_dbuf_free is an optional early release of one owner.

: GkMemImpl { i dptr GpuKit kit }

The owner of one pooled block: its device pointer and the kit whose pool it came from (held, so the kit — and its device — outlive it).

% Drop GkMemImpl { @ drop GkMemImpl x → v { ( _gk_pool_give ( _GpuKit_ptr . x kit ) . x dptr ) } }

: GkMem { s ctl }

@ GkMem_share GkMem h → GkMem

@ GkMem_drop sink GkMem h → v

@ gk_buf_none i dtype → GkBuf

A buffer that is not there: what a failed allocation returns, and the placeholder for an absent operand (a missing bias, mask, …).

@ gk_buf_view GkBuf b i off i n → GkBuf

n elements of b starting at element off, same dtype. The view holds b's block: it stays valid however long the view lives.

@ gk_buf_view_as GkBuf b i boff i n i dtype → GkBuf

n elements of type dtype at BYTE offset boff into b — a view that reinterprets (an f32 block read as i64 indices, say). Holds b's block like gk_buf_view.

@ gk_buf_wrap i dptr i n i dtype → GkBuf

A buffer over device memory something else owns (another package's allocation, a raw pointer out of a C library). It owns nothing: the caller keeps the memory alive for as long as the GkBuf is used.

@ _gk_tname i dtype → s

@ _gk_pfx i dtype → s

Kernel-name prefix per dtype so the cache never mixes element types.

@ _gk_scal i dtype f v → i

A scalar kernel argument holding value v in the buffer dtype's width: f32 bit pattern for GK_F32, f64 bits otherwise (the arg cell is 8 bytes; the kernel reads the leading 4 for a float parameter).

@ gk_buf_ok GkBuf b → b

@ gk_buf_len GkBuf b → i

@ gk_buf_dtype GkBuf b → i

@ gk_buf_esz GkBuf b → i

Bytes per element. Needed by anyone addressing a sub-range of a buffer (the device pointer is byte-addressed), which is why it is public rather than the file-private __gk_esz it wraps.

@ gk_dbuf_new GpuKit kit__h i n i dtype → GkBuf

@ gk_dbuf_free sink GkBuf b → v

Let go of b now rather than at the end of its owner's scope.

@ gk_dbuf_upload GpuKit kit GkBuf b ( Vec f ) src → b

Host f64 vector → device (converted to the buffer's element type: f32 rounds, i64 truncates like a C cast). Copies min(len(src), b.n) elements; missing elements upload as 0. The direct (non-converting) f64 path is taken only when src actually holds b.n elements — a short vector goes through the padding stage, never out of bounds.

@ gk_dbuf_download GpuKit kit GkBuf b ( Vec f ) dst → b

Device → host f64 vector (converted from the buffer's element type). dst must already hold at least b.n elements (fails closed otherwise); the first b.n are overwritten in place.

@ gk_dbuf_upload_raw GpuKit kit GkBuf b * u src → b

Upload bytes that are ALREADY in the buffer's element type: no conversion, no staging allocation, one memcpy to the device. src must address b.n elements of that type and stay valid for the call — a memory-mapped file is exactly the intended source.

The converting entry point above is the general one, but for a checkpoint whose storage dtype already matches the device buffer it walks every element through f64 twice (widen on read, narrow on upload) and allocates two host buffers the size of the tensor to do it. On a 4.6 GB model that is most of the load.

@ gk_dbuf_upload_i GpuKit kit GkBuf b ( Vec i ) src → b

Exact i64 host view (GK_I64 buffers only — index tensors must never make a round trip through f64). Same length contracts as the f64 view.

@ gk_dbuf_download_i GpuKit kit GkBuf b ( Vec i ) dst → b

@ gk_arg_dev GkBuf b → i

Compile-cached launch over pre-built args (device pointers via gk_arg_dev, scalars via gpu_arg_i64/i32/f32). Syncs before returning unless autosync is off (below).

: ~ b g_gk_autosync T

Autosync: gk_run_dev (and every gkd_* kernel on top of it) normally syncs the device after each launch. An executor that chains hundreds of launches can turn this off, launch away, and gk_sync once at the end — the CUDA stream serialises kernels, and downloads synchronise implicitly; the CPU backend runs launches synchronously either way.

@ gk_autosync b on → v

@ gk_sync GpuKit kit__h → b

: ~ b g_gk_prof F

── Per-kernel profiling ────────────────────────────────────────────── "The model is slow" is not actionable; "78% of the frame is in three kernels" is. With profiling on, every gk_run_dev launch is bracketed by a CUDA event pair (the kit's own, made on first use) and the device time is accumulated into the kernel's cache slot, so gk_prof_report ranks the kernels by where the time actually went. Events measure the GPU, not the launch — a host clock around an async launch measures neither.

It is not free: each launch waits for its own end event, which serialises the stream. Turn it on to find the hot kernel, off to time the program.

@ gk_prof GpuKit kit__h b on → v

@ gk_prof_on → b

@ gk_prof_reset GpuKit kit__h → v

Drop every accumulated count — e.g. after a warm-up frame, so the numbers describe the steady state and not the compiles.

@ gk_prof_total GpuKit kit__h → i

Total device time accumulated so far, in nanoseconds.

@ gk_prof_report GpuKit kit__h → v

Kernels by device time, slowest first, with the share of the total.

@ gk_run_dev GpuKit kit__h s src s name i grid i block ( Vec i ) args → b

@ gkd_ew GpuKit kit s opname s op GkBuf o GkBuf a GkBuf b → b

@ gkd_add GpuKit kit GkBuf o GkBuf a GkBuf b → b

@ gkd_sub GpuKit kit GkBuf o GkBuf a GkBuf b → b

@ gkd_mul GpuKit kit GkBuf o GkBuf a GkBuf b → b

@ gkd_div GpuKit kit GkBuf o GkBuf a GkBuf b → b

@ _gk_vi ( Vec i ) v i k → i

@ gkd_ew_bc GpuKit kit s opname s op GkBuf o GkBuf a GkBuf b ( Vec i ) odims ( Vec i ) astr ( Vec i ) bstr → b

out over odims; a and b read through their stride tables (entries ≥ 0, 0 = broadcast). Float dtypes only.

@ gkd_map GpuKit kit s mapname s expr GkBuf o GkBuf a → b

@ gkd_relu GpuKit kit GkBuf o GkBuf a → b

@ gkd_sigmoid GpuKit kit GkBuf o GkBuf a → b

@ gkd_exp GpuKit kit GkBuf o GkBuf a → b

@ gkd_tanh GpuKit kit GkBuf o GkBuf a → b

@ gkd_sqrt GpuKit kit GkBuf o GkBuf a → b

@ gkd_log GpuKit kit GkBuf o GkBuf a → b

: i GKD_RT 8

Rows / columns of C one CPU-backend thread owns. See __gkd_mm_tiled.

: i GKD_CT 32

@ _gkd_ceil i x i q → i

@ _gkd_mac i dtype s dst s x s y → String

The multiply-accumulate step, spelled the way each element type's contract requires: F64 with explicit round-to-nearest intrinsics (NVRTC's fmad contraction would otherwise fuse and change the bits), F32 plain (its contract is true-float32 semantics, which the verified model goldens pin). dst and the product operands are C expressions.

@ _gkd_gemm_tiled s tn i dtype → String

The register-tiled body with a GEMM epilogue: alpha, beta and an optional per-column bias. Kept separate from __gkd_mm_tiled rather than parameterised, because the inner loop is what has to stay simple enough to vectorise and the epilogue runs once per output.

Wants B as [K,N], i.e. transb=0. That is the layout the tiling exists for — B+t*N+cb is 32 contiguous elements reused by all 8 rows of the tile. Staging a transb=1 weight into it at CALL time was measured and is a LOSS: the transpose is a full memory round trip over a cold weight that is then used once, and the model came out 2.4x slower even though an isolated benchmark said 1.5x faster (the benchmark kept the weight in L3 across repeats). Transpose at load time instead.

: i GKD_BM 64

Shared-memory tiled matmul body for the CUDA backend — the body behind gkd_gemm and gkd_bmm both.

One thread per output element (the shape this file shipped with) is bandwidth-bound: every MAC costs a fresh global load of B, so a 4090 runs it at ~3 TFLOP/s, about 4% of what its FP32 units can do. Here a 256-thread block owns a 64x64 tile of Y, stages 64x16 of A and 16x64 of B into shared memory, and each thread keeps a 4x4 register tile. B is read once per tile instead of once per output row, and the arithmetic intensity goes from 1 MAC per load to 32. Measured on a 4090 at 783x1024x1024 f32: 3.1 -> 11.5 TFLOP/s, and 0.6 -> 11.1 with B transposed (the per-element kernel reads a transposed B along the wrong axis entirely).

The accumulation order is unchanged: each output still sums t = 0..K-1 ascending, only blocked. Out-of-range staging slots hold a zero, so a ragged tail adds exact 0.0 terms — the value is bit-identical to the per-element kernel, which is what the verified model goldens pin. Nothing here reassociates.

transb is baked into the body (and into the kernel name) rather than branched on per k: with B as [N,K] the staging loop walks k contiguously per column, so transb=1 loads coalesced too — which is why, unlike the CPU path, this one does not care which layout the weight arrives in.

batched picks the bmm signature (per-batch A/B strides as/bs, a bare store) over the gemm one (alpha/beta and a per-column bias). The tile: a block owns GKD_BM x GKD_BN of Y and each of its 256 threads a GKD_TM x GKD_TN register block, stepping K in GKD_BK.

Why 128x64x16 with an 8x4 register tile and not the obvious 64x64 with 4x4: at 4x4 a thread reads 8 floats from shared per k step to do 16 MACs, and Ada's 128 B/clk of shared bandwidth cannot keep 128 FMA lanes fed at that ratio — the kernel stalls on shared, not on math. 8x4 halves the ratio to 12 reads per 32 MACs. Going the whole way to 128x128/8x8 halves it again, but M here is ~800 rows and N is often 1024, which is 56 blocks — under half of a 4090's 128 SMs, so the wider tile loses more to idle SMs than it gains per SM.

: i GKD_BN 64

: i GKD_BK 16

: i GKD_TM 4

: i GKD_TN 4

@ _gkd_smem_body s tn i dtype i transb i batched i big → String

big swaps the 64x64/4x4 tile for 128x64 with an 8x4 register tile — the shape the comment above argues for and gkd_gemm now picks when M is deep enough to keep every SM busy at it: 8x4 halves the shared-reads-per-MAC ratio (12 per 32 against 8 per 16), which is the bound the 64x64 tile runs into on Ada. Same staging, same ascending-k accumulation, so the value is still bit-identical to the per-element kernel — only the blocking changes.

@ gkd_matmul GpuKit kit GkBuf c GkBuf a GkBuf b i m i k i n → b

@ gkd_bmm GpuKit kit GkBuf y GkBuf a GkBuf b i batch i m i kk i n i abatch i bbatch → b

@ gkd_gather GpuKit kit GkBuf y GkBuf d GkBuf ix i outer i axin i inner i nidx → b

@ gkd_scatter GpuKit kit GkBuf y GkBuf d GkBuf ix GkBuf u i outer i ax i inner i nidx → b

y = d with y[outer, ix[g], inner] = u[outer, g, inner]. When several indices collide the surviving write is unspecified (ONNX leaves this undefined too). y may be d itself for in-place update.

@ gkd_softmax_rows GpuKit kit GkBuf o GkBuf a i rows i cols → b

@ gkd_sum GpuKit kit GkBuf a → ?f

---

kernels.nu

gpukit/kernels.nu — a small library of ready-made f64 kernels over gk_run.

These cover the operations ML code writes over and over. They generate the CUDA-C source, cache it on the kit (by a fixed entry name), and marshal through gk_run — so a caller never writes a kernel or a device buffer for the common cases.

Numerics: f is a C double, so every buffer is float64. Elementwise ops and gk_matmul_f accumulate in the same order a sequential host loop would, so they are bit-identical across the CUDA / CPU / pure backends and to a naive host implementation. gk_reduce_sum_f / gk_dot_f combine per-thread partial sums (a parallel order), so they match a host sum to rounding but are not guaranteed bit-identical to a strictly sequential accumulation.

API

@ gk_add_f GpuKit kit ( Vec f ) out ( Vec f ) a ( Vec f ) b → b

@ gk_sub_f GpuKit kit ( Vec f ) out ( Vec f ) a ( Vec f ) b → b

@ gk_mul_f GpuKit kit ( Vec f ) out ( Vec f ) a ( Vec f ) b → b

@ gk_div_f GpuKit kit ( Vec f ) out ( Vec f ) a ( Vec f ) b → b

@ gk_map_f GpuKit kit s kname ( Vec f ) out ( Vec f ) in s expr → b

── Unary map: out[i] = <expr>, with x bound to in[i] ──────────────── kname is the CUDA entry name (a valid C identifier, stable per expr so caching works). expr is C over the local double x, e.g. x*x, 1.0/(1.0+exp(-x)), x>0.0?x:0.0.

: i GK_MM_RT 8 // rows of C per thread (CPU tile)

: i GK_MM_CT 32 // cols of C per thread (CPU tile)

@ gk_matmul_f GpuKit kit ( Vec f ) c ( Vec f ) a ( Vec f ) b i m i k i n → b

@ _gk_zeros i n → ( Vec f )

@ _gk_partial_threads i n → i

@ gk_reduce_sum_f GpuKit kit ( Vec f ) x → ?f

Sum of x. None on a device error.

@ gk_dot_f GpuKit kit ( Vec f ) a ( Vec f ) b → ?f

Dot product a·b (equal lengths assumed). None on a device error.

---

devops.nu

gpukit/devops.nu — the NN operator family over device-resident GkBufs.

These are the kernels a CNN / transformer forward pass needs beyond the arithmetic core in dev.nu: gemm, conv/pool, normalisation, activation, axis softmax, axis data movement (concat/slice/permute/resize/expand), reductions and argmax/index selection. The kernel bodies are lifted from packages/onnx's original f32 operator set and generalised over the element type, keeping the arithmetic ORDER identical — so the GK_F32 instantiation is bit-for-bit the kernels the onnx/objdet/yoloe models were verified with, and GK_F64 gets the same operators in double.

Conventions:

• every wrapper validates buffers, dtypes and sizes and fails closed

(returns F) — a wrong shape never reaches the device;

• layout metadata is explicit (NCHW dims, axis views (outer,ax,inner));

batch N is assumed 1 for the conv family, group 1;

• arithmetic ops require a float dtype; pure data movement

(copy/slice/permute/resize/expand) accepts GK_I64 too;

• index outputs (argmax) are GK_I64.

API

@ gkd_gemm GpuKit kit GkBuf y GkBuf a GkBuf b GkBuf c i hasb i m i n i k f alpha f beta i transb → b

@ gkd_conv2d GpuKit kit GkBuf y GkBuf x GkBuf w GkBuf bias i hasb i cin i h i wd i cout i kh i kw i oh i ow i ph i pw i sh i sw → b

@ gkd_conv2d_dil GpuKit kit GkBuf y GkBuf x GkBuf w GkBuf bias i hasb i cin i h i wd i cout i kh i kw i oh i ow i ph i pw i sh i sw i dh i dw → b

── 2-D DILATED convolution, NCHW, batch 1, group 1 ────────────────── The atrous variant (U²-Net's RSU4F, DeepLab-family heads): tap (r, s) reads input row iy = oy·sh − ph + r·dh. gkd_conv2d is the dh=dw=1 special case and keeps its specialised fast kernel; this one is a plain one-thread-per-output body, because dilated layers in the models that need them are a handful of small feature maps, not the hot path.

@ gkd_convtranspose2d GpuKit kit GkBuf y GkBuf x GkBuf w GkBuf bias i hasb i cin i h i wd i cout i kh i kw i oh i ow i ph i pw i sh i sw → b

@ gkd_maxpool2d GpuKit kit GkBuf y GkBuf x i c i h i wd i kh i kw i oh i ow i sh i sw i ph i pw → b

@ gkd_batchnorm GpuKit kit GkBuf y GkBuf x GkBuf sc GkBuf bb GkBuf mean GkBuf var i c i hw f eps → b

@ gkd_leakyrelu GpuKit kit GkBuf y GkBuf x f alpha → b

@ gkd_clip GpuKit kit GkBuf y GkBuf x f lo f hi → b

@ gkd_erf GpuKit kit GkBuf y GkBuf x → b

@ gkd_layernorm GpuKit kit GkBuf y GkBuf x GkBuf sc GkBuf bi i outer i ax f eps → b

: i GKD_ATTN_SMEM 33792

── Fused scaled-dot-product attention ────────────────────────────────

O[h,i,:] = softmax_j( scale * Q[h,i,:]·K[h,j,:] ) · V[h,j,:]

Q is [heads, n, hd], K and V are [heads, nkv, hd], O is [heads, n, hd].

The composed form — bmm, scale, softmax, bmm — has to MATERIALISE the [heads, n, nkv] score matrix, and that is the whole cost. For a streaming aggregator holding 72 frames of 783 tokens, one block's scores are 2.8 GB, written once by the first bmm, read and written four times by the softmax and read once more by the second bmm. The arithmetic is a rounding error next to the traffic, and the buffer alone bounds how long a sequence fits in memory.

So the scores are never written. A block owns BQ queries, walks the keys in tiles, and keeps a running row maximum and denominator, rescaling its accumulator when the maximum moves — the online softmax of Milakov & Gimelshein, which is what FlashAttention and torch's own SDPA do. K and V are read once each; nothing of size n*nkv exists.

This is the one op in this file that is NOT bit-identical to its composed equivalent. The sum over keys is blocked and rescaled rather than run left to right, so the low bits differ — as they do between torch's SDPA and its own naive path. It is the same value to within f32 rounding, and closer to the reference than the composed form is, but a golden taken with gkd_bmm + gkd_softmax_ax will not reproduce bit for bit. Whether gkd_attention will run for this head width on this backend — so a caller can decide, before it allocates, whether it needs the [heads, n, nkv] score buffer the composed fallback would want. At a 72-frame window of 783 tokens that buffer is 2.8 GB, and it is the difference between a long sequence fitting and not. Shared-memory budget one attention block may take: a third of the 100 KB an Ada SM can carve out, so three blocks stay resident.

@ gkd_attention_ok GpuKit kit i hd → b

@ gkd_attention GpuKit kit GkBuf o GkBuf q GkBuf k GkBuf v

Unmasked: every key is visible to every query.

@ gkd_attention_masked GpuKit kit GkBuf o GkBuf q GkBuf k GkBuf v GkBuf mask

With a per-key additive bias: mask holds (heads/hpb) rows of nkv floats, head h reading row h/hpb, and the bias is added to the score before the softmax — 0 keeps a key, -1e30 drops it. That is how a padded batch says "these positions are not there": the key never reaches the running maximum and contributes exactly zero to the denominator, so the row is the row it would have been without the padding. Pass a zero GkBuf (dptr 0) for mask to run unmasked — the kernel is then character-for-character the one gkd_attention has always generated, so nothing already compiled is invalidated.

@ gkd_attention_batch GpuKit kit GkBuf o GkBuf q GkBuf k GkBuf v GkBuf mask

The fused attention over a BATCH of sequences in the layout a transformer's linear layers actually produce: Q, K, V and O are [batch·n, heads·hd] — token rows, heads interleaved within the row — NOT the [heads, n, hd] gkd_attention wants. A caller running batch=1 used two gkd_perm round trips per tensor per block to get in and out of head-major layout; a caller running a real batch would also need a 4-D permute whose dims gkd_perm bakes into the kernel name, i.e. an NVRTC compile per (batch, length) pair. This entry removes both: the kernel reads the interleaved layout directly (row stride heads·hd, head offset h·hd — the loads along d stay contiguous and coalesced exactly as before), and no permuted copy of Q/K/V/O ever exists.

mask is one additive row of nkv floats PER BATCH ELEMENT ([batch, nkv]; 0 keeps a key, -1e30 drops it), or a zero GkBuf to run unmasked. Same tile, same online softmax, same accumulation order as gkd_attention_masked — a sequence embedded at batch b is bit-identical to the same sequence permuted through the head-major entry. gkd_attention_ok answers for this entry too (the tile constraint is the same).

@ gkd_softmax_ax GpuKit kit GkBuf y GkBuf x i outer i ax i inner → b

@ gkd_copy_ax GpuKit kit GkBuf dst GkBuf src i outer i src_ax i inner i dst_ax i off → b

Copy src into a concat output along an axis viewed as (outer, src_ax, inner); the slot starts at off in the output's axis of size dst_ax.

@ gkd_slice_ax GpuKit kit GkBuf dst GkBuf src i outer i sz i inner i src_ax i soff → b

Extract a contiguous axis slice into a fresh tensor: viewing src as (outer, src_ax, inner), dst[o,a,i] = src[o, soff+a, i] for a in [0,sz).

@ gkd_perm GpuKit kit GkBuf y GkBuf x ( Vec i ) dims ( Vec i ) perm → b

General N-D (≤6) transpose: output axis k reads input axis perm[k]. dims are the INPUT dims; both vecs hold ndim (≤6) entries and perm must be a permutation of 0..ndim−1.

@ gkd_resize_bilinear GpuKit kit GkBuf y GkBuf x i c i h i wd i oh i ow i align → b

Nearest-neighbour upsample by integer factors (NCHW): Y[c,oy,ox] = X[c, oy/sh, ox/sw]. ── Bilinear resize, NCHW (batch folded into C) ──────────────────────

The resampler every decoder upsamples with, and the one gkd_resize_nn is not. align picks the coordinate convention, and the two are NOT interchangeable:

align_corners=1 : src = dst · (in−1)/(out−1) endpoints pinned align_corners=0 : src = (dst + 0.5)·in/out − 0.5, clamped

A single output pixel (out == 1) maps to src 0 under align_corners, which is what torch does rather than dividing by zero.

Edge taps are clamped, so a coordinate that lands on the last row or column interpolates with itself rather than reading past the end.

@ gkd_resize_nn GpuKit kit GkBuf y GkBuf x i c i h i wd i oh i ow i sh i sw → b

@ gkd_expandlast GpuKit kit GkBuf y GkBuf x i outer i rep → b

Expand the last axis: input (outer,1) → (outer,rep), Y[o,r] = X[o].

@ gkd_reducel2 GpuKit kit GkBuf y GkBuf x i outer i ax → b

L2 norm along the last axis: input (outer, ax) → output (outer).

@ gkd_argmax GpuKit kit GkBuf y GkBuf x i outer i ax → b

ArgMax along the last axis: input (outer, ax) in any element type → GK_I64 indices (outer). Ties resolve to the first maximum.

@ gkd_eos_gather GpuKit kit GkBuf y GkBuf data GkBuf tok i bsz i l i d → b

EOS read-out (CLIP): given per-row token ids tok [B,L] (GK_I64) and features data [B,L,D], select each row's max-id token's features → Y [B,D]. One thread per output element.