rlx_pkg/
auto_tier.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
5//! Auto-tier heuristics and optional name allow-lists.
6
7use crate::tier::StorageTier;
8use crate::write::PackedWeight;
9use std::collections::HashSet;
10
11/// Options for [`apply_auto_tier`].
12#[derive(Debug, Clone)]
13pub struct AutoTierOptions {
14    /// Names forced to hot (exact match).
15    pub hot_names: HashSet<String>,
16    /// If true, names containing these substrings become warm.
17    pub warm_substrings: Vec<String>,
18    /// Tensors larger than this (bytes) default to warm unless hot-listed.
19    pub warm_min_bytes: usize,
20}
21
22impl Default for AutoTierOptions {
23    fn default() -> Self {
24        Self {
25            hot_names: HashSet::new(),
26            warm_substrings: vec![
27                "expert".into(),
28                "experts".into(),
29                "lora".into(),
30                "ffn_gate_exps".into(),
31                "ffn_up_exps".into(),
32                "ffn_down_exps".into(),
33            ],
34            warm_min_bytes: 16 << 20, // 16 MiB
35        }
36    }
37}
38
39/// Mutate `weights` tiers in place using heuristics / allow-lists.
40///
41/// Does **not** replace GGUF-class quantization — warm is for cold/redundant
42/// host blobs, not a substitute for Q4_K.
43pub fn apply_auto_tier(weights: &mut [PackedWeight], opts: &AutoTierOptions) {
44    for w in weights.iter_mut() {
45        if opts.hot_names.contains(&w.name) {
46            w.tier = StorageTier::Hot;
47            continue;
48        }
49        let lower = w.name.to_ascii_lowercase();
50        if opts
51            .warm_substrings
52            .iter()
53            .any(|s| lower.contains(&s.to_ascii_lowercase()))
54        {
55            w.tier = StorageTier::Warm;
56            continue;
57        }
58        if w.data.len() >= opts.warm_min_bytes && w.tier == StorageTier::Hot {
59            w.tier = StorageTier::Warm;
60        }
61    }
62}
rlx_pkg/auto_tier.rs

rlx_pkg/
auto_tier.rs
