Trait GpuBackend 

pub trait GpuBackend: Send + Sync {
    // Required methods
    fn alloc(&self, bytes: usize) -> Result<DevicePtr>;
    fn alloc_managed(&self, bytes: usize) -> Result<DevicePtr>;
    fn free(&self, ptr: DevicePtr) -> Result<()>;
    fn copy_h2d(&self, src: &[u8], dst: DevicePtr) -> Result<()>;
    fn copy_d2h(&self, src: DevicePtr, dst: &mut [u8]) -> Result<()>;
    fn copy_d2d(
        &self,
        src: DevicePtr,
        dst: DevicePtr,
        bytes: usize,
    ) -> Result<()>;
    fn launch(
        &self,
        func: KernelHandle,
        grid: [u32; 3],
        block: [u32; 3],
        shared_mem: u32,
        stream: u64,
        params: &mut [*mut c_void],
    ) -> Result<()>;
    fn synchronize(&self, stream: u64) -> Result<()>;
    fn default_stream(&self) -> u64;
    fn kernel(&self, module: &str, func_name: &str) -> Result<KernelHandle>;
    fn has_module(&self, module: &str) -> bool;
    fn op_cache(&self) -> &OpCache;
    fn memset(&self, ptr: DevicePtr, value: u8, bytes: usize) -> Result<()>;
    fn memset_async(
        &self,
        ptr: DevicePtr,
        value: u8,
        bytes: usize,
        stream: u64,
    ) -> Result<()>;
    fn total_memory(&self) -> Result<usize>;
    fn free_memory(&self) -> Result<usize>;
    fn sm_count(&self) -> Result<u32>;

