TensorRT RAG batch inference fails with embedding dimension mismatch after model change

Problem Description

A production inference server that runs a Retrieval‑Augmented Generation (RAG) pipeline with TensorRT started throwing runtime errors after the passage encoder checkpoint was upgraded. The server processes thousands of queries per second in batches of up to 128 on a single GPU. Within seconds of deployment the following errors appeared in the logs:

TensorRT Runtime Error: Dimension mismatch between input binding 'input_ids' (expected [batch, seq_len]) and provided tensor [batch, seq_len, 768].
Assertion failed: engine->bindingDimensions(bindingIndex) == inputDims
  -- shape incompatibility between query_encoder_output (batch, 768) and passage_encoder_output (batch, 1024).
Failed to enqueue TensorRT execution context: Engine binding dimension mismatch for 'embedding_output' (expected 768, got 1024).

Symptoms observed in the monitoring system:

• Spike in inference_latency_ms to >5 s for the RAG endpoint.
• Increased engine_enqueue_failures_total counter.
• Batch jobs of 4‑8 k queries/sec abort after the first few seconds.

Root Cause Analysis

The RAG architecture consists of two independent encoders:

• Query encoder – produces a [batch, D_q] embedding.
• Passage encoder – produces a [batch, N_doc, D_p] embedding.

TensorRT engines are built with static binding dimensions unless dynamic shapes are explicitly declared (TensorRT Developer Guide – Dynamic Shapes). The original engine was built assuming D_q = D_p = 768, matching the facebook/rag-token-base checkpoint.

During a model upgrade the passage encoder checkpoint was switched to a 1024‑dimensional variant (facebook/rag-token-nq style). The ONNX export reflected the new dimension, but the TensorRT engine was not rebuilt. Consequently the execution context still expected a 768‑dimensional embedding for the passage encoder while the runtime supplied a 1024‑dimensional tensor.

Because the query encoder remained unchanged, its output shape still matched the engine binding (768). The mismatch triggered the binding validation step described in the TensorRT Runtime Execution Context – Binding Management, resulting in the observed errors.

Investigation and Debugging

1. Verify engine bindings

import tensorrt as trt

logger = trt.Logger(trt.Logger.INFO)
runtime = trt.Runtime(logger)
engine = runtime.deserialize_cuda_engine(open("rag_engine.trt", "rb").read())
for i in range(engine.num_bindings):
    name = engine.get_binding_name(i)
    dims = engine.get_binding_dimensions(i)
    print(f"{i}: {name} -> {dims}")

Expected output (original engine):

0: input_ids -> ( -1, 128 )
1: attention_mask -> ( -1, 128 )
2: embedding_output -> ( -1, 768 )

Actual output after upgrade showed the same binding dimensions, confirming the engine was stale.

2. Inspect ONNX export

python -c "
import torch
from transformers import RagTokenizer, RagRetriever, RagSequenceForGeneration
model = RagSequenceForGeneration.from_pretrained('facebook/rag-token-nq')
model.save_pretrained('rag_nq')
"
onnxsim rag_nq/model.onnx rag_nq/simplified.onnx

The simplified ONNX graph revealed the passage encoder output shape (batch, N_doc, 1024), confirming the dimension change.

3. Check dynamic shape profile

The TensorRT builder was configured with a static profile:

builder = trt.Builder(logger)
profile = builder.create_optimization_profile()
profile.set_shape("input_ids", (1, 128), (32, 128), (64, 128))
builder.add_optimization_profile(profile)

Documentation (TensorRT Python API – Building and Serializing Engines) states that any dimension not covered by a profile must be static. The new 1024‑dimensional output was never part of a profile, causing the engine to reject the tensor.

4. Correlate with community reports

• GitHub issue #14892 describes the same failure after updating a RAG checkpoint.
• HuggingFace issue #3021 notes that the ONNX parser propagates the new embedding size, but the engine must be rebuilt.
• Stack Overflow question 78543219 shows identical error messages.

Resolution

1. Rebuild the TensorRT engine with updated dimensions

Two approaches are viable:

1. Static binding with new dimension – rebuild the engine after the checkpoint change.
2. Dynamic output dimension – expose the passage encoder output as a dynamic shape and provide a profile that covers both 768 and 1024.

Static rebuild (simplest)

# rebuild_engine.py
import tensorrt as trt
import onnx
import os

logger = trt.Logger(trt.Logger.INFO)
builder = trt.Builder(logger)
config = builder.create_builder_config()
config.set_flag(trt.BuilderFlag.FP16)  # keep mixed‑precision if desired

