How the Llama 2 Inference Loop Handles Variable-Length Prompt Batches

The Llama 2 inference loop in llama/generation.py pads all prompts to a common tensor size and uses a boolean mask to ensure padding tokens never influence generation, enabling efficient parallel processing of variable-length batches.

The meta-llama/llama repository implements a sophisticated batching strategy in its text generation pipeline. When processing multiple requests with the Llama 2 inference loop, the system must reconcile different prompt lengths while maintaining autoregressive correctness. This article examines the exact padding, masking, and token replacement mechanisms used in llama/generation.py to handle variable-length prompt batches without sacrificing computational efficiency.

The Batching Challenge in Text Generation

Text generation models process sequences as tensors, requiring uniform dimensions across the batch dimension. When prompts vary in length—such as a short question alongside a long document excerpt—naive truncation or separate processing wastes compute or breaks parallelism. The Llama 2 implementation solves this through strategic padding and attention masking.

Inside the Llama.generate Method

The core inference logic resides in llama/generation.py within the Llama.generate method (lines 58–210). This method accepts prompt_tokens: List[List[int]], where each inner list represents a tokenized prompt of arbitrary length. The method orchestrates tensor allocation, the forward pass loop, and post-processing to ensure that variable-length prompts generate correctly within a single batched CUDA kernel launch.

Step-by-Step: Padding and Masking Logic

Calculating Batch Dimensions

Before allocating tensors, the method determines the batch size and length extremes to establish the working tensor dimensions:

bsz = len(prompt_tokens)
min_prompt_len = min(len(t) for t in prompt_tokens)
max_prompt_len = max(len(t) for t in prompt_tokens)
total_len = min(params.max_seq_len, max_gen_len + max_prompt_len)

(Source: llama/generation.py, lines 58–66)

These calculations ensure the output tensor accommodates the longest prompt plus requested generation tokens, capped by the model's maximum sequence length.

Padding with the Tokenizer's pad_id

The code allocates a CUDA tensor filled with the tokenizer's pad_id and copies each prompt into its respective row, left-aligning the sequences:

tokens = torch.full((bsz, total_len), pad_id, dtype=torch.long, device="cuda")
for k, t in enumerate(prompt_tokens):
    tokens[k, :len(t)] = torch.tensor(t, dtype=torch.long, device="cuda")

(Source: llama/generation.py, lines 67–71)

This left-padding approach ensures that trailing positions contain only padding tokens, which will be masked during generation.

Creating the Input Text Mask

To prevent the model from attending to padding positions, the method creates a boolean mask identifying real prompt tokens:

input_text_mask = tokens != pad_id

(Source: llama/generation.py, lines 75–76)

This mask becomes critical during the generation loop to distinguish between prompt tokens that should be preserved and positions requiring new generation.

The Generation Loop with Variable Lengths

Selective Token Updates with torch.where

The generation loop begins at min_prompt_len and iterates to total_len. At each position cur_pos, the model produces logits for the next token. The key innovation is the selective update using torch.where to handle variable prompt lengths within the batch:

next_token = torch.where(
    input_text_mask[:, cur_pos],   # keep original token while still in prompt

    tokens[:, cur_pos],
    next_token                    # replace with sampled token after prompt

)

(Source: llama/generation.py, lines 94–99)

This ensures that for sequences where cur_pos falls within the original prompt, the existing token is preserved, while positions beyond the prompt length receive newly sampled tokens. Because the mask is applied per-sample, each sequence in the batch advances independently.

Per-Sample Early Stopping

The loop tracks completion status individually for each batch element to avoid unnecessary computation once shorter sequences finish:

eos_reached = (~input_text_mask[:, cur_pos]) & (next_token == tokenizer.eos_id)
if eos_reached.all():
    break

(Source: llama/generation.py, lines 107–112)

The eos_reached tensor updates only for positions outside the original prompt (where generation occurs), allowing the loop to exit early once all sequences have generated an end-of-sequence token, even if they started at different times.

Post-Processing and Decoding

After the loop terminates, Llama.generate processes the output tensor to remove padding and optionally the EOS token. According to lines 176–210 in llama/generation.py, the method iterates through each batch row, slicing from the start of the sequence up to the first padding token or EOS, depending on the echo flag. This yields clean text outputs regardless of the original prompt lengths or padding artifacts.

Practical Implementation Example

The following example demonstrates variable-length batch processing using the official Llama 2 API:

from llama import Llama

# Initialize generator

generator = Llama.build(
    ckpt_dir="checkpoints",
    tokenizer_path="tokenizer.model",
    max_seq_len=512,
    max_batch_size=8
)

# Variable-length prompts

prompts = [
    "Explain quantum entanglement.",
    "Write a detailed analysis of climate change impacts on coastal ecosystems, including specific species examples.",
    "What is 2+2?"
]

# Tokenize with BOS, no EOS

prompt_tokens = [
    generator.tokenizer.encode(p, bos=True, eos=False) 
    for p in prompts
]

# Generate - padding and masking handled internally

results = generator.text_completion(
    prompts=prompts,
    max_gen_len=50,
    temperature=0.7,
    top_p=0.9
)

for result in results:
    print(result["generation"])

This code executes batched inference across three prompts of significantly different lengths. The Llama.generate method automatically applies the padding and masking logic described above, allowing the GPU to process all sequences in parallel while maintaining correct autoregressive generation for each individual prompt.

Summary

• The Llama 2 inference loop in llama/generation.py handles variable-length prompts by padding all sequences to the maximum combined length of prompts plus generation tokens.
• A boolean mask (input_text_mask) distinguishes real prompt tokens from padding, ensuring the model never attends to padded positions.
• During generation, torch.where selectively updates only positions beyond each prompt's original length, preserving pre-filled prompt tokens while sampling new ones.
• Per-sample early stopping via eos_reached allows the loop to terminate once all variable-length sequences complete, optimizing inference time.

Frequently Asked Questions

Does Llama 2 use left-padding or right-padding for variable-length batches?

Llama 2 uses left-padding (pre-padding) in the Llama.generate method. The code allocates a tensor filled with pad_id and copies each prompt starting at index 0, leaving trailing positions as padding. This left-alignment ensures that new tokens are generated at consistent positions relative to the prompt end across the batch.

How does the inference loop prevent the model from using padding tokens as input?

The loop creates a boolean mask input_text_mask = tokens != pad_id that identifies real prompt tokens. During generation, the torch.where operation uses this mask to preserve original prompt tokens and only replace padding positions with newly sampled tokens. Because the mask ensures padding tokens are never written back into the active generation positions, they never influence subsequent forward passes.

Can the Llama 2 inference loop handle a batch where some prompts are already complete while others continue generating?

Yes, the implementation tracks completion status per sample using the eos_reached tensor. The loop continues until eos_reached.all() returns True, meaning it only breaks once every sequence in the batch has generated an EOS token. This allows shorter prompts to reach EOS and stop contributing new tokens while longer prompts continue generating, all within the same batched CUDA kernel launches.

What happens to the padding in the final output returned to the user?

After the generation loop completes, the Llama.generate method post-processes each sequence to remove padding tokens. According to lines 176–210 in llama/generation.py, the code slices each output row up to the first padding token or EOS (depending on the echo flag), ensuring the final decoded strings contain only the generated text without padding artifacts.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}