update docs for deploying via trt-llm

michaelfeil · michaelfeil · commit dd4b565a4bce · 2025-03-24T18:22:45.000Z
diff --git a/11-embeddings-reranker-classification-tensorrt/BEI-allenai-llama-3.1-tulu-3-8b-reward-model-fp8/README.md b/11-embeddings-reranker-classification-tensorrt/BEI-allenai-llama-3.1-tulu-3-8b-reward-model-fp8/README.md
@@ -102,7 +102,9 @@ environment_variables: {}
 external_package_dirs: []
 model_metadata:
   example_model_input:
-    inputs: Baseten is a fast inference provider
+    inputs:
+    - - Baseten is a fast inference provider
+    - - Classify this separately.
     raw_scores: true
     truncate: true
     truncation_direction: Right
diff --git a/11-embeddings-reranker-classification-tensorrt/BEI-allenai-llama-3.1-tulu-3-8b-reward-model-fp8/config.yaml b/11-embeddings-reranker-classification-tensorrt/BEI-allenai-llama-3.1-tulu-3-8b-reward-model-fp8/config.yaml
@@ -3,7 +3,9 @@ environment_variables: {}
 external_package_dirs: []
 model_metadata:
   example_model_input:
-    inputs: Baseten is a fast inference provider
+    inputs:
+    - - Baseten is a fast inference provider
+    - - Classify this separately.
     raw_scores: true
     truncate: true
     truncation_direction: Right
diff --git a/11-embeddings-reranker-classification-tensorrt/BEI-baseten-example-meta-llama-3-70b-instructforsequenceclassification-fp8/README.md b/11-embeddings-reranker-classification-tensorrt/BEI-baseten-example-meta-llama-3-70b-instructforsequenceclassification-fp8/README.md
@@ -102,7 +102,9 @@ environment_variables: {}
 external_package_dirs: []
 model_metadata:
   example_model_input:
-    inputs: Baseten is a fast inference provider
+    inputs:
+    - - Baseten is a fast inference provider
+    - - Classify this separately.
     raw_scores: true
     truncate: true
     truncation_direction: Right
diff --git a/11-embeddings-reranker-classification-tensorrt/BEI-baseten-example-meta-llama-3-70b-instructforsequenceclassification-fp8/config.yaml b/11-embeddings-reranker-classification-tensorrt/BEI-baseten-example-meta-llama-3-70b-instructforsequenceclassification-fp8/config.yaml
@@ -3,7 +3,9 @@ environment_variables: {}
 external_package_dirs: []
 model_metadata:
   example_model_input:
-    inputs: Baseten is a fast inference provider
+    inputs:
+    - - Baseten is a fast inference provider
+    - - Classify this separately.
     raw_scores: true
     truncate: true
     truncation_direction: Right
diff --git a/11-embeddings-reranker-classification-tensorrt/BEI-mixedbread-ai-mxbai-rerank-base-v2-reranker-fp8/README.md b/11-embeddings-reranker-classification-tensorrt/BEI-mixedbread-ai-mxbai-rerank-base-v2-reranker-fp8/README.md
@@ -102,7 +102,9 @@ environment_variables: {}
 external_package_dirs: []
 model_metadata:
   example_model_input:
-    inputs: Baseten is a fast inference provider
+    inputs:
+    - - Baseten is a fast inference provider
+    - - Classify this separately.
     raw_scores: true
     truncate: true
     truncation_direction: Right
diff --git a/11-embeddings-reranker-classification-tensorrt/BEI-mixedbread-ai-mxbai-rerank-base-v2-reranker-fp8/config.yaml b/11-embeddings-reranker-classification-tensorrt/BEI-mixedbread-ai-mxbai-rerank-base-v2-reranker-fp8/config.yaml
@@ -3,7 +3,9 @@ environment_variables: {}
 external_package_dirs: []
 model_metadata:
   example_model_input:
-    inputs: Baseten is a fast inference provider
+    inputs:
+    - - Baseten is a fast inference provider
+    - - Classify this separately.
     raw_scores: true
     truncate: true
     truncation_direction: Right
