qwen-embedding-trt

Convert Qwen3-Embedding HuggingFace checkpoints to TensorRT engines for high-throughput inference.

Why this exists

Qwen3-Embedding ships as a Qwen3ForCausalLM checkpoint with a custom pooling/normalization step. Neither path supports it cleanly out of the box:

• TensorRT-LLM (NVIDIA's optimized LLM serving stack) supports Qwen3 only as a generation model — its standard outputs are vocab-size logits, not the hidden states an embedding model pools from.
• Standard TensorRT (the generic ONNX → engine path) doesn't have a built-in Qwen architecture and trips on a few transformers internals during ONNX export.

This repo bridges that gap: a small wrapper around AutoModel that produces last-token-pooled, L2-normalized embeddings, plus the workarounds needed to get it through ONNX export and into a deployable TRT engine. Output is a single model.plan file you can serve with Triton's tensorrt_plan platform.

What it does

1. Loads a Qwen3 HF checkpoint with AutoModel (backbone only — drops the LM head)
2. Wraps the backbone with last-token pooling + L2 normalization
3. Exports to ONNX (with workarounds for torch.vmap masking and FP16/BF16 quirks)
4. Builds a TRT engine via trtexec
5. Validates by comparing engine output to a BF16 PyTorch baseline (cosine ≥ 0.99)

Validated on Qwen3-Embedding-0.6B and Qwen3-Embedding-4B; the same pipeline should work for the 8B variant unchanged.

Quickstart

# Build the image once
docker build -t qwen-embedding-trt .

# 1. Get HF weights
# Option A — from a local directory (recommended if HuggingFace Hub is unavailable):
export WEIGHTS=/path/to/your/Qwen3-Embedding-0.6B

# Option B — download from HuggingFace Hub:
# huggingface-cli download Qwen/Qwen3-Embedding-0.6B --local-dir ./weights/
# export WEIGHTS=$(pwd)/weights

# 2. Export to ONNX
docker run --rm --gpus all -v "$(pwd)":/work -v "${WEIGHTS}":/weights:ro -w /work \
  qwen-embedding-trt \
  python3 scripts/01_export_onnx.py --weights /weights --out ./model.onnx

# 3. Build TRT engine
docker run --rm --gpus all -v "$(pwd)":/work -w /work \
  qwen-embedding-trt \
  bash scripts/02_build_engine.sh ./model.onnx ./model.plan

# 4. Capture validation baseline (compares BF16 HF output to engine output)
docker run --rm --gpus all -v "$(pwd)":/work -v "${WEIGHTS}":/weights:ro -w /work \
  qwen-embedding-trt \
  python3 scripts/03_validate.py --weights /weights --baseline-only

Pass criterion: mean cosine ≥ 0.99 vs the BF16 HF reference. See Validation for the full Triton validation flow.

TRT version lock: The .plan file is tied to the TRT version inside this image. Your Triton serving container must use the same tritonserver:25.12-py3 tag, or the engine will fail to load with Error Code 6: API Usage Error.

What it doesn't do

• Doesn't ship a serving container — load the resulting .plan into any Triton instance using the tensorrt_plan platform
• Doesn't support non-Qwen3 architectures (the masking patch is Qwen3-specific)
• Doesn't support FP16 (Qwen attention overflows FP16 → NaN; BF16 is required)
• Doesn't support sequence lengths > 512 by default — change MAX_SEQ_LEN in the wrapper if you need longer

How it works

The non-obvious bits, in order of how much time they cost us:

1. attn_implementation="eager" (mandatory for ONNX)

Flash Attention has no ONNX op representation. Eager attention lowers to standard primitives (matmul, softmax) that ONNX can express. This costs ~3-5× inference speed vs Flash, but is a hard prerequisite for the ONNX path.

2. Monkey-patching transformers' vmap-based mask builder

transformers >= 4.51 builds the causal mask via _vmap_for_bhqkv which calls torch.vmap. Both ONNX exporters (TorchScript and dynamo) fail on vmap with AssertionError: ProxyTorchDispatchMode not registered. We replace create_causal_mask and create_sliding_window_causal_mask with broadcast-only equivalents using only ops ONNX supports. See src/qwen_embedding_trt/wrapper.py.

3. dynamo=False to force TorchScript ONNX path

PyTorch 2.5's torch.onnx.export defaults to dynamo, which emits at opset 21 and asks onnxscript to downconvert to opset 17. That downconversion fails with a C-API assertion in onnxscript 0.7 and silently writes a 10 MB stub. TorchScript emits opset 17 natively — slower at export time but actually produces the model.

4. ONNX external data format (>2 GB models)

A naive single-file ONNX export of a 4B model produces an 8 GB protobuf, which protobuf can't reload. We convert to external data format (.onnx descriptor + .onnx.data sidecar). TRT's OnnxParser.parse_from_file handles both natively.

5. AutoModel, not AutoModelForCausalLM

AutoModelForCausalLM includes the LM head, so its outputs are vocab-size logits. AutoModel returns just the backbone, so we get last_hidden_state directly. This is the key change vs naive Qwen serving recipes.

Validation

The validation script compares engine output to a BF16 HuggingFace baseline. There are two modes:

Baseline-only (no running server needed — run this after engine build):

docker run --rm --gpus all -v "$(pwd)":/work -w /work \
  qwen-embedding-trt \
  python3 scripts/03_validate.py --weights ./weights --baseline-only
# Saves baseline.npz to the current directory

Triton validation (validates against a live Triton server):

python3 scripts/03_validate.py \
  --weights ./weights \
  --triton-url http://localhost:8000 \
  --model-name qwen_embedding

The script tokenizes the test inputs locally and calls Triton's native KServe v2 HTTP API (/v2/models/{model_name}/infer) directly — no adapter layer needed.

Tests

pip install -r requirements-dev.txt
pytest                  # unit + onnxruntime smoke tests, CPU only
pytest -m gpu           # full e2e against a real Qwen checkpoint, requires GPU

CI runs the CPU tests on every PR. The GPU test is documented as a manual pre-release check.

License

Apache 2.0. See LICENSE.

qwen-embedding-trt