Fix ROCm backend (target MIGraphX EP) and add CUDA support
All checks were successful
Mirror to GitHub / mirror (push) Successful in 1s
release / build (push) Successful in 1m50s

ROCm silently fell back to CPU: the code targeted ONNX Runtime's classic
ROCMExecutionProvider, but distro ROCm-enabled ONNX Runtime builds (e.g.
Arch's onnxruntime-rocm) are commonly compiled with --use_migraphx instead,
and registration failures were invisible since breadmill never installed a
tracing subscriber. Switches the rocm feature to target MIGraphX, adds a
default tracing subscriber so EP registration success/failure is always
visible, and fixes a real crash where MIGraphX's output sequence padding
could index the attention mask out of bounds during mean-pooling.

Also adds a CUDA backend (--cuda / backend = "cuda") mirroring the same
ort execution-provider pattern, for NVIDIA hardware.

Version bump: 0.1.0 -> 0.2.0.
This commit is contained in:
Breadway 2026-07-03 21:58:39 +08:00
parent d3843f3131
commit 2618a33fd5
10 changed files with 227 additions and 96 deletions

View file

@ -24,17 +24,27 @@ const TOKENIZER_URL: &str =
"https://huggingface.co/nomic-ai/nomic-embed-text-v1.5/resolve/main/tokenizer.json";
fn main() {
// Surfaces ort's EP-registration warnings/errors (e.g. a GPU EP silently
// falling back to CPU) by default, without requiring RUST_LOG to be set.
tracing_subscriber::fmt()
.with_env_filter(
tracing_subscriber::EnvFilter::try_from_default_env()
.unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("warn,ort=info")),
)
.init();
let raw_args: Vec<String> = std::env::args().collect();
// Extract global flags before command dispatch.
let use_npu = raw_args.iter().any(|a| a == "--npu");
let use_rocm = raw_args.iter().any(|a| a == "--rocm");
let use_cuda = raw_args.iter().any(|a| a == "--cuda");
// Build a view of argv without backend flags for command matching.
let args: Vec<&str> = raw_args
.iter()
.skip(1)
.filter(|a| a.as_str() != "--npu" && a.as_str() != "--rocm")
.filter(|a| a.as_str() != "--npu" && a.as_str() != "--rocm" && a.as_str() != "--cuda")
.map(|s| s.as_str())
.collect();
@ -49,7 +59,7 @@ fn main() {
}
}
Some("--reindex") | Some("reindex") => {
if let Err(e) = run_daemon(true, use_npu, use_rocm) {
if let Err(e) = run_daemon(true, use_npu, use_rocm, use_cuda) {
eprintln!("breadmill: {}", e);
std::process::exit(1);
}
@ -66,7 +76,7 @@ fn main() {
cli_status();
}
None | Some("serve") | Some("--serve") => {
if let Err(e) = run_daemon(false, use_npu, use_rocm) {
if let Err(e) = run_daemon(false, use_npu, use_rocm, use_cuda) {
eprintln!("breadmill: {}", e);
std::process::exit(1);
}
@ -74,7 +84,7 @@ fn main() {
Some(cmd) => {
eprintln!("breadmill: unknown command: {}", cmd);
eprintln!(
"usage: breadmill [serve|reindex|fetch-model|query <text>|status] [--npu|--rocm] [--version]"
"usage: breadmill [serve|reindex|fetch-model|query <text>|status] [--npu|--rocm|--cuda] [--version]"
);
std::process::exit(1);
}
@ -83,7 +93,7 @@ fn main() {
// ---- Daemon -----------------------------------------------------------------
fn run_daemon(force_reindex: bool, use_npu: bool, use_rocm: bool) -> Result<(), String> {
fn run_daemon(force_reindex: bool, use_npu: bool, use_rocm: bool, use_cuda: bool) -> Result<(), String> {
let config = breadsearch_shared::Config::load();
let state_dir = breadsearch_shared::state_dir();
let cache_dir = breadsearch_shared::cache_dir();
@ -101,6 +111,9 @@ fn run_daemon(force_reindex: bool, use_npu: bool, use_rocm: bool) -> Result<(),
} else if use_rocm || config.model.backend == "rocm" {
eprintln!("breadmill: ROCm backend selected");
Backend::Rocm
} else if use_cuda || config.model.backend == "cuda" {
eprintln!("breadmill: CUDA backend selected");
Backend::Cuda
} else {
Backend::Cpu
};