metrale_gpu_runtime/cuda_backend/
alloc_ledger.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: The device-allocation ledger's entries and reports: size and
4//! allocating call site of every live allocation.
5//! 2026-10-01: Also [`LedgerProbe`], which reads the ledger without the backend, and
6//! [`chunk_slack`]: the device memory the driver holds beyond the ledger for small allocations.
7//!
8//! Owner: gpu-runtime (CUDA backend).
9//! Invariants: none beyond the types.
10
11use super::MetraleCudaBackend;
12
13/// 2026-09-25: One live device allocation: its size in bytes and the call
14/// site of `GpuBackend::alloc` or `alloc_managed` that made it (both are
15/// `#[track_caller]` on the trait and on the CUDA impl).
16#[derive(Clone, Copy)]
17pub(super) struct AllocRecord {
18    pub(super) bytes: usize,
19    pub(super) site: &'static std::panic::Location<'static>,
20}
21
22/// 2026-10-01: A read handle on one backend's ledger, for a reader that cannot hold the
23/// backend, which moves into the model it builds. It reads the map
24/// [`MetraleCudaBackend::live_bytes`] reads, so the two always agree.
25#[derive(Clone)]
26pub struct LedgerProbe(
27    std::sync::Arc<parking_lot::Mutex<std::collections::HashMap<u64, AllocRecord>>>,
28);
29
30impl LedgerProbe {
31    /// 2026-10-01: Total bytes on the ledger now.
32    pub fn live_bytes(&self) -> usize {
33        ledger_bytes(&self.0.lock())
34    }
35}
36
37fn ledger_bytes(ledger: &std::collections::HashMap<u64, AllocRecord>) -> usize {
38    ledger.values().map(|r| r.bytes).sum()
39}
40
41/// 2026-10-01: The chunk the driver packs allocations smaller than itself into, on GB10: such
42/// allocations come back at 512 B to 64 KiB alignment inside 2 MiB-aligned ranges, while larger
43/// ones start 2 MiB-aligned (nsys device-memory events of two serves, 2026-10-01). A chunk that
44/// holds any live small allocation stays resident whole.
45pub const SMALL_ALLOC_CHUNK: u64 = 2 << 20;
46
47/// 2026-10-01: The bytes of the `chunk`-sized ranges that hold live allocations smaller than
48/// `chunk`, less those allocations' bytes: memory the driver keeps resident that no ledger entry
49/// accounts for. `allocs` is `(device address, bytes)`; allocations of `chunk` bytes or more are
50/// not packed and count nothing. On GB10 this is 52 MiB for the dense 27B (2,274 small
51/// allocations in 270 chunks) and 429 MiB for Nemotron-3-Nano (24,456 in 2,033 chunks), which
52/// matches the latter's driver use above the other reserve terms to within 11 MiB.
53pub fn chunk_slack(allocs: impl Iterator<Item = (u64, usize)>, chunk: u64) -> usize {
54    let mut chunks = std::collections::BTreeSet::new();
55    let mut small = 0usize;
56    for (addr, bytes) in allocs.filter(|&(_, b)| (b as u64) < chunk && b > 0) {
57        small += bytes;
58        let last = addr + bytes as u64 - 1;
59        chunks.extend(addr / chunk..=last / chunk);
60    }
61    (chunks.len() as u64 * chunk).saturating_sub(small as u64) as usize
62}
63
64impl MetraleCudaBackend {
65    /// 2026-09-25: Enter an allocation in the ledger. `site` is the caller of
66    /// `GpuBackend::alloc` or `alloc_managed`.
67    pub(crate) fn record_alloc(
68        &self,
69        ptr: crate::gpu::DevicePtr,
70        bytes: usize,
71        site: &'static std::panic::Location<'static>,
72    ) {
73        self.live_allocs
74            .lock()
75            .insert(ptr.0, AllocRecord { bytes, site });
76    }
77
78    /// 2026-10-02: Every live allocation: address, bytes and call site (`file:line`).
79    pub fn live_allocations(&self) -> Vec<(crate::gpu::DevicePtr, usize, String)> {
80        self.live_allocs
81            .lock()
82            .iter()
83            .map(|(&a, r)| {
84                let site = format!("{}:{}", r.site.file(), r.site.line());
85                (crate::gpu::DevicePtr(a), r.bytes, site)
86            })
87            .collect()
88    }
89
90    pub(crate) fn forget_alloc(&self, ptr: crate::gpu::DevicePtr) {
91        self.live_allocs.lock().remove(&ptr.0);
92    }
93
94    /// 2026-09-25: Total bytes on the ledger.
95    pub fn live_bytes(&self) -> usize {
96        ledger_bytes(&self.live_allocs.lock())
97    }
98
99    /// 2026-10-01: `chunk_slack` of the live allocations, at `SMALL_ALLOC_CHUNK`.
100    pub fn chunk_slack_bytes(&self) -> usize {
101        let ledger = self.live_allocs.lock();
102        chunk_slack(ledger.iter().map(|(&a, r)| (a, r.bytes)), SMALL_ALLOC_CHUNK)
103    }
104
105    /// 2026-10-01: A [`LedgerProbe`] on this backend's ledger.
106    pub fn ledger_probe(&self) -> LedgerProbe {
107        LedgerProbe(std::sync::Arc::clone(&self.live_allocs))
108    }
109
110    /// 2026-09-25: A text report of the ledger: the total, then up to `top_n`
111    /// call sites of at least `min_mb` MiB, largest first, then (if any are
112    /// left) one line summing the other sites, then up to `top_n` source
113    /// files, largest first.
114    pub fn alloc_report(&self, top_n: usize, min_mb: usize) -> String {
115        use std::collections::HashMap;
116        let mut by_site: HashMap<String, (usize, usize)> = HashMap::new();
117        let mut total = 0usize;
118        for rec in self.live_allocs.lock().values() {
119            total += rec.bytes;
120            let key = format!("{}:{}", rec.site.file(), rec.site.line());
121            let e = by_site.entry(key).or_insert((0, 0));
122            e.0 += rec.bytes;
123            e.1 += 1;
124        }
125        let mut rows: Vec<(String, usize, usize)> =
126            by_site.into_iter().map(|(k, v)| (k, v.0, v.1)).collect();
127        rows.sort_by(|a, b| b.1.cmp(&a.1));
128
129        let mut out = format!(
130            "GPU allocation ledger: {:.2} GB live across {} sites\n",
131            total as f64 / 1e9,
132            rows.len()
133        );
134        let mut shown = 0usize;
135        let mut folded_bytes = 0usize;
136        let mut folded_sites = 0usize;
137        for (site, bytes, count) in rows {
138            if shown < top_n && bytes >= min_mb * 1024 * 1024 {
139                out.push_str(&format!(
140                    "  {:>9.1} MB  x{:<5} {}\n",
141                    bytes as f64 / (1024.0 * 1024.0),
142                    count,
143                    site
144                ));
145                shown += 1;
146            } else {
147                folded_bytes += bytes;
148                folded_sites += 1;
149            }
150        }
151        if folded_sites > 0 {
152            out.push_str(&format!(
153                "  {:>9.1} MB  across {} smaller sites\n",
154                folded_bytes as f64 / (1024.0 * 1024.0),
155                folded_sites
156            ));
157        }
158
159        // 2026-09-25: Per-file rollup: a file that allocates from many lines,
160        // each under the cut above, still shows as one entry here.
161        let mut by_file: HashMap<&str, (usize, usize)> = HashMap::new();
162        for rec in self.live_allocs.lock().values() {
163            let e = by_file.entry(rec.site.file()).or_insert((0, 0));
164            e.0 += rec.bytes;
165            e.1 += 1;
166        }
167        let mut frows: Vec<(&str, usize, usize)> =
168            by_file.into_iter().map(|(k, v)| (k, v.0, v.1)).collect();
169        frows.sort_by(|a, b| b.1.cmp(&a.1));
170        out.push_str("  ── by file ──\n");
171        for (file, bytes, count) in frows.into_iter().take(top_n) {
172            out.push_str(&format!(
173                "  {:>9.1} MB  x{:<5} {}\n",
174                bytes as f64 / (1024.0 * 1024.0),
175                count,
176                file
177            ));
178        }
179        out
180    }
181}
182
183#[cfg(test)]
184#[path = "alloc_ledger_tests.rs"]
185mod tests;
metrale_gpu_runtime/cuda_backend/alloc_ledger.rs

metrale_gpu_runtime/cuda_backend/
alloc_ledger.rs
