Mistral Architecture › Pre-fill and Chunking
Language AI HandbookAll chapters Mistral Architecture › Pre-fill and ChunkingOne section · Sections in this chapter6 / 14 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
When processing a prompt before generation begins, standard transformers compute attention for all tokens at once. For long prompts, this can create memory spikes that exceed available capacity. Mistral introduces pre-fill chunking to address this.
The idea is simple: instead of processing the entire prompt in a single forward pass, split it into chunks of size (the window size). Process each chunk sequentially, using the rolling KV cache to maintain context between chunks. This bounds memory usage during pre-fill to the same level as during generation.
def chunked_prefill_memory(
prompt_length, window_size, hidden_size, num_heads, dtype_bytes=2
):
"""
Calculate peak memory during chunked pre-fill.
Instead of processing all tokens at once, we process
chunks of size window_size sequentially.
"""
chunk_size = min(prompt_length, window_size)
# Attention matrix: chunk_size x chunk_size per head
attention_mem = chunk_size * chunk_size * num_heads * dtype_bytes
# Activations for current chunk: chunk_size x hidden_size
activation_mem = chunk_size * hidden_size * dtype_bytes
return attention_mem + activation_mem
def standard_prefill_memory(
prompt_length, hidden_size, num_heads, dtype_bytes=2
):
"""Calculate memory for standard (non-chunked) pre-fill."""
# Full attention matrix: n x n per head
attention_mem = prompt_length * prompt_length * num_heads * dtype_bytes
# Activations: n x hidden_size
activation_mem = prompt_length * hidden_size * dtype_bytes
return attention_mem + activation_mem
prompt_lengths = [1024, 2048, 4096, 8192, 16384]
window_size = 4096
hidden_size = 4096
num_heads = 32Pre-fill Memory Comparison:
Prompt Length Standard Chunked Reduction
----------------------------------------------------------
1,024 72.0 MB 72.0 MB 0.0%
2,048 272.0 MB 272.0 MB 0.0%
4,096 1056.0 MB 1056.0 MB 0.0%
8,192 4160.0 MB 1056.0 MB 74.6%
16,384 16512.0 MB 1056.0 MB 93.6%For an 8K token prompt, chunked pre-fill uses 75% less memory than standard pre-fill. This reduction enables processing of longer prompts on the same hardware, or processing multiple requests concurrently.
