rlx_orpheus/backbone/
options.rs1use rlx_llama32::MetalGgufPrefillMode;
7use rlx_runtime::Device;
8
9#[derive(Debug, Clone)]
11pub struct BackboneLoadOptions {
12 pub metal_prefill: MetalGgufPrefillMode,
14 pub use_fast_kv: Option<bool>,
17 pub memory_efficient: bool,
20}
21
22impl Default for BackboneLoadOptions {
23 fn default() -> Self {
24 Self::synthesis()
25 }
26}
27
28impl BackboneLoadOptions {
29 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 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 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 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 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}