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 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 alloc_report(&self, _top_n: usize, _min_mb: usize) -> Option<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 debug_sync_kernels(&self) -> bool { ... }
    fn kernel_registry(&self) -> Option<Arc<AtlasRegistry>> { ... }
    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 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<()> { ... }
}

GPU backend trait — SBIO IORouter for all CUDA operations.

Implementations: AtlasCudaBackend (production), MockGpuBackend (tests).

Required Methods

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

Allocate bytes of device memory.

#[track_caller] so the CUDA backend’s ledger records WHICH code asked for the memory. It must stay on the trait declaration as well as the impl: nearly every caller goes through &dyn GpuBackend, and without it here the vtable would attribute every allocation in the process to the one line inside the backend.

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

Allocate managed (unified) memory. On GB10, this allows over-subscribing physical GPU memory — Linux pages overflow to NVMe swap automatically. Managed memory is slower than device memory but avoids OOM.

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

Free device memory.

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

Copy from host to device.

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

Copy from device to host.

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

Copy device to device.

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

Launch a kernel on the given CUDA stream.

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

Synchronize a CUDA stream (blocks until all work completes).

fn default_stream(&self) -> u64

Get the default stream handle.

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

Look up a kernel function by module and function name.

#[track_caller] on the DECLARATION is what makes the caller location survive the &dyn GpuBackend vtable — every lookup in Atlas goes through dynamic dispatch, so without it the audit can only ever name the backend’s own line. The location is what turns an unresolved-lookup report from a name list into a work item.

fn op_cache(&self) -> &OpCache

This backend’s memoized kernel handles and scratch allocations.

Required rather than defaulted: an op that memoizes a KernelHandle or a DevicePtr anywhere else is caching something that belongs to this backend’s model, and a default would let a new backend forget.

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

Set device memory to a byte value (synchronous — waits for completion).

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

Set device memory to a byte value on the given stream (async — does not wait).

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

Total device memory in bytes.

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

Free device memory in bytes.

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

Number of streaming multiprocessors (CUDA SMs / HIP CUs) on the device.

Queried from the driver, never assumed: dispatch rules that ask “does this grid still fill the machine?” are wrong on every part whose SM count differs from the one they were tuned on. Callers must resolve it ONCE at construction and keep the value, not call it per launch.

Provided Methods

fn sweep_unreleased(&self) -> usize

Free every allocation this backend made that nobody released, and report how many there were.

The teardown backstop. Enumerating owners does not scale: the loaders fuse weights into fresh allocations owned by layer structs, which no pool releases — measured at 15.3 GB per cycle on a 27B, linear over six cycles. A backend is created per model, so its outstanding set IS that model’s leak.

Default 0: a backend that does not track allocations has nothing to sweep, which is honest for the mock and for Metal.

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

Live device bytes this backend has allocated and not freed, if it tracks them. None for backends with no ledger (mock/CPU).

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

Attribution of live device memory by allocating call site, biggest first. None for backends with no ledger.

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

Synchronous device-to-host copy ordered after work on stream.

Unlike copy_d2h (which uses the default stream and only orders against work already on the default stream), this method enqueues the copy on stream. CUDA serializes the copy after any prior kernel launches on stream, so the bytes read are guaranteed to reflect post-kernel state.

Required when reading bytes that were just written by kernels on a non-default stream — e.g. high_speed_swap_offload_new_blocks reading WHT+quantize output bytes.

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

Typed-args kernel launch.

CUDA’s default impl packs args into u64 slots and forwards to launch(). The Metal backend overrides this to map each KernelArg::Buffer to setBuffer:offset:atIndex: and each KernelArg::Bytes to setBytes:length:atIndex:.

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

Whether stream is inside an active CUDA-graph capture. Telemetry taps MUST check this before any sync/D2H on a potentially-captured stream — those calls invalidate the capture (CUDA 901) and wedge the serve. Default false (backends without capture, or without a query API, never capture through this trait’s eager paths).

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