diff --git a/11-embeddings-reranker-classification-tensorrt/BEI-mixedbread-ai-mxbai-rerank-large-v2-reranker-fp8/README.md b/11-embeddings-reranker-classification-tensorrt/BEI-mixedbread-ai-mxbai-rerank-large-v2-reranker-fp8/README.md
@@ -102,7 +102,9 @@ environment_variables: {}
 external_package_dirs: []
 model_metadata:
   example_model_input:
-    inputs: Baseten is a fast inference provider
+    inputs:
+    - - Baseten is a fast inference provider
+    - - Classify this separately.
     raw_scores: true
     truncate: true
     truncation_direction: Right
diff --git a/11-embeddings-reranker-classification-tensorrt/BEI-mixedbread-ai-mxbai-rerank-large-v2-reranker-fp8/config.yaml b/11-embeddings-reranker-classification-tensorrt/BEI-mixedbread-ai-mxbai-rerank-large-v2-reranker-fp8/config.yaml
@@ -3,7 +3,9 @@ environment_variables: {}
 external_package_dirs: []
 model_metadata:
   example_model_input:
-    inputs: Baseten is a fast inference provider
+    inputs:
+    - - Baseten is a fast inference provider
+    - - Classify this separately.
     raw_scores: true
     truncate: true
     truncation_direction: Right
diff --git a/11-embeddings-reranker-classification-tensorrt/BEI-papluca-xlm-roberta-base-language-detection-classification/README.md b/11-embeddings-reranker-classification-tensorrt/BEI-papluca-xlm-roberta-base-language-detection-classification/README.md
@@ -101,7 +101,9 @@ environment_variables: {}
 external_package_dirs: []
 model_metadata:
   example_model_input:
-    inputs: Baseten is a fast inference provider
+    inputs:
+    - - Baseten is a fast inference provider
+    - - Classify this separately.
     raw_scores: true
     truncate: true
     truncation_direction: Right
diff --git a/11-embeddings-reranker-classification-tensorrt/BEI-papluca-xlm-roberta-base-language-detection-classification/config.yaml b/11-embeddings-reranker-classification-tensorrt/BEI-papluca-xlm-roberta-base-language-detection-classification/config.yaml
@@ -3,7 +3,9 @@ environment_variables: {}
 external_package_dirs: []
 model_metadata:
   example_model_input:
-    inputs: Baseten is a fast inference provider
+    inputs:
+    - - Baseten is a fast inference provider
+    - - Classify this separately.
     raw_scores: true
     truncate: true
     truncation_direction: Right
diff --git a/11-embeddings-reranker-classification-tensorrt/BEI-samlowe-roberta-base-go_emotions-classification/README.md b/11-embeddings-reranker-classification-tensorrt/BEI-samlowe-roberta-base-go_emotions-classification/README.md
@@ -101,7 +101,9 @@ environment_variables: {}
 external_package_dirs: []
 model_metadata:
   example_model_input:
-    inputs: Baseten is a fast inference provider
+    inputs:
+    - - Baseten is a fast inference provider
+    - - Classify this separately.
     raw_scores: true
     truncate: true
     truncation_direction: Right
diff --git a/11-embeddings-reranker-classification-tensorrt/BEI-samlowe-roberta-base-go_emotions-classification/config.yaml b/11-embeddings-reranker-classification-tensorrt/BEI-samlowe-roberta-base-go_emotions-classification/config.yaml
@@ -3,7 +3,9 @@ environment_variables: {}
 external_package_dirs: []
 model_metadata:
   example_model_input:
-    inputs: Baseten is a fast inference provider
+    inputs:
+    - - Baseten is a fast inference provider
+    - - Classify this separately.
     raw_scores: true
     truncate: true
     truncation_direction: Right
diff --git a/11-embeddings-reranker-classification-tensorrt/BEI-skywork-skywork-reward-llama-3.1-8b-v0.2-reward-model-fp8/README.md b/11-embeddings-reranker-classification-tensorrt/BEI-skywork-skywork-reward-llama-3.1-8b-v0.2-reward-model-fp8/README.md
@@ -102,7 +102,9 @@ environment_variables: {}
 external_package_dirs: []
 model_metadata:
   example_model_input:
