use std::path::{Path, PathBuf};
use std::sync::Mutex;
use anyhow::{Context, Result, bail};
use rlx_core::validate_standard_device;
use rlx_core::weight_loader::GgufLoader;
use rlx_llama_base::LlamaBaseConfig;
use rlx_llama32::{Llama32Generator, Llama32Runner, Llama32RunnerBuilder, MetalGgufPrefillMode};
use crate::backbone::BackboneLoadOptions;
use crate::device::lm_kv_decode_supported;
use rlx_qwen3::{SampleOpts, apply_repetition_penalty, sample_token_at};
use rlx_qwen35::encode_prompt_from_gguf;
use rlx_runtime::Device;
use rlx_runtime::{
llama_decode_bucket_compile_peak_bytes, llama_decode_oneshot_compile_peak_bytes,
memory_headroom_bytes, process_rss_bytes, soft_memory_budget_bytes, would_exceed_soft_budget,
};
use crate::runner::GenerationConfig;
use crate::tokens::{
STOP_TOKEN_ID, custom_token_id_to_code, is_snac_slot_token, mask_logits_for_snac_slot,
use_snac_logit_mask_for,
};
pub const DEFAULT_N_CTX: u32 = 2048;
pub const DEFAULT_COMPILE_SEQ_CAP: u32 = 256;
pub const SYNTHESIS_COMPILE_SEQ_CAP: u32 = 128;
enum DecodeStrategy {
Bucket { max_past: usize },
Dynamic { cache_capacity: usize },
}
fn shared_ram_mode() -> bool {
if matches!(
std::env::var("ORPHEUS_LOW_MEM").ok().as_deref(),
Some("1") | Some("true") | Some("TRUE")
) {
return true;
}
soft_memory_budget_bytes()
.is_some_and(|budget| budget < llama_decode_bucket_compile_peak_bytes() / 2)
}
pub(crate) fn low_mem_mode() -> bool {
shared_ram_mode()
|| memory_headroom_bytes().is_some_and(|h| h < llama_decode_oneshot_compile_peak_bytes())
}
fn log_soft_memory_budget() {
if let Some(budget) = soft_memory_budget_bytes() {
let rss = process_rss_bytes().unwrap_or(0);
let headroom = budget.saturating_sub(rss);
eprintln!(
"[orpheus/backbone] soft RAM budget {:.1} GiB (RSS {:.1} GiB, headroom {:.1} GiB)",
budget as f64 / 1024.0 / 1024.0 / 1024.0,
rss as f64 / 1024.0 / 1024.0 / 1024.0,
headroom as f64 / 1024.0 / 1024.0 / 1024.0,
);
}
}
fn synthesis_decode_plan(compile_cap: usize, opts: &BackboneLoadOptions) -> DecodeStrategy {
let max_past = decode_bucket_max_past(compile_cap);
let want_bucket = env_flag("ORPHEUS_BUCKET_DECODE") == Some(true)
|| (!opts.memory_efficient && env_flag("ORPHEUS_BUCKET_DECODE") != Some(false));
if want_bucket {
if would_exceed_soft_budget(llama_decode_bucket_compile_peak_bytes()) {
eprintln!(
"[orpheus/backbone] bucket decode denied (headroom < {:.1} GiB) — dynamic decode",
llama_decode_bucket_compile_peak_bytes() as f64 / 1024.0 / 1024.0 / 1024.0,
);
} else {
eprintln!(
"[orpheus/backbone] bucket decode enabled (max_past={max_past}) — set ORPHEUS_BUCKET_DECODE=0 to save RAM"
);
return DecodeStrategy::Bucket { max_past };
}
} else if env_flag("ORPHEUS_BUCKET_DECODE") == Some(false) {
eprintln!("[orpheus/backbone] ORPHEUS_BUCKET_DECODE=0 — dynamic decode");
}
let cache_capacity = if low_mem_mode() || opts.memory_efficient {
1
} else {
2
};
DecodeStrategy::Dynamic { cache_capacity }
}
fn compile_seq_cap_for(n_ctx: u32, opts: &BackboneLoadOptions) -> usize {
let base = if env_flag("ORPHEUS_HIGH_MEM") == Some(true) {
DEFAULT_COMPILE_SEQ_CAP
} else if low_mem_mode() {
64
} else if opts.memory_efficient {
SYNTHESIS_COMPILE_SEQ_CAP
} else {
DEFAULT_COMPILE_SEQ_CAP
};
std::env::var("ORPHEUS_COMPILE_SEQ_CAP")
.ok()
.and_then(|s| s.parse::<u32>().ok())
.map(|c| c.min(n_ctx) as usize)
.unwrap_or((base.min(n_ctx)) as usize)
}
fn prefill_cache_capacity() -> usize {
std::env::var("ORPHEUS_PREFILL_CACHE")
.ok()
.and_then(|s| s.parse().ok())
.unwrap_or(if low_mem_mode() { 1 } else { 2 })
}
fn decode_bucket_max_past(compile_cap: usize) -> usize {
std::env::var("ORPHEUS_DECODE_CACHE_CAP")
.ok()
.and_then(|s| s.parse().ok())
.unwrap_or(compile_cap.min(64))
}
fn env_flag(name: &str) -> Option<bool> {
match std::env::var(name).ok().as_deref() {
Some("1") | Some("true") | Some("TRUE") => Some(true),
Some("0") | Some("false") | Some("FALSE") => Some(false),
_ => None,
}
}
fn sample_opts(cfg: &GenerationConfig) -> SampleOpts {
if cfg.greedy || std::env::var("ORPHEUS_GREEDY").ok().as_deref() == Some("1") {
return SampleOpts::greedy();
}
SampleOpts::temperature(cfg.temperature, cfg.seed)
.with_top_p(cfg.top_p)
.with_top_k(cfg.top_k)
}
fn use_compile_caches() -> bool {
env_flag("ORPHEUS_COMPILE_CACHE") == Some(true) && !low_mem_mode()
}
fn metal_f32_kv_enabled() -> bool {
env_flag("ORPHEUS_METAL_KV") == Some(true)
}
fn use_packed_lm(device: Device) -> bool {
if let Some(v) = env_flag("ORPHEUS_PACKED_LM") {
return v;
}
if env_flag("ORPHEUS_FAST_LM") == Some(true) {
return false;
}
if matches!(device, Device::Metal) && metal_f32_kv_enabled() {
return false;
}
matches!(device, Device::Cpu) && !lm_kv_decode_supported(device)
}
fn use_fast_kv(device: Device, opts: &BackboneLoadOptions) -> bool {
if let Some(v) = opts.use_fast_kv {
return v;
}
if env_flag("ORPHEUS_FAST_LM") == Some(false) {
return false;
}
if env_flag("ORPHEUS_METAL_KV") == Some(false) {
return false;
}
lm_kv_decode_supported(device) && !use_packed_lm(device)
}
enum LmEngine {
Packed(Box<Llama32Runner>),
Kv(Box<Llama32Generator>),
}
fn push_orpheus_stream_token(
tok: u32,
stream_index: &mut usize,
on_code: &mut impl FnMut(i32) -> Result<()>,
) -> Result<bool> {
if tok == STOP_TOKEN_ID {
return Ok(true);
}
if is_snac_slot_token(tok, *stream_index) {
if let Some(code) = custom_token_id_to_code(tok, *stream_index) {
if code > 0 {
on_code(code)?;
}
}
*stream_index += 1;
}
Ok(false)
}
fn adjust_orpheus_logits(
logits: &mut [f32],
stream_index: usize,
token_counts: &std::collections::HashMap<u32, u32>,
repetition_penalty: f32,
apply_penalty: bool,
lm_device: Device,
lm_decode_on_cpu: bool,
) {
if use_snac_logit_mask_for(lm_device, lm_decode_on_cpu) {
mask_logits_for_snac_slot(logits, stream_index);
}
if apply_penalty {
apply_repetition_penalty(logits, token_counts, repetition_penalty);
}
}
pub struct BackboneModel {
engine: Mutex<LmEngine>,
weights: PathBuf,
n_ctx: u32,
lm_device: Device,
lm_decode_on_cpu: bool,
pub seed: Option<u32>,
}
impl BackboneModel {
pub fn weights_path(&self) -> &Path {
&self.weights
}
pub fn load(path: &Path, n_ctx: u32) -> Result<Self> {
Self::load_on(path, n_ctx, Device::Cpu)
}
pub fn load_on(path: &Path, n_ctx: u32, device: Device) -> Result<Self> {
Self::load_on_with(path, n_ctx, device, BackboneLoadOptions::for_device(device))
}
pub fn load_on_with(
path: &Path,
n_ctx: u32,
device: Device,
opts: BackboneLoadOptions,
) -> Result<Self> {
validate_standard_device("orpheus", device)?;
let path_str = path
.to_str()
.ok_or_else(|| anyhow::anyhow!("non-utf8 weights path"))?;
let compile_cap = compile_seq_cap_for(n_ctx, &opts);
log_soft_memory_budget();
let lm_device = device;
let lm_decode_on_cpu = matches!(lm_device, Device::Cpu)
|| matches!(lm_device, Device::Gpu | Device::Vulkan)
|| (lm_device == Device::Metal
&& matches!(
opts.metal_prefill.resolve(),
MetalGgufPrefillMode::CpuF32
| MetalGgufPrefillMode::Auto
| MetalGgufPrefillMode::PackedGguf,
));
let engine = if use_fast_kv(lm_device, &opts) {
let mut gguf = GgufLoader::from_file(path_str).context("open GGUF for KV generator")?;
if gguf.architecture() != "llama" {
bail!(
"rlx-orpheus: expected llama GGUF at {}; got `{}`",
path.display(),
gguf.architecture()
);
}
let cfg = rlx_llama32::llama32_cfg_from_gguf(gguf.file())?;
let mut generator = Llama32Generator::from_loader_at_mode(
cfg,
&mut gguf,
lm_device,
path,
if lm_device == Device::Cpu {
MetalGgufPrefillMode::CpuF32
} else {
opts.metal_prefill.resolve()
},
)?
.with_compile_seq_cap(compile_cap);
let decode = synthesis_decode_plan(compile_cap, &opts);
let use_bucket = matches!(decode, DecodeStrategy::Bucket { .. });
generator = match decode {
DecodeStrategy::Bucket { max_past } => generator.with_decode_cache(max_past),
DecodeStrategy::Dynamic { cache_capacity } => {
generator.with_dynamic_decode_cache(cache_capacity)
}
};
if use_bucket && use_compile_caches() {
generator = generator.with_prefill_cache(prefill_cache_capacity());
}
let decode_on_cpu = lm_decode_on_cpu;
eprintln!(
"[orpheus/backbone] kv-cache LM on {lm_device:?} {} (prefill={:?}, decode={}, compile_cap={compile_cap}, decode={})",
path.display(),
if lm_device == Device::Cpu {
MetalGgufPrefillMode::CpuF32
} else {
opts.metal_prefill.resolve()
},
if decode_on_cpu { "cpu" } else { "device" },
if use_bucket { "bucket" } else { "dynamic" },
);
LmEngine::Kv(Box::new(generator))
} else {
let base = LlamaBaseConfig::from_gguf_path(path)
.with_context(|| format!("parse GGUF {}", path.display()))?;
if base.arch != "llama" {
bail!(
"rlx-orpheus: expected llama GGUF at {}; got `{}`",
path.display(),
base.arch
);
}
let runner = Llama32RunnerBuilder::default()
.weights(path)
.max_seq(n_ctx as usize)
.device(lm_device)
.packed_weights(use_packed_lm(lm_device))
.stream(false)
.sample(SampleOpts::greedy())
.build()
.context("build Llama32Runner for Orpheus")?;
eprintln!(
"[orpheus/backbone] {} on {device:?} {}",
if use_packed_lm(device) {
"packed GGUF prefill"
} else {
"f32 runner"
},
path.display()
);
LmEngine::Packed(Box::new(runner))
};
Ok(Self {
engine: Mutex::new(engine),
weights: path.to_path_buf(),
n_ctx,
lm_device,
lm_decode_on_cpu,
seed: Some(GenerationConfig::default().seed as u32),
})
}
pub fn generate_codes(&self, prompt: &str, cfg: &GenerationConfig) -> Result<Vec<i32>> {
let prompt_ids = encode_prompt_from_gguf(&self.weights, prompt)
.with_context(|| format!("tokenize prompt for {}", self.weights.display()))?;
self.generate_codes_from_prompt(&prompt_ids, cfg)
}
pub fn generate_codes_from_prompt(
&self,
prompt_ids: &[u32],
cfg: &GenerationConfig,
) -> Result<Vec<i32>> {
let mut codes = Vec::new();
self.generate_codes_from_prompt_streaming(prompt_ids, cfg, |c| {
codes.push(c);
Ok(())
})?;
Ok(codes)
}
pub fn generate_codes_from_prompt_streaming(
&self,
prompt_ids: &[u32],
cfg: &GenerationConfig,
mut on_code: impl FnMut(i32) -> Result<()>,
) -> Result<()> {
if prompt_ids.len() as u32 > self.n_ctx {
bail!(
"prompt too long: {} tokens > n_ctx={}",
prompt_ids.len(),
self.n_ctx
);
}
let lm_device = self.lm_device;
let lm_decode_on_cpu = self.lm_decode_on_cpu;
let mut engine = self
.engine
.lock()
.map_err(|e| anyhow::anyhow!("backbone lock poisoned: {e}"))?;
let mut stream_index = 0usize;
match &mut *engine {
LmEngine::Kv(generator) => {
let sample = sample_opts(cfg);
let apply_penalty = cfg.repetition_penalty > 1.0;
let mut token_counts = std::collections::HashMap::<u32, u32>::new();
generator.prefill(prompt_ids);
eprintln!(
"[orpheus/backbone] generating up to {} tokens (prompt len={})",
cfg.max_new_tokens,
prompt_ids.len()
);
for step in 0..cfg.max_new_tokens {
let penalty = cfg.repetition_penalty;
let slot_ix = stream_index;
let next = generator
.step_cached_adjust(sample, step as u64, |logits| {
adjust_orpheus_logits(
logits,
slot_ix,
&token_counts,
penalty,
apply_penalty,
lm_device,
lm_decode_on_cpu,
);
})
.context("Orpheus LM cached decode step")?;
if step > 0 && step % 7 == 0 {
eprintln!(
"[orpheus/backbone] step {}/{} ({} tokens in history)",
step,
cfg.max_new_tokens,
generator.tokens().len()
);
}
if push_orpheus_stream_token(next, &mut stream_index, &mut on_code)? {
break;
}
*token_counts.entry(next).or_insert(0) += 1;
}
if std::env::var("ORPHEUS_DEBUG_TOKENS").ok().as_deref() == Some("1") {
let generated = generator.tokens()[prompt_ids.len()..].to_vec();
eprintln!(
"[orpheus/backbone] raw tokens ({}): {:?}",
generated.len(),
&generated[..generated.len().min(16)]
);
}
}
LmEngine::Packed(runner) => {
let sample = sample_opts(cfg);
let mut history = prompt_ids.to_vec();
let mut token_counts = std::collections::HashMap::<u32, u32>::new();
for step in 0..cfg.max_new_tokens {
let slot_ix = stream_index;
let mut logits = runner.predict_logits(&history)?;
adjust_orpheus_logits(
&mut logits,
slot_ix,
&token_counts,
cfg.repetition_penalty,
true,
lm_device,
lm_decode_on_cpu,
);
let next = sample_token_at(&logits, sample, step as u64) as u32;
if push_orpheus_stream_token(next, &mut stream_index, &mut on_code)? {
break;
}
*token_counts.entry(next).or_insert(0) += 1;
history.push(next);
}
if std::env::var("ORPHEUS_DEBUG_TOKENS").ok().as_deref() == Some("1") {
let generated = history[prompt_ids.len()..].to_vec();
eprintln!(
"[orpheus/backbone] raw tokens ({}): {:?}",
generated.len(),
&generated[..generated.len().min(16)]
);
}
}
}
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::backbone::BackboneLoadOptions;
use crate::tokens::STOP_TOKEN_ID;
#[test]
fn synthesis_defaults_to_dynamic_decode() {
let _guard = EnvGuard::unset("ORPHEUS_BUCKET_DECODE");
let plan = synthesis_decode_plan(128, &BackboneLoadOptions::synthesis());
assert!(matches!(plan, DecodeStrategy::Dynamic { .. }));
}
#[test]
fn reference_parity_allows_bucket_when_env_set() {
let _guard = EnvGuard::set("ORPHEUS_BUCKET_DECODE", "1");
let plan = synthesis_decode_plan(128, &BackboneLoadOptions::reference_parity());
let _ = plan;
}
struct EnvGuard {
key: &'static str,
prev: Option<String>,
}
impl EnvGuard {
fn set(key: &'static str, val: &str) -> Self {
let prev = std::env::var(key).ok();
unsafe { std::env::set_var(key, val) };
Self { key, prev }
}
fn unset(key: &'static str) -> Self {
let prev = std::env::var(key).ok();
unsafe { std::env::remove_var(key) };
Self { key, prev }
}
}
impl Drop for EnvGuard {
fn drop(&mut self) {
match &self.prev {
Some(v) => unsafe { std::env::set_var(self.key, v) },
None => unsafe { std::env::remove_var(self.key) },
}
}
}
#[test]
fn stop_token_yields_no_codes() {
let mut codes = Vec::new();
let tokens = [STOP_TOKEN_ID];
let mut stream_index = 0usize;
for &tok in &tokens {
let stop = push_orpheus_stream_token(tok, &mut stream_index, &mut |c| {
codes.push(c);
Ok(())
})
.expect("push token");
if stop {
break;
}
}
assert!(codes.is_empty());
}
}