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