This configuration builds an inference engine to serve Mistral 7B on an H100 GPU. It is very similar to the configuration for any other Mistral model, including fine-tuned variants.

Setup

See the end-to-end engine builder tutorial prerequisites for full setup instructions.

Make sure you have accessed the gated model on Hugging Face and set your hf_access_token in your Baseten workspace secrets.

Please upgrade to the latest version of Truss with pip install --upgrade truss before following this example.

pip install --upgrade truss
truss init mistral-7b-trt-llm
cd mistral-7b-trt-llm
rm model/model.py

Configuration

This configuration is optimized for low latency, with a batch size of 8. It applies post-training quantization to fp8 for further speed gains.

config.yaml
model_name: Mistral Engine 7B v0.2 6
python_version: py39
resources:
  accelerator: H100:1
  use_gpu: true
secrets:
  hf_access_token: "set token in baseten workspace"
trt_llm:
  build:
    base_model: mistral
    checkpoint_repository:
      repo: mistralai/Mistral-7B-Instruct-v0.2
      source: HF
    max_input_len: 8192
    quantization_type: fp8

Deployment

truss push --publish --trusted

Usage

call_model.py
import requests
import os

# Model ID for production deployment
model_id = ""
# Read secrets from environment variables
baseten_api_key = os.environ["BASETEN_API_KEY"]

# Call model endpoint
resp = requests.post(
    f"https://model-{model_id}.api.baseten.co/production/predict",
    headers={"Authorization": f"Api-Key {baseten_api_key}"},
    json={
      "messages": [{"role": "user", "content": "How awesome is TensorRT-LLM?"}],
      "max_tokens": 1024
    },
    stream=True
)

# Print the generated tokens as they get streamed
for content in resp.iter_content():
    print(content.decode("utf-8"), end="", flush=True)
prompt
string

The input text prompt to guide the language model’s generation.

One of prompt XOR messages is required.

messages
List[Dict]

A list of dictionaries representing the message history, typically used in conversational contexts.

One of prompt XOR messages is required.

max_tokens
int

The maximum number of tokens to generate in the output. Controls the length of the generated text.

beam_width
int
default: "1"

The number of beams used in beam search. Maximum of 1.

repetition_penalty
float

A penalty applied to repeated tokens to discourage the model from repeating the same words or phrases.

presence_penalty
float

A penalty applied to tokens already present in the prompt to encourage the generation of new topics.

temperature
float

Controls the randomness of the output. Lower values make the output more deterministic, while higher values increase randomness.

length_penalty
float

A penalty applied to the length of the generated sequence to control verbosity. Higher values make the model favor shorter outputs.

end_id
int

The token ID that indicates the end of the generated sequence.

pad_id
int

The token ID used for padding sequences to a uniform length.

runtime_top_k
int

Limits the sampling pool to the top k tokens, ensuring the model only considers the most likely tokens at each step.

runtime_top_p
float

Applies nucleus sampling to limit the sampling pool to a cumulative probability p, ensuring only the most likely tokens are considered.