cubecl_runtime/config/
memory.rs

1use super::logger::{LogLevel, LoggerConfig};
2
3/// Configuration for memory settings in `CubeCL`.
4///
5/// Unknown fields are rejected so a leftover `pools` entry — the pools are the
6/// memory management's to lay out, not a setting — or a misspelled option is a
7/// load error rather than a silently dropped setting.
8#[derive(Clone, Debug, serde::Serialize, serde::Deserialize, Default)]
9#[serde(deny_unknown_fields)]
10pub struct MemoryConfig {
11    /// Logger configuration for memory-related logs, using specific log levels.
12    #[serde(default)]
13    pub logger: LoggerConfig<MemoryLogLevel>,
14    /// Configuration for persistent memory pools.
15    #[serde(default)]
16    pub persistent_memory: PersistentMemory,
17}
18
19/// Configuration options for persistent memory pools in `CubeCL` runtimes.
20#[derive(Clone, Debug, serde::Serialize, serde::Deserialize, Default)]
21pub enum PersistentMemory {
22    /// Persistent memory is enabled but used only when explicitly specified.
23    #[default]
24    #[serde(rename = "enabled")]
25    Enabled,
26    /// Like `enabled`, and automatic allocations whose size matches an
27    /// existing persistent bucket are also served from the persistent pool.
28    ///
29    /// A good heuristic for training, where the recurring allocations are
30    /// weight-shaped and updated within the same pool. Less suited to
31    /// inference: activations that happen to match a weight size get pulled
32    /// into exact-sized persistent slices.
33    #[serde(rename = "size-match")]
34    SizeMatch,
35    /// Persistent memory is disabled, allowing dynamic allocations.
36    #[serde(rename = "disabled")]
37    Disabled,
38    /// Persistent memory is enforced, preventing dynamic allocations.
39    ///
40    /// # Warning
41    ///
42    /// Enforcing persistent memory may cause out-of-memory errors if tensors of varying sizes are used.
43    #[serde(rename = "enforced")]
44    Enforced,
45}
46
47/// Log levels for memory-related events in `CubeCL`.
48#[derive(Default, Clone, Copy, Debug, serde::Serialize, serde::Deserialize)]
49pub enum MemoryLogLevel {
50    /// No memory-related logging.
51    #[default]
52    #[serde(rename = "disabled")]
53    Disabled,
54    /// Logs basic memory events, such as creating memory pages and manually cleaning memory.
55    #[serde(rename = "basic")]
56    Basic,
57    /// Logs detailed memory information.
58    #[serde(rename = "full")]
59    Full,
60}
61
62impl LogLevel for MemoryLogLevel {}
63
64#[cfg(all(test, feature = "std"))]
65mod tests {
66    use super::*;
67
68    #[test]
69    fn pools_rejected_in_config_files() {
70        // Pool layouts are a programmatic setting; a leftover `pools` entry in
71        // a config file must be a load error, not a silently ignored setting.
72        assert!(toml::from_str::<MemoryConfig>("pools = \"sub-slices\"").is_err());
73        assert!(
74            toml::from_str::<MemoryConfig>("[[pools]]\ntype = \"sliced\"\npage_size = \"1MiB\"\n")
75                .is_err()
76        );
77    }
78}
cubecl_runtime/config/memory.rs

cubecl_runtime/config/
memory.rs
