cubecl_runtime/config/
memory.rs

1use super::logger::{LogLevel, LoggerConfig};
2use super::size::MemorySize;
3use alloc::vec::Vec;
4
5/// Configuration for memory settings in `CubeCL`.
6///
7/// Unknown fields are rejected so a leftover `pools` entry (now a programmatic
8/// setting, see [`MemoryPoolsConfig`]) or a misspelled option is a load error
9/// rather than a silently dropped setting.
10#[derive(Clone, Debug, serde::Serialize, serde::Deserialize, Default)]
11#[serde(deny_unknown_fields)]
12pub struct MemoryConfig {
13    /// Logger configuration for memory-related logs, using specific log levels.
14    #[serde(default)]
15    pub logger: LoggerConfig<MemoryLogLevel>,
16    /// Configuration for persistent memory pools.
17    #[serde(default)]
18    pub persistent_memory: PersistentMemory,
19}
20
21/// A pool layout override for a runtime's **main GPU** memory: a preset or an
22/// explicit list of pool entries, tried in order at allocation time (the first
23/// pool that accepts an allocation's size serves it).
24///
25/// This is a **programmatic** setting, deliberately not a config-file one —
26/// pool layouts are dynamic (e.g. resized per model just before a load) and
27/// must not freeze at startup. Apply it with
28/// [`install_memory_pools`](crate::client::ComputeClient::install_memory_pools):
29/// it rebuilds the calling stream's pools in place and becomes the layout for
30/// streams created afterwards. Auxiliary pools (pinned CPU, staging, uniforms)
31/// are never affected.
32#[derive(Clone, Debug, PartialEq)]
33pub enum MemoryPoolsConfig {
34    /// A named preset matching the runtime-level presets.
35    Preset(MemoryPoolsPreset),
36    /// An explicit pool list, mirroring
37    /// [`MemoryPoolOptions`](crate::memory_management::MemoryPoolOptions).
38    Explicit(Vec<MemoryPoolConfig>),
39}
40
41/// The presets of [`MemoryConfiguration`](crate::memory_management::MemoryConfiguration).
42#[derive(Clone, Copy, Debug, PartialEq, Eq)]
43pub enum MemoryPoolsPreset {
44    /// The runtime's `SubSlices` preset: a ladder of size-bucketed pools that
45    /// sub-slice large pages.
46    SubSlices,
47    /// The runtime's `ExclusivePages` preset: one page per allocation, in
48    /// exponentially spaced size buckets.
49    ExclusivePages,
50}
51
52/// One pool entry; mirrors [`MemoryPoolOptions`](crate::memory_management::MemoryPoolOptions)
53/// and [`PoolType`](crate::memory_management::PoolType).
54///
55/// Sizes are aligned up to the device alignment.
56#[derive(Clone, Debug, PartialEq)]
57pub enum MemoryPoolConfig {
58    /// Every allocation gets its own page
59    /// ([`PoolType::ExclusivePages`](crate::memory_management::PoolType::ExclusivePages)).
60    Exclusive {
61        /// Largest allocation this pool accepts. `0` is valid: it makes a pool
62        /// dedicated to zero-sized (sub-alignment) allocations.
63        max_alloc_size: MemorySize,
64        /// Period (in parent allocation count) after which unused pages are
65        /// deallocated. `None` never deallocates.
66        dealloc_period: Option<u64>,
67    },
68    /// Every allocation is its own device allocation, reused by exact size
69    /// ([`PoolType::Direct`](crate::memory_management::PoolType::Direct)).
70    ///
71    /// Wastes only alignment padding, at the cost of a device allocation per
72    /// distinct size rather than per page.
73    Direct {
74        /// Reserved bytes above which the pool returns free slices to the
75        /// driver, releasing just enough for the allocation that crossed it.
76        ///
77        /// A watermark, not a budget: an allocation that still does not fit
78        /// once everything free is gone is served anyway. `None` never
79        /// reclaims on its own, leaving it to an explicit cleanup.
80        reclaim_at: Option<MemorySize>,
81    },
82    /// Allocations are slices of larger pages
83    /// ([`PoolType::SlicedPages`](crate::memory_management::PoolType::SlicedPages)).
84    Sliced {
85        /// Size of each page.
86        page_size: MemorySize,
87        /// Largest slice this pool accepts. Defaults to `page_size`.
88        max_slice_size: Option<MemorySize>,
89        /// Hard cap on the pool's total reserved bytes: exceeding it is an
90        /// error instead of silent growth. `None` grows unbounded. Note:
91        /// runtimes that create one memory management per stream (CUDA, HIP)
92        /// apply the cap per stream.
93        max_pool_size: Option<MemorySize>,
94        /// Period (in parent allocation count) after which unused pages are
95        /// deallocated. `None` never deallocates.
96        dealloc_period: Option<u64>,
97    },
98}
99
100/// Configuration options for persistent memory pools in `CubeCL` runtimes.
101#[derive(Clone, Debug, serde::Serialize, serde::Deserialize, Default)]
102pub enum PersistentMemory {
103    /// Persistent memory is enabled but used only when explicitly specified.
104    #[default]
105    #[serde(rename = "enabled")]
106    Enabled,
107    /// Like `enabled`, and automatic allocations whose size matches an
108    /// existing persistent bucket are also served from the persistent pool.
109    ///
110    /// A good heuristic for training, where the recurring allocations are
111    /// weight-shaped and updated within the same pool. Less suited to
112    /// inference: activations that happen to match a weight size get pulled
113    /// into exact-sized persistent slices.
114    #[serde(rename = "size-match")]
115    SizeMatch,
116    /// Persistent memory is disabled, allowing dynamic allocations.
117    #[serde(rename = "disabled")]
118    Disabled,
119    /// Persistent memory is enforced, preventing dynamic allocations.
120    ///
121    /// # Warning
122    ///
123    /// Enforcing persistent memory may cause out-of-memory errors if tensors of varying sizes are used.
124    #[serde(rename = "enforced")]
125    Enforced,
126}
127
128/// Log levels for memory-related events in `CubeCL`.
129#[derive(Default, Clone, Copy, Debug, serde::Serialize, serde::Deserialize)]
130pub enum MemoryLogLevel {
131    /// No memory-related logging.
132    #[default]
133    #[serde(rename = "disabled")]
134    Disabled,
135    /// Logs basic memory events, such as creating memory pages and manually cleaning memory.
136    #[serde(rename = "basic")]
137    Basic,
138    /// Logs detailed memory information.
139    #[serde(rename = "full")]
140    Full,
141}
142
143impl LogLevel for MemoryLogLevel {}
144
145#[cfg(all(test, feature = "std"))]
146mod tests {
147    use super::*;
148
149    #[test]
150    fn pools_rejected_in_config_files() {
151        // Pool layouts are a programmatic setting; a leftover `pools` entry in
152        // a config file must be a load error, not a silently ignored setting.
153        assert!(toml::from_str::<MemoryConfig>("pools = \"sub-slices\"").is_err());
154        assert!(
155            toml::from_str::<MemoryConfig>("[[pools]]\ntype = \"sliced\"\npage_size = \"1MiB\"\n")
156                .is_err()
157        );
158    }
159}
cubecl_runtime/config/memory.rs

cubecl_runtime/config/
memory.rs
