Fix ROCm backend (target MIGraphX EP) and add CUDA support
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
parent
d3843f3131
commit
2618a33fd5
10 changed files with 227 additions and 96 deletions
171
Cargo.lock
generated
171
Cargo.lock
generated
|
|
@ -25,7 +25,7 @@ checksum = "b169f7a6d4742236a0a00c541b845991d0ac43e546831af1249753ab4c3aa3a0"
|
|||
dependencies = [
|
||||
"cfg-if",
|
||||
"cipher",
|
||||
"cpufeatures 0.2.17",
|
||||
"cpufeatures",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -111,15 +111,6 @@ dependencies = [
|
|||
"generic-array",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "block-buffer"
|
||||
version = "0.12.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d2f6c7dbe95a6ed67ad9f18e57daf93a2f034c524b99fd2b76d18fdfeb6660aa"
|
||||
dependencies = [
|
||||
"hybrid-array",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "block-padding"
|
||||
version = "0.3.3"
|
||||
|
|
@ -142,7 +133,7 @@ dependencies = [
|
|||
|
||||
[[package]]
|
||||
name = "breadmill"
|
||||
version = "0.1.0"
|
||||
version = "0.2.0"
|
||||
dependencies = [
|
||||
"breadsearch-shared",
|
||||
"hex",
|
||||
|
|
@ -154,8 +145,9 @@ dependencies = [
|
|||
"rusqlite",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"sha2 0.11.0",
|
||||
"sha2",
|
||||
"tokenizers",
|
||||
"tracing-subscriber",
|
||||
"ureq 2.12.1",
|
||||
"usearch",
|
||||
"zip",
|
||||
|
|
@ -163,7 +155,7 @@ dependencies = [
|
|||
|
||||
[[package]]
|
||||
name = "breadsearch"
|
||||
version = "0.1.0"
|
||||
version = "0.2.0"
|
||||
dependencies = [
|
||||
"bread-theme",
|
||||
"breadsearch-shared",
|
||||
|
|
@ -174,7 +166,7 @@ dependencies = [
|
|||
|
||||
[[package]]
|
||||
name = "breadsearch-shared"
|
||||
version = "0.1.0"
|
||||
version = "0.2.0"
|
||||
dependencies = [
|
||||
"serde",
|
||||
"serde_json",
|
||||
|
|
@ -315,7 +307,7 @@ version = "0.4.4"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "773f3b9af64447d2ce9850330c473515014aa235e6a783b02db81ff39e4a3dad"
|
||||
dependencies = [
|
||||
"crypto-common 0.1.7",
|
||||
"crypto-common",
|
||||
"inout",
|
||||
]
|
||||
|
||||
|
|
@ -383,12 +375,6 @@ dependencies = [
|
|||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "const-oid"
|
||||
version = "0.10.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "a6ef517f0926dd24a1582492c791b6a4818a4d94e789a334894aa15b0d12f55c"
|
||||
|
||||
[[package]]
|
||||
name = "constant_time_eq"
|
||||
version = "0.3.1"
|
||||
|
|
@ -420,15 +406,6 @@ dependencies = [
|
|||
"libc",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "cpufeatures"
|
||||
version = "0.3.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "8b2a41393f66f16b0823bb79094d54ac5fbd34ab292ddafb9a0456ac9f87d201"
|
||||
dependencies = [
|
||||
"libc",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "crc"
|
||||
version = "3.4.0"
|
||||
|
|
@ -497,15 +474,6 @@ dependencies = [
|
|||
"typenum",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "crypto-common"
|
||||
version = "0.2.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ce6e4c961d6cd6c9a86db418387425e8bdeaf05b3c8bc1411e6dca4c252f1453"
|
||||
dependencies = [
|
||||
"hybrid-array",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "cxx"
|
||||
version = "1.0.194"
|
||||
|
|
@ -688,22 +656,11 @@ version = "0.10.7"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292"
|
||||
dependencies = [
|
||||
"block-buffer 0.10.4",
|
||||
"crypto-common 0.1.7",
|
||||
"block-buffer",
|
||||
"crypto-common",
|
||||
"subtle",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "digest"
|
||||
version = "0.11.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f1dd6dbb5841937940781866fa1281a1ff7bd3bf827091440879f9994983d5c2"
|
||||
dependencies = [
|
||||
"block-buffer 0.12.1",
|
||||
"const-oid",
|
||||
"crypto-common 0.2.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "dirs"
|
||||
version = "5.0.1"
|
||||
|
|
@ -779,7 +736,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
|||
checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb"
|
||||
dependencies = [
|
||||
"libc",
|
||||
"windows-sys 0.52.0",
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -1350,7 +1307,7 @@ version = "0.12.1"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "6c49c37c09c17a53d937dfbb742eb3a961d65a994e6bcdcf37e7399d0cc8ab5e"
|
||||
dependencies = [
|
||||
"digest 0.10.7",
|
||||
"digest",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -1375,15 +1332,6 @@ version = "1.10.1"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "6dbf3de79e51f3d586ab4cb9d5c3e2c14aa28ed23d180cf89b4df0454a69cc87"
|
||||
|
||||
[[package]]
|
||||
name = "hybrid-array"
|
||||
version = "0.4.12"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9155a582abd142abc056962c29e3ce5ff2ad5469f4246b537ed42c5deba857da"
|
||||
dependencies = [
|
||||
"typenum",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "icu_collections"
|
||||
version = "2.2.0"
|
||||
|
|
@ -1624,6 +1572,12 @@ dependencies = [
|
|||
"libc",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "lazy_static"
|
||||
version = "1.5.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe"
|
||||
|
||||
[[package]]
|
||||
name = "libc"
|
||||
version = "0.2.186"
|
||||
|
|
@ -1708,7 +1662,7 @@ dependencies = [
|
|||
"nom_locate",
|
||||
"rand",
|
||||
"rangemap",
|
||||
"sha2 0.10.9",
|
||||
"sha2",
|
||||
"stringprep",
|
||||
"thiserror 2.0.18",
|
||||
"ttf-parser",
|
||||
|
|
@ -1758,6 +1712,15 @@ version = "0.2.2"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "670fdfda89751bc4a84ac13eaa63e205cf0fd22b4c9a5fbfa085b63c1f1d3a30"
|
||||
|
||||
[[package]]
|
||||
name = "matchers"
|
||||
version = "0.2.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d1525a2a28c7f4fa0fc98bb91ae755d1e2d1505079e05539e35bc876b5d65ae9"
|
||||
dependencies = [
|
||||
"regex-automata",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "matrixmultiply"
|
||||
version = "0.3.10"
|
||||
|
|
@ -1775,7 +1738,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
|||
checksum = "d89e7ee0cfbedfc4da3340218492196241d89eefb6dab27de5df917a6d2e78cf"
|
||||
dependencies = [
|
||||
"cfg-if",
|
||||
"digest 0.10.7",
|
||||
"digest",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -1924,6 +1887,15 @@ dependencies = [
|
|||
"windows-sys 0.48.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "nu-ansi-term"
|
||||
version = "0.50.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5"
|
||||
dependencies = [
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "num-complex"
|
||||
version = "0.4.6"
|
||||
|
|
@ -2104,7 +2076,7 @@ version = "0.12.2"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f8ed6a7761f76e3b9f92dfb0a60a6a6477c61024b775147ff0973a02653abaf2"
|
||||
dependencies = [
|
||||
"digest 0.10.7",
|
||||
"digest",
|
||||
"hmac",
|
||||
]
|
||||
|
||||
|
|
@ -2416,7 +2388,7 @@ dependencies = [
|
|||
"errno",
|
||||
"libc",
|
||||
"linux-raw-sys",
|
||||
"windows-sys 0.52.0",
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -2587,8 +2559,8 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
|||
checksum = "e3bf829a2d51ab4a5ddf1352d8470c140cadc8301b2ae1789db023f01cedd6ba"
|
||||
dependencies = [
|
||||
"cfg-if",
|
||||
"cpufeatures 0.2.17",
|
||||
"digest 0.10.7",
|
||||
"cpufeatures",
|
||||
"digest",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -2598,19 +2570,17 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
|||
checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283"
|
||||
dependencies = [
|
||||
"cfg-if",
|
||||
"cpufeatures 0.2.17",
|
||||
"digest 0.10.7",
|
||||
"cpufeatures",
|
||||
"digest",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "sha2"
|
||||
version = "0.11.0"
|
||||
name = "sharded-slab"
|
||||
version = "0.1.7"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "446ba717509524cb3f22f17ecc096f10f4822d76ab5c0b9822c5f9c284e825f4"
|
||||
checksum = "f40ca3c46823713e0d4209592e8d6e826aa57e928f09752619fc696c499637f6"
|
||||
dependencies = [
|
||||
"cfg-if",
|
||||
"cpufeatures 0.3.0",
|
||||
"digest 0.11.3",
|
||||
"lazy_static",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -2758,7 +2728,7 @@ dependencies = [
|
|||
"getrandom 0.3.4",
|
||||
"once_cell",
|
||||
"rustix",
|
||||
"windows-sys 0.52.0",
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -2810,6 +2780,15 @@ dependencies = [
|
|||
"syn",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "thread_local"
|
||||
version = "1.1.9"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f60246a4944f24f6e018aa17cdeffb7818b76356965d03b07d6a9886e8962185"
|
||||
dependencies = [
|
||||
"cfg-if",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "time"
|
||||
version = "0.3.51"
|
||||
|
|
@ -2997,6 +2976,36 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
|||
checksum = "db97caf9d906fbde555dd62fa95ddba9eecfd14cb388e4f491a66d74cd5fb79a"
|
||||
dependencies = [
|
||||
"once_cell",
|
||||
"valuable",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tracing-log"
|
||||
version = "0.2.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ee855f1f400bd0e5c02d150ae5de3840039a3f54b025156404e34c23c03f47c3"
|
||||
dependencies = [
|
||||
"log",
|
||||
"once_cell",
|
||||
"tracing-core",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tracing-subscriber"
|
||||
version = "0.3.23"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "cb7f578e5945fb242538965c2d0b04418d38ec25c79d160cd279bf0731c8d319"
|
||||
dependencies = [
|
||||
"matchers",
|
||||
"nu-ansi-term",
|
||||
"once_cell",
|
||||
"regex-automata",
|
||||
"sharded-slab",
|
||||
"smallvec",
|
||||
"thread_local",
|
||||
"tracing",
|
||||
"tracing-core",
|
||||
"tracing-log",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -3167,6 +3176,12 @@ version = "1.0.4"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be"
|
||||
|
||||
[[package]]
|
||||
name = "valuable"
|
||||
version = "0.1.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ba73ea9cf16a25df0c8caa16c51acb937d5712a8429db78a3ee29d5dcacd3a65"
|
||||
|
||||
[[package]]
|
||||
name = "vcpkg"
|
||||
version = "0.2.15"
|
||||
|
|
|
|||
|
|
@ -80,3 +80,7 @@ Start from `breadbox/breadbox/src/main.rs`. **Reuse verbatim:** the gtk4-layer-s
|
|||
- nomic prefixes + mean-pool + normalize must match between index and query or recall collapses.
|
||||
- `ort` linking: prefer the crate's downloaded/bundled ONNX Runtime to avoid version skew with Arch's `onnxruntime`.
|
||||
- Office formats (docx/odt) are best-effort in v1; md/txt/org/pdf are the reliable path.
|
||||
- GPU EPs (ROCm/CUDA) fail to register silently at the ONNX Runtime level and fall back to CPU — always check
|
||||
startup logs for `Successfully registered` before trusting a GPU build is actually accelerating. See
|
||||
[README: GPU backend notes](README.md#gpu-backend-notes) for the MIGraphX-vs-ROCMExecutionProvider distinction
|
||||
and the per-shape JIT-compile-and-cache behavior that matters for interactive query latency.
|
||||
|
|
|
|||
50
README.md
50
README.md
|
|
@ -22,13 +22,27 @@ Optional features:
|
|||
| Feature | What it adds |
|
||||
|---------|-------------|
|
||||
| `npu` | AMD XDNA NPU via VitisAI ONNX Runtime EP (requires Ryzen AI SDK) |
|
||||
| `rocm` | AMD iGPU via ROCm ONNX Runtime EP |
|
||||
| `rocm` | AMD iGPU via the MIGraphX ONNX Runtime EP (ROCm-backed) |
|
||||
| `cuda` | NVIDIA GPU via the CUDA ONNX Runtime EP |
|
||||
|
||||
```
|
||||
# NPU build
|
||||
cargo build --release -p breadmill --features npu
|
||||
|
||||
# ROCm (AMD iGPU/dGPU) build
|
||||
cargo build --release -p breadmill --features rocm
|
||||
|
||||
# CUDA (NVIDIA GPU) build
|
||||
cargo build --release -p breadmill --features cuda
|
||||
```
|
||||
|
||||
`rocm`/`cuda`/`npu` all use `ort`'s `load-dynamic` mode: at runtime, breadmill
|
||||
dlopens whatever `libonnxruntime.so` the dynamic linker resolves (or
|
||||
`ORT_DYLIB_PATH` if set). GPU acceleration only works if that ONNX Runtime
|
||||
build actually has the matching execution provider compiled in — breadmill
|
||||
logs a clear `Successfully registered` / `not enabled in this build` line for
|
||||
this at startup (see [GPU backend notes](#gpu-backend-notes) below).
|
||||
|
||||
## Setup
|
||||
|
||||
**1. Fetch the embedding model** (~550 MB, downloaded once from Hugging Face):
|
||||
|
|
@ -89,6 +103,7 @@ breadmill status
|
|||
# Backend flags (requires the matching Cargo feature)
|
||||
breadmill --npu
|
||||
breadmill --rocm
|
||||
breadmill --cuda
|
||||
```
|
||||
|
||||
## Config
|
||||
|
|
@ -109,7 +124,7 @@ snippet_len = 200 # max characters in result snippet
|
|||
[model]
|
||||
name = "nomic-embed-text-v1.5"
|
||||
dim = 768
|
||||
backend = "cpu" # "cpu", "npu", or "rocm"
|
||||
backend = "cpu" # "cpu", "npu", "rocm", or "cuda"
|
||||
```
|
||||
|
||||
`roots` and `excludes` support `~/` expansion. The index respects `.gitignore` files found during the walk.
|
||||
|
|
@ -124,6 +139,37 @@ Set `backend = "npu"` in config (or pass `--npu`) when running a build compiled
|
|||
4. `/etc/vaip_config.json`
|
||||
5. `/opt/xilinx/vaip_config.json`
|
||||
|
||||
### GPU backend notes
|
||||
|
||||
Both `rocm` and `cuda` need a system ONNX Runtime that was actually built with
|
||||
the matching execution provider — the crate's own downloaded binary is CPU-only.
|
||||
Point `ORT_DYLIB_PATH` at one, or install a distro package that provides
|
||||
`libonnxruntime.so` with the EP baked in and let the dynamic linker find it.
|
||||
|
||||
**ROCm (`--rocm` / `backend = "rocm"`)** targets ONNX Runtime's **MIGraphX**
|
||||
execution provider, not the classic `ROCMExecutionProvider`. Distro
|
||||
ROCm-enabled ONNX Runtime packages (e.g. Arch's `onnxruntime-rocm`) are
|
||||
commonly built with `--use_migraphx` rather than `--use_rocm`, so this is the
|
||||
EP that's actually available in practice; the classic ROCm EP needs a bespoke
|
||||
`--use_rocm` build most distros don't package. Startup logs a
|
||||
`Successfully registered `MIGraphXExecutionProvider`` line when it's really
|
||||
active — check for it if in doubt, since a failed GPU EP registration falls
|
||||
back to CPU silently at the ONNX Runtime level (breadmill's own log line is
|
||||
only a statement of intent, not a confirmation).
|
||||
|
||||
MIGraphX JIT-compiles the model per distinct input sequence length and caches
|
||||
the compiled kernel to disk (each compile takes ~60–120s and produces a
|
||||
~500MB `.mxr` file). Set `ORT_MIGRAPHX_MODEL_CACHE_PATH=/path/to/cache` so
|
||||
that cost is paid once per shape instead of on every daemon restart. Because
|
||||
query text length varies, expect an occasional multi-second stall the first
|
||||
time a new token length is seen — fine for background document indexing,
|
||||
noticeable for interactive query embedding.
|
||||
|
||||
**CUDA (`--cuda` / `backend = "cuda"`)** targets the standard
|
||||
`CUDAExecutionProvider` and needs a CUDA-enabled ONNX Runtime + a working
|
||||
CUDA/cuDNN install. Unverified on real NVIDIA hardware in this repo — only
|
||||
compile-checked, since development happened on an AMD-only machine.
|
||||
|
||||
## Runtime paths
|
||||
|
||||
| Purpose | Path |
|
||||
|
|
|
|||
|
|
@ -1,13 +1,19 @@
|
|||
[package]
|
||||
name = "breadmill"
|
||||
version = "0.1.0"
|
||||
version = "0.2.0"
|
||||
edition = "2021"
|
||||
license = "MIT"
|
||||
|
||||
[features]
|
||||
default = []
|
||||
npu = ["ort/vitis", "ort/load-dynamic"]
|
||||
rocm = ["ort/rocm", "ort/load-dynamic"]
|
||||
# "rocm" targets the MIGraphX execution provider, not ONNX Runtime's classic
|
||||
# ROCMExecutionProvider (--use_rocm build). Distro ROCm-enabled ONNX Runtime
|
||||
# packages (e.g. Arch's onnxruntime-rocm) are commonly built with --use_migraphx
|
||||
# instead; the classic ROCm EP needs a bespoke --use_rocm build that ships
|
||||
# libonnxruntime_providers_rocm.so, which most distros don't package.
|
||||
rocm = ["ort/migraphx", "ort/load-dynamic"]
|
||||
cuda = ["ort/cuda", "ort/load-dynamic"]
|
||||
|
||||
[[bin]]
|
||||
name = "breadmill"
|
||||
|
|
@ -26,6 +32,10 @@ breadsearch-shared = { path = "../breadsearch-shared" }
|
|||
ort = { version = "2.0.0-rc.12", default-features = false, features = ["std", "tracing", "download-binaries", "tls-native", "copy-dylibs", "api-23"] }
|
||||
tokenizers = "0"
|
||||
|
||||
# Surfaces ort's own EP-registration tracing (e.g. a GPU EP silently failing to
|
||||
# register and falling back to CPU) as visible log output instead of nowhere.
|
||||
tracing-subscriber = { version = "0.3", features = ["env-filter"] }
|
||||
|
||||
# Vector index
|
||||
usearch = "2"
|
||||
|
||||
|
|
|
|||
|
|
@ -19,8 +19,14 @@ pub enum Backend {
|
|||
/// AMD XDNA NPU via the VitisAI ONNX Runtime execution provider.
|
||||
/// `cache_dir` is used to store the compiled NPU model between runs.
|
||||
Npu { cache_dir: PathBuf },
|
||||
/// AMD iGPU via the ROCm ONNX Runtime execution provider.
|
||||
/// AMD iGPU via the MIGraphX ONNX Runtime execution provider (ROCm-backed).
|
||||
/// Distro ROCm ONNX Runtime builds (e.g. Arch's onnxruntime-rocm) are
|
||||
/// commonly compiled with `--use_migraphx`, not `--use_rocm`, so this
|
||||
/// targets `MIGraphXExecutionProvider` rather than the classic
|
||||
/// `ROCMExecutionProvider`.
|
||||
Rocm,
|
||||
/// NVIDIA GPU via the CUDA ONNX Runtime execution provider.
|
||||
Cuda,
|
||||
}
|
||||
|
||||
pub struct OrtEmbedder {
|
||||
|
|
@ -108,11 +114,14 @@ impl OrtEmbedder {
|
|||
let actual_seq = shape[1] as usize;
|
||||
let actual_dim = shape[2] as usize;
|
||||
|
||||
// Mean-pool over non-padding positions
|
||||
// Mean-pool over non-padding positions. Some execution providers (e.g.
|
||||
// MIGraphX) pad the output sequence dimension for kernel efficiency, so
|
||||
// actual_seq can exceed mask.len() — only positions covered by our own
|
||||
// attention mask are meaningful, so cap the loop at whichever is shorter.
|
||||
let mut result = vec![0.0f32; actual_dim];
|
||||
let mut count = 0usize;
|
||||
|
||||
for t in 0..actual_seq {
|
||||
for t in 0..actual_seq.min(mask.len()) {
|
||||
if mask[t] > 0 {
|
||||
for d in 0..actual_dim {
|
||||
result[d] += data[t * actual_dim + d];
|
||||
|
|
@ -155,6 +164,7 @@ fn configure_eps(builder: SessionBuilder, backend: &Backend) -> Result<SessionBu
|
|||
Backend::Cpu => Ok(builder),
|
||||
Backend::Npu { cache_dir } => npu_session(builder, cache_dir),
|
||||
Backend::Rocm => rocm_session(builder),
|
||||
Backend::Cuda => cuda_session(builder),
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -195,14 +205,19 @@ fn build_vitis_ep(cache_dir: &Path) -> Result<ort::ep::ExecutionProviderDispatch
|
|||
.build())
|
||||
}
|
||||
|
||||
// ---- ROCm EP (AMD iGPU) -----------------------------------------------------
|
||||
// ---- MIGraphX EP (AMD iGPU, ROCm-backed) -------------------------------------
|
||||
|
||||
#[cfg(feature = "rocm")]
|
||||
fn rocm_session(builder: SessionBuilder) -> Result<SessionBuilder, String> {
|
||||
eprintln!("breadmill: using ROCm execution provider (device 0)");
|
||||
eprintln!("breadmill: using MIGraphX execution provider (device 0)");
|
||||
eprintln!(
|
||||
"breadmill: note — check the log line above/below for \"Successfully registered \
|
||||
`MIGraphXExecutionProvider`\"; if it's missing, the ONNX Runtime in use wasn't built \
|
||||
with MIGraphX support and inference silently fell back to CPU"
|
||||
);
|
||||
builder
|
||||
.with_execution_providers([
|
||||
ort::execution_providers::ROCmExecutionProvider::default().build(),
|
||||
ort::ep::MIGraphX::default().with_device_id(0).build(),
|
||||
ort::ep::CPU::default().build(),
|
||||
])
|
||||
.map_err(|e| e.to_string())
|
||||
|
|
@ -214,6 +229,30 @@ fn rocm_session(builder: SessionBuilder) -> Result<SessionBuilder, String> {
|
|||
Ok(builder)
|
||||
}
|
||||
|
||||
// ---- CUDA EP (NVIDIA GPU) ----------------------------------------------------
|
||||
|
||||
#[cfg(feature = "cuda")]
|
||||
fn cuda_session(builder: SessionBuilder) -> Result<SessionBuilder, String> {
|
||||
eprintln!("breadmill: using CUDA execution provider (device 0)");
|
||||
eprintln!(
|
||||
"breadmill: note — check the log line above/below for \"Successfully registered \
|
||||
`CUDAExecutionProvider`\"; if it's missing, the ONNX Runtime in use wasn't built \
|
||||
with CUDA support and inference silently fell back to CPU"
|
||||
);
|
||||
builder
|
||||
.with_execution_providers([
|
||||
ort::ep::CUDA::default().with_device_id(0).build(),
|
||||
ort::ep::CPU::default().build(),
|
||||
])
|
||||
.map_err(|e| e.to_string())
|
||||
}
|
||||
|
||||
#[cfg(not(feature = "cuda"))]
|
||||
fn cuda_session(builder: SessionBuilder) -> Result<SessionBuilder, String> {
|
||||
eprintln!("breadmill: CUDA backend requested but not compiled in (rebuild with --features cuda); using CPU");
|
||||
Ok(builder)
|
||||
}
|
||||
|
||||
/// Locate the VitisAI EP config file required by the AMD Ryzen AI SDK.
|
||||
///
|
||||
/// Search order:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
};
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
[package]
|
||||
name = "breadsearch-shared"
|
||||
version = "0.1.0"
|
||||
version = "0.2.0"
|
||||
edition = "2021"
|
||||
license = "MIT"
|
||||
|
||||
|
|
|
|||
|
|
@ -129,7 +129,7 @@ pub struct ModelConfig {
|
|||
pub name: String,
|
||||
#[serde(default = "default_dim")]
|
||||
pub dim: usize,
|
||||
/// Compute backend: "cpu" or "npu" (VitisAI/XDNA).
|
||||
/// Compute backend: "cpu", "npu" (VitisAI/XDNA), "rocm" (MIGraphX/AMD GPU), or "cuda" (NVIDIA GPU).
|
||||
#[serde(default = "default_backend")]
|
||||
pub backend: String,
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
[package]
|
||||
name = "breadsearch"
|
||||
version = "0.1.0"
|
||||
version = "0.2.0"
|
||||
edition = "2021"
|
||||
license = "MIT"
|
||||
|
||||
|
|
|
|||
|
|
@ -5,6 +5,10 @@ After=default.target
|
|||
|
||||
[Service]
|
||||
Type=simple
|
||||
# Uncomment if built with --features rocm: persists MIGraphX's compiled-kernel
|
||||
# cache across restarts (each new sequence length otherwise costs a ~60-120s
|
||||
# recompile). See README.md#gpu-backend-notes.
|
||||
#Environment=ORT_MIGRAPHX_MODEL_CACHE_PATH=%h/.cache/breadsearch/migraphx-cache
|
||||
ExecStart=%h/.cargo/bin/breadmill
|
||||
Restart=on-failure
|
||||
RestartSec=5
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue