spark_runtime/cuda_backend/
alloc_ledger.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! The device-allocation ledger: what is live, how big it is, and which call
4//! site asked for it.
5//!
6//! On GB10 every `cuMemAlloc` consumes host RAM, so an allocation outside the
7//! util pledge is how the box ends up in swap — and before this existed the
8//! KV budget inferred "Atlas-own" bytes from a free-memory delta, which counts
9//! a co-tenant's pages as ours. The ledger replaces that inference with a
10//! measurement.
11//!
12//! Split from `cuda_backend.rs` for the 500-LoC cap, which the file crossed
13//! when this landed. Exact piecewise copy, with one correction: the doc
14//! comment above `record_alloc` described the `registry` field, not the
15//! method, on main.
16
17use super::AtlasCudaBackend;
18
19/// One live device allocation: how big it is and who asked for it.
20///
21/// `site` is the CALLER of `GpuBackend::alloc`, not this file, because both
22/// allocating methods are `#[track_caller]`. Costs one pointer per live
23/// allocation and nothing at all per kernel launch — allocation is a
24/// load-time event, not a hot-path one.
25#[derive(Clone, Copy)]
26pub(super) struct AllocRecord {
27    pub(super) bytes: usize,
28    pub(super) site: &'static std::panic::Location<'static>,
29}
30
31impl AtlasCudaBackend {
32    /// Enter an allocation in the ledger. `site` is the CALLER of
33    /// `GpuBackend::alloc` (both allocating methods are `#[track_caller]`),
34    /// which is what makes `alloc_report` name a file rather than this one.
35    pub(crate) fn record_alloc(
36        &self,
37        ptr: crate::gpu::DevicePtr,
38        bytes: usize,
39        site: &'static std::panic::Location<'static>,
40    ) {
41        self.live_allocs
42            .lock()
43            .insert(ptr.0, AllocRecord { bytes, site });
44    }
45
46    pub(crate) fn forget_alloc(&self, ptr: crate::gpu::DevicePtr) {
47        self.live_allocs.lock().remove(&ptr.0);
48    }
49
50    /// Total live device bytes this backend has allocated and not freed.
51    pub fn live_bytes(&self) -> usize {
52        self.live_allocs.lock().values().map(|r| r.bytes).sum()
53    }
54
55    /// Human-readable attribution of live device memory, biggest site first.
56    ///
57    /// Aggregated by allocating call site rather than by pointer: one site
58    /// looping over 48 SSM layers is one line reading 9.7 GB across 48
59    /// allocations, which is the shape that makes an over-sized pool obvious.
60    /// Sites below `min_mb` are folded into a remainder line so the report
61    /// stays readable while still summing to the true total.
62    pub fn alloc_report(&self, top_n: usize, min_mb: usize) -> String {
63        use std::collections::HashMap;
64        let mut by_site: HashMap<String, (usize, usize)> = HashMap::new();
65        let mut total = 0usize;
66        for rec in self.live_allocs.lock().values() {
67            total += rec.bytes;
68            let key = format!("{}:{}", rec.site.file(), rec.site.line());
69            let e = by_site.entry(key).or_insert((0, 0));
70            e.0 += rec.bytes;
71            e.1 += 1;
72        }
73        let mut rows: Vec<(String, usize, usize)> =
74            by_site.into_iter().map(|(k, v)| (k, v.0, v.1)).collect();
75        rows.sort_by(|a, b| b.1.cmp(&a.1));
76
77        let mut out = format!(
78            "GPU allocation ledger: {:.2} GB live across {} sites\n",
79            total as f64 / 1e9,
80            rows.len()
81        );
82        let mut shown = 0usize;
83        let mut folded_bytes = 0usize;
84        let mut folded_sites = 0usize;
85        for (site, bytes, count) in rows {
86            if shown < top_n && bytes >= min_mb * 1024 * 1024 {
87                out.push_str(&format!(
88                    "  {:>9.1} MB  x{:<5} {}\n",
89                    bytes as f64 / (1024.0 * 1024.0),
90                    count,
91                    site
92                ));
93                shown += 1;
94            } else {
95                folded_bytes += bytes;
96                folded_sites += 1;
97            }
98        }
99        if folded_sites > 0 {
100            out.push_str(&format!(
101                "  {:>9.1} MB  across {} smaller sites\n",
102                folded_bytes as f64 / (1024.0 * 1024.0),
103                folded_sites
104            ));
105        }
106
107        // Per-FILE rollup. The by-site view above has a blind spot: a
108        // subsystem that allocates many distinct buffers from many distinct
109        // lines is split into pieces small enough to fall below the cut and
110        // vanish into the remainder. The vision encoder is exactly that shape
111        // — ~16 buffers (scores, probs, qr/kr/vt, merge, rope, ...) each from
112        // its own line — so ~2.2 GB was invisible in a top-12 by site while
113        // being the fourth-largest consumer in the process. Rolling up by
114        // file costs one more pass over the same map and makes a subsystem
115        // legible as a subsystem.
116        let mut by_file: HashMap<&str, (usize, usize)> = HashMap::new();
117        for rec in self.live_allocs.lock().values() {
118            let e = by_file.entry(rec.site.file()).or_insert((0, 0));
119            e.0 += rec.bytes;
120            e.1 += 1;
121        }
122        let mut frows: Vec<(&str, usize, usize)> =
123            by_file.into_iter().map(|(k, v)| (k, v.0, v.1)).collect();
124        frows.sort_by(|a, b| b.1.cmp(&a.1));
125        out.push_str("  ── by file ──\n");
126        for (file, bytes, count) in frows.into_iter().take(top_n) {
127            out.push_str(&format!(
128                "  {:>9.1} MB  x{:<5} {}\n",
129                bytes as f64 / (1024.0 * 1024.0),
130                count,
131                file
132            ));
133        }
134        out
135    }
136}
spark_runtime/cuda_backend/alloc_ledger.rs

spark_runtime/cuda_backend/
alloc_ledger.rs