A55 diagnostic: read every allocation’s trailing guard band back and report the ones a kernel wrote past. Returns the violation count. Ok(0) when ATLAS_REDZONE is unset or the backend has no red zones — every backend but CUDA.

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

A55 bisection: poison guard bands [lo, hi) with 0xEE and the rest with 0x00. Layout-preserving by construction — nothing is allocated, moved or resized.

fn debug_sync_kernels(&self) -> bool

Synchronise the stream after every kernel launch, so an asynchronous illegal-address fault is reported at the kernel that caused it rather than at a later sync. Resolved once when the backend is built; read on the launch path, which is why it is not a per-launch getenv.

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

This backend’s model-scoped kernel modules, for the few callers that need the registry itself rather than a kernel handle — resolving a __device__ symbol, for instance. None on backends that have no such concept, which is why it is an accessor rather than a downcast.

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

Async host-to-device copy: src may be dropped or overwritten the moment this returns.

That is what the ~90 call sites in spark-model rely on — nearly all hand over a stack array or local Vec that dies at the end of the statement — and it used to hold only by accident. See crate::pinned_hosts for why, and for how the CUDA backend now MAKES the promise true (page-locked source ⇒ it buys the ordering that the pageable path gets from the driver for free) instead of inheriting it.

Use GpuBackend::copy_h2d_async_retained when the source outlives the next synchronisation and the extra ordering is not wanted.

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

Async host-to-device copy for a source the CALLER keeps alive.

src must remain valid, and must not be rewritten, until the next synchronisation point on stream. In exchange it never inserts an implicit sync — what makes a batched scatter out of one pinned staging blob (N enqueues + one synchronize) worth doing; see GpuBackend::copy_d2h_async for the measured shape. The name marks, greppably, every site making a promise the compiler cannot check.

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

Async device-to-host copy (no stream synchronization).

The counterpart of GpuBackend::copy_h2d_async, and the ONLY D2H primitive usable for a batched gather: copy_d2h and copy_d2h_on_stream both cuStreamSynchronize INSIDE the call, so an N-chunk gather pays N full stream drains. Measured cost of that shape: the SSM snapshot spill moved 66,846,720 B as 60 blocking copy_d2h calls in ~400 ms (~165 MB/s), while the mirror-image scatter (copy_h2d_async ×60 + ONE synchronize) moved the same bytes through the same host buffer in ~28 ms.

Lifetime requirement (same as copy_h2d_async): the destination buffer must remain valid, and must not be read or re-used, until the next synchronization point on this stream.

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

Async device-to-device copy (no stream synchronization).

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

Strided device-to-device 2D (pitched) copy: height rows of width_bytes, source rows spaced by src_pitch, dest rows by dst_pitch. Default = per-row copy_d2d_async loop; the CUDA backend overrides with ONE cudaMemcpy2DAsync (replaces the per-token Z-copy loop = up to num_tokens×num_ssm_layers launches/forward).

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

Begin capturing CUDA operations on stream into a graph.

All kernel launches and async copies on this stream between begin_capture and end_capture are recorded (not executed). The stream must NOT be the legacy default stream (handle 0).

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

End capture and return an instantiated graph ready for replay.

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

Replay all operations captured in the graph on stream.

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

Destroy an instantiated graph, freeing resources.

fn abort_capture_if_active(&self, _stream: u64)

Best-effort: if stream is mid graph-capture, end that capture so the stream returns to normal mode (discarding any partial graph). Call this on an error path that unwound out of a begin_capture/end_capture region (e.g. a fold refuse bailed mid-capture) — otherwise the stream is left recording and every subsequent op fails with STREAM_CAPTURE_UNSUPPORTED, bricking the server. No-op if not capturing.

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

Free device memory as the DRIVER reports it, with no host leg.

free_memory is max(cuMemGetInfo, MemAvailable) (ANOMALIES A73), so it cannot separate driver-committed device memory from reclaimable host page cache — which is exactly the separation a per-request leak measurement needs. Default falls back to free_memory for backends that have no distinct driver leg.

fn live_alloc_count(&self) -> usize

Live (allocated, not yet freed) device allocations on this backend.

A COUNT, not bytes: it answers “did this request hand back every buffer it took?” without an allocator-size ledger. Default 0 = not tracked.

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

Create a new CUDA stream (for overlapping work).

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

Bind the CUDA context to the current thread.

Must be called on any thread that uses GPU operations (alloc, launch, etc.) if it’s different from the thread that created the backend.

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

Create a CUDA event (for inter-stream synchronization).

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

Record an event on a stream (marks a point in the stream’s work).

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

Make a stream wait for an event (GPU-side sync, CPU does not block).

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

Block the calling host thread until all work already recorded against the event — e.g. an async D2H copy issued on the graph stream followed by record_event, then event_synchronize right before the host dereferences the destination pinned buffer. Cheaper than synchronize(stream) when the stream has work beyond the event you care about: this only waits for the recorded point, not for everything subsequently enqueued.

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

Destroy an event.

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

Device-side alias of a page-locked host pointer from Self::alloc_host_pinned (cuMemHostGetDevicePointer). On UMA parts (GB10) this lets a KERNEL write results directly into host-visible memory, eliminating the copy-engine op for tiny readbacks entirely. Default: unsupported.

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

Allocate page-locked (pinned) host memory for efficient async H2D.

On DGX Spark (UMA/LPDDR5X), pinned memory enables true async DMA without internal CUDA staging overhead. Small metadata buffers should be packed into a single pinned region and copied in one call.

Returns a raw pointer to bytes of page-locked host memory. Caller must call free_host_pinned to release.

The returned region is ZEROED. Callers pack these buffers with alignment padding between fields and then form a &[u8] over the whole packed range for one copy_h2d; a slice over a never-written byte is UB no matter what the device later does with it. Every implementation must uphold this — cuMemAllocHost_v2 and newBufferWithLength do not zero on their own and their wrappers memset explicitly.

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

Free page-locked host memory previously allocated by alloc_host_pinned.

Trait Implementations

impl ModelResource<dyn GpuBackend> for BufferArena

Release every buffer this arena owns.

The destructure below is exhaustive on purpose — no ... A buffer added to BufferArena without a matching free is a leak that only shows up as the next model failing to fit, so the compiler is made to refuse the addition instead. If this line stops compiling, the fix is to free the new field, not to add a wildcard.

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

Human name, for the teardown report and for attributing a failure.

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

Release everything this owns. Must be idempotent: the host calls it, and a Drop backstop may call it again.

impl ModelResource<dyn GpuBackend> for OpCache

Release the scratch allocations.

The kernel handles are not freed here: they are module-scoped and die with the AtlasRegistry the backend holds, which cuda_host::release unloads once every handle to it is gone. Freeing them here would be a double-unload.

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

Human name, for the teardown report and for attributing a failure.

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

Release everything this owns. Must be idempotent: the host calls it, and a Drop backstop may call it again.

impl ModelResource<dyn GpuBackend> for PagedKvCache

Release both pools of every layer.

Each layer allocates its K and V pools separately, so freeing per layer is correct. The block bookkeeping (free_blocks, block_ref_counts) is host state indexing into those pools — cleared with them so a released cache cannot hand out a block into freed memory.

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

Human name, for the teardown report and for attributing a failure.

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

Release everything this owns. Must be idempotent: the host calls it, and a Drop backstop may call it again.

impl ModelResource<dyn GpuBackend> for WeightStore

Release every weight tensor.

Safe to free per-entry because the loaders allocate per-tensor: the fast path calls gpu.alloc(meta.len) once per tensor before inserting it (fast_weights/mod.rs:360-388), and no loader inserts an .offset() view of a shared block into this map. (Fused per-expert views DO exist — see weight_loader/step3p7.rs:93 — but they live in the layer structs that own the fused allocation, not here, so this cannot double-free them.)

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

Human name, for the teardown report and for attributing a failure.

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

Release everything this owns. Must be idempotent: the host calls it, and a Drop backstop may call it again.

Implementors

impl GpuBackend for AtlasCudaBackend