Serve embeddings with BEI

Baseten Embeddings Inference is Baseten’s solution for production grade inference on embedding, classification and reranking models using TensorRT-LLM. With Baseten Embeddings Inference you get the following benefits:

Lowest-latency inference across any embedding solution (vLLM, SGlang, Infinity, TEI, Ollama)¹
Highest-throughput inference across any embedding solution (vLLM, SGlang, Infinity, TEI, Ollama) - thanks to XQA kernels, FP8 and dynamic batching.²
High parallelism: up to 1400 client embeddings per second
Cached model weights for fast vertical scaling and high availability - no Hugging Face hub dependency at runtime
Ahead-of-time compilation, memory allocation and fp8 post-training quantization

Get started with embedding models:

Embedding models are LLMs without a lm_head for language generation. Typical architectures that are supported for embeddings are LlamaModel, BertModel, RobertaModel or Gemma2Model, and contain the safetensors, config, tokenizer and sentence-transformer config files. A good example is the repo BAAI/bge-multilingual-gemma2. To deploy a model for embeddings, set the following config in your local directory.

config.yaml

model_name: BEI-Linq-Embed-Mistral
resources:
  accelerator: RTX_PRO_6000:2
  use_gpu: true
trt_llm:
  build:
    base_model: encoder
    checkpoint_repository:
      # for a different model, change the repo to e.g. to "Salesforce/SFR-Embedding-Mistral"
      # "BAAI/bge-en-icl" or "BAAI/bge-m3"
      repo: "Linq-AI-Research/Linq-Embed-Mistral"
      revision: main
      source: HF
    # only Llama, Mistral and Qwen Models support quantization.
    # others, use: "quantization_type: no_quant"
    quantization_type: fp8

The embedding and reranking examples use 2 NVIDIA RTX PRO 6000 GPUs per replica. BEI loads the complete model on each GPU and distributes batches across them. Each GPU must fit the model and its batch; GPU memory is not combined. Use RTX_PRO_6000 for 1 GPU. Leave tensor_parallel_count unset. With config.yaml in your local directory, deploy the model to Baseten and promote it to the production environment:

baseten model push
baseten model deployment promote --model-id <model-id> --deployment-id <deployment-id>

uvx truss push --promote

Deployed embedding models are OpenAI compatible without any additional settings. You may use the client code below to consume the model.

from openai import OpenAI
import os

client = OpenAI(
    api_key=os.environ['BASETEN_API_KEY'],
    # add the deployment URL
    base_url="https://model-xxxxxx.api.baseten.co/environments/production/sync/v1"
)

embedding = client.embeddings.create(
    input=["Baseten Embeddings are fast.", "Embed this sentence!"],
    model="not-required"
)

Example deployment of classification, reranking, and classification models

Besides embedding models, BEI deploys high-throughput rerank and classification models. You can identify suitable architectures by their ForSequenceClassification suffix in the Hugging Face repo. The use-case for these models is either Reward Modeling, Reranking documents in RAG or tasks like content moderation.

model_name: BEI-mixedbread-rerank-large-v2-fp8
resources:
  accelerator: RTX_PRO_6000:2
  cpu: '1'
  memory: 10Gi
  use_gpu: true
trt_llm:
  build:
    base_model: encoder
    checkpoint_repository:
      repo: michaelfeil/mxbai-rerank-large-v2-seq
      revision: main
      source: HF
    # only Llama, Mistral and Qwen Models support quantization
    quantization_type: fp8

As OpenAI does not offer reranking or classification, we are sending a simple request to the endpoint. Depending on the model, you might want to apply a specific prompt template first.

import requests
import os

headers = {
    f"Authorization": f"Bearer {os.environ['BASETEN_API_KEY']}"
}

# model specific prompt for mixedbread's reranker v2.
prompt = (
  "<|endoftext|><|im_start|>system\nYou are Qwen, created by Alibaba Cloud. You are a helpful assistant.\n<|im_end|>\n<|im_start|>user\n"
  "query: {query} \ndocument: {doc} \nYou are a search relevance expert who evaluates how well documents match search queries. For each query-document pair, carefully analyze the semantic relationship between them, then provide your binary relevance judgment (0 for not relevant, 1 for relevant).\nRelevance:<|im_end|>\n<|im_start|>assistant\n"
).format(query="What is Baseten?",doc="Baseten is a fast inference provider.")

requests.post(
    headers=headers,
    url="https://model-xxxxxx.api.baseten.co/environments/production/sync/predict",
    json={
        "inputs": prompt,
        "raw_scores": True,
    }
)