    // Provided methods
    fn sweep_unreleased(&self) -> usize { ... }
    fn live_bytes(&self) -> Option<usize> { ... }
    fn chunk_slack_bytes(&self) -> Option<usize> { ... }
    fn alloc_report(&self, _top_n: usize, _min_mb: usize) -> Option<String> { ... }
    fn live_allocations(&self) -> Option<Vec<(DevicePtr, usize, String)>> { ... }
    fn copy_d2h_on_stream(
        &self,
        src: DevicePtr,
        dst: &mut [u8],
        stream: u64,
    ) -> Result<()> { ... }
    fn launch_typed(
        &self,
        func: KernelHandle,
        grid: [u32; 3],
        block: [u32; 3],
        shared_mem: u32,
        stream: u64,
        args: &[KernelArg<'_>],
    ) -> Result<()> { ... }
    fn stream_is_capturing(&self, _stream: u64) -> bool { ... }
    fn scan_redzones(&self) -> Result<usize> { ... }
    fn poison_redzones(&self, _lo: usize, _hi: usize) -> Result<()> { ... }
    fn kernel_n_tile(&self, kernel: KernelHandle) -> Result<u32> { ... }
    fn kernel_casts_a_to_e4m3(&self, kernel: KernelHandle) -> bool { ... }
    fn debug_sync_kernels(&self) -> bool { ... }
    fn kernel_registry(&self) -> Option<Arc<MetraleRegistry>> { ... }
    fn copy_h2d_async(
        &self,
        src: &[u8],
        dst: DevicePtr,
        _stream: u64,
    ) -> Result<()> { ... }
    fn copy_h2d_async_retained(
        &self,
        src: &[u8],
        dst: DevicePtr,
        stream: u64,
    ) -> Result<()> { ... }
    fn copy_d2h_async(
        &self,
        src: DevicePtr,
        dst: &mut [u8],
        _stream: u64,
    ) -> Result<()> { ... }
    fn copy_d2d_async(
        &self,
        src: DevicePtr,
        dst: DevicePtr,
        bytes: usize,
        _stream: u64,
    ) -> Result<()> { ... }
    fn copy_d2d_2d_async(
        &self,
        src: DevicePtr,
        src_pitch: usize,
        dst: DevicePtr,
        dst_pitch: usize,
        width_bytes: usize,
        height: usize,
        stream: u64,
    ) -> Result<()> { ... }
    fn begin_capture(&self, _stream: u64) -> Result<()> { ... }
    fn end_capture(&self, _stream: u64) -> Result<GraphHandle> { ... }
    fn launch_graph(&self, _graph: GraphHandle, _stream: u64) -> Result<()> { ... }
    fn destroy_graph(&self, _graph: GraphHandle) -> Result<()> { ... }
    fn abort_capture_if_active(&self, _stream: u64) { ... }
    fn device_free_memory(&self) -> Result<usize> { ... }
    fn live_alloc_count(&self) -> usize { ... }
    fn create_stream(&self) -> Result<u64> { ... }
    fn bind_to_thread(&self) -> Result<()> { ... }
    fn create_event(&self) -> Result<u64> { ... }
    fn record_event(&self, _event: u64, _stream: u64) -> Result<()> { ... }
    fn stream_wait_event(&self, _stream: u64, _event: u64) -> Result<()> { ... }
    fn event_synchronize(&self, _event: u64) -> Result<()> { ... }
    fn event_query(&self, _event: u64) -> Result<bool> { ... }
    fn destroy_event(&self, _event: u64) -> Result<()> { ... }
    fn host_ptr_to_device(&self, _host: *mut u8) -> Result<DevicePtr> { ... }
    fn alloc_host_pinned(&self, bytes: usize) -> Result<*mut u8> { ... }
    fn free_host_pinned(&self, ptr: *mut u8, bytes: usize) -> Result<()> { ... }
    fn alloc_arena(&self, bytes: usize) -> Result<DevicePtr> { ... }
    fn free_arena(&self, ptr: DevicePtr) -> Result<()> { ... }
}

2026-09-25: The device interface the model code calls. Implementations: MetraleCudaBackend (cuda feature), MetalGpuBackend (metal feature) and MockGpuBackend (tests and the test-utils feature).

Required Methods

fn alloc(&self, bytes: usize) -> Result<DevicePtr>

2026-09-25: Allocate bytes of device memory.

#[track_caller] is on the declaration as well as the CUDA impl, so the allocation ledger records the caller’s location through &dyn GpuBackend.

fn alloc_managed(&self, bytes: usize) -> Result<DevicePtr>

2026-09-25: Allocate managed memory: cuMemAllocManaged on CUDA, plain alloc on Metal and the mock.

fn free(&self, ptr: DevicePtr) -> Result<()>

2026-09-25: Free device memory. A null pointer is a no-op.

fn copy_h2d(&self, src: &[u8], dst: DevicePtr) -> Result<()>

2026-09-25: Copy from host to device; returns when the copy is done.

fn copy_d2h(&self, src: DevicePtr, dst: &mut [u8]) -> Result<()>

2026-09-25: Copy from device to host; returns when the copy is done.

fn copy_d2d(&self, src: DevicePtr, dst: DevicePtr, bytes: usize) -> Result<()>

2026-09-25: Copy device to device; returns when the copy is done.

fn launch( &self, func: KernelHandle, grid: [u32; 3], block: [u32; 3], shared_mem: u32, stream: u64, params: &mut [*mut c_void], ) -> Result<()>

2026-09-25: Launch a kernel with untyped parameter pointers. The Metal backend refuses this; use launch_typed.

fn synchronize(&self, stream: u64) -> Result<()>

2026-09-25: Block until all work queued on stream has completed.

fn default_stream(&self) -> u64

2026-09-25: The default stream handle.

fn kernel(&self, module: &str, func_name: &str) -> Result<KernelHandle>

2026-09-25: Look up a kernel function by module and function name.

#[track_caller] is on the declaration so the CUDA backend’s kernel audit records the caller’s location through &dyn GpuBackend.

fn has_module(&self, module: &str) -> bool

2026-09-25: Whether module is loaded in this backend. Ask before looking up a kernel that only some targets carry: the CUDA backend’s kernel audit records every failed lookup.

fn op_cache(&self) -> &OpCache

2026-09-25: This backend’s memoized kernel handles and scratch allocations. It has no default, so every backend owns its own cache.

fn memset(&self, ptr: DevicePtr, value: u8, bytes: usize) -> Result<()>

2026-09-25: Set device memory to a byte value; returns when done.

fn memset_async( &self, ptr: DevicePtr, value: u8, bytes: usize, stream: u64, ) -> Result<()>

2026-09-25: Set device memory to a byte value, enqueued on stream on CUDA. Metal and the mock write at once.

fn total_memory(&self) -> Result<usize>

2026-09-25: Total device memory in bytes.

fn free_memory(&self) -> Result<usize>

2026-09-25: Free device memory in bytes. On an integrated GPU the CUDA backend reports the larger of the driver’s figure and host MemAvailable.

fn sm_count(&self) -> Result<u32>

2026-09-25: Number of streaming multiprocessors on the device. The CUDA backend asks the driver, the mock returns 48, and Metal returns an error. Resolve it once at construction, not per launch.

Provided Methods

fn sweep_unreleased(&self) -> usize

2026-09-25: Free every allocation still on this backend’s ledger and return how many there were. The default returns 0 and frees nothing; only the CUDA backend overrides it.

fn live_bytes(&self) -> Option<usize>

2026-09-25: Live device bytes this backend has allocated and not freed: Some on CUDA and the mock, None (the default) on Metal.

fn chunk_slack_bytes(&self) -> Option<usize>

2026-10-01: Device memory the driver holds for this backend’s small allocations beyond their bytes (cuda_backend::alloc_ledger::chunk_slack): Some on CUDA, None (the default) on a backend without an allocation ledger.

fn alloc_report(&self, _top_n: usize, _min_mb: usize) -> Option<String>

2026-09-25: Live device memory by allocating call site; None (the default) on every backend but CUDA.

fn live_allocations(&self) -> Option<Vec<(DevicePtr, usize, String)>>

2026-10-02: Every live allocation on the ledger: its address, bytes and allocating call site (file:line); None (the default) on every backend but CUDA.

fn copy_d2h_on_stream( &self, src: DevicePtr, dst: &mut [u8], stream: u64, ) -> Result<()>

2026-09-25: Blocking device-to-host copy ordered after the work already on stream. copy_d2h orders only against the default stream; use this to read what kernels on another stream just wrote. The CUDA backend enqueues the copy on stream and waits; the default synchronises stream and then calls copy_d2h.

fn launch_typed( &self, func: KernelHandle, grid: [u32; 3], block: [u32; 3], shared_mem: u32, stream: u64, args: &[KernelArg<'_>], ) -> Result<()>

2026-09-25: Launch a kernel with typed arguments.

The default packs them with pack_kernel_args and calls launch. The Metal backend binds each KernelArg::Buffer with setBuffer:offset:atIndex: and each KernelArg::Bytes with setBytes:length:atIndex:; the mock records them.

fn stream_is_capturing(&self, _stream: u64) -> bool

2026-09-25: Whether stream is inside an active CUDA-graph capture. Check it before a sync or D2H on a stream that may be capturing: those invalidate the capture (CUDA_ERROR_STREAM_CAPTURE_INVALIDATED, 901). The CUDA backend reports a failed query as capturing, and metrale_scale builds always report false; the default is false.

fn scan_redzones(&self) -> Result<usize>

2026-09-25: Read every allocation’s trailing guard band back and return how many no longer hold the fill byte. Ok(0) when METRALE_REDZONE is unset, and on every backend but CUDA.

fn poison_redzones(&self, _lo: usize, _hi: usize) -> Result<()>

2026-09-25: Refill the guard bands of allocations with index in [lo, hi) with 0xEE and the rest with 0x00; nothing is allocated, moved or resized. A no-op when METRALE_REDZONE is unset, and on every backend but CUDA.

fn kernel_n_tile(&self, kernel: KernelHandle) -> Result<u32>

2026-09-29: The output-column (N) tile width kernel publishes as the extern "C" __device__ unsigned int <entry>_n_tile of its module, read when the handle was resolved. A launch whose grid covers N in tiles must size it from this, never from a width assumed at the call site: one entry name can carry different tiles in different target trees. An error for a kernel that publishes none, a zero width included.

fn kernel_casts_a_to_e4m3(&self, kernel: KernelHandle) -> bool

2026-09-29: Whether kernel casts its BF16 A operand to E4M3 with no scale, as its module declares with extern "C" __device__ unsigned int <entry>_a_e4m3. Values past 448 saturate there; launchers check the range in debug builds (ops::check_e4m3_activation_range in model-layers).

fn debug_sync_kernels(&self) -> bool

2026-09-25: Whether KernelLaunch::launch synchronises after every launch, so an asynchronous fault is reported at the kernel that caused it. The CUDA backend reads METRALE_DEBUG_SYNC_KERNELS=1 once, when it is built; the default is false.

fn kernel_registry(&self) -> Option<Arc<MetraleRegistry>>

2026-09-25: This backend’s kernel module registry, for callers that need more than a kernel handle. None (the default) on every backend but CUDA.

fn copy_h2d_async(&self, src: &[u8], dst: DevicePtr, _stream: u64) -> Result<()>

2026-09-25: Host-to-device copy on stream; src may be dropped or overwritten as soon as this returns.

The CUDA backend synchronises stream after the enqueue when src is page-locked memory registered in crate::pinned_hosts. The default is the blocking copy_h2d. Use GpuBackend::copy_h2d_async_retained when the source outlives the next synchronisation.

fn copy_h2d_async_retained( &self, src: &[u8], dst: DevicePtr, stream: u64, ) -> Result<()>

2026-09-25: Host-to-device copy on stream for a source the caller keeps alive: src must stay valid and unchanged until the next synchronisation point on stream. The CUDA backend adds no synchronisation; the default is copy_h2d_async.

fn copy_d2h_async( &self, src: DevicePtr, dst: &mut [u8], _stream: u64, ) -> Result<()>

2026-09-25: Device-to-host copy enqueued on stream without waiting. dst must stay valid, and must not be read or reused, until the next synchronisation point on stream. On CUDA, copy_d2h and copy_d2h_on_stream each wait inside the call, so this is the D2H to use for many chunks followed by one synchronize. The default, used by Metal, is the blocking copy_d2h; the mock copies at once and counts the call.

fn copy_d2d_async( &self, src: DevicePtr, dst: DevicePtr, bytes: usize, _stream: u64, ) -> Result<()>

2026-09-25: Device-to-device copy enqueued on stream. The default is the blocking copy_d2d.

fn copy_d2d_2d_async( &self, src: DevicePtr, src_pitch: usize, dst: DevicePtr, dst_pitch: usize, width_bytes: usize, height: usize, stream: u64, ) -> Result<()>

2026-09-25: Pitched device-to-device copy: height rows of width_bytes, source rows src_pitch apart and destination rows dst_pitch apart. The default loops copy_d2d_async per row; the CUDA backend issues one cudaMemcpy2DAsync.

fn begin_capture(&self, _stream: u64) -> Result<()>

2026-09-25: Begin capturing the work enqueued on stream into a graph (cuStreamBeginCapture on CUDA); captured work is recorded, not run. The default does nothing.

fn end_capture(&self, _stream: u64) -> Result<GraphHandle>

2026-09-25: End the capture and return the instantiated graph. The default returns GraphHandle(0).

fn launch_graph(&self, _graph: GraphHandle, _stream: u64) -> Result<()>

2026-09-25: Replay a captured graph on stream.

fn destroy_graph(&self, _graph: GraphHandle) -> Result<()>

2026-09-25: Destroy an instantiated graph.

fn abort_capture_if_active(&self, _stream: u64)

2026-09-25: Best effort: end any capture active on stream and discard the partial graph, so the stream runs work again. Call it on an error path that left a begin_capture/end_capture region early. The default does nothing.

fn device_free_memory(&self) -> Result<usize>

2026-09-25: Free device memory as the driver reports it (cuMemGetInfo on CUDA), never substituted by host memory. The default is free_memory.

fn live_alloc_count(&self) -> usize

2026-09-25: Number of live device allocations on this backend: a count, not bytes. The default, 0, means not tracked (Metal).

fn create_stream(&self) -> Result<u64>

2026-09-25: Create a stream. The default returns 0, the default stream.

fn bind_to_thread(&self) -> Result<()>

2026-09-25: Make the backend’s CUDA context current on this thread. Call it on any thread other than the one that built the backend before using the GPU. The default does nothing.

fn create_event(&self) -> Result<u64>

2026-09-25: Create an event for ordering work across streams. The default returns 0.

fn record_event(&self, _event: u64, _stream: u64) -> Result<()>

2026-09-25: Record event at the current point of stream’s work.

fn stream_wait_event(&self, _stream: u64, _event: u64) -> Result<()>

2026-09-25: Make stream wait for event on the device; the host does not block.

fn event_synchronize(&self, _event: u64) -> Result<()>

2026-09-25: Block the host until the work recorded before event has completed. Unlike synchronize(stream), it does not wait for work enqueued after the event.

fn event_query(&self, _event: u64) -> Result<bool>

2026-09-25: Whether the work recorded before event has completed; never blocks. The default returns true.

fn destroy_event(&self, _event: u64) -> Result<()>

2026-09-25: Destroy an event.

fn host_ptr_to_device(&self, _host: *mut u8) -> Result<DevicePtr>

2026-09-25: The device address of a page-locked host pointer from Self::alloc_host_pinned (cuMemHostGetDevicePointer_v2 on CUDA), so a kernel can write into host-visible memory. The default returns an error.

fn alloc_host_pinned(&self, bytes: usize) -> Result<*mut u8>

2026-09-25: Allocate bytes of host memory for H2D staging: page-locked on CUDA (cuMemAllocHost_v2), a shared MTLBuffer on Metal, and a 64-byte aligned heap block by default. Release it with free_host_pinned.

The region is zeroed on every backend, so a caller may form a &[u8] over all of it, padding included.

fn free_host_pinned(&self, ptr: *mut u8, bytes: usize) -> Result<()>

2026-09-25: Free memory from alloc_host_pinned; bytes must be the allocated size.

fn alloc_arena(&self, bytes: usize) -> Result<DevicePtr>

2026-09-25: Device memory for a weight arena, freed with free_arena. The CUDA backend allocates it off the allocation ledger, whose live_bytes factory::build counts as this process’s own memory when it sizes the KV cache. The default is the ledgered alloc.

fn free_arena(&self, ptr: DevicePtr) -> Result<()>

Trait Implementations

impl ModelResource<dyn GpuBackend> for BufferArena

2026-09-25: The destructure in release names every field, with no .., so a field added to BufferArena does not compile until it is handled here. Free the new buffer rather than adding a wildcard: a missed free is a leak that shows only as the next model failing to fit.

fn label(&self) -> &'static str

2026-09-25: Human name, used to attribute a failure in Teardown::release_all’s error.

fn release(&mut self, gpu: &dyn GpuBackend) -> Result<()>

2026-09-25: Release everything this owns. Must be idempotent.

impl ModelResource<dyn GpuBackend> for OpCache

2026-09-25: Frees every scratch allocation and empties the scratch map.

Kernel handles need no release here: they point into the modules of the backend’s MetraleRegistry, which unloads them when it is dropped or passed to cuda_host::release.

fn label(&self) -> &'static str

2026-09-25: Human name, used to attribute a failure in Teardown::release_all’s error.

fn release(&mut self, gpu: &dyn GpuBackend) -> Result<()>

2026-09-25: Release everything this owns. Must be idempotent.

Implementors

impl GpuBackend for MetraleCudaBackend