Skip to main content

rlx_orpheus/backbone/
options.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3
4//! Load-time options for the Orpheus GGUF backbone.
5
6use rlx_llama32::MetalGgufPrefillMode;
7use rlx_runtime::Device;
8
9/// Options when constructing [`crate::backbone::BackboneModel`].
10#[derive(Debug, Clone)]
11pub struct BackboneLoadOptions {
12    /// Metal GGUF prefill strategy for the KV generator path.
13    pub metal_prefill: MetalGgufPrefillMode,
14    /// Force incremental KV decode vs packed runner. `None` uses env
15    /// (`ORPHEUS_METAL_KV`, `ORPHEUS_PACKED_LM`, …).
16    pub use_fast_kv: Option<bool>,
17    /// When true (default for [`Self::synthesis`]), use dynamic decode and skip
18    /// the bucket compile ladder unless `ORPHEUS_BUCKET_DECODE=1`.
19    pub memory_efficient: bool,
20}
21
22impl Default for BackboneLoadOptions {
23    fn default() -> Self {
24        Self::synthesis()
25    }
26}
27
28impl BackboneLoadOptions {
29    /// Low-RAM TTS: CPU F32 GGUF prefill + CPU decode on Metal hosts (reference path).
30    pub fn synthesis() -> Self {
31        Self {
32            metal_prefill: MetalGgufPrefillMode::CpuF32,
33            use_fast_kv: Some(true),
34            memory_efficient: true,
35        }
36    }
37
38    /// GPU TTS on Metal / CUDA / ROCm: on-device prefill + decode (mmap GGUF).
39    /// wgpu / Vulkan use CPU GGUF prefill + decode (parity) with SNAC logit mask.
40    pub fn gpu(device: Device) -> Self {
41        Self {
42            metal_prefill: metal_prefill_for_device(device),
43            use_fast_kv: Some(true),
44            memory_efficient: true,
45        }
46    }
47
48    /// Pick GPU path when an accelerator is available and low-memory mode is off.
49    pub fn for_device(device: Device) -> Self {
50        if super::rlx::low_mem_mode() || matches!(device, Device::Cpu) {
51            Self::synthesis()
52        } else if matches!(
53            device,
54            Device::Metal | Device::Cuda | Device::Rocm | Device::Gpu | Device::Vulkan
55        ) {
56            Self::gpu(device)
57        } else {
58            Self::synthesis()
59        }
60    }
61
62    /// Reference parity / throughput: bucket decode allowed when RAM budget permits.
63    pub fn reference_parity() -> Self {
64        Self {
65            metal_prefill: MetalGgufPrefillMode::CpuF32,
66            use_fast_kv: Some(true),
67            memory_efficient: false,
68        }
69    }
70
71    /// Experimental: packed Metal KV + CPU F32 logits.
72    pub fn fast_prefill() -> Self {
73        Self {
74            metal_prefill: MetalGgufPrefillMode::PackedGguf,
75            use_fast_kv: Some(true),
76            memory_efficient: true,
77        }
78    }
79
80    pub fn with_memory_efficient(mut self, enabled: bool) -> Self {
81        self.memory_efficient = enabled;
82        self
83    }
84
85    pub fn with_metal_prefill(mut self, mode: MetalGgufPrefillMode) -> Self {
86        self.metal_prefill = mode;
87        self
88    }
89
90    pub fn with_fast_kv(mut self, enabled: bool) -> Self {
91        self.use_fast_kv = Some(enabled);
92        self
93    }
94}
95
96fn metal_prefill_for_device(device: Device) -> MetalGgufPrefillMode {
97    if let Ok(s) = std::env::var("ORPHEUS_METAL_PREFILL") {
98        if let Some(m) = MetalGgufPrefillMode::parse(&s) {
99            return m;
100        }
101    }
102    match device {
103        Device::Metal => MetalGgufPrefillMode::MetalF32,
104        Device::Gpu | Device::Vulkan => MetalGgufPrefillMode::CpuF32,
105        _ => MetalGgufPrefillMode::Auto,
106    }
107}
108
109#[cfg(test)]
110mod tests {
111    use super::*;
112
113    #[test]
114    fn gpu_on_metal_uses_metal_f32_prefill() {
115        let opts = BackboneLoadOptions::gpu(Device::Metal);
116        assert_eq!(opts.metal_prefill, MetalGgufPrefillMode::MetalF32);
117    }
118
119    #[test]
120    fn for_device_wgpu_uses_cpu_gguf_path() {
121        let opts = BackboneLoadOptions::for_device(Device::Gpu);
122        assert_eq!(opts.metal_prefill, MetalGgufPrefillMode::CpuF32);
123        assert_eq!(opts.use_fast_kv, Some(true));
124    }
125}