Benchmarks and performance optimizations

Embedding models on BEI are fast, and offer currently the fastest implementation for embeddings across all open-source and closed-source providers. The team behind the implementation is the authors of infinity. We recommend using fp8 quantization for Llama, Mistral, and Qwen2 models on L4 or newer (L4, H100, H200, and B200). Quality difference between fp8 and bfloat16 is often negligible: embedding models often retain >99% cosine similarity between both precisions, and reranking models retain the ranking order despite a difference in the retained output. For more details, check out the technical launch post.

The team at Baseten has additional options for sharing cached model weights across replicas - for fast horizontal scaling. Please contact us to enable this option.

Deploy custom or fine-tuned models on BEI

We support the deployment of the below models, as well as all finetuned variants of these models (same architecture & customized weights). The following repositories are supported - this list is not exhaustive.

Model Repository	Architecture	Function
`Salesforce/SFR-Embedding-Mistral`	MistralModel	embedding
`BAAI/bge-m3`	BertModel	embedding
`BAAI/bge-multilingual-gemma2`	Gemma2Model	embedding
`mixedbread-ai/mxbai-embed-large-v1`	BertModel	embedding
`BAAI/bge-large-en-v1.5`	BertModel	embedding
`allenai/Llama-3.1-Tulu-3-8B-RM`	LlamaForSequenceClassification	classifier
`ncbi/MedCPT-Cross-Encoder`	BertForSequenceClassification	reranker/classifier
`SamLowe/roberta-base-go_emotions`	XLMRobertaForSequenceClassification	classifier
`mixedbread/mxbai-rerank-large-v2-seq`	Qwen2ForSequenceClassification	reranker/classifier
`BAAI/bge-en-icl`	LlamaModel	embedding
`BAAI/bge-reranker-v2-m3`	BertForSequenceClassification	reranker/classifier
`Skywork/Skywork-Reward-Llama-3.1-8B-v0.2`	LlamaForSequenceClassification	classifier
`Snowflake/snowflake-arctic-embed-l`	BertModel	embedding
`nomic-ai/nomic-embed-code`	Qwen2Model	embedding

¹ measured on H100-HBM3 (bert-large-335M, for BAAI/bge-en-icl: 9ms) ² measured on H100-HBM3 (leading model architecture on MTEB, MistralModel-7B)

Deploy a reranker on two L4s

Create config.yaml to deploy the ModernBERT reranker with BEI-Bert:

config.yaml

model_name: modernbert-reranker
resources:
  accelerator: L4:2
  use_gpu: true
trt_llm:
  inference_stack: v1
  build:
    base_model: encoder_bert
    checkpoint_repository:
      source: HF
      repo: Alibaba-NLP/gte-reranker-modernbert-base
    quantization_type: no_quant
  runtime:
    webserver_default_route: /rerank

Run baseten model push from this directory. Set accelerator: L4:4 for four GPUs. Each GPU loads a complete model and consumes batches from a shared queue. Set checkpoint_repository.revision to the commit you validated for reproducible deployments. For small embedding, reranking, and decision models, start with two or four L4s or RTX 6000 Ada GPUs where available. Select RTX resources with your Baseten representative. Measure throughput and p95 latency with concurrent traffic. See multi-GPU sizing for per-GPU memory requirements.

Text model request examples

Each example requires a deployment of its named checkpoint. Set the model’s production URL and your API key:

export MODEL_URL="https://model-<MODEL_ID>.api.baseten.co/environments/production/sync"
export BASETEN_API_KEY="<API_KEY>"

See the model matrix for endpoint and precision requirements. For Laya, use the typed decision deployment.

Qwen3 embeddings

With Qwen/Qwen3-Embedding-0.6B deployed, include the retrieval instruction in queries. Embed documents without the query instruction.

curl --fail-with-body "$MODEL_URL/v1/embeddings" \
  -H "Authorization: Bearer $BASETEN_API_KEY" \
  -H 'Content-Type: application/json' \
  -d '{"model":"Qwen/Qwen3-Embedding-0.6B","input":["Instruct: Given a web search query, retrieve relevant passages that answer the query\nQuery: What causes tides?"]}'

Query and document prompts

With michaelfeil/embeddinggemma-300m deployed, use its query and document prompts:

curl --fail-with-body "$MODEL_URL/v1/embeddings" \
  -H "Authorization: Bearer $BASETEN_API_KEY" \
  -H 'Content-Type: application/json' \
  -d '{"model":"michaelfeil/embeddinggemma-300m","input":["task: search result | query: What causes tides?","title: Tides | text: The Moon and Sun exert gravitational forces on the oceans."]}'

