wehub-resource-sync
491 lines
21 KiB
Python
491 lines
21 KiB
Python
"""Asynchronous embedding inference engine for encoder and decoder models."""
|
||
|
||
import asyncio
|
||
import concurrent.futures
|
||
import json
|
||
import os
|
||
from typing import List, Literal, Optional, Tuple, Union # noqa: UP035
|
||
|
||
import numpy as np
|
||
import tvm
|
||
from tvm import relax
|
||
from tvm.runtime import Device
|
||
from tvm_ffi import Shape
|
||
|
||
from mlc_llm.serve import engine_utils
|
||
from mlc_llm.support.auto_device import detect_device
|
||
from mlc_llm.tokenizers import Tokenizer
|
||
|
||
|
||
class AsyncEmbeddingEngine:
|
||
"""Asynchronous embedding inference engine.
|
||
|
||
Supports both encoder models (BERT-style) and decoder-only embedding models
|
||
(e.g. Qwen3-Embeddings). Uses a ThreadPoolExecutor for background inference
|
||
so that the asyncio event loop is not blocked.
|
||
|
||
Parameters
|
||
----------
|
||
model : str
|
||
Path to the model weight directory.
|
||
|
||
model_lib : str
|
||
Path to the compiled model library (.so/.dylib file).
|
||
|
||
device : Union[str, Device]
|
||
Device string, e.g. "auto", "cuda:0", "metal".
|
||
|
||
pooling_strategy : Optional[str]
|
||
Pooling strategy: "cls" (first token), "mean" (masked average),
|
||
or "last" (last token). If None, auto-detected based on model type:
|
||
encoder -> "cls", decoder -> "last".
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
model: str,
|
||
model_lib: str,
|
||
device: Union[str, Device] = "auto",
|
||
*,
|
||
pooling_strategy: Optional[str] = None,
|
||
) -> None:
|
||
# Reuse existing utility: device detection
|
||
self.device = detect_device(device) if isinstance(device, str) else device
|
||
# Reuse existing utility: tokenizer
|
||
self.tokenizer = Tokenizer(model)
|
||
|
||
# Load TVM module, metadata, and params via engine_utils helpers
|
||
ex = tvm.runtime.load_module(model_lib)
|
||
vm = relax.VirtualMachine(ex, device=self.device)
|
||
self._mod = vm.module
|
||
self._metadata = json.loads(self._mod["_metadata"]())
|
||
self._params = engine_utils.load_embedding_params(model, self.device, self._metadata)
|
||
|
||
# Detect model type and set pooling strategy
|
||
self.embedding_metadata = engine_utils.get_embedding_metadata(self._metadata)
|
||
if self.embedding_metadata:
|
||
self.model_type = self.embedding_metadata["model_type"]
|
||
self.pooling_strategy = self.embedding_metadata["pooling_strategy"]
|
||
self.normalize = self.embedding_metadata["normalize"]
|
||
else:
|
||
self.model_type = engine_utils.detect_embedding_model_type(self._mod)
|
||
self.pooling_strategy = "cls" if self.model_type == "encoder" else "last"
|
||
self.normalize = True
|
||
# Allow caller to override pooling strategy
|
||
if pooling_strategy:
|
||
self.pooling_strategy = pooling_strategy
|
||
|
||
# Initialize model-type-specific functions
|
||
if self.model_type == "encoder":
|
||
self._init_encoder(model)
|
||
else:
|
||
self._init_decoder(model)
|
||
|
||
# Background thread pool (1 worker = serialized GPU inference)
|
||
self._executor = concurrent.futures.ThreadPoolExecutor(
|
||
max_workers=1, thread_name_prefix="embedding"
|
||
)
|
||
self._terminated = False
|
||
|
||
def _init_encoder(self, model: str) -> None:
|
||
"""Initialize encoder (BERT-style) model functions and special tokens."""
|
||
self._prefill_func = self._mod["prefill"]
|
||
self._cls_token_id: Optional[int] = None
|
||
self._sep_token_id: Optional[int] = None
|
||
tok_config_path = os.path.join(model, "tokenizer_config.json")
|
||
if os.path.exists(tok_config_path):
|
||
with open(tok_config_path, encoding="utf-8") as f:
|
||
tok_config = json.load(f)
|
||
# Try added_tokens_decoder first (newer HF format)
|
||
added = tok_config.get("added_tokens_decoder", {})
|
||
for tid, info in added.items():
|
||
if info.get("content") == tok_config.get("cls_token"):
|
||
self._cls_token_id = int(tid)
|
||
if info.get("content") == tok_config.get("sep_token"):
|
||
self._sep_token_id = int(tid)
|
||
# Fallback: encode the special token strings via tokenizer
|
||
if self._cls_token_id is None and tok_config.get("cls_token"):
|
||
ids = list(self.tokenizer.encode(tok_config["cls_token"]))
|
||
if len(ids) == 1:
|
||
self._cls_token_id = ids[0]
|
||
if self._sep_token_id is None and tok_config.get("sep_token"):
|
||
ids = list(self.tokenizer.encode(tok_config["sep_token"]))
|
||
if len(ids) == 1:
|
||
self._sep_token_id = ids[0]
|
||
|
||
def _init_decoder(self, model: str) -> None:
|
||
"""Initialize decoder (Qwen3-Embeddings style) model functions."""
|
||
# Prefer tokenizer post-processing (HF-style) for terminal/pooling token handling.
|
||
# Only fall back to manual EOS append when tokenizer does not define a post-processor
|
||
# that actually appends a token at the end of the sequence.
|
||
self._decoder_tokenizer_appends_eos = False
|
||
tokenizer_json_path = os.path.join(model, "tokenizer.json")
|
||
if os.path.exists(tokenizer_json_path):
|
||
with open(tokenizer_json_path, encoding="utf-8") as f:
|
||
tokenizer_json = json.load(f)
|
||
post_proc = tokenizer_json.get("post_processor")
|
||
if post_proc is not None:
|
||
# Check if the post-processor actually appends a special token at the end
|
||
# (e.g. TemplateProcessing with "$A <|endoftext|>"). We verify by encoding
|
||
# a test string and checking if the last token is a known special token.
|
||
test_tokens = list(self.tokenizer.encode("test"))
|
||
if len(test_tokens) > 0:
|
||
vocab = tokenizer_json.get("added_tokens", [])
|
||
special_ids = {t["id"] for t in vocab if t.get("special", False)}
|
||
if test_tokens[-1] in special_ids:
|
||
self._decoder_tokenizer_appends_eos = True
|
||
|
||
# Read EOS token from config — fallback only when tokenizer does not auto-append.
|
||
self._decoder_eos_token_id: Optional[int] = None
|
||
config_path = os.path.join(model, "mlc-chat-config.json")
|
||
if os.path.exists(config_path):
|
||
with open(config_path, encoding="utf-8") as f:
|
||
chat_config = json.load(f)
|
||
eos = chat_config.get("eos_token_id")
|
||
if isinstance(eos, list):
|
||
self._decoder_eos_token_id = eos[0]
|
||
elif isinstance(eos, int):
|
||
self._decoder_eos_token_id = eos
|
||
|
||
self._embed_func = self._mod["embed"]
|
||
self._prefill_to_hidden_func = self._mod["prefill_to_last_hidden_states"]
|
||
self._batch_prefill_to_hidden_func = self._mod["batch_prefill_to_last_hidden_states"]
|
||
if self._mod.implements_function("create_tir_paged_kv_cache"):
|
||
self._create_kv_cache_func = self._mod["create_tir_paged_kv_cache"]
|
||
elif self._mod.implements_function("create_flashinfer_paged_kv_cache"):
|
||
self._create_kv_cache_func = self._mod["create_flashinfer_paged_kv_cache"]
|
||
else:
|
||
raise RuntimeError("Cannot find KV cache creation function in model library.")
|
||
self._kv_state_add_sequence = tvm.get_global_func("vm.builtin.kv_state_add_sequence")
|
||
self._kv_state_remove_sequence = tvm.get_global_func("vm.builtin.kv_state_remove_sequence")
|
||
self._kv_state_begin_forward = tvm.get_global_func("vm.builtin.kv_state_begin_forward")
|
||
self._kv_state_end_forward = tvm.get_global_func("vm.builtin.kv_state_end_forward")
|
||
self._nd_reshape = tvm.get_global_func("vm.builtin.reshape")
|
||
|
||
def embed(self, inputs: List[str]) -> Tuple[List[List[float]], int]: # noqa: UP006
|
||
"""Compute embeddings for a list of input strings (synchronous).
|
||
|
||
Parameters
|
||
----------
|
||
inputs : List[str]
|
||
The input strings to embed.
|
||
|
||
Returns
|
||
-------
|
||
embeddings : List[List[float]]
|
||
The L2-normalized embedding vectors.
|
||
total_tokens : int
|
||
Total number of tokens processed.
|
||
"""
|
||
if self.model_type == "encoder":
|
||
return self._embed_encoder(inputs)
|
||
return self._embed_decoder(inputs)
|
||
|
||
async def async_embed(self, inputs: List[str]) -> Tuple[List[List[float]], int]: # noqa: UP006
|
||
"""Compute embeddings asynchronously in a background thread.
|
||
|
||
This method does not block the asyncio event loop.
|
||
|
||
Parameters
|
||
----------
|
||
inputs : List[str]
|
||
The input strings to embed.
|
||
|
||
Returns
|
||
-------
|
||
embeddings : List[List[float]]
|
||
The L2-normalized embedding vectors.
|
||
total_tokens : int
|
||
Total number of tokens processed.
|
||
"""
|
||
loop = asyncio.get_running_loop()
|
||
return await loop.run_in_executor(self._executor, self.embed, inputs)
|
||
|
||
def _embed_encoder(
|
||
self,
|
||
inputs: List[str], # noqa: UP006
|
||
) -> Tuple[List[List[float]], int]: # noqa: UP006
|
||
"""Encoder model embedding (BERT-style).
|
||
|
||
Processes each input individually to avoid batch padding artifacts.
|
||
|
||
Encoder uses bidirectional attention, so chunked prefill is NOT possible
|
||
(each token must attend to all other tokens in the full sequence).
|
||
Inputs exceeding prefill_chunk_size are truncated.
|
||
|
||
(Additional Strategy)
|
||
TODO: For better long-text support, implement sliding window + mean pooling:
|
||
1. Split text into overlapping windows of prefill_chunk_size (stride=chunk/2)
|
||
2. Encode each window independently
|
||
3. Mean-pool all window embeddings → final embedding → L2 normalize
|
||
This preserves information from the full text at the cost of N× compute.
|
||
""" # noqa: RUF002
|
||
embeddings: List[List[float]] = [] # noqa: UP006
|
||
total_tokens = 0
|
||
prefill_chunk = self._metadata.get("prefill_chunk_size", 512)
|
||
|
||
for text in inputs:
|
||
tokens = list(self.tokenizer.encode(text))
|
||
# Add [CLS] and [SEP] if needed
|
||
if self._cls_token_id is not None and (
|
||
len(tokens) == 0 or tokens[0] != self._cls_token_id
|
||
):
|
||
tokens = [self._cls_token_id, *tokens]
|
||
if self._sep_token_id is not None and (
|
||
len(tokens) == 0 or tokens[-1] != self._sep_token_id
|
||
):
|
||
tokens = [*tokens, self._sep_token_id]
|
||
|
||
# Truncate to compiled buffer limit (keep [CLS] at start, [SEP] at end)
|
||
if len(tokens) > prefill_chunk:
|
||
tokens = tokens[:prefill_chunk]
|
||
if self._sep_token_id is not None:
|
||
tokens[-1] = self._sep_token_id
|
||
|
||
seq_len = len(tokens)
|
||
total_tokens += seq_len
|
||
|
||
token_ids = np.array([tokens], dtype=np.int32) # [1, seq_len]
|
||
attention_mask: np.ndarray = np.ones((1, seq_len), dtype=np.int32) # [1, seq_len]
|
||
|
||
tokens_tvm = tvm.runtime.tensor(token_ids, device=self.device)
|
||
mask_tvm = tvm.runtime.tensor(attention_mask, device=self.device)
|
||
|
||
output = self._prefill_func(tokens_tvm, mask_tvm, self._params)
|
||
# .numpy() copies to CPU, escaping TVM workspace buffer reuse across calls.
|
||
output_np = output.numpy() # [1, seq_len, hidden_size]
|
||
|
||
# Pooling
|
||
if self.pooling_strategy == "cls":
|
||
pooled = output_np[0, 0, :]
|
||
elif self.pooling_strategy == "mean":
|
||
pooled = output_np[0].mean(axis=0)
|
||
else: # "last"
|
||
pooled = output_np[0, -1, :]
|
||
|
||
# L2 normalize
|
||
pooled = pooled.astype(np.float32)
|
||
if self.normalize:
|
||
norm = np.linalg.norm(pooled)
|
||
if norm > 1e-12:
|
||
pooled = pooled / norm
|
||
|
||
embeddings.append(pooled.tolist())
|
||
|
||
return embeddings, total_tokens
|
||
|
||
def _embed_decoder(self, inputs: List[str]) -> Tuple[List[List[float]], int]: # noqa: UP006
|
||
"""Decoder model embedding with batch prefill optimization.
|
||
|
||
When total tokens fit within prefill_chunk_size, all inputs are processed
|
||
in a single batch forward pass using shared KV cache. Otherwise, falls back
|
||
to sequential chunked prefill per input.
|
||
"""
|
||
# Read KV cache config from metadata
|
||
prefill_chunk = self._metadata.get("prefill_chunk_size", 2048)
|
||
max_seq_len = self._metadata.get("context_window_size", 32768)
|
||
if max_seq_len == -1:
|
||
max_seq_len = self._metadata.get("sliding_window_size", -1)
|
||
assert max_seq_len > 0, f"max_seq_len must be positive, got {max_seq_len}"
|
||
support_sliding = int(self._metadata.get("sliding_window_size", -1) != -1)
|
||
|
||
# Tokenize all inputs. Prefer tokenizer post-processor output. If absent (older models),
|
||
# fall back to appending eos_token_id when missing.
|
||
token_lists: List[List[int]] = [] # noqa: UP006
|
||
for text in inputs:
|
||
tokens = list(self.tokenizer.encode(text))
|
||
if (
|
||
not self._decoder_tokenizer_appends_eos
|
||
and self._decoder_eos_token_id is not None
|
||
and (len(tokens) == 0 or tokens[-1] != self._decoder_eos_token_id)
|
||
):
|
||
tokens.append(self._decoder_eos_token_id)
|
||
if len(tokens) > max_seq_len:
|
||
tokens = tokens[:max_seq_len]
|
||
token_lists.append(tokens)
|
||
|
||
total_tokens = sum(len(t) for t in token_lists)
|
||
|
||
# Fast path: all tokens fit in one prefill chunk → batch forward
|
||
if total_tokens <= prefill_chunk and all(len(t) > 0 for t in token_lists):
|
||
return self._batch_embed_decoder(
|
||
token_lists, total_tokens, max_seq_len, prefill_chunk, support_sliding
|
||
)
|
||
|
||
# Greedy sub-batching: pack texts into sub-batches that fit within
|
||
# prefill_chunk, preserving input order. Oversize texts (single text
|
||
# exceeding prefill_chunk) fall back to sequential chunked prefill.
|
||
sub_batches = self._build_sub_batches(token_lists, prefill_chunk)
|
||
all_embeddings: List[List[float]] = [] # noqa: UP006
|
||
for batch_type, batch, batch_total in sub_batches:
|
||
if batch_type == "batch":
|
||
embs, _ = self._batch_embed_decoder(
|
||
batch, batch_total, max_seq_len, prefill_chunk, support_sliding
|
||
)
|
||
else:
|
||
embs, _ = self._sequential_embed_decoder(
|
||
batch, batch_total, max_seq_len, prefill_chunk, support_sliding
|
||
)
|
||
all_embeddings.extend(embs)
|
||
|
||
return all_embeddings, total_tokens
|
||
|
||
@staticmethod
|
||
def _build_sub_batches(
|
||
token_lists: List[List[int]], # noqa: UP006
|
||
prefill_chunk: int,
|
||
) -> List[Tuple[Literal["batch", "sequential"], List[List[int]], int]]: # noqa: UP006
|
||
"""Partition token lists into sub-batches that fit within prefill_chunk.
|
||
|
||
Each sub-batch is a tuple of (mode, token_lists, total_token_count).
|
||
Empty token lists are skipped to avoid invalid batch processing.
|
||
"""
|
||
sub_batches: List[Tuple[Literal["batch", "sequential"], List[List[int]], int]] = [] # noqa: UP006
|
||
current_batch: List[List[int]] = [] # noqa: UP006
|
||
current_tokens = 0
|
||
|
||
for tokens in token_lists:
|
||
if not tokens:
|
||
continue
|
||
token_len = len(tokens)
|
||
is_oversized = token_len > prefill_chunk
|
||
if current_batch and (is_oversized or current_tokens + token_len > prefill_chunk):
|
||
sub_batches.append(("batch", current_batch, current_tokens))
|
||
current_batch, current_tokens = [], 0
|
||
if is_oversized:
|
||
sub_batches.append(("sequential", [tokens], token_len))
|
||
else:
|
||
current_batch.append(tokens)
|
||
current_tokens += token_len
|
||
if current_batch:
|
||
sub_batches.append(("batch", current_batch, current_tokens))
|
||
|
||
return sub_batches
|
||
|
||
def _batch_embed_decoder(
|
||
self,
|
||
token_lists: List[List[int]], # noqa: UP006
|
||
total_tokens: int,
|
||
max_seq_len: int,
|
||
prefill_chunk: int,
|
||
support_sliding: int,
|
||
) -> Tuple[List[List[float]], int]: # noqa: UP006
|
||
"""Batch prefill: process all inputs in a single forward pass."""
|
||
batch_size = len(token_lists)
|
||
|
||
# Create KV cache for the entire batch
|
||
kv_cache = self._create_kv_cache_func(
|
||
Shape([batch_size]),
|
||
Shape([max_seq_len]),
|
||
Shape([prefill_chunk]),
|
||
Shape([16]),
|
||
Shape([support_sliding]),
|
||
)
|
||
|
||
# Register all sequences
|
||
seq_ids = list(range(batch_size))
|
||
seq_lens = [len(t) for t in token_lists]
|
||
for sid in seq_ids:
|
||
self._kv_state_add_sequence(kv_cache, sid)
|
||
|
||
# Begin forward with all sequences at once
|
||
self._kv_state_begin_forward(kv_cache, Shape(seq_ids), Shape(seq_lens))
|
||
|
||
# Concatenate all tokens → embed → batch prefill
|
||
all_tokens = []
|
||
for tokens in token_lists:
|
||
all_tokens.extend(tokens)
|
||
token_ids = tvm.runtime.tensor(np.array(all_tokens, dtype=np.int32), device=self.device)
|
||
all_embed = self._embed_func(token_ids, self._params)
|
||
all_embed = self._nd_reshape(all_embed, Shape([1, total_tokens, all_embed.shape[-1]]))
|
||
|
||
hidden_states, _ = self._batch_prefill_to_hidden_func(all_embed, kv_cache, self._params)
|
||
# .numpy() copies to CPU, escaping TVM workspace buffer reuse across calls.
|
||
# (torch.from_dlpack is zero-copy and hits aliasing bugs on 2nd+ invocation.)
|
||
hidden_np = hidden_states.numpy()
|
||
self._kv_state_end_forward(kv_cache)
|
||
for sid in seq_ids:
|
||
self._kv_state_remove_sequence(kv_cache, sid)
|
||
|
||
# Extract last token hidden state per sequence
|
||
embeddings: List[List[float]] = [] # noqa: UP006
|
||
offset = 0
|
||
for tokens in token_lists:
|
||
last_pos = offset + len(tokens) - 1
|
||
pooled = hidden_np[0, last_pos, :].astype(np.float32)
|
||
if self.normalize:
|
||
norm = np.linalg.norm(pooled)
|
||
if norm > 1e-12:
|
||
pooled = pooled / norm
|
||
embeddings.append(pooled.tolist())
|
||
offset += len(tokens)
|
||
|
||
return embeddings, total_tokens
|
||
|
||
def _sequential_embed_decoder(
|
||
self,
|
||
token_lists: List[List[int]], # noqa: UP006
|
||
total_tokens: int,
|
||
max_seq_len: int,
|
||
prefill_chunk: int,
|
||
support_sliding: int,
|
||
) -> Tuple[List[List[float]], int]: # noqa: UP006
|
||
"""Sequential chunked prefill: process each input independently."""
|
||
embeddings: List[List[float]] = [] # noqa: UP006
|
||
|
||
for tokens in token_lists:
|
||
if len(tokens) == 0:
|
||
continue
|
||
|
||
# Create KV cache for this single sequence
|
||
kv_cache = self._create_kv_cache_func(
|
||
Shape([1]),
|
||
Shape([max_seq_len]),
|
||
Shape([prefill_chunk]),
|
||
Shape([16]),
|
||
Shape([support_sliding]),
|
||
)
|
||
self._kv_state_add_sequence(kv_cache, 0)
|
||
|
||
# Process tokens in chunks
|
||
hidden = None
|
||
for chunk_start in range(0, len(tokens), prefill_chunk):
|
||
chunk_end = min(chunk_start + prefill_chunk, len(tokens))
|
||
chunk_tokens = tokens[chunk_start:chunk_end]
|
||
chunk_len = len(chunk_tokens)
|
||
|
||
token_ids = tvm.runtime.tensor(
|
||
np.array(chunk_tokens, dtype=np.int32), device=self.device
|
||
)
|
||
chunk_embed = self._embed_func(token_ids, self._params)
|
||
chunk_embed = self._nd_reshape(
|
||
chunk_embed, Shape([1, chunk_len, chunk_embed.shape[-1]])
|
||
)
|
||
self._kv_state_begin_forward(kv_cache, Shape([0]), Shape([chunk_len]))
|
||
hidden, kv_cache = self._prefill_to_hidden_func(chunk_embed, kv_cache, self._params)
|
||
# .numpy() copies to CPU, escaping TVM buffer aliasing.
|
||
hidden_np = hidden.numpy()
|
||
self._kv_state_end_forward(kv_cache)
|
||
|
||
self._kv_state_remove_sequence(kv_cache, 0)
|
||
|
||
pooled = hidden_np[0, -1, :] if hidden_np.ndim == 3 else hidden_np[-1, :]
|
||
pooled = pooled.astype(np.float32)
|
||
if self.normalize:
|
||
norm = np.linalg.norm(pooled)
|
||
if norm > 1e-12:
|
||
pooled = pooled / norm
|
||
embeddings.append(pooled.tolist())
|
||
|
||
return embeddings, total_tokens
|
||
|
||
def terminate(self) -> None:
|
||
"""Terminate the engine and clean up the thread pool."""
|
||
if getattr(self, "_terminated", True):
|
||
return
|
||
self._terminated = True
|
||
self._executor.shutdown(wait=False)
|
||
|
||
def __del__(self):
|
||
self.terminate()
|