Struct CostModel 

pub struct CostModel {
    pub a: f64,
    pub b: f64,
    pub max_workspace_bytes: usize,
}

Quadratic-aware workspace cost model for ONNX attention inference.

BGE-M3 uses multi-head attention whose intermediate tensor footprint scales as O(batch * seq^2) (attention score matrix) plus O(batch * seq) (FFN intermediates, projection matrices). The total peak workspace is approximately:

peak ≈ a * (batch * seq) + b * (batch * seq^2)

where a (bytes/token-position) captures the FFN / projection contribution and b (bytes/token-position^2) captures the attention contribution.

At sequence length 512 attention is small relative to FFN, so a linear approximation works. At 8192, b * N^2 dominates by ~16×, so using only a would under-budget by that same factor.

Coefficients are derived at startup by crate::probe or set conservatively from compile-time defaults when measurement is unavailable.

Fields

a: f64

Bytes per token-position (linear term: FFN intermediates, projections).

b: f64

Bytes per token-position-squared (quadratic term: attention scores).

max_workspace_bytes: usize

Maximum workspace bytes available per worker for a single session.run() call.

Implementations

impl CostModel

pub const CONSERVATIVE_A: f64 = 16_384.0

Conservative static defaults calibrated so a (16, 512) chunk lands at ~140 MB workspace — matching the old static budget at the previous default BGE_M3_ONNX_BATCH_SIZE = 16, MAX_SEQ_LENGTH = 512.

These are used when the probe cannot run (no ORT, no model, macOS without cgroup support) or when BGE_M3_DISABLE_AUTO_BUDGET is set.

Formula check: 16 KiB/token × 16 × 512 + 8 B/token² × 16 × 512² = 16384 × 8192 + 8 × 16 × 262144 = 134 217 728 + 33 554 432 = 167 772 160 ≈ 160 MB per chunk (workers run sequentially inside one worker).

pub const CONSERVATIVE_B: f64 = 8.0

Conservative quadratic coefficient (bytes per token-position squared).

pub const DEFAULT_MAX_WORKSPACE: usize

Default maximum workspace per worker when memory cannot be detected.

2 GiB is conservatively safe for the Fargate 28 GiB task with 7 workers (28 GB * 0.7 safety / 7 workers ≈ 2.8 GB); we round down for headroom.

pub fn conservative(max_workspace_bytes: usize) -> Self

Constructs a CostModel with conservative defaults and the given workspace ceiling.

pub fn chunk_cost(&self, count: usize, max_seq: usize) -> u128

Estimated peak workspace (bytes) for a single session.run() call with count texts and max_seq as the padded sequence length.

Uses saturating arithmetic on u128 to avoid overflow at large inputs.

pub fn fits(&self, count: usize, max_seq: usize) -> bool

Returns true if the chunk fits within the workspace budget.

Trait Implementations

impl Clone for CostModel

fn clone(&self) -> CostModel

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl Debug for CostModel

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.

impl Copy for CostModel