1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
//! `EmbedEngine` — the device-agnostic embedding entry point: GPU when a wgpu adapter exists,
//! the native CPU encoder otherwise.
//!
//! wgpu ships no software fallback adapter (llvmpipe-class drivers are external system
//! installs), so "works on machines without a GPU" is carried entirely by the `encoder_cpu`
//! runtime. Device policy (force-CPU, adapter selection) is the CALLER's decision — this module
//! reads no environment variables; adapters and CLIs translate their config into `auto()` vs
//! `cpu()`.
use crate::GpuCtx;
use crate::encoder::EncoderGpu;
use crate::encoder_cpu::CpuEncoder;
use crate::encoder_weights::{EncBatch, EncoderConfig};
use crate::pooling::EmbedOut;
use anyhow::Result;
use std::path::Path;
/// A loaded embedding model on whichever device was available.
// One engine is constructed per model for a process lifetime and moved at most once — the
// variant-size asymmetry (weights live inline in CpuEncoder) has no runtime cost worth a Box.
#[allow(clippy::large_enum_variant)]
pub enum EmbedEngine {
/// GPU path: owns the wgpu context and the prebaked-plan executor.
Gpu {
/// The wgpu device/queue the encoder dispatches on.
ctx: GpuCtx,
/// The prebaked whole-model executor.
enc: EncoderGpu,
},
/// Native CPU path (`gemm` + rayon) — also the parity oracle.
Cpu(CpuEncoder),
}
impl EmbedEngine {
/// Load on the best available device: any wgpu adapter (Metal/Vulkan/DX12/GL) first, else
/// the CPU encoder. GPU-side failures during LOAD fall through to CPU with a warning on
/// stderr — a machine with a broken driver must still embed.
pub fn auto(dir: &Path, max_tokens: usize) -> Result<Self> {
// (CrossEncoder scoring heads used to be pinned here: the GPU pool kernels implement
// Mean/LastToken/PerToken and not the pooler+classifier head, so `/score` — the
// ontology-verifier's hot path — ran CPU-only BY CONSTRUCTION and measured 2.07× torch-CPU.
// The head is now a GPU plan arm: CLS → pooler GEMM (tanh) → classifier GEMM. The pin is
// gone; a CrossEncoder checkpoint takes the same GPU-or-fall-back-to-CPU path as any other.)
match GpuCtx::new() {
Ok(ctx) => match EncoderGpu::load(&ctx, dir, max_tokens) {
Ok(enc) => Ok(Self::Gpu { ctx, enc }),
Err(e) => {
eprintln!(
"inferencelayer: GPU encoder load failed ({e:#}); falling back to CPU"
);
Ok(Self::Cpu(CpuEncoder::load(dir)?))
}
},
Err(_) => Ok(Self::Cpu(CpuEncoder::load(dir)?)),
}
}
/// Load on the CPU unconditionally (deployment policy, contention avoidance, oracle runs).
pub fn cpu(dir: &Path) -> Result<Self> {
Ok(Self::Cpu(CpuEncoder::load(dir)?))
}
/// Embed a ragged batch on whichever device this engine holds.
pub fn encode(&mut self, batch: &EncBatch) -> Result<EmbedOut> {
match self {
Self::Gpu { ctx, enc } => enc.encode(ctx, batch),
Self::Cpu(enc) => enc.encode(batch),
}
}
/// Per-token hidden states `[T, hidden]` (pre-pooling), on whichever device this engine holds.
/// Heads that consume token states rather than a pooled vector — GLiNER's span head — run on
/// this, which is what lets them ride the GPU instead of being pinned to the CPU encoder.
pub fn forward_hidden(&mut self, batch: &EncBatch) -> Result<Vec<f32>> {
match self {
Self::Gpu { ctx, enc } => enc.forward_hidden(ctx, batch),
Self::Cpu(enc) => enc.forward_hidden(batch),
}
}
/// The parsed model configuration.
pub fn config(&self) -> &EncoderConfig {
match self {
Self::Gpu { enc, .. } => enc.config(),
Self::Cpu(enc) => enc.config(),
}
}
/// Human-readable device tag: `"<Backend>/<adapter> (<precision>)"` or `"cpu"`.
/// The wgpu context this engine holds, when it is on the GPU. Downstream heads (GLiNER's span
/// head) reuse it rather than acquiring a SECOND adapter — two contexts would mean two devices,
/// two copies of every weight, and a needless round-trip between them.
pub fn gpu_ctx(&self) -> Option<&GpuCtx> {
match self {
Self::Gpu { ctx, .. } => Some(ctx),
Self::Cpu(_) => None,
}
}
/// The GPU device AND encoder, when running on one — the pair a task head needs to dispatch its
/// own GEMMs. Same rationale as `gpu_ctx`: reuse this device rather than acquiring a second
/// adapter. `None` on the CPU path, which is the signal to use the CPU kernels.
pub fn gpu_parts(&self) -> Option<(&GpuCtx, &EncoderGpu)> {
match self {
Self::Gpu { ctx, enc } => Some((ctx, enc)),
Self::Cpu(_) => None,
}
}
pub fn device(&self) -> String {
match self {
Self::Gpu { ctx, enc } => format!("{} ({})", ctx.backend, enc.precision()),
Self::Cpu(_) => "cpu".to_string(),
}
}
}