metrale_gpu_runtime/
cuda_backend.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: The CUDA `GpuBackend`: [`MetraleCudaBackend`], the driver API
4//! declarations its submodules call, and its device-allocation ledger.
5//!
6//! Owner: gpu-runtime (CUDA backend).
7//! Invariants:
8//! - A pointer that `alloc` or `alloc_managed` returns is on `live_allocs`
9//!   until `free` or `sweep_unreleased` takes it off; `free` takes it off
10//!   before calling `cuMemFree_v2`, whatever that call returns.
11//! - `alloc_arena` allocations are never on `live_allocs`.
12
13use std::ffi::c_void;
14
15use anyhow::{Result, bail};
16use std::sync::Arc;
17
18use crate::registry::MetraleRegistry;
19
20pub mod arch_preflight;
21mod fault_probe;
22mod gpu_copy;
23mod gpu_impl;
24mod gpu_impl_graph;
25mod n_tile;
26mod redzone;
27pub mod tensormap;
28pub(crate) use redzone::{
29    ALLOC_SEQ, RedZone, redzone_bytes, redzone_fill, redzone_min_idx, redzone_trace_idx,
30};
31mod memory;
32pub use memory::{cuda_free_memory_bytes, spawn_oom_watchdog};
33pub(crate) use memory::{device_is_integrated, effective_free_bytes};
34
35unsafe extern "C" {
36    pub(super) fn cuMemAlloc_v2(dptr: *mut u64, bytesize: usize) -> i32;
37    pub(super) fn cuMemFree_v2(dptr: u64) -> i32;
38    pub(super) fn cuMemcpyHtoDAsync_v2(
39        dst: u64,
40        src: *const c_void,
41        bytes: usize,
42        stream: u64,
43    ) -> i32;
44    pub(super) fn cuMemcpyDtoHAsync_v2(
45        dst: *mut c_void,
46        src: u64,
47        bytes: usize,
48        stream: u64,
49    ) -> i32;
50    pub(super) fn cuMemcpyDtoDAsync_v2(dst: u64, src: u64, bytes: usize, stream: u64) -> i32;
51    pub(super) fn cuStreamSynchronize(stream: u64) -> i32;
52    pub(super) fn cuStreamQuery(stream: u64) -> i32;
53    pub(super) fn cuMemHostGetDevicePointer_v2(
54        dptr: *mut u64,
55        host: *mut std::ffi::c_void,
56        flags: u32,
57    ) -> i32;
58    pub(super) fn cuMemGetInfo_v2(free: *mut usize, total: *mut usize) -> i32;
59    /// 2026-09-25: The device of the calling thread's current context.
60    pub(super) fn cuCtxGetDevice(device: *mut i32) -> i32;
61    /// 2026-09-25: The `CUdevice` for a device ordinal, with no current
62    /// context needed on the calling thread; `arch_preflight` uses it for
63    /// that reason.
64    pub(super) fn cuDeviceGet(device: *mut i32, ordinal: i32) -> i32;
65    pub(super) fn cuDeviceGetAttribute(pi: *mut i32, attrib: u32, dev: i32) -> i32;
66    pub(super) fn cuMemsetD8Async(dst: u64, value: u8, n: usize, stream: u64) -> i32;
67    /// 2026-09-25: Synchronous variants, used only by the red-zone diagnostic
68    /// (`redzone.rs`).
69    pub(super) fn cuMemsetD8_v2(dst: u64, value: u8, n: usize) -> i32;
70    pub(super) fn cuMemcpyDtoH_v2(dst: *mut c_void, src: u64, bytes: usize) -> i32;
71    pub(super) fn cuStreamBeginCapture(hStream: u64, mode: u32) -> i32;
72    // 2026-09-25: Not declared under `metrale_scale`; there
73    // `stream_is_capturing` reports false without asking the driver.
74    #[cfg(not(metrale_scale))]
75    pub(super) fn cuStreamIsCapturing(hStream: u64, captureStatus: *mut u32) -> i32;
76    pub(super) fn cuStreamEndCapture(hStream: u64, phGraph: *mut u64) -> i32;
77    // 2026-09-25: Graph instantiate: `cuGraphInstantiate` under `metrale_scale`
78    // (build.rs sets it when METRALE_TARGET_HW starts with `strix`),
79    // `cuGraphInstantiateWithFlags` otherwise. Both take
80    // `(CUgraphExec*, CUgraph, u64)`.
81    #[cfg(not(metrale_scale))]
82    pub(super) fn cuGraphInstantiateWithFlags(
83        phGraphExec: *mut u64,
84        hGraph: u64,
85        flags: u64,
86    ) -> i32;
87    #[cfg(metrale_scale)]
88    pub(super) fn cuGraphInstantiate(phGraphExec: *mut u64, hGraph: u64, flags: u64) -> i32;
89    pub(super) fn cuGraphLaunch(hGraphExec: u64, hStream: u64) -> i32;
90    pub(super) fn cuGraphExecDestroy(hGraphExec: u64) -> i32;
91    pub(super) fn cuGraphDestroy(hGraph: u64) -> i32;
92    fn cuCtxGetCurrent(pctx: *mut u64) -> i32;
93    pub(super) fn cuCtxSetCurrent(ctx: u64) -> i32;
94    pub(super) fn cuStreamCreate(phStream: *mut u64, flags: u32) -> i32;
95    pub(super) fn cuMemAllocHost_v2(pp: *mut *mut c_void, bytesize: usize) -> i32;
96    pub(super) fn cuMemFreeHost(p: *mut c_void) -> i32;
97    pub(super) fn cuMemAllocManaged(dptr: *mut u64, bytesize: usize, flags: u32) -> i32;
98    pub(super) fn cuEventCreate(phEvent: *mut u64, flags: u32) -> i32;
99    pub(super) fn cuEventRecord(hEvent: u64, hStream: u64) -> i32;
100    pub(super) fn cuStreamWaitEvent(hStream: u64, hEvent: u64, flags: u32) -> i32;
101    pub(super) fn cuEventSynchronize(hEvent: u64) -> i32;
102    pub(super) fn cuEventQuery(hEvent: u64) -> i32;
103    pub(super) fn cuEventDestroy_v2(hEvent: u64) -> i32;
104}
105
106/// 2026-09-25: The CUDA `GpuBackend`, holding one model's kernel modules.
107///
108/// The modules unload when the last `Arc<MetraleRegistry>` handle drops
109/// (`MetraleRegistry`'s `Drop`); this struct holds one of those handles.
110pub struct MetraleCudaBackend {
111    /// 2026-09-25: This model's kernel modules. `registry()` and
112    /// `kernel_registry()` hand out clones of the `Arc`.
113    registry: Arc<MetraleRegistry>,
114    /// 2026-09-25: `METRALE_DEBUG_SYNC_KERNELS=1`: sync after every launch.
115    /// Read once, in `new`.
116    debug_sync_kernels: bool,
117    /// 2026-09-25: This model's kernel handles and op scratch, dropped with
118    /// the backend.
119    op_cache: crate::op_cache::OpCache,
120    /// 2026-09-25: Every device allocation `alloc` and `alloc_managed` made
121    /// and nobody freed, keyed by pointer, with its size and allocating call
122    /// site (`AllocRecord`).
123    ///
124    /// Model teardown ends with `sweep_unreleased`, which frees what is left:
125    /// allocations no `ModelResource` released, such as weights the loaders
126    /// fused into layer-owned buffers (`TransformerModel::release_pools`).
127    ///
128    /// Process-lifetime workspaces are not on it: the cuBLASLt, CUTLASS and
129    /// FlashInfer workspaces (`ctx` in `cublaslt.rs` and `cutlass.rs`,
130    /// `workspaces` in `flashinfer.rs`) call `cuMemAlloc_v2` directly, so a sweep
131    /// cannot free memory a static still points at.
132    ///
133    /// 2026-10-01: Shared, so a [`LedgerProbe`] can read it after the backend has moved into a
134    /// model.
135    live_allocs: Arc<parking_lot::Mutex<std::collections::HashMap<u64, AllocRecord>>>,
136    /// 2026-09-25: The trailing guard band of every live allocation padded by
137    /// `alloc` (`METRALE_REDZONE=<bytes>` set, creation index at least
138    /// `METRALE_REDZONE_MIN_IDX`). A zone covers
139    /// `[user_ptr + user_bytes, user_ptr + user_bytes + pad_bytes)` and is
140    /// filled with `METRALE_REDZONE_FILL` when allocated.
141    /// [`MetraleCudaBackend::scan_redzones`] reads the zones back and reports
142    /// the ones that changed: a write past the end of the buffer.
143    redzones: parking_lot::Mutex<Vec<RedZone>>,
144    /// 2026-09-25: The process CUDA host's stream (`MetraleRegistry::raw_stream`).
145    default_stream: u64,
146    /// 2026-09-25: The context current on the thread that ran `new`, after the
147    /// registry load; `bind_to_thread` makes it current on another thread.
148    cuda_ctx: u64,
149}
150
151impl MetraleCudaBackend {
152    /// 2026-09-25: Load `ptx_modules` on GPU `ordinal` and build the backend.
153    ///
154    /// The module set (a `TargetPtxSet`'s `modules`, e.g. from
155    /// `metrale_kernels::ptx_for_model`) is loaded afresh for each call; the
156    /// CUDA context and stream are the process host's, shared by every
157    /// backend (`cuda_host::host`). Fails if no context is current after the
158    /// load.
159    pub fn new(ordinal: usize, ptx_modules: &[(&'static str, &'static [u8])]) -> Result<Self> {
160        // 2026-09-25: Reset the per-run telemetry before any kernel lookup,
161        // so the kernel audit lists only this backend's lookups.
162        metrale_telemetry::run_metrics::reset_for_new_run();
163        let registry = MetraleRegistry::load(ordinal, ptx_modules)
164            .map_err(|e| anyhow::anyhow!("MetraleRegistry load failed: {e}"))?;
165        let default_stream = registry.raw_stream();
166        crate::timing::install_for_stream(default_stream);
167
168        let mut cuda_ctx: u64 = 0;
169        let status = unsafe { cuCtxGetCurrent(&mut cuda_ctx) };
170        if status != 0 || cuda_ctx == 0 {
171            bail!("cuCtxGetCurrent failed: status {status}, ctx {cuda_ctx:#x}");
172        }
173
174        tracing::info!(
175            "MetraleCudaBackend initialized on GPU {ordinal} with {} PTX modules",
176            ptx_modules.len()
177        );
178
179        Ok(Self {
180            live_allocs: Arc::new(parking_lot::Mutex::new(std::collections::HashMap::new())),
181            redzones: parking_lot::Mutex::new(Vec::new()),
182            registry,
183            debug_sync_kernels: std::env::var("METRALE_DEBUG_SYNC_KERNELS").as_deref() == Ok("1"),
184            op_cache: crate::op_cache::OpCache::new(),
185            default_stream,
186            cuda_ctx,
187        })
188    }
189
190    /// 2026-09-25: Number of allocations on the ledger; one lock, no copy.
191    pub(crate) fn live_alloc_len(&self) -> usize {
192        self.live_allocs.lock().len()
193    }
194
195    /// 2026-09-25: Drain the ledger and free every allocation on it. Returns
196    /// how many there were; when there were any, one warning gives their
197    /// total size and the five largest allocating call sites.
198    ///
199    /// `TransformerModel::release_pools` calls it last, after every
200    /// `ModelResource::release`. A freed pointer is already off the ledger
201    /// (`free` calls `forget_alloc` first), so nothing is freed twice.
202    pub fn sweep_unreleased(&self) -> usize {
203        let outstanding: Vec<(u64, AllocRecord)> = self.live_allocs.lock().drain().collect();
204        let count = outstanding.len();
205        if count > 0 {
206            let bytes: usize = outstanding.iter().map(|(_, r)| r.bytes).sum();
207            // 2026-09-25: Aggregated by call site, so a pool allocated per
208            // layer from one line is one entry.
209            let mut by_site: std::collections::HashMap<String, (usize, usize)> =
210                std::collections::HashMap::new();
211            for (_, r) in &outstanding {
212                let e = by_site
213                    .entry(format!("{}:{}", r.site.file(), r.site.line()))
214                    .or_insert((0, 0));
215                e.0 += r.bytes;
216                e.1 += 1;
217            }
218            let mut rows: Vec<_> = by_site.into_iter().collect();
219            rows.sort_by(|a, b| b.1.0.cmp(&a.1.0));
220            let top: Vec<String> = rows
221                .iter()
222                .take(5)
223                .map(|(site, (b, n))| {
224                    format!("{site} ({:.1} MB x{n})", *b as f64 / (1024.0 * 1024.0))
225                })
226                .collect();
227            tracing::warn!(
228                "sweep: {count} allocation(s) totalling {:.2} GB had no owner; \
229                 largest sites: {}",
230                bytes as f64 / 1e9,
231                top.join(", ")
232            );
233        }
234        for (raw, _) in outstanding {
235            // 2026-09-25: Not `free`: the ledger is already drained, and a
236            // failed free is logged so the loop still frees the rest.
237            let status = unsafe { cuMemFree_v2(raw) };
238            if status != 0 && !crate::registry::is_teardown_noop(status) {
239                tracing::warn!("sweep: cuMemFree failed for {raw:#x}: status {status}");
240            }
241        }
242        count
243    }
244
245    pub fn registry(&self) -> &Arc<MetraleRegistry> {
246        &self.registry
247    }
248
249    pub(crate) fn debug_sync_kernels(&self) -> bool {
250        self.debug_sync_kernels
251    }
252}
253
254/// 2026-09-25: Dropping the backend runs `sweep_unreleased`.
255///
256/// A load that fails part-way builds no model, so no teardown frees what it
257/// had allocated; this does. After a model's teardown the ledger is already
258/// drained (`TransformerModel::release_pools`), and the sweep frees nothing.
259impl Drop for MetraleCudaBackend {
260    fn drop(&mut self) {
261        let swept = self.sweep_unreleased();
262        if swept > 0 {
263            // 2026-09-25: The warning names both possible causes, an
264            // abandoned load or an allocation with no registered owner; the
265            // count alone cannot tell them apart.
266            tracing::warn!(
267                "backend drop reclaimed {swept} allocation(s) that no owner released — \
268                 expected if a load was abandoned part-way, otherwise an unregistered owner"
269            );
270        }
271    }
272}
273
274/// 2026-09-25: [`effective_free_bytes`] for a poll, where the integrated
275/// query may have failed. `None` counts as discrete, so an unanswered query
276/// can only lower the reading, never add host memory to it.
277pub(crate) fn polled_free_bytes(
278    cu_free: usize,
279    mem_available: Option<usize>,
280    integrated: Option<bool>,
281) -> usize {
282    effective_free_bytes(cu_free, mem_available, integrated.unwrap_or(false))
283}
284
285/// 2026-09-25: `MemAvailable` from `/proc/meminfo`, in bytes. `None` when the
286/// file cannot be read or has no parsable `MemAvailable` line.
287fn system_available_memory_bytes() -> Option<usize> {
288    let contents = std::fs::read_to_string("/proc/meminfo").ok()?;
289    for line in contents.lines() {
290        if line.starts_with("MemAvailable:") {
291            let kb: usize = line.split_whitespace().nth(1)?.parse().ok()?;
292            return Some(kb * 1024);
293        }
294    }
295    None
296}
297
298#[path = "cuda_backend/alloc_ledger.rs"]
299mod alloc_ledger;
300use alloc_ledger::AllocRecord;
301pub use alloc_ledger::LedgerProbe;
302
303#[cfg(test)]
304mod tests;
metrale_gpu_runtime/cuda_backend.rs

metrale_gpu_runtime/
cuda_backend.rs
