Implementation:Pytorch Serve Llama2 Generate

Knowledge Sources	Pytorch_Serve
Domains	LLM_Serving, Text_Generation, Sampling
Last Updated	2026-02-13 18:52 GMT

Overview

Llama2_Generate provides generation utilities for the Llama2 model in TorchServe. It implements nucleus (top-p) sampling, autoregressive text generation, single-turn text completion, and multi-turn chat completion. The module defines typed data structures for dialog messages, completion predictions, and chat predictions.

Description

The generate.py module contains four core functions and several type definitions that together form the generation pipeline for Llama2 models. It supports both text completion (single prompt in, single text out) and chat completion (multi-turn dialog in, response out) use cases.

Key Components

• sample_top_p(): Implements nucleus sampling (top-p sampling), which selects from the smallest set of tokens whose cumulative probability exceeds the threshold p
• generate(): Core autoregressive generation loop that produces tokens one at a time, feeding each generated token back as input
• text_completion(): Wraps generate() to provide a simple text-in, text-out interface for single-turn prompts
• chat_completion(): Formats multi-turn dialog messages into the Llama2 chat template and generates a response
• Type Definitions: Role, Message, CompletionPrediction, ChatPrediction, and Dialog provide structured types for the generation pipeline

Usage

from examples.large_models.tp_llama.generate import (
    sample_top_p,
    generate,
    text_completion,
    chat_completion,
    Role,
    Message,
    CompletionPrediction,
    ChatPrediction,
    Dialog,
)

Code Reference

Source Location

File	Lines	Repository
examples/large_models/tp_llama/generate.py	L1-346	pytorch/serve
examples/large_models/tp_llama/generate.py	L43-66	sample_top_p() function
examples/large_models/tp_llama/generate.py	L69-176	generate() function
examples/large_models/tp_llama/generate.py	L179-231	text_completion() function
examples/large_models/tp_llama/generate.py	L234-346	chat_completion() function

Signature

# Type definitions
Role = Literal["system", "user", "assistant"]


class Message(TypedDict):
    role: Role
    content: str


class CompletionPrediction(TypedDict, total=False):
    generation: str
    tokens: list[str]
    logprobs: list[float]


class ChatPrediction(TypedDict, total=False):
    generation: Message
    tokens: list[str]
    logprobs: list[float]


Dialog = list[Message]


def sample_top_p(probs: torch.Tensor, p: float) -> torch.Tensor:
    """
    Perform nucleus (top-p) sampling on a probability distribution.

    Sorts probabilities in descending order, computes the cumulative sum,
    and masks out tokens whose cumulative probability exceeds the threshold p.
    Then samples from the filtered distribution.

    Args:
        probs (torch.Tensor): Probability distribution tensor of shape (batch, vocab).
        p (float): Cumulative probability threshold (0.0 < p <= 1.0).

    Returns:
        torch.Tensor: Sampled token index of shape (batch, 1).
    """
    ...


def generate(
    model,
    prompt_tokens: list[list[int]],
    max_gen_len: int,
    temperature: float = 0.6,
    top_p: float = 0.9,
    logprobs: bool = False,
    echo: bool = False,
) -> tuple:
    """
    Autoregressive token generation for a batch of prompts.

    Iteratively generates tokens by feeding the model output back as input.
    Supports temperature scaling, top-p sampling, and optional log-probability
    collection.

    Args:
        model: The Llama2 model instance.
        prompt_tokens (list[list[int]]): Batch of tokenized prompts.
        max_gen_len (int): Maximum number of new tokens to generate.
        temperature (float): Sampling temperature (default 0.6).
        top_p (float): Nucleus sampling threshold (default 0.9).
        logprobs (bool): Whether to collect token log-probabilities.
        echo (bool): Whether to include prompt tokens in output.

    Returns:
        tuple: (generated_tokens, generated_logprobs) where generated_tokens
               is a list of token lists and generated_logprobs is optionally
               a list of log-probability lists.
    """
    ...


def text_completion(
    model,
    tokenizer,
    prompts: list[str],
    max_gen_len: int,
    temperature: float = 0.6,
    top_p: float = 0.9,
    logprobs: bool = False,
    echo: bool = False,
) -> list[CompletionPrediction]:
    """
    Generate text completions for a batch of string prompts.

    Tokenizes each prompt, calls generate(), and decodes the output tokens
    back to strings.

    Args:
        model: The Llama2 model instance.
        tokenizer: SentencePiece tokenizer.
        prompts (list[str]): List of text prompts.
        max_gen_len (int): Maximum tokens to generate per prompt.
        temperature (float): Sampling temperature.
        top_p (float): Nucleus sampling threshold.
        logprobs (bool): Whether to return log-probabilities.
        echo (bool): Whether to include the prompt in the output.

    Returns:
        list[CompletionPrediction]: List of completion results with
            'generation' text and optional 'tokens' and 'logprobs'.
    """
    ...


