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    /// CUDA-specific memory settings. Ignored by other backends.
18    #[serde(default)]
19    pub cuda: CudaMemoryConfig,
20}
21
22/// Memory settings for the CUDA runtime (`[memory.cuda]` in `cubecl.toml`).
23#[derive(Clone, Debug, Default, serde::Serialize, serde::Deserialize)]
24#[serde(deny_unknown_fields)]
25pub struct CudaMemoryConfig {
26    /// Allocator backing `CubeCL`'s pools, selected when the runtime initializes.
27    #[serde(default)]
28    pub allocator: CudaAllocator,
29}
30
31/// CUDA allocator used to obtain backing memory for `CubeCL`'s pools.
32#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
33#[serde(rename_all = "lowercase")]
34pub enum CudaAllocator {
35    /// Allocate and free synchronously, avoiding an additional CUDA memory pool.
36    /// Pool growth and cleanup may synchronize GPU work.
37    #[default]
38    Sync,
39    /// Use CUDA's stream-ordered allocator, which may reserve extra backing memory
40    /// but can preserve concurrency during pool growth and cleanup. Falls back to
41    /// synchronous allocation if an asynchronous allocation fails.
42    Async,
43}
44
45/// Configuration options for persistent memory pools in `CubeCL` runtimes.
46#[derive(Clone, Debug, serde::Serialize, serde::Deserialize, Default)]
47pub enum PersistentMemory {
48    /// Persistent memory is enabled but used only when explicitly specified.
49    #[default]
50    #[serde(rename = "enabled")]
51    Enabled,
52    /// Like `enabled`, and automatic allocations whose size matches an
53    /// existing persistent bucket are also served from the persistent pool.
54    ///
55    /// A good heuristic for training, where the recurring allocations are
56    /// weight-shaped and updated within the same pool. Less suited to
57    /// inference: activations that happen to match a weight size get pulled
58    /// into exact-sized persistent slices.
59    #[serde(rename = "size-match")]
60    SizeMatch,
61    /// Persistent memory is disabled, allowing dynamic allocations.
62    #[serde(rename = "disabled")]
63    Disabled,
64    /// Persistent memory is enforced, preventing dynamic allocations.
65    ///
66    /// # Warning
67    ///
68    /// Enforcing persistent memory may cause out-of-memory errors if tensors of varying sizes are used.
69    #[serde(rename = "enforced")]
70    Enforced,
71}
72
73/// Log levels for memory-related events in `CubeCL`.
74#[derive(Default, Clone, Copy, Debug, serde::Serialize, serde::Deserialize)]
75pub enum MemoryLogLevel {
76    /// No memory-related logging.
77    #[default]
78    #[serde(rename = "disabled")]
79    Disabled,
80    /// Logs basic memory events, such as creating memory pages and manually cleaning memory.
81    #[serde(rename = "basic")]
82    Basic,
83    /// Logs detailed memory information.
84    #[serde(rename = "full")]
85    Full,
86}
87
88impl LogLevel for MemoryLogLevel {}
89
90#[cfg(all(test, feature = "std"))]
91mod tests {
92    use super::*;
93
94    #[test]
95    fn cuda_allocator_config() {
96        use crate::config::CubeClRuntimeConfig;
97
98        let config: CubeClRuntimeConfig = toml::from_str("").unwrap();
99        assert_eq!(config.memory.cuda.allocator, CudaAllocator::Sync);
100
101        let config: CubeClRuntimeConfig =
102            toml::from_str("[memory.cuda]\nallocator = \"async\"").unwrap();
103        assert_eq!(config.memory.cuda.allocator, CudaAllocator::Async);
104    }
105
106    #[test]
107    fn pools_rejected_in_config_files() {
108        // Pool layouts are a programmatic setting; a leftover `pools` entry in
109        // a config file must be a load error, not a silently ignored setting.
110        assert!(toml::from_str::<MemoryConfig>("pools = \"sub-slices\"").is_err());
111        assert!(
112            toml::from_str::<MemoryConfig>("[[pools]]\ntype = \"sliced\"\npage_size = \"1MiB\"\n")
113                .is_err()
114        );
115    }
116}
cubecl_runtime/config/memory.rs

cubecl_runtime/config/
memory.rs