For nomic-ai/modernbert-embed-base, use the same request shape with search_query: and search_document: prefixes. For BAAI/bge-m3, send text without a prefix. Change the model value to match your deployment.

Rank documents

With Alibaba-NLP/gte-reranker-modernbert-base deployed:

curl --fail-with-body "$MODEL_URL/rerank" \
  -H "Authorization: Bearer $BASETEN_API_KEY" \
  -H 'Content-Type: application/json' \
  -d '{"query":"What causes tides?","texts":["The Moon and Sun exert gravitational forces on the oceans.","Clouds form when water vapor condenses."]}'

Sequence classification

With SamLowe/roberta-base-go_emotions deployed through encoder_bert and webserver_default_route: /predict:

curl --fail-with-body "$MODEL_URL/predict" \
  -H "Authorization: Bearer $BASETEN_API_KEY" \
  -H 'Content-Type: application/json' \
  -d '{"inputs":["Thank you, that solved my problem!"]}'

Gemma4 and Qwen3 MoE sequence classifiers use this request shape when the checkpoint includes trained classification weights and label configuration. A base language-model checkpoint doesn’t provide that classifier head.

Typed decisions

Deploy michaelfeil/laya-typed-decisions with BEI-Bert by setting base_model: encoder_bert in the trt_llm configuration. The checkpoint uses BF16 and its custom decision head; set quantization_type: no_quant and call the explicit /sync/v1/systemone route. Create config.yaml:

config.yaml

model_name: laya-typed-decisions
resources:
  accelerator: L4:2
  use_gpu: true
trt_llm:
  inference_stack: v1
  build:
    base_model: encoder_bert
    checkpoint_repository:
      source: HF
      repo: michaelfeil/laya-typed-decisions
    quantization_type: no_quant
    max_num_tokens: 8192

Both GPUs load the complete model. Use L4 for one GPU or L4:4 for four. Baseten manages startup, ports, and health checks through the BEI-Bert deployment. Use the explicit sync URL below for decisions; leave webserver_default_route unset. Deploy and promote it:

baseten model push
baseten model deployment promote --model-id "<MODEL_ID>" --deployment-id "<DEPLOYMENT_ID>"
export MODEL_URL="https://model-<MODEL_ID>.api.baseten.co/environments/production/sync"
export BASETEN_API_KEY="<API_KEY>"
curl --fail-with-body "$MODEL_URL/health" \
  -H "Authorization: Bearer $BASETEN_API_KEY"

Wait for deployment readiness before sending requests. Submit several questions about one state:

curl --fail-with-body "$MODEL_URL/v1/systemone" \
  -H "Authorization: Bearer $BASETEN_API_KEY" \
  -H 'Content-Type: application/json' \
  -d '{
    "state": "I was billed twice. Please refund the extra charge.",
    "questions": {
      "team": {
        "type": "choice",
        "instructions": "Which team should handle this request?",
        "criteria": {"billing": "billing and payments", "support": "technical support"}
      },
      "refund": {
        "type": "noul",
        "instructions": "The customer requests a refund."
      },
      "urgency": {
        "type": "score",
        "instructions": "How urgent is this request?",
        "criteria": ["low", "medium", "high"]
      }
    }
  }'

Laya accepts plain text in state. Serialize structured records into text before sending them. Answers include the selected choice, a boolean probability, or an expected zero-based score, plus confidence and action probability fields.

Examples

Models

Engines

Custom Docker servers

Custom Python models

Chains

API workflows

Baseten for Model Labs

Serve embeddings with BEI

Get started with embedding models:

Example deployment of classification, reranking, and classification models

Benchmarks and performance optimizations

Deploy custom or fine-tuned models on BEI

Deploy a reranker on two L4s

Text model request examples

Qwen3 embeddings

Query and document prompts

Rank documents

Sequence classification

Typed decisions

​Get started with embedding models:

​Example deployment of classification, reranking, and classification models

​Benchmarks and performance optimizations

​Deploy custom or fine-tuned models on BEI

​Deploy a reranker on two L4s

​Text model request examples

​Qwen3 embeddings

​Query and document prompts

​Rank documents

​Sequence classification

​Typed decisions

Get started with embedding models:

Example deployment of classification, reranking, and classification models

Benchmarks and performance optimizations

Deploy custom or fine-tuned models on BEI

Deploy a reranker on two L4s

Text model request examples

Qwen3 embeddings

Query and document prompts

Rank documents

Sequence classification

Typed decisions