Add base_model=encoder_torch for BEI-torch backend

truss/base/constants.py

@@ -27,6 +27,10 @@
 BEI_TRTLLM_CLIENT_BATCH_SIZE = 128
 BEI_MAX_CONCURRENCY_TARGET_REQUESTS = 2048
 BEI_REQUIRED_MAX_NUM_TOKENS = 16384
+# BEI-torch (vLLM) rejects --max-batch-tokens below this; long-context embedders
+# (Nemotron-3-Embed at 32k) need the headroom. User max_num_tokens is clamped
+# upward to this floor, matching the ENCODER/ENCODER_BERT behavior.
+BEI_TORCH_REQUIRED_MAX_NUM_TOKENS = 32768
 
 TRTLLM_MIN_MEMORY_REQUEST_GI = 10
 HF_MODELS_API_URL = "https://huggingface.co/api/models"

truss/base/trt_llm_config.py

@@ -50,6 +50,11 @@ class TrussTRTLLMModel(str, Enum):
     # the encoder_bert setting will specfically optimize for thoughput and cold-start latency of small models (<4B parameters)
     # supports also splade and colbert style models or ModernBert.
     ENCODER_BERT = "encoder_bert"
+    # BEI torch backend (text-embeddings-router + vLLM). Serves embedding, reranker, and
+    # classification models whose architecture the TRT-LLM `encoder` path cannot compile
+    # (e.g. LlamaBidirectional / Ministral3-Embed, Gemma2Embedding, jina-v3). No engine
+    # build step; the checkpoint is served directly by the torch backend.
+    ENCODER_TORCH = "encoder_torch"
     # Decoder will launch the backend that is optimized for decoder only models such as LLama3ForCausalLM, Qwen3MoeForCausalLM etc.
     DECODER = "decoder"
     # a ERROR will be raised if you push one of the below models. Don't use
@@ -365,7 +370,31 @@ def uses_lora(self) -> bool:
 
     def _bei_specfic_migration(self):
         """performs embedding specfic optimizations (no kv-cache, high batch size)"""
-        if self.base_model == TrussTRTLLMModel.ENCODER:
+        if self.base_model == TrussTRTLLMModel.ENCODER_TORCH:
+            # BEI-torch has no engine build, so build-time quantization is not applicable.
+            # Users wanting quantized weights must point checkpoint_repository at a
+            # pre-quantized HF checkpoint.
+            if self.quantization_type != TrussTRTLLMQuantizationType.NO_QUANT:
+                raise ValueError(
+                    f"base_model=encoder_torch does not support build-time quantization; "
+                    f"you set build.quantization_type={self.quantization_type.value}. "
+                    "Point checkpoint_repository at a pre-quantized HF checkpoint instead."
+                )
+            # Speculative decoding is a decoder concept; embedding/reranker serving via
+            # BEI-torch has no draft/target loop.
+            if self.speculator is not None:
+                raise ValueError(
+                    "base_model=encoder_torch does not support speculative decoding; "
+                    "remove trt_llm.build.speculator."
+                )
+            # LoRA adapters are a TRT-LLM build-time concept; vLLM serves the
+            # checkpoint as-is, so adapters would be silently ignored.
+            if self.lora_adapters is not None:
+                raise ValueError(
+                    "base_model=encoder_torch does not support lora_adapters; "
+                    "remove trt_llm.build.lora_adapters."
+                )
+        elif self.base_model == TrussTRTLLMModel.ENCODER:
             # Encoder specific settings
             if self.max_seq_len:
                 logger.info(
@@ -572,6 +601,7 @@ class VersionsOverrides(PydanticTrTBaseModel):
     briton_version: Optional[str] = None
     bei_version: Optional[str] = None
     bei_bert_version: Optional[str] = None
+    bei_torch_version: Optional[str] = None
     v2_llm_version: Optional[str] = None
 
     @model_validator(mode="before")
@@ -581,6 +611,7 @@ def version_must_start_with_number(cls, data):
             "briton_version",
             "bei_version",
             "bei_bert_version",
+            "bei_torch_version",
         ]:
             v = data.get(field)
             if v is not None and (not v or not v[0].isdigit()):
@@ -596,6 +627,10 @@ class ImageVersions(PydanticTrTBaseModel):
     # INTERNAL
     bei_image: str
     beibert_image: str
+    # Base image for the BEI torch backend (text-embeddings-router + vLLM). Backend
+    # inserts a resolved image reference; the default here is a placeholder so
+    # local unit tests can construct ImageVersions without wiring core-product.
+    bei_torch_image: str = "baseten/bei:torch-dev-placeholder"
     briton_image: str
     v2_llm_image: str
 
