This configuration builds an inference engine to serve Llama 3.1 8B on an A100 GPU. It is very similar to the configuration for any other Llama 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 llama-3-1-8b-trt-llm
cd llama-3-1-8b-trt-llm
rm model/model.py

​
Configuration

This is a well-rounded configuration that balances latency and throughput. It supports long sequence lengths for multi-step chat and has a batch size of 32 as that’s reasonable for an eight-billion-parameter model on an A100 GPU. The model is served unquantized in fp16.

config.yaml
model_name: Llama 3.1 8B Engine
resources:
  accelerator: A100
secrets:
  hf_access_token: "set token in baseten workspace"
trt_llm:
  build:
    base_model: llama
    checkpoint_repository:
      repo: meta-llama/Llama-3.1-8B-Instruct
      source: HF
    max_seq_len: 8192

​
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.

Build your first LLM engineMistral with TensorRT-LLM