-    inputs: Baseten is a fast inference provider
+    inputs:
+    - - Baseten is a fast inference provider
+    - - Classify this separately.
     raw_scores: true
     truncate: true
     truncation_direction: Right
diff --git a/11-embeddings-reranker-classification-tensorrt/BEI-skywork-skywork-reward-llama-3.1-8b-v0.2-reward-model-fp8/config.yaml b/11-embeddings-reranker-classification-tensorrt/BEI-skywork-skywork-reward-llama-3.1-8b-v0.2-reward-model-fp8/config.yaml
@@ -3,7 +3,9 @@ environment_variables: {}
 external_package_dirs: []
 model_metadata:
   example_model_input:
-    inputs: Baseten is a fast inference provider
+    inputs:
+    - - Baseten is a fast inference provider
+    - - Classify this separately.
     raw_scores: true
     truncate: true
     truncation_direction: Right
diff --git a/11-embeddings-reranker-classification-tensorrt/Briton-meta-llama-llama-3.1-8b-instruct-with-speculative-lookahead-decoding-fp8/README.md b/11-embeddings-reranker-classification-tensorrt/Briton-meta-llama-llama-3.1-8b-instruct-with-speculative-lookahead-decoding-fp8/README.md
@@ -165,6 +165,7 @@ trt_llm:
       repo: meta-llama/Llama-3.1-8B-Instruct
       revision: main
       source: HF
+    max_batch_size: 64
     max_seq_len: 131072
     num_builder_gpus: 4
     plugin_configuration:
diff --git a/11-embeddings-reranker-classification-tensorrt/Briton-meta-llama-llama-3.1-8b-instruct-with-speculative-lookahead-decoding-fp8/config.yaml b/11-embeddings-reranker-classification-tensorrt/Briton-meta-llama-llama-3.1-8b-instruct-with-speculative-lookahead-decoding-fp8/config.yaml
@@ -30,6 +30,7 @@ trt_llm:
       repo: meta-llama/Llama-3.1-8B-Instruct
       revision: main
       source: HF
+    max_batch_size: 64
     max_seq_len: 131072
     num_builder_gpus: 4
     plugin_configuration:
diff --git a/11-embeddings-reranker-classification-tensorrt/Briton-meta-llama-llama-3.3-70b-instruct-speculative-with-1b-external-draft-fp8/README.md b/11-embeddings-reranker-classification-tensorrt/Briton-meta-llama-llama-3.3-70b-instruct-speculative-with-1b-external-draft-fp8/README.md
@@ -164,6 +164,7 @@ trt_llm:
       repo: meta-llama/Llama-3.3-70B-Instruct
       revision: main
       source: HF
+    max_batch_size: 64
     max_seq_len: 131072
     plugin_configuration:
       use_fp8_context_fmha: true
diff --git a/11-embeddings-reranker-classification-tensorrt/Briton-meta-llama-llama-3.3-70b-instruct-speculative-with-1b-external-draft-fp8/config.yaml b/11-embeddings-reranker-classification-tensorrt/Briton-meta-llama-llama-3.3-70b-instruct-speculative-with-1b-external-draft-fp8/config.yaml
@@ -29,6 +29,7 @@ trt_llm:
       repo: meta-llama/Llama-3.3-70B-Instruct
       revision: main
       source: HF
+    max_batch_size: 64
     max_seq_len: 131072
     plugin_configuration:
       use_fp8_context_fmha: true
diff --git a/11-embeddings-reranker-classification-tensorrt/Briton-qwen-qwen2.5-7b-instruct-with-speculative-lookahead-decoding-fp8/README.md b/11-embeddings-reranker-classification-tensorrt/Briton-qwen-qwen2.5-7b-instruct-with-speculative-lookahead-decoding-fp8/README.md
@@ -164,6 +164,7 @@ trt_llm:
       repo: Qwen/Qwen2.5-7B-Instruct
       revision: main
       source: HF
