rlx_runtime/resource_budget.rs
1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3// SPDX-License-Identifier: MIT OR Apache-2.0
4// RLX — versatile ML compiler + runtime.
5//! **Centralized resource budget for large-model inference** — the single source
6//! of truth for how much a model may hold resident, applied uniformly across ALL
7//! large models (DeepSeek-V4, Kimi-K3, Llama4, GLM-MoE, …). It unifies the two
8//! knobs that were previously scattered:
9//! * **RAM budget** — how much weight memory may be resident
10//! ([`crate::memory_estimate::soft_memory_budget_bytes`]).
11//! * **Experts-per-time** — how many MoE experts stay resident per layer (the
12//! offload/paging cache size; feeds [`crate::expert_pool::ExpertPoolConfig`]
13//! and `CompiledGraph::set_moe_resident_experts`).
14//!
15//! Model crates read one [`ResourceBudget`] (from env, or overridden by a model's
16//! config) and hand it to the runtime — they never re-implement budgeting. Env:
17//! `RLX_MAX_RAM_BYTES` (bytes; else `RLX_SOFT_MEMORY_BUDGET_BYTES` / physical RAM)
18//! and `RLX_MAX_RESIDENT_EXPERTS` (per layer; else derived from the RAM budget).
19
20use crate::expert_pool::{ExpertPoolConfig, ExpertRefreshPolicy};
21use crate::memory_estimate::soft_memory_budget_bytes;
22
23/// How much a model may hold resident. `None` fields mean "unbounded / derive".
24#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
25pub struct ResourceBudget {
26 /// Max resident weight RAM (bytes). `None` → the physical-RAM-derived soft
27 /// budget ([`soft_memory_budget_bytes`]).
28 pub max_ram_bytes: Option<usize>,
29 /// Max MoE experts resident PER LAYER. `None` → derive from `max_ram_bytes`
30 /// (or all experts if RAM is unbounded).
31 pub max_resident_experts: Option<usize>,
32}
33
34impl ResourceBudget {
35 /// Unbounded: full RAM, all experts resident.
36 pub const UNBOUNDED: Self = Self {
37 max_ram_bytes: None,
38 max_resident_experts: None,
39 };
40
41 /// Read from the environment. `RLX_MAX_RAM_BYTES` and `RLX_MAX_RESIDENT_EXPERTS`
42 /// override; absent fields fall back to derivation at query time.
43 pub fn from_env() -> Self {
44 let max_ram_bytes =
45 rlx_ir::env::var("RLX_MAX_RAM_BYTES").and_then(|v| v.trim().parse::<usize>().ok());
46 let max_resident_experts = rlx_ir::env::var("RLX_MAX_RESIDENT_EXPERTS")
47 .and_then(|v| v.trim().parse::<usize>().ok());
48 Self {
49 max_ram_bytes,
50 max_resident_experts,
51 }
52 }
53
54 /// The effective RAM ceiling: explicit `max_ram_bytes`, else the soft
55 /// physical-RAM budget (`None` only when RAM is unknown and unset).
56 pub fn effective_ram_bytes(&self) -> Option<usize> {
57 self.max_ram_bytes.or_else(soft_memory_budget_bytes)
58 }
59
60 /// How many experts to keep resident PER LAYER: honor an explicit
61 /// `max_resident_experts`, else derive from the RAM ceiling by reserving
62 /// `backbone_bytes` (attention/norms/embed) and dividing the remainder by
63 /// `bytes_per_expert`. Always in `1..=num_experts`.
64 pub fn resident_experts(
65 &self,
66 num_experts: usize,
67 bytes_per_expert: usize,
68 backbone_bytes: usize,
69 ) -> usize {
70 if let Some(e) = self.max_resident_experts {
71 return e.clamp(1, num_experts.max(1));
72 }
73 match self.effective_ram_bytes() {
74 Some(ram) => {
75 let for_experts = ram.saturating_sub(backbone_bytes);
76 (for_experts / bytes_per_expert.max(1)).clamp(1, num_experts.max(1))
77 }
78 None => num_experts,
79 }
80 }
81
82 /// Build an [`ExpertPoolConfig`] honoring this budget (offload kicks in when the
83 /// resident count is below `num_experts`).
84 pub fn expert_pool_config(
85 &self,
86 num_experts: usize,
87 bytes_per_expert: usize,
88 backbone_bytes: usize,
89 refresh: ExpertRefreshPolicy,
90 ) -> ExpertPoolConfig {
91 ExpertPoolConfig::new(
92 num_experts,
93 self.resident_experts(num_experts, bytes_per_expert, backbone_bytes),
94 refresh,
95 )
96 }
97
98 /// Whether weight offload/streaming should engage: the model doesn't fit the
99 /// RAM ceiling, or an explicit expert cap is below the expert count.
100 pub fn needs_offload(&self, model_bytes: usize, num_experts: usize) -> bool {
101 self.max_resident_experts.is_some_and(|e| e < num_experts)
102 || self
103 .effective_ram_bytes()
104 .is_some_and(|ram| model_bytes > ram)
105 }
106}
107
108#[cfg(test)]
109mod tests {
110 use super::*;
111
112 #[test]
113 fn explicit_expert_cap_wins() {
114 let b = ResourceBudget {
115 max_ram_bytes: Some(1 << 40),
116 max_resident_experts: Some(8),
117 };
118 assert_eq!(b.resident_experts(256, 1 << 20, 0), 8);
119 assert!(b.needs_offload(1 << 30, 256));
120 }
121
122 #[test]
123 fn derive_experts_from_ram() {
124 // 20 GB RAM, 4 GB backbone, 1 GB/expert → 16 experts fit.
125 let b = ResourceBudget {
126 max_ram_bytes: Some(20 * (1 << 30)),
127 max_resident_experts: None,
128 };
129 assert_eq!(b.resident_experts(256, 1 << 30, 4 * (1 << 30)), 16);
130 }
131
132 #[test]
133 fn clamps_and_unbounded() {
134 // Tiny RAM → at least 1 expert.
135 let tiny = ResourceBudget {
136 max_ram_bytes: Some(1),
137 max_resident_experts: None,
138 };
139 assert_eq!(tiny.resident_experts(256, 1 << 30, 0), 1);
140 // UNBOUNDED still honors the physical-RAM soft budget (a safety floor), so
141 // with tiny per-expert bytes everything fits → all experts.
142 assert_eq!(ResourceBudget::UNBOUNDED.resident_experts(256, 1, 0), 256);
143 // A model that comfortably fits needs no offload under UNBOUNDED.
144 assert!(!ResourceBudget::UNBOUNDED.needs_offload(1 << 20, 256));
145 }
146}