@@ -618,7 +653,11 @@ def model_post_init(self, __context):
             self.runtime.enable_chunked_context
             and (
                 self.build.base_model
-                not in (TrussTRTLLMModel.ENCODER, TrussTRTLLMModel.ENCODER_BERT)
+                not in (
+                    TrussTRTLLMModel.ENCODER,
+                    TrussTRTLLMModel.ENCODER_BERT,
+                    TrussTRTLLMModel.ENCODER_TORCH,
+                )
             )
             and not (
                 self.build.plugin_configuration.use_paged_context_fmha
@@ -636,7 +675,11 @@ def model_post_init(self, __context):
         if (
             self.runtime.webserver_default_route is None
             and self.build.base_model
-            in (TrussTRTLLMModel.ENCODER, TrussTRTLLMModel.ENCODER_BERT)
+            in (
+                TrussTRTLLMModel.ENCODER,
+                TrussTRTLLMModel.ENCODER_BERT,
+                TrussTRTLLMModel.ENCODER_TORCH,
+            )
             and not ENGINE_BUILDER_TRUSS_RUNTIME_MIGRATION
         ):
             if hf_cfg is not None:
@@ -672,9 +715,11 @@ def model_post_init(self, __context):
                         f"but you set `trt_llm.build.base_model` to `decoder`. "
                         f"Please set it to `encoder_bert`."
                     )
-            if (
-                "ForCausalLM" in arch
-                and self.build.base_model != TrussTRTLLMModel.DECODER
+            if "ForCausalLM" in arch and self.build.base_model not in (
+                TrussTRTLLMModel.DECODER,
+                # encoder_torch is the sanctioned way to serve a causal-arch checkpoint
+                # as an embedding/reranker via bidirectional-attention override in BEI-torch.
+                TrussTRTLLMModel.ENCODER_TORCH,
             ):
                 logger.warning(
                     f"Your model architecture {arch} indicates a CausalLM based model. "
@@ -889,7 +934,7 @@ def trt_llm_common_validation(config: "TrussConfig"):
             pass
         else:
             raise ValueError(
-                "TRT-LLM is not supported on CUDA_COMPUTE_75 (T4) and CUDA_COMPUTE_70 (V100) GPUs. \n"
+                "TRT-LLM and BEI-torch are not supported on CUDA_COMPUTE_75 (T4) and CUDA_COMPUTE_70 (V100) GPUs. \n"
                 "the lowest supported CUDA compute capability is CUDA_COMPUTE_80 (A100) or A10G (CUDA_COMPUTE_86)"
             )
     elif trt_llm_config.build.quantization_type in [
@@ -947,6 +992,7 @@ def trt_llm_validation_v1(config: "TrussConfig") -> "TrussConfig":
     if trt_llm_config_v1.build.base_model not in [
         TrussTRTLLMModel.ENCODER,
         TrussTRTLLMModel.ENCODER_BERT,
+        TrussTRTLLMModel.ENCODER_TORCH,
     ]:
         current_tags = config.model_metadata.get("tags", [])
         if (

truss/config.schema.json

@@ -1097,6 +1097,7 @@
             "briton_version": null,
             "bei_version": null,
             "bei_bert_version": null,
+            "bei_torch_version": null,
             "v2_llm_version": null
           }
         }
@@ -1129,6 +1130,7 @@
             "briton_version": null,
             "bei_version": null,
             "bei_bert_version": null,
+            "bei_torch_version": null,
             "v2_llm_version": null
           }
         }
@@ -1569,6 +1571,7 @@
       "enum": [
         "encoder",
         "encoder_bert",
+        "encoder_torch",
         "decoder",
         "palmyra",
         "qwen",
@@ -1775,6 +1778,18 @@
           "default": null,
           "title": "Bei Bert Version"
         },
+        "bei_torch_version": {
+          "anyOf": [
+            {
+              "type": "string"
+            },
+            {
+              "type": "null"
+            }
+          ],
+          "default": null,
+          "title": "Bei Torch Version"
+        },
         "v2_llm_version": {
           "anyOf": [
             {

truss/contexts/docker_build_setup.py

@@ -46,6 +46,12 @@ def _fill_trt_llm_versions(
         ):
             print(f"Using BEI image: {image_versions.bei_image}")
             tr.set_base_image(image_versions.bei_image, "/usr/bin/python3")
+        elif (
+            tr.spec.config.trt_llm.build.base_model
+            == trt_llm_config.TrussTRTLLMModel.ENCODER_TORCH
+        ):
+            print(f"Using BEI torch image: {image_versions.bei_torch_image}")
+            tr.set_base_image(image_versions.bei_torch_image, "/usr/bin/python3")
         elif (
             tr.spec.config.trt_llm.build.base_model
             == trt_llm_config.TrussTRTLLMModel.ENCODER_BERT

truss/contexts/image_builder/serving_image_builder.py

@@ -27,6 +27,7 @@
     BASE_SERVER_REQUIREMENTS_TXT_FILENAME,
     BEI_MAX_CONCURRENCY_TARGET_REQUESTS,
     BEI_REQUIRED_MAX_NUM_TOKENS,
+    BEI_TORCH_REQUIRED_MAX_NUM_TOKENS,
     BEI_TRTLLM_CLIENT_BATCH_SIZE,
     CHAINS_CODE_DIR,
     CONSTRAINTS_TXT_FILENAME,
@@ -610,6 +611,56 @@ def prepare_trtllm_bei_encoder_build_dir(self, build_dir: Path):
         )
         copy_tree_path(DOCKER_SERVER_TEMPLATES_DIR, build_dir, ignore_patterns=[])
 
+    def prepare_bei_encoder_torch_build_dir(self, build_dir: Path):
+        """prepares the build directory for an ENCODER_TORCH model to launch the BEI torch backend
+        (text-embeddings-router + vLLM). No engine build; the checkpoint is served directly."""
+        config = self._spec.config
+        assert (
+            config.trt_llm
+            and config.trt_llm.build
+            and config.trt_llm.build.base_model == TrussTRTLLMModel.ENCODER_TORCH
+        ), (
+            "prepare_bei_encoder_torch_build_dir should only be called for ENCODER_TORCH tensorrt-llm model"
+        )
+        assert isinstance(config.trt_llm.root, TRTLLMConfigurationV1), (
+            "prepare_bei_encoder_torch_build_dir should only be called for inference_stack v1 tensorrt-llm model"
+        )
+        trt_llm_config: TRTLLMConfigurationV1 = config.trt_llm.root
+        # Torch backend has no TRT-LLM 32-cap; 256 is a safe upper bound for KV-less
+        # embedding servers before per-request scheduling overhead dominates.
+        runtime_max_batch_size = min(trt_llm_config.build.max_batch_size, 256)
+        # vLLM rejects --max-batch-tokens below BEI_TORCH_REQUIRED_MAX_NUM_TOKENS;
+        # long-context embedders (Nemotron 32k) need the headroom.
+        runtime_max_batch_tokens = max(
+            trt_llm_config.build.max_num_tokens, BEI_TORCH_REQUIRED_MAX_NUM_TOKENS
+        )
+        port = 7997
+        start_command = " ".join(
+            [
+                "truss-transfer-cli /tmp/bei-model && text-embeddings-router --model-id /tmp/bei-model",
+                f"--port {port}",
+                f"--max-batch-requests {runtime_max_batch_size}",
+                f"--max-batch-tokens {runtime_max_batch_tokens}",
+                # sentences per JSON payload; capped for request-based autoscaling.
+                f"--max-client-batch-size {BEI_TRTLLM_CLIENT_BATCH_SIZE}",
+                # concurrent requests before returning 429.
+                # https://docs.baseten.co/performance/concurrency#concurrency-target
+                f"--max-concurrent-requests {BEI_MAX_CONCURRENCY_TARGET_REQUESTS}",
+                # stricter validation in newer BEI-torch images requires this to
+                # silently trim overlong inputs instead of 4xx-ing the request.
+                "--auto-truncate",
+            ]
+        )
+        self._spec.config.docker_server = DockerServer(
+            start_command=f"/bin/sh -c '{start_command}'",
+            server_port=port,
+            predict_endpoint=trt_llm_config.runtime.webserver_default_route
+            or "/v1/embeddings",
+            readiness_endpoint="/health",
+            liveness_endpoint="/health",
+        )
+        copy_tree_path(DOCKER_SERVER_TEMPLATES_DIR, build_dir, ignore_patterns=[])
+
     def prepare_trtllm_bert_encoder_build_dir(self, build_dir: Path):
         """prepares the build directory for a trtllm ENCODER model to launch a Baseten Embeddings Inference (BEI) server"""
         config = self._spec.config
@@ -665,7 +716,11 @@ def prepare_trtllm_decoder_build_dir(self, build_dir: Path):
             config.trt_llm
             and config.trt_llm.build
             and config.trt_llm.build.base_model
-            not in [TrussTRTLLMModel.ENCODER, TrussTRTLLMModel.ENCODER_BERT]
+            not in [
+                TrussTRTLLMModel.ENCODER,
+                TrussTRTLLMModel.ENCODER_BERT,
+                TrussTRTLLMModel.ENCODER_TORCH,
+            ]
         ), (
             "prepare_trtllm_decoder_build_dir should only be called for decoder tensorrt-llm model"
         )
@@ -730,6 +785,8 @@ def prepare_image_build_dir(
                 elif config.trt_llm.build.base_model == TrussTRTLLMModel.ENCODER_BERT:
                     # Run the specific encoder_bert build
                     self.prepare_trtllm_bert_encoder_build_dir(build_dir=build_dir)
+                elif config.trt_llm.build.base_model == TrussTRTLLMModel.ENCODER_TORCH:
+                    self.prepare_bei_encoder_torch_build_dir(build_dir=build_dir)
                 else:
                     self.prepare_trtllm_decoder_build_dir(build_dir=build_dir)
 

truss/tests/conftest.py

@@ -816,6 +816,30 @@ def trtllm_config_encoder(default_config) -> Dict[str, Any]:
     return trtllm_config
 
 
+@pytest.fixture
+def trtllm_config_encoder_torch(default_config) -> Dict[str, Any]:
+    trtllm_config = default_config
+    trtllm_config["resources"] = {
+        "accelerator": Accelerator.H100.value,
+        "cpu": "1",
+        "memory": "30Gi",
+        "use_gpu": True,
+        "node_count": 1,
+    }
+    trtllm_config["trt_llm"] = {
+        "build": {
+            "base_model": "encoder_torch",
+            "checkpoint_repository": {
+                "source": "HF",
+                "repo": "nvidia/Nemotron-3-Embed-8B-BF16",
+            },
+            "max_num_tokens": 32768,
+        },
+        "runtime": {},
+    }
+    return trtllm_config
+
+
 @pytest.fixture
 def deprecated_trtllm_config(default_config) -> Dict[str, Any]:
     trtllm_config = default_config

truss/tests/trt_llm/test_trt_llm_config.py

@@ -116,6 +116,40 @@ def test_trt_llm_encoder(trtllm_config_encoder):
     assert config.build.plugin_configuration.use_paged_context_fmha is False
 
 
+def test_trt_llm_encoder_torch_accepts_no_quant(trtllm_config_encoder_torch):
+    config = TRTLLMConfigurationV1(**trtllm_config_encoder_torch["trt_llm"])
+    assert config.build.base_model.value == "encoder_torch"
+    assert config.build.quantization_type.value == "no_quant"
+    assert config.build.max_num_tokens == 32768
+
+
+def test_trt_llm_encoder_torch_rejects_build_time_quantization(
+    trtllm_config_encoder_torch,
+):
+    trtllm_config_encoder_torch["trt_llm"]["build"]["quantization_type"] = "fp8"
+    with pytest.raises(pydantic.ValidationError, match="encoder_torch"):
+        TRTLLMConfigurationV1(**trtllm_config_encoder_torch["trt_llm"])
+
+
+def test_trt_llm_encoder_torch_rejects_speculator(trtllm_config_encoder_torch):
+    trtllm_config_encoder_torch["trt_llm"]["build"]["speculator"] = {
+        "speculative_decoding_mode": "LOOKAHEAD_DECODING",
+        "lookahead_windows_size": 4,
+        "lookahead_ngram_size": 3,
+        "lookahead_verification_set_size": 2,
+    }
+    with pytest.raises(pydantic.ValidationError, match="speculator"):
+        TRTLLMConfigurationV1(**trtllm_config_encoder_torch["trt_llm"])
+
+
+def test_trt_llm_encoder_torch_rejects_lora_adapters(trtllm_config_encoder_torch):
+    trtllm_config_encoder_torch["trt_llm"]["build"]["lora_adapters"] = {
+        "adapter1": {"source": "HF", "repo": "nvidia/Nemotron-3-Embed-8B-BF16-lora"}
+    }
+    with pytest.raises(pydantic.ValidationError, match="lora_adapters"):
+        TRTLLMConfigurationV1(**trtllm_config_encoder_torch["trt_llm"])
+
+
 def test_trt_llm_encoder_autoconfig(trtllm_config_encoder):
     trt_llm_config = TRTLLMConfigurationV1(**trtllm_config_encoder["trt_llm"])
     try:

truss/trt_llm/config_checks.py

@@ -45,6 +45,7 @@ def add_openai_tag(tr: TrussHandle) -> str:
         if trt_llm_config.build.base_model in [
             TrussTRTLLMModel.ENCODER,
             TrussTRTLLMModel.ENCODER_BERT,
+            TrussTRTLLMModel.ENCODER_TORCH,
         ]:
             return ("", False)
         # only briton requires openai-compatible tag, all others don't care about the openai tag