def chat_completion(
    model,
    tokenizer,
    dialogs: list[Dialog],
    max_gen_len: int,
    temperature: float = 0.6,
    top_p: float = 0.9,
    logprobs: bool = False,
) -> list[ChatPrediction]:
    """
    Generate chat responses for multi-turn dialogs.

    Formats each dialog into the Llama2 chat template with special tokens
    ([INST], [/INST], <<SYS>>, <</SYS>>), generates a response, and
    returns structured ChatPrediction results.

    Args:
        model: The Llama2 model instance.
        tokenizer: SentencePiece tokenizer.
        dialogs (list[Dialog]): List of multi-turn dialogs.
        max_gen_len (int): Maximum tokens to generate per response.
        temperature (float): Sampling temperature.
        top_p (float): Nucleus sampling threshold.
        logprobs (bool): Whether to return log-probabilities.

    Returns:
        list[ChatPrediction]: List of chat predictions with
            'generation' Message and optional 'tokens' and 'logprobs'.
    """
    ...

Import

import torch
from typing import Literal, TypedDict

from examples.large_models.tp_llama.generate import (
    sample_top_p,
    generate,
    text_completion,
    chat_completion,
)

I/O Contract

Function	Input	Output	Notes
sample_top_p(probs, p)	probs: torch.Tensor (batch, vocab); p: float	torch.Tensor sampled token index (batch, 1)	Nucleus sampling; filters tokens by cumulative probability
generate(model, prompt_tokens, max_gen_len, ...)	Model, tokenized prompts (list[list[int]]), generation params	tuple: (token lists, optional logprob lists)	Core autoregressive loop; supports batched generation
text_completion(model, tokenizer, prompts, ...)	Model, tokenizer, text prompts (list[str]), generation params	list[CompletionPrediction]	Tokenizes, generates, decodes; single-turn text completion
chat_completion(model, tokenizer, dialogs, ...)	Model, tokenizer, dialogs (list[Dialog]), generation params	list[ChatPrediction]	Formats Llama2 chat template; multi-turn conversation

Llama2 Chat Template Format

<s>[INST] <<SYS>>
{system_message}
<</SYS>>

{user_message_1} [/INST] {assistant_response_1} </s><s>[INST] {user_message_2} [/INST]

Usage Examples

Example 1: Text completion with a single prompt

from examples.large_models.tp_llama.generate import text_completion

prompts = [
    "The theory of relativity states that",
    "In machine learning, a transformer is",
]

results = text_completion(
    model=llama_model,
    tokenizer=tokenizer,
    prompts=prompts,
    max_gen_len=128,
    temperature=0.6,
    top_p=0.9,
)

for prompt, result in zip(prompts, results):
    print(f"Prompt: {prompt}")
    print(f"Completion: {result['generation']}")
    print()

Example 2: Multi-turn chat completion

from examples.large_models.tp_llama.generate import chat_completion, Dialog

dialogs: list[Dialog] = [
    [
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": "What is TorchServe?"},
    ],
    [
        {"role": "system", "content": "You are a coding expert."},
        {"role": "user", "content": "How do I serve a model with TorchServe?"},
        {"role": "assistant", "content": "You can use the torch-model-archiver..."},
        {"role": "user", "content": "Can you show me the config file?"},
    ],
]

results = chat_completion(
    model=llama_model,
    tokenizer=tokenizer,
    dialogs=dialogs,
    max_gen_len=256,
    temperature=0.6,
    top_p=0.9,
)

for dialog, result in zip(dialogs, results):
    print(f"Last user message: {dialog[-1]['content']}")
    print(f"Response: {result['generation']['content']}")
    print()

Example 3: Nucleus sampling behavior

import torch
from examples.large_models.tp_llama.generate import sample_top_p

# Example probability distribution over a vocabulary of 5 tokens
probs = torch.tensor([[0.4, 0.3, 0.15, 0.1, 0.05]])

# With p=0.9, tokens with cumulative prob < 0.9 are kept
# Sorted: [0.4, 0.3, 0.15, 0.1, 0.05]
# Cumulative: [0.4, 0.7, 0.85, 0.95, 1.0]
# Tokens 0, 1, 2 are kept (cumulative 0.85 < 0.9), token 3 is the boundary
sampled = sample_top_p(probs, p=0.9)
print(f"Sampled token index: {sampled.item()}")

Related Pages

• Principle:Pytorch_Serve_LLM_Text_Generation -- The principle of LLM text generation with sampling strategies
• Implementation:Pytorch_Serve_Llama2_Model -- The Llama2 model architecture used by these generation functions
• Implementation:Pytorch_Serve_GptHandler -- GPT Fast handler with related generation logic