Embeddings (Encoder-Only Models)#
trtllm-serve can serve encoder-only models (BERT-style classifiers, reward
models, text-embedding models) through an OpenAI-compatible POST /v1/embeddings
endpoint with native dynamic batching — coalescing many independent concurrent
requests into a single forward pass for high throughput, the way the NVIDIA Triton
Inference Server dynamic batcher does.
This replaces the need to run a separate Triton Inference Server in front of an
encoder model: point your existing OpenAI-style embeddings client at trtllm-serve
and it works unchanged.
Quick start#
Launch an embeddings server with the embeddings subcommand:
trtllm-serve embeddings <hf_model_or_path> \
--max_batch_size 32 \
--max_queue_delay 0.005 \
--max_queue_size 2048 \
--host 0.0.0.0 --port 8000
Send a request with any OpenAI-compatible client or curl:
curl http://localhost:8000/v1/embeddings \
-H "Content-Type: application/json" \
-d '{"model": "<model>", "input": ["hello world", "foo bar"]}'
from openai import OpenAI
client = OpenAI(base_url="http://localhost:8000/v1", api_key="tensorrt_llm")
resp = client.embeddings.create(model="<model>", input=["hello world", "foo bar"])
for item in resp.data:
print(item.index, len(item.embedding))
The response is the standard OpenAI embeddings shape:
{
"object": "list",
"data": [
{"object": "embedding", "index": 0, "embedding": [ ... ]},
{"object": "embedding", "index": 1, "embedding": [ ... ]}
],
"model": "<model>",
"usage": {"prompt_tokens": 8, "total_tokens": 8}
}
Request fields#
The endpoint accepts the standard OpenAI /v1/embeddings fields:
Field |
Type |
Notes |
|---|---|---|
|
str |
Model name. |
|
str | list[str] | list[int] | list[list[int]] |
Text(s) or pre-tokenized token-id list(s). |
|
|
|
|
int (optional) |
Matryoshka output size. Only supported by Matryoshka-trained text-embedding models; rejected with |
|
str (optional) |
Ignored; accepted for compatibility. |
|
bool (default |
TRT-LLM extension. Encoder models such as BERT generally need their special tokens (e.g. |
There are no required TRT-LLM-specific request fields — existing OpenAI-compatible
embeddings clients work by pointing at the trtllm-serve URL.
Dynamic batching#
A lightweight in-server batcher coalesces concurrent requests in front of the encoder forward pass. It exposes three knobs that mirror the Triton dynamic batcher:
|
Behavior |
Triton equivalent |
|---|---|---|
|
Upper bound on the number of requests fused into one forward pass. A batch reaching this size is dispatched immediately. |
maximum / |
|
Hold window: how long an incoming request waits for others to join its batch before dispatch. |
|
|
Maximum number of in-flight queued requests. Further requests are rejected with HTTP 429 (backpressure). |
|
A batch is dispatched as soon as any of these fires: it reaches --max_batch_size,
adding the next request would exceed the engine’s --max_num_tokens budget, or the
--max_queue_delay hold window elapses.
Migrating from the Triton Inference Server dynamic batcher#
If you currently serve an encoder model with the Triton inflight_batcher_llm backend
and a config.pbtxt dynamic_batching { ... } block, map the settings directly:
Triton |
|
|---|---|
|
|
|
|
|
|
Adopt the same values you tuned in Triton as a starting point, then adjust for your latency/throughput budget.
Error handling#
Condition |
HTTP status |
|---|---|
Input longer than |
400 |
Request queue full ( |
429 |
Invalid request body |
400 |
Embedding responses are unary (non-streaming).
Output semantics and scope#
The endpoint is model-output-agnostic: it returns whatever per-request vector the model emits, serialized into the OpenAI embeddings schema.
Classifier / reward models (e.g. a BERT sequence classifier): the returned vector is the model’s class-logit / score vector (
[num_labels]).Text-embedding models — the Qwen3-Embedding family (
Qwen3-Embedding-0.6B,-4B,-8B) is supported. These ship as aQwen3ForCausalLMdecoder plus a sentence-transformers pooling pipeline; the embeddings server detects this and serves the L2-normalized last-token hidden state (a[hidden_size]sentence-embedding vector — 1024 / 2560 / 4096 respectively), with no extra flags. A configurable pooling method (CLS / mean) for other sentence-transformers backbones remains a follow-up.
Notes:
The embeddings path uses the synchronous
llm.encode()fast path (EncoderExecutor): a single forward pass per batch, no KV cache, sampler, or decode loop.One encoder model per server instance. Generation and embedding modes are not mixed in one server.
Single-GPU per server. The encode path runs in-process and does not use the multi-GPU worker proxy, so the
embeddingscommand does not expose tensor/pipeline parallelism flags (if a--configfile sets them, startup fails with a clear error). To scale out, see Scaling out across GPUs below.A single in-server worker drives the GPU (no
num_workersknob): the GPU serializes forwards and the underlying executor is not safe for concurrent calls. Increase throughput with--max_batch_size/--max_queue_delay, not more workers.
Scaling out across GPUs#
Embedding / encoder-only models are usually small and fit comfortably on a single GPU.
The recommended way to use more GPUs is therefore data parallelism: run one
single-GPU trtllm-serve embeddings instance per GPU and put a load balancer in front
of them. There is no cross-GPU communication, so throughput scales close to linearly
with the number of replicas.
# One replica per GPU (8x B200 example), each on its own port.
for i in $(seq 0 7); do
CUDA_VISIBLE_DEVICES=$i trtllm-serve embeddings <model> --port $((8000 + i)) &
done
# Then point any HTTP load balancer (nginx, k8s Service, etc.) at ports 8000-8007.
Tensor / pipeline parallelism (sharding a single model across GPUs) is only needed
for an embedding model too large to fit on one GPU — uncommon for encoder-only models.
It is not yet supported by the embeddings command and is planned as a follow-up.
Relationship to llm.encode()#
The server reuses the existing Python llm.encode() API
(LLM(..., encode_only=True)) under the hood; the only addition is the async
coalescing layer plus the HTTP surface. The synchronous llm.encode() API continues
to work unchanged for direct Python callers.
Encoder CUDA graphs#
Encoder CUDA graphs capture the encoder forward pass once per input shape at startup and replay it for each batch, removing the per-kernel launch overhead of eager execution. The gain is largest for small encoders and small batches, where launching kernels, not running them, dominates the latency.
Enable them with an EncodeCudaGraphConfig on an encode_only=True LLM:
from tensorrt_llm import LLM
from tensorrt_llm.llmapi import EncodeCudaGraphConfig
llm = LLM(
model="<model>",
encode_only=True,
cuda_graph_config=EncodeCudaGraphConfig(
batch_sizes=[1, 2, 4, 8],
num_tokens=[128, 256, 512, 1024],
seq_lens=[64, 128],
enable_padding=True,
),
)
outputs = llm.encode(["hello world", "foo bar"])
For trtllm-serve embeddings, put the same settings in the --config YAML:
cuda_graph_config:
batch_sizes: [1, 2, 4, 8]
num_tokens: [128, 256, 512, 1024]
seq_lens: [64, 128]
enable_padding: true
Field |
Meaning |
|---|---|
|
Number of requests per batch to capture graphs for. |
|
Total token counts per batch (all requests’ tokens, packed) to capture graphs for. |
|
Longest-request lengths to capture graphs for. |
|
Round each batch up to the nearest captured shape. Recommended; without it only exact matches use a graph. |
|
Alternative to listing |
A graph is captured at startup for every feasible combination of these buckets. At
runtime a batch uses the graph for its (batch size, total tokens, longest request)
shape, after rounding up if enable_padding is set. A batch with no matching graph,
such as one larger than the largest bucket, runs eagerly: the result is the same, just
not accelerated. More buckets cover more traffic at the cost of startup time and graph
memory.
Encoder CUDA graphs require the TRTLLM attention backend; with other backends every
batch runs eagerly. Multi-item scoring batches also run eagerly. For encoder-decoder
models, see encoder_cuda_graph_config in
Use encoder-decoder models with the PyTorch backend.
Extra model inputs#
llm.encode() passes extra keyword arguments through to the model’s forward(), for
example token_type_ids for BERT. With encoder CUDA graphs, declare each tensor
argument in extra_model_inputs so the graph is captured with it:
import torch
from transformers import AutoTokenizer
from tensorrt_llm import LLM
from tensorrt_llm.llmapi import EncodeCudaGraphConfig, EncodeExtraInputSpec
llm = LLM(
model="<bert_model>",
encode_only=True,
cuda_graph_config=EncodeCudaGraphConfig(
batch_sizes=[1, 2, 4, 8],
num_tokens=[128, 256, 512, 1024],
seq_lens=[64, 128],
enable_padding=True,
extra_model_inputs=[
EncodeExtraInputSpec(name="token_type_ids", shape=("num_tokens",), dtype="int32"),
],
),
)
# Sentence pairs: the tokenizer marks the second sentence of each pair as segment 1.
pairs = [("What is TensorRT-LLM?", "An inference library."),
("Is it fast?", "Yes.")]
tokenizer = AutoTokenizer.from_pretrained("<bert_model>")
encoded = tokenizer([q for q, _ in pairs], [a for _, a in pairs])
prompts = [{"prompt_token_ids": ids} for ids in encoded["input_ids"]]
# Packed: every prompt's values concatenated in prompt order, no padding.
token_type_ids = torch.tensor(
[t for row in encoded["token_type_ids"] for t in row], dtype=torch.int32)
outputs = llm.encode(prompts, token_type_ids=token_type_ids)
Each EncodeExtraInputSpec has a name (the forward() argument), a dtype, and a
shape with exactly one symbolic dimension; the other dimensions are fixed integers:
Symbolic dimension |
Size at call time |
Example |
|---|---|---|
|
Total tokens in the batch, packed in prompt order |
|
|
Number of prompts, one row each |
per-request features: |
Rules:
Every tensor argument passed to
llm.encode()must be declared; an undeclared tensor raisesValueError. A non-tensor argument is allowed but runs that call eagerly.Every declared input must be passed on every call, with the declared dtype and shape. To omit one, pass zeros.
These checks apply only to calls that can use a graph. Multi-item scoring batches and non-
TRTLLMattention backends always run eagerly, so their arguments pass through unchecked.Pass tensors where they already live, on the host or on the device.
input_ids,position_ids,seq_lens,multi_item_part_lens,attn_metadataandreturn_context_logitsare managed by the runtime and cannot be declared.Extra model inputs are available through the Python
llm.encode()API only.trtllm-serve embeddingsdoes not support them yet and rejects a--configthat declaresextra_model_inputsat startup. They are also not supported for encoder-decoder models or encoders that take fixed-shape feature tensors.