# How the Attention Mask Functions in Llama 2's Forward Pass

> Discover how the Llama 2 attention mask combines padding and causal masks to enforce constraints and ignore tokens in the forward pass. Optimize your LLM understanding.

- Repository: [Meta Llama/llama](https://github.com/meta-llama/llama)
- Tags: internals
- Published: 2026-03-05

---

**The Llama 2 attention mask combines a binary padding mask with a pre-computed causal triangular mask inside `LlamaModel.forward`, then adds this combined mask to raw attention logits before the softmax operation to enforce causal constraints and ignore padding tokens.**

The `meta-llama/llama` repository implements a **causal self-attention** mechanism that relies on a sophisticated masking strategy to handle variable-length batches and prevent future-token leakage. Understanding how the attention mask functions in the Llama 2 forward pass is essential for custom inference implementations or fine-tuning workflows. The core logic resides in [`llama/model.py`](https://github.com/meta-llama/llama/blob/main/llama/model.py), where the model constructs and applies a combined mask through additive masking rather than multiplicative approaches.

## Building the Combined Attention Mask

### Padding Mask Creation

Inside `LlamaModel.forward`, the workflow begins by distinguishing valid tokens from padding. The implementation converts `input_ids` into a binary mask using `(input_ids != self.pad_token_id).to(dtype)`, resulting in a tensor where real tokens hold a value of `1` and padding tokens hold `0`. This operation creates the foundation for the **padding mask**, ensuring the model does not waste computation on padded positions added to equalize sequence lengths within a batch.

### Dimension Expansion for Broadcasting

The binary mask undergoes critical reshaping to facilitate broadcasting across attention heads. The code expands the mask from shape `(batch, seq_len)` to `[batch, 1, 1, seq_len]`, allowing it to align with the four-dimensional query-key score matrix during attention calculations. This expansion occurs via tensor slicing operations that prepare the mask for efficient combination with the causal component.

### Integrating the Causal Mask

Llama 2 utilizes a pre-computed lower-triangular causal mask, typically stored as `self.bias` with shape `[1, 1, seq_len, seq_len]`. The forward pass adds this causal mask to the expanded padding mask, producing a combined tensor. Wherever attention should be blocked—whether due to padding or future-token constraints—the sum yields `float("-inf")` (or the PyTorch constant `torch.finfo(dtype).min`). This **additive combination** generates the final mask that drives the softmax behavior.

## Applying the Mask in LlamaAttention

Within each `LlamaAttention` block, the raw attention scores—computed as `query @ key.transpose(-2, -1) / sqrt(d_k)`—receive the combined mask through simple addition. Because the softmax function interprets large negative values as zero probability, positions marked with `-inf` effectively disappear from the attention distribution according to the analysis of [`llama/model.py`](https://github.com/meta-llama/llama/blob/main/llama/model.py) around line 260. This **additive masking** approach proves efficient because the same mask broadcasts across all attention heads without requiring separate tensor operations per head or per layer.

## Practical Implementation Examples

When using the library directly, mask construction happens automatically inside `LlamaModel.forward`, but understanding the manual process aids custom implementations:

```python
import torch
from llama import LlamaForCausalLM, LlamaTokenizer

# Load model & tokenizer (assume weights are already downloaded)

tokenizer = LlamaTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf")
model = LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")

# Example inputs – note the trailing padding token (id 0)

input_ids = torch.tensor([[1, 5, 23, 0, 0]])   # shape: (batch, seq_len)

# 1️⃣ Build the attention mask (same logic as the library)

pad_id = tokenizer.pad_token_id
attention_mask = (input_ids != pad_id).to(model.dtype)   # 1 for real tokens, 0 for pads

attention_mask = attention_mask[:, None, None, :]        # (batch, 1, 1, seq_len)

# 2️⃣ Forward pass – the mask is automatically combined with the causal mask

outputs = model(input_ids=input_ids, attention_mask=attention_mask)
logits = outputs.logits

```

For custom inference loops that bypass high-level APIs, you can manually replicate the mask combination logic:

```python
seq_len = input_ids.size(1)
device = input_ids.device

# Causal lower-triangular mask (1 = keep, 0 = mask)

causal_mask = torch.tril(torch.ones(seq_len, seq_len, device=device)).unsqueeze(0).unsqueeze(0)

# Final additive mask: -inf where we want to block attention

mask = (1.0 - attention_mask) * -1e9          # pads -> -inf

combined_mask = causal_mask + mask             # shape (1, 1, seq_len, seq_len)

# Inside an attention block (simplified)

scores = q @ k.transpose(-2, -1) / math.sqrt(d_k)
scores = scores + combined_mask
weights = torch.softmax(scores, dim=-1)

```

These snippets demonstrate that the **attention mask is not a separate tensor handed to the attention heads**; it is merged with the causal mask and added to the logits before softmax, exactly as implemented in the `meta-llama/llama` source code.

## Why the Mask Architecture Matters

- **Correct Language Modeling**: The causal component ensures autoregressive training integrity by preventing tokens from attending to future positions, eliminating the possibility of data leakage during training.
- **Batch Efficiency**: The padding mask enables processing of variable-length sequences within a single tensor, eliminating wasted computation on pad tokens and maximizing GPU utilization.
- **Numerical Stability**: Using `-inf` or `torch.finfo(dtype).min` prevents overflow during softmax computation, maintaining gradient stability throughout deep transformer stacks.

## Summary

- The Llama 2 forward pass constructs a **combined attention mask** in [`llama/model.py`](https://github.com/meta-llama/llama/blob/main/llama/model.py) that merges padding and causal constraints into a single tensor.
- The mask expands to shape `[batch, 1, 1, seq_len]` to broadcast efficiently across query-key pairs and attention heads.
- **Additive masking** adds the combined mask to attention logits before softmax, setting blocked positions to negative infinity to zero out their probability mass.
- This approach handles variable-length batches efficiently while enforcing strict causal attention boundaries required for autoregressive generation.
- The implementation leverages `self.bias` (the causal mask) alongside pad-token identification to create the final blocking tensor applied in `LlamaAttention`.

## Frequently Asked Questions

### What is the difference between the padding mask and causal mask in Llama 2?

The padding mask identifies which positions contain actual tokens versus padding (based on `pad_token_id`), while the causal mask enforces the autoregressive property by blocking attention to future tokens through a lower-triangular matrix. Llama 2 combines these into a single additive mask that handles both constraints simultaneously before the softmax operation.

### Why does Llama 2 use additive masking instead of multiplicative masking?

Additive masking proves more numerically stable and computationally efficient for transformer architectures. By adding large negative values (`-inf` or `torch.finfo(dtype).min`) to attention logits before softmax, the model achieves zero probability for blocked positions without requiring separate multiplication operations that could introduce floating-point precision issues or additional memory overhead.

### Where in the codebase is the attention mask actually applied to the attention scores?

The mask application occurs within the attention computation in [`llama/model.py`](https://github.com/meta-llama/llama/blob/main/llama/model.py), specifically where the raw attention scores (the query-key dot products scaled by `sqrt(d_k)`) are added to the combined mask tensor immediately before the softmax normalization step. This location corresponds to the internal `LlamaAttention` implementation within the model file.

### Can the same attention mask be reused across different layers in Llama 2?

Yes, the combined attention mask is constructed once in `LlamaModel.forward` and reused across all transformer layers in the stack. Its broadcastable shape `[batch, 1, 1, seq_len]` allows efficient sharing across all attention heads and layers without consuming excessive memory or requiring per-layer mask recomputation.