# Load updated ONNX
onnx_model_path = "rag_nq/simplified.onnx"
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
parser = trt.OnnxParser(network, logger)
with open(onnx_model_path, "rb") as f:
    parser.parse(f.read())

# Create a profile that matches the new output dimension
profile = builder.create_optimization_profile()
profile.set_shape("input_ids", (1, 128), (32, 128), (128, 128))
builder.add_optimization_profile(profile)

engine = builder.build_engine(network, config)
with open("rag_engine_v2.trt", "wb") as f:
    f.write(engine.serialize())

After rebuilding, the binding inspection shows:

2: embedding_output -> ( -1, 1024 )

Dynamic output dimension (future‑proof)

# dynamic_engine.py
profile = builder.create_optimization_profile()
profile.set_shape("input_ids", (1, 128), (64, 128), (128, 128))
profile.set_shape("embedding_output", (1, 768), (64, 768), (128, 1024))  # allow both sizes
builder.add_optimization_profile(profile)

By declaring embedding_output as a dynamic tensor, the same engine can serve both 768‑ and 1024‑dimensional passage encoders, eliminating the need for a rebuild on every checkpoint change.

2. Update the inference server to reload the new engine

# inference_server.py (excerpt)
def load_engine(path):
    with open(path, "rb") as f:
        return runtime.deserialize_cuda_engine(f.read())

engine = load_engine("rag_engine_v2.trt")
context = engine.create_execution_context()
# Ensure the new context is used for all incoming batches

3. Clear stale engine cache

If the server caches engines by model name, purge the cache before redeploying:

rm -rf /var/cache/trt_engine/*

Verification

1. Run a single‑batch sanity check:

python run_inference.py --batch-size 1 --model rag_engine_v2.trt

Expected log snippet:

[INFO] Execution context created – binding dimensions validated.
[INFO] Inference completed in 12.3 ms (batch=1)

Execute a high‑throughput load test (e.g., hey or custom load generator) with batch size 128:

hey -c 64 -n 10000 -m POST -d query_payload.json http://inference-server/rag

Metrics to confirm:

• inference_latency_ms stable around 15‑20 ms per batch.
• No increase in engine_enqueue_failures_total.
• GPU utilization remains within expected range (≈70 %).

Inspect TensorRT logs for binding validation messages:

[TensorRT] Binding dimensions validated: input_ids=(128,128), embedding_output=(128,1024)
[TensorRT] Enqueue succeeded for batch size 128

Prevention and Best Practices

• Version‑aware engine rebuilds: Automate engine regeneration whenever a model checkpoint changes. Include the checkpoint hash in the engine filename.
• Dynamic shape profiles for RAG: Follow the TensorRT Model Optimization Guide – Variable Input Dimensions for RAG and declare both embedding_output and input_ids as dynamic.
• Engine cache invalidation: Implement a cache‑key that incorporates model version, precision mode, and max sequence length. On mismatch, discard the stale engine.
• Continuous integration test: Add a regression test that builds an engine from the current ONNX, runs a batch inference, and asserts that output shapes match the expected dimensions.
• Monitoring: Alert on engine_enqueue_failures_total and on any spike in inference_latency_ms > 100 ms for the RAG endpoint.

Related Topic Hub: Model Serving Troubleshooting Hub

FAQ

1. Why does the error appear only under peak load?
   During peak load the server processes larger batches (e.g., 128). The static binding for embedding_output was sized for 768 dimensions; small batches sometimes fit within the internal buffer, but larger batches expose the mismatch, triggering the validation error.
2. Can I keep the old engine and just reshape the output tensor?
   No. TensorRT validates binding dimensions against the engine’s compiled metadata before kernel launch. A reshape without rebuilding violates the contract and results in the same assertion failure.
3. Is mixed‑precision (FP16) a factor in the mismatch?
   Mixed‑precision alone is not the cause, but if the query encoder is built FP16 while the passage encoder remains FP32, the engine may generate separate kernels with differing expected output sizes, amplifying shape‑validation errors when batch size grows.
4. How do I know which dimensions are dynamic in an existing engine?
   Use the C++/Python API to query engine.get_binding_dimensions(). Dimensions set to -1 are dynamic. The profile object can be inspected via engine.get_optimization_profile() for min/opt/max shapes.
5. Will enabling TensorRT’s implicit batch mode avoid this issue?
   Implicit batch mode is deprecated and does not support dynamic output dimensions. Switching to explicit batch with proper profiles is the recommended approach for RAG models.