+    max_batch_size: 64
     max_seq_len: 32768
     num_builder_gpus: 4
     quantization_config:
diff --git a/11-embeddings-reranker-classification-tensorrt/Briton-qwen-qwen2.5-7b-instruct-with-speculative-lookahead-decoding-fp8/config.yaml b/11-embeddings-reranker-classification-tensorrt/Briton-qwen-qwen2.5-7b-instruct-with-speculative-lookahead-decoding-fp8/config.yaml
@@ -29,6 +29,7 @@ trt_llm:
       repo: Qwen/Qwen2.5-7B-Instruct
       revision: main
       source: HF
+    max_batch_size: 64
     max_seq_len: 32768
     num_builder_gpus: 4
     quantization_config:
diff --git a/11-embeddings-reranker-classification-tensorrt/Briton-qwen-qwq-32b-reasoning-with-speculative-fp8/README.md b/11-embeddings-reranker-classification-tensorrt/Briton-qwen-qwq-32b-reasoning-with-speculative-fp8/README.md
@@ -164,6 +164,7 @@ trt_llm:
       repo: Qwen/QwQ-32B
       revision: main
       source: HF
+    max_batch_size: 64
     max_seq_len: 40960
     num_builder_gpus: 4
     quantization_config:
diff --git a/11-embeddings-reranker-classification-tensorrt/Briton-qwen-qwq-32b-reasoning-with-speculative-fp8/config.yaml b/11-embeddings-reranker-classification-tensorrt/Briton-qwen-qwq-32b-reasoning-with-speculative-fp8/config.yaml
@@ -29,6 +29,7 @@ trt_llm:
       repo: Qwen/QwQ-32B
       revision: main
       source: HF
+    max_batch_size: 64
     max_seq_len: 40960
     num_builder_gpus: 4
     quantization_config:
diff --git a/11-embeddings-reranker-classification-tensorrt/templating/generate_templates.py b/11-embeddings-reranker-classification-tensorrt/templating/generate_templates.py
@@ -460,7 +460,10 @@ class Predictor(Task):
     model_metadata: dict = field(
         default_factory=lambda: dict(
             example_model_input={
-                "inputs": "Baseten is a fast inference provider",
+                "inputs": [
+                    ["Baseten is a fast inference provider"],
+                    ["Classify this separately."],
+                ],
                 "raw_scores": True,
                 "truncate": True,
                 "truncation_direction": "Right",
@@ -491,7 +494,7 @@ class Predictor(Task):
     headers=headers,
     url="https://model-xxxxxx.api.baseten.co/environments/production/sync/predict",
     json={
-        "inputs": "Baseten is a fast inference provider",
+        "inputs": [["Baseten is a fast inference provider", ["classify this separately."]],
         "raw_scores": True,
         "truncate": True,
         "truncation_direction": "Right"
@@ -501,10 +504,18 @@ class Predictor(Task):
 Returns:
 ```json
 [
-  {
-    "label": "excitement",
-    "score": 0.99
-  }
+  [
+    {
+        "label": "excitement",
+        "score": 0.99
+    }
+  ],
+  [
+    {
+        "label": "excitement",
+        "score": 0.01
+    }
+  ]
 ]
 ```
 Important, this is different from the `predict` route that you usually call. (https://model-xxxxxx.api.baseten.co/environments/production/predict), it contains an additional `sync` before that.
@@ -1015,6 +1026,7 @@ def llamalike_lookahead(
         lookahead_ngram_size=5,
         lookahead_verification_set_size=5,
     )
+    config.build.max_batch_size = 64
     return config
 
 
@@ -1034,6 +1046,7 @@ def llamalike_spec_dec(
             source=CheckpointSource.HF,
         ),
     )
+    config.build.max_batch_size = 64
     # the draft model and its kv cache live in the free memory of the target model's left-over KV cache
     config.runtime.kv_cache_free_gpu_mem_fraction = 0.45
 
