rlx_orpheus/
device.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3
4//! Device selection for Orpheus synthesis (LM backbone + SNAC decoder backends).
5
6use anyhow::{Result, bail};
7use rlx_runtime::{Device, is_available};
8
9/// Resolved execution targets for Orpheus (LM device + SNAC backend).
10#[derive(Debug, Clone, Copy, PartialEq, Eq)]
11pub struct OrpheusRuntimeDevice {
12    /// RLX GGUF backbone device (Metal on Apple Silicon by default).
13    pub lm: Device,
14    /// When true, SNAC runs on native RLX CoreML ([`Device::Ane`]); LM uses [`Self::lm`].
15    pub snac_coreml: bool,
16}
17
18impl OrpheusRuntimeDevice {
19    pub fn snac_load_options(self) -> crate::decoder::SnacLoadOptions {
20        crate::decoder::SnacLoadOptions {
21            coreml: self.snac_coreml,
22        }
23    }
24}
25
26fn env_flag(name: &str) -> Option<bool> {
27    match std::env::var(name).ok().as_deref() {
28        Some("1") | Some("true") | Some("TRUE") => Some(true),
29        Some("0") | Some("false") | Some("FALSE") => Some(false),
30        _ => None,
31    }
32}
33
34/// Whether incremental KV LM decode is supported for Orpheus on `device`.
35///
36/// MLX is opt-in (`ORPHEUS_MLX_KV=1`) — default off until rlx-mlx Orpheus decode is stable.
37pub fn lm_kv_decode_supported(device: Device) -> bool {
38    match device {
39        Device::Mlx => env_flag("ORPHEUS_MLX_KV") == Some(true),
40        Device::Metal | Device::Cuda | Device::Rocm | Device::Gpu | Device::Vulkan => true,
41        Device::Cpu => env_flag("ORPHEUS_FAST_LM") == Some(true),
42        _ => false,
43    }
44}
45
46/// Best accelerator for Orpheus LM inference.
47///
48/// Prefers **Metal** on Apple Silicon (on-device prefill + decode).
49/// MLX is skipped unless [`lm_kv_decode_supported`] (opt-in via `ORPHEUS_MLX_KV=1`).
50pub fn preferred_synth_device() -> Device {
51    if is_available(Device::Metal) {
52        Device::Metal
53    } else if is_available(Device::Mlx) && lm_kv_decode_supported(Device::Mlx) {
54        Device::Mlx
55    } else if is_available(Device::Cuda) {
56        Device::Cuda
57    } else if is_available(Device::Rocm) {
58        Device::Rocm
59    } else if is_available(Device::Gpu) {
60        Device::Gpu
61    } else if is_available(Device::Vulkan) {
62        Device::Vulkan
63    } else {
64        Device::Cpu
65    }
66}
67
68/// Device for integration tests; honors `ORPHEUS_SYNTH_DEVICE` then [`preferred_synth_device`].
69#[cfg(feature = "llama")]
70pub fn synth_device_for_tests() -> Device {
71    if let Ok(s) = std::env::var("ORPHEUS_SYNTH_DEVICE") {
72        if let Ok(d) = rlx_cli::parse_device(s.trim()) {
73            return d;
74        }
75    }
76    preferred_synth_device()
77}
78
79#[cfg(feature = "llama")]
80pub fn parse_orpheus_device(s: &str) -> Result<OrpheusRuntimeDevice> {
81    resolve_orpheus_device(s)
82}
83
84#[cfg(feature = "llama")]
85pub fn resolve_orpheus_device(s: &str) -> Result<OrpheusRuntimeDevice> {
86    let s = s.trim();
87    if s.eq_ignore_ascii_case("auto") {
88        return Ok(OrpheusRuntimeDevice {
89            lm: preferred_synth_device(),
90            snac_coreml: false,
91        });
92    }
93    if s.eq_ignore_ascii_case("coreml") {
94        let lm = preferred_synth_device();
95        if !matches!(lm, Device::Metal | Device::Mlx | Device::Cpu) {
96            bail!("coreml SNAC path expects an Apple host with Metal or MLX; got lm={lm:?}");
97        }
98        return Ok(OrpheusRuntimeDevice {
99            lm,
100            snac_coreml: true,
101        });
102    }
103    Ok(OrpheusRuntimeDevice {
104        lm: rlx_cli::parse_device(s)?,
105        snac_coreml: false,
106    })
107}
108
109#[cfg(test)]
110mod tests {
111    use super::*;
112
113    #[test]
114    fn mlx_kv_off_by_default() {
115        assert!(!lm_kv_decode_supported(Device::Mlx));
116    }
117
118    #[test]
119    fn metal_kv_on_by_default() {
120        assert!(lm_kv_decode_supported(Device::Metal));
121    }
122
123    #[test]
124    fn wgpu_kv_on_by_default() {
125        assert!(lm_kv_decode_supported(Device::Gpu));
126        assert!(lm_kv_decode_supported(Device::Vulkan));
127    }
128
129    #[test]
130    fn resolve_coreml_sets_snac_flag() {
131        let rt = resolve_orpheus_device("coreml").unwrap();
132        assert!(rt.snac_coreml);
133    }
134}
rlx_orpheus/device.rs

rlx_orpheus/
device.rs
