bge_m3_embedding_server/embedder/
dual.rs

1// Copyright (c) 2026 J. Patrick Fulton
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! Paired dense + sparse embedding pipeline (one forward pass per chunk).
16
17use anyhow::Result;
18use ort::value::TensorRef;
19
20use super::error::ort_err;
21use super::jit_guard::{self, TrtJitGuard};
22use super::math::{normalize_l2, seq_len_distribution, sparse_maxpool, sparse_project};
23use super::tokenize::{build_chunk_arrays, tokenize_no_pad};
24use super::types::{DualEmbedding, EmbedStats, SparseEmbedding};
25use crate::binpack::{CostModel, bin_pack};
26use crate::config::ModelVariant;
27
28/// Produces paired dense + sparse embeddings using **one** `session.run()` per chunk.
29///
30/// Both projections are derived from the same forward pass:
31/// - **FP32**: extracts both `sentence_embedding` (dense) and `token_embeddings`
32///   (sparse base) from the model's dual outputs.
33/// - **FP16/INT8**: extracts dense from the CLS token (position 0) of
34///   `last_hidden_state`, and sparse from the full hidden states of the same
35///   tensor. This avoids a second forward pass.
36///
37/// Numerically equivalent to calling [`super::dense::embed_dense`] and
38/// [`super::sparse::embed_sparse`] separately, within FP rounding tolerance.
39///
40/// `guard` is the in-band TRT JIT admission guard; see
41/// [`super::dense::embed_dense`] for the refusal semantics. This route (the
42/// fused dual-output graph at `seq=8192`) is the one most prone to the
43/// pathological autotuner allocation the guard prevents.
44#[allow(clippy::cast_possible_truncation, clippy::too_many_lines)]
45pub(super) fn embed_both(
46    session: &mut ort::session::Session,
47    tokenizer: &tokenizers::Tokenizer,
48    texts: &[String],
49    cost_model: &CostModel,
50    model_variant: ModelVariant,
51    guard: Option<&TrtJitGuard>,
52) -> Result<(Vec<DualEmbedding>, EmbedStats)> {
53    let (weight, bias) = crate::weights::sparse_linear();
54    let weight_view = weight.view();
55
56    let tokenize_start = std::time::Instant::now();
57    let encodings = tokenize_no_pad(tokenizer, texts)?;
58    let seq_lens: Vec<usize> = encodings.iter().map(|e| e.get_ids().len()).collect();
59    let tokenize_ms = u64::try_from(tokenize_start.elapsed().as_millis()).unwrap_or(u64::MAX);
60
61    let seq_dist = seq_len_distribution(&seq_lens);
62    let total_token_positions: usize = seq_lens.iter().sum();
63    let chunks = bin_pack(&seq_lens, cost_model);
64    jit_guard::guard_chunks(guard, &chunks, &seq_lens).map_err(anyhow::Error::new)?;
65
66    let mut all_dual: Vec<Option<DualEmbedding>> = (0..texts.len()).map(|_| None).collect();
67
68    let mut max_chunk_seq: usize = 0;
69    let mut inference_ms: u64 = 0;
70
71    for (chunk_idx, chunk_indices) in chunks.iter().enumerate() {
72        let chunk_max = chunk_indices
73            .iter()
74            .map(|&i| seq_lens[i])
75            .max()
76            .unwrap_or(1)
77            .max(1);
78
79        max_chunk_seq = max_chunk_seq.max(chunk_max);
80
81        let (ids_array, mask_array) = build_chunk_arrays(&encodings, chunk_indices, chunk_max)?;
82
83        let ids_tensor = TensorRef::from_array_view(ids_array.view()).map_err(ort_err)?;
84        let mask_tensor = TensorRef::from_array_view(mask_array.view()).map_err(ort_err)?;
85
86        let chunk_start = std::time::Instant::now();
87        let outputs = {
88            let _span = tracing::debug_span!(
89                "chunk",
90                chunk_idx,
91                batch = chunk_indices.len(),
92                max_seq = chunk_max
93            )
94            .entered();
95            session
96                .run(ort::inputs! {
97                    "input_ids" => ids_tensor,
98                    "attention_mask" => mask_tensor,
99                })
100                .map_err(ort_err)?
101        };
102        let chunk_ms = u64::try_from(chunk_start.elapsed().as_millis()).unwrap_or(u64::MAX);
103        inference_ms = inference_ms.saturating_add(chunk_ms);
104        tracing::debug!(
105            chunk_idx,
106            batch = chunk_indices.len(),
107            max_seq = chunk_max,
108            elapsed_ms = chunk_ms,
109            "both chunk inference complete"
110        );
111
112        // Extract dense + token-level hidden states from the same outputs.
113        // FP32: separate sentence_embedding + token_embeddings outputs.
114        // FP16/INT8: derive dense (CLS) and sparse-base from last_hidden_state.
115        let (dense_emb, token_emb) = match model_variant {
116            ModelVariant::Fp32 => {
117                let dense = outputs["sentence_embedding"]
118                    .try_extract_array::<f32>()
119                    .map_err(ort_err)?
120                    .to_owned();
121                let tokens = outputs["token_embeddings"]
122                    .try_extract_array::<f32>()
123                    .map_err(ort_err)?
124                    .to_owned();
125                (dense, tokens)
126            }
127            ModelVariant::Fp16 | ModelVariant::Int8 => {
128                let lhs = outputs["last_hidden_state"]
129                    .try_extract_array::<f32>()
130                    .map_err(ort_err)?;
131                let dense = lhs.index_axis(ndarray::Axis(1), 0).to_owned();
132                let tokens = lhs.to_owned();
133                (dense, tokens)
134            }
135        };
136
137        for (chunk_pos, &orig_idx) in chunk_indices.iter().enumerate() {
138            // Dense: CLS row, L2-normalized.
139            let dense_row = dense_emb.index_axis(ndarray::Axis(0), chunk_pos);
140            let mut dense_vec = dense_row
141                .as_slice()
142                .expect("dense embedding should be contiguous")
143                .to_vec();
144            normalize_l2(&mut dense_vec);
145
146            // Sparse: project each token's hidden state, then max-pool.
147            let enc = &encodings[orig_idx];
148            let ids = enc.get_ids();
149            let mask = enc.get_attention_mask();
150            let batch_hidden = token_emb.index_axis(ndarray::Axis(0), chunk_pos);
151
152            let scores: Vec<f32> = (0..ids.len())
153                .map(|j| {
154                    let hidden = batch_hidden.index_axis(ndarray::Axis(0), j);
155                    let hidden_slice = hidden
156                        .as_slice()
157                        .expect("hidden state should be contiguous");
158                    sparse_project(hidden_slice, &weight_view, *bias)
159                })
160                .collect();
161
162            let (indices, values) = sparse_maxpool(ids, mask, &scores);
163
164            all_dual[orig_idx] = Some(DualEmbedding {
165                dense: dense_vec,
166                sparse: SparseEmbedding { indices, values },
167            });
168        }
169    }
170
171    let stats = EmbedStats {
172        chunks: chunks.len(),
173        max_chunk_seq,
174        total_token_positions,
175        tokenize_ms,
176        inference_ms,
177        seq_len_min: seq_dist.min,
178        seq_len_max: seq_dist.max,
179        seq_len_mean: seq_dist.mean,
180        seq_len_p95: seq_dist.p95,
181    };
182
183    Ok((
184        all_dual
185            .into_iter()
186            .map(|d| d.expect("every slot must be filled"))
187            .collect(),
188        stats,
189    ))
190}
bge_m3_embedding_server/embedder/dual.rs

bge_m3_embedding_server/embedder/
dual.rs
