CUDA Driver API FFI Bindings #
Minimal bindings to the CUDA Driver API for PTX JIT compilation and kernel execution. No nvcc or CUDA runtime needed — only the NVIDIA driver (libcuda.so).
Usage: cuDriverInit let dev ← cuDeviceGet 0 let ctx ← cuCtxCreate dev let mod ← cuModuleLoadData ptxString let func ← cuModuleGetFunction mod "main" let buf ← cuMalloc 1024 cuLaunchKernel func (gridX, gridY, gridZ) (blockX, blockY, blockZ) 0 #[buf]
Equations
Instances For
Equations
Instances For
Equations
Instances For
Equations
Instances For
Equations
Instances For
Driver initialization #
Context #
Module (PTX JIT) #
Variant that takes raw bytes (cubin / fatbin / pre-compiled PTX blobs containing non-UTF-8 bytes). Skips disk-cache and JIT paths; the input must already be in a form the driver accepts directly.
Raise a kernel's dynamic shared-memory limit (required for kernels
requesting > 48 KB of smem). Wraps cuFuncSetAttribute with
CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES.
Utilities #
Memory-mapped file I/O #
Copy a slice from mmapped memory to a Lean ByteArray (for metadata).
Copy mmapped data directly to GPU buffer (zero Lean-side copy).
Fast file read: mmap + memcpy. Faster than IO.FS.readBinFile for large files.
Persistent mmap with GC-managed lifetime #
MMappedFile is an opaque Lean external object backed by a C-side
mmap region. Its finalizer (registered on the external class at
registration time) calls munmap when the last Lean reference is
dropped — matching llama.cpp's unique_ptr<llama_mmap> RAII pattern.
Unlike the legacy mmapFile above, this path never copies the file
into a ByteArray. Tensor weights are kept as (mmap, offset, size)
triples and uploaded to the GPU via cuMemcpyHtoDAsync directly from
the mmap region — saving ~5 GB of Lean-heap allocation and enabling
async H2D (the mmap region stays alive until all in-flight copies
finish, because the stream sync point still holds a Lean reference).
Opaque handle to a memory-mapped file. GC-finalised with
munmap.
Instances For
Copy a slice of the mmap into a fresh Lean ByteArray. Use only
for metadata / parsing; tensor weights should use
cuMemcpyHtoDFromMmap instead.
H2D copy from mmap region directly to GPU. When stream = 0,
issues the sync cuMemcpyHtoD; otherwise async on the given
stream. Caller must keep the MMappedFile referenced until the
stream has synchronised (Lean GC automatically does this via any
structure that holds both the GPU buffer handle and the mmap).
Pin a sub-range of an mmap region and map it into CUDA's unified VA
space. Returns the device-side pointer that kernels can ld.global
directly — the driver pulls pages over PCIe on demand. Same trick
llama.cpp uses for tok_embd_per_layer (its getrows kernel reads the
host pointer of the CPU_Mapped buffer). offset must be 4 KB-aligned;
size is rounded up to a multiple of 4 KB. Cost: ~200 ms/GB on the
first call, once at model load — no per-token cuMemcpy afterwards.
Memory #
Kernel launch #
Metadata-free launch descriptor table (Option B+) #
These register kernel launch metadata (func, grid/block, arg ptrs) in
C-owned storage, then fire by numeric id. The Lean side holds only a
USize id per kernel and makes exactly one FFI call per launch — no
Array USize allocation, no ByteArray for params, no closure for
buffer resolution. See docs/llama-fusion-analysis/53-metadata-free- forward-design.md.
Reset the descriptor pool. Call once when rebuilding an inference schedule (not per token).
Register a descriptor with immutable grid/block and a snapshot of buffer pointers. Returns the descriptor id.
Rebind one buffer pointer in an existing descriptor (used only for sites whose buffers genuinely change between calls).
Fast launch via descriptor id. This is the Lean-heap-free hot
path. stream=0 = default stream.
Fused rebind-all-args + launch. Used when the descriptor's
buffer pointers may differ between calls (e.g. same logical call
site across 42 layers with per-layer weights). Writes each arg
into the C-owned arg_storage in place — still zero Lean heap
past the FFI boundary. args.size must equal the descriptor's
n_args.
Raw-bytes launch for external PTX (e.g. llama.cpp) whose kernels take
mixed-type args (uint3 structs, f16 scalars, ...). argBytes holds
packed arg values; argOffsets[i] = byte offset of the i-th arg.
CUDA receives void** where each entry points at argBytes + offset.
Stream-aware variant of cuLaunchKernelRaw. Required when the caller
is inside cuStreamBeginCapture — launches on the default stream are
NOT captured into the graph, so their execution ordering diverges from
the captured hesper kernels on graph replay, producing garbage output.
stream = 0 (default stream) matches the old behaviour.
L2 cache persistence (Ampere+ / Ada / Hopper, CC ≥ 8.0) #
Set the persisting-L2 limit on the current context. The value is
clamped to the device's MAX_PERSISTING_L2_CACHE_SIZE attribute.
Returns the effective limit the driver applied.
Install a persisting access-policy window on the default (null) stream.
Subsequent launches on this stream will prefer keeping
[ptr, ptr+size) in L2 across kernel boundaries.
Evict all persisting lines from L2 and return the cache to its default non-persistent behaviour.
CUDA Graphs #
Capture a sequence of kernel launches once, then replay the whole graph per decode token with a single driver call. llama.cpp uses this to amortise the ~1.2 µs/dispatch host overhead that dominates our ~10 ms/tok host budget. See docs/llama-fusion-analysis/12-complete-cuda-flow.md §5 for the llama.cpp reference flow.
Opaque handles are size_t on the C side; we marshall as USize.
cudaStream_t handle (null stream ≡ 0).
Equations
Instances For
cudaGraph_t handle (capture product).
Equations
Instances For
cudaGraphExec_t handle (instantiated graph ready to launch).
Equations
Instances For
Create a dedicated non-blocking stream on which subsequent launches can be captured. Call once at context init.
Create a blocking (default-synchronising) stream. Ops on this
stream implicitly sync with the null stream, so readBuffer
(which goes through the null stream) observes prior work
without an explicit cuStreamSynchronize.
Begin graph capture on stream. All cuLaunchKernelOnStream
launches between this call and cuStreamEndCapture are recorded
into the produced cudaGraph_t.
End capture and return the captured graph.
Turn a captured graph into an executable that the driver can replay
with one launch. Must be called before the first cuGraphLaunch.
Replay a previously-instantiated graph on stream. Returns when
the kernels have been submitted (they still run asynchronously).
Launch a kernel onto a specific stream — needed so we can capture
the launches into a graph (the default cuLaunchKernel goes to the
null stream, which is not captureable).
Synchronise a stream. Required after cuGraphLaunch when the
caller wants to read back results.
Host→device memcpy on an explicit stream. Needed so writes
issued DURING stream capture get recorded as memcpy nodes in the
resulting graph (rather than forcing a sync). The driver captures
the (src-host-ptr, dst-device-ptr, size) triple; on each replay
the memcpy re-reads src from host memory, so subsequent tokens'
values flow through without re-capturing.
Pinned host memory (staging buffers for CUDA Graphs) #
ByteArray is Lean-GC'd and its address is not stable across
writeBufferOffset calls. CUDA Graph capture records the pointer,
so replay against a freed ByteArray tombstones with
CUDA_ERROR_ILLEGAL_ADDRESS. The correct source for capturable
writes is page-locked (pinned) host memory allocated via
cuMemHostAlloc and held for the whole session. llama.cpp uses the
same trick for its per-token scalar uploads.
Allocate size bytes of pinned host memory. Returns the host
virtual-address as a USize. The memory survives until
cuMemFreeHost is called (or the process exits).
Allocate size bytes of pinned host-mapped memory and return
(hostPtr, devPtr). The device pointer aliases the same physical
memory through CUDA's unified VA space, so kernels can st.global
into it directly and the host can read the result with no driver
call once the producing stream has been synchronised. Used to
eliminate the per-token cuMemcpyDtoH(4 byte) argmax bubble that
drains the stream implicitly (~9.8 ms/tok on graphs-OFF).
Read a UInt32 from a pinned host pointer with no driver call.
Caller must have synchronised the stream that wrote it
(cuStreamSynchronize once is enough — independent of how many
kernels updated the slot).
Write a small scalar (≤ 8 bytes) into a pinned host buffer. The data is plain Lean bytes; the C++ side memcpys them into the pinned region. No GPU involvement. Use before every graph launch to update a captured memcpy node's source.
Host→device memcpy where the source is a pinned host pointer (stable across the session). Safe to use inside stream capture; the graph records the host pointer, and every replay reads the current contents.
Fused: write src bytes to pinned slot (hostPtr+offset) then
immediately queue an async H2D copy to (dst) on stream. One
FFI crossing instead of two (cuWritePinned + cuMemcpyHtoDFromPinned),
eliminating ~half the Lean→C boundary overhead on hot per-scalar
write sites.