xbe.at

NLP and transformers / BERT and language models

Optimizing Transformer Models with KV Cache and Advanced Techniques

10 min read, 14 sections

Contents
  1. Introduction to KV Cache
  2. Implementing KV Cache in Attention Mechanism
  3. Benefits of KV Cache
  4. Introduction to Flash Attention
  5. Memory Efficiency of Flash Attention
  6. Implementing vLLM for Efficient Inference
  7. Continuous Batching in vLLM
  8. PagedAttention in vLLM
  9. Optimizing Transformer Inference with vLLM
  10. Implementing Efficient Token Generation
  11. Real-life Example: Text Completion
  12. Real-life Example: Dialogue Generation
  13. Comparing Standard Attention and Flash Attention
  14. Additional Resources

Introduction to KV Cache

KV Cache, short for Key-Value Cache, is a technique used to optimize the performance of transformer models. It stores precomputed key and value tensors from previous time steps, reducing redundant computations in autoregressive generation tasks.

import torch

class KVCache:
    def __init__(self, max_seq_len, num_layers, num_heads, head_dim):
        self.max_seq_len = max_seq_len
        self.num_layers = num_layers
        self.num_heads = num_heads
        self.head_dim = head_dim
        self.cache = {
            'k': torch.zeros(num_layers, max_seq_len, num_heads, head_dim),
            'v': torch.zeros(num_layers, max_seq_len, num_heads, head_dim)
        }

    def update(self, layer, position, k, v):
        self.cache['k'][layer, position] = k
        self.cache['v'][layer, position] = v

    def get(self, layer, start, end):
        return self.cache['k'][layer, start:end], self.cache['v'][layer, start:end]

Implementing KV Cache in Attention Mechanism

Let's integrate the KV Cache into a simplified attention mechanism to demonstrate its usage in practice.

def attention_with_kv_cache(query, kv_cache, layer, position):
    # Retrieve cached keys and values
    keys, values = kv_cache.get(layer, 0, position + 1)
    
    # Compute attention scores
    scores = torch.matmul(query, keys.transpose(-2, -1)) / math.sqrt(query.size(-1))
    
    # Apply softmax to get attention weights
    weights = torch.softmax(scores, dim=-1)
    
    # Compute weighted sum of values
    output = torch.matmul(weights, values)
    
    return output

# Usage example
query = torch.rand(1, 1, 64)  # Assume 1 head with dimension 64
kv_cache = KVCache(max_seq_len=100, num_layers=1, num_heads=1, head_dim=64)
layer = 0
position = 10

result = attention_with_kv_cache(query, kv_cache, layer, position)
print(result.shape)  # Expected output: torch.Size([1, 1, 64])

Benefits of KV Cache

KV Cache significantly improves the efficiency of transformer models during inference, especially for autoregressive tasks. By storing precomputed key and value tensors, it eliminates the need to recompute them for each new token, resulting in faster generation times and reduced computational overhead.

def generate_with_kv_cache(model, input_ids, max_length):
    kv_cache = KVCache(max_seq_len=max_length, num_layers=model.num_layers,
                       num_heads=model.num_heads, head_dim=model.head_dim)
    
    for i in range(max_length - len(input_ids)):
        # Forward pass with KV Cache
        logits = model(input_ids, kv_cache)
        
        # Select the next token
        next_token = torch.argmax(logits[:, -1, :])
        input_ids = torch.cat([input_ids, next_token.unsqueeze(0)], dim=-1)
        
        if next_token == model.eos_token_id:
            break
    
    return input_ids

# Usage example (assuming 'model' is a pre-trained transformer)
input_ids = torch.tensor([1, 2, 3])  # Start of sequence
generated_ids = generate_with_kv_cache(model, input_ids, max_length=50)
print(generated_ids)

Introduction to Flash Attention

Flash Attention is an optimized attention algorithm that reduces memory usage and increases computational efficiency. It achieves this by splitting the input sequence into blocks and processing them iteratively, avoiding the need to store the full attention matrix in memory.

import torch
import math

def flash_attention(q, k, v, block_size=256):
    batch_size, seq_len, num_heads, head_dim = q.shape
    scale = 1 / math.sqrt(head_dim)
    
    output = torch.zeros_like(q)
    normalizer = torch.zeros((batch_size, num_heads, seq_len, 1))
    
    for i in range(0, seq_len, block_size):
        block_q = q[:, i:i+block_size]
        
        for j in range(0, seq_len, block_size):
            block_k = k[:, j:j+block_size]
            block_v = v[:, j:j+block_size]
            
            scores = torch.matmul(block_q, block_k.transpose(-2, -1)) * scale
            block_probs = torch.softmax(scores, dim=-1)
            
            output[:, i:i+block_size] += torch.matmul(block_probs, block_v)
            normalizer[:, :, i:i+block_size] += block_probs.sum(dim=-1, keepdim=True)
    
    return output / normalizer

# Usage example
batch_size, seq_len, num_heads, head_dim = 2, 1024, 8, 64
q = torch.randn(batch_size, seq_len, num_heads, head_dim)
k = torch.randn(batch_size, seq_len, num_heads, head_dim)
v = torch.randn(batch_size, seq_len, num_heads, head_dim)

result = flash_attention(q, k, v)
print(result.shape)  # Expected output: torch.Size([2, 1024, 8, 64])

Memory Efficiency of Flash Attention

Flash Attention significantly reduces memory usage compared to standard attention mechanisms. By processing the input in blocks, it avoids storing the full attention matrix, which can be prohibitively large for long sequences.

def compare_memory_usage(seq_len, num_heads, head_dim):
    batch_size = 1
    
    # Standard attention memory usage
    q = torch.randn(batch_size, seq_len, num_heads, head_dim)
    k = torch.randn(batch_size, seq_len, num_heads, head_dim)
    v = torch.randn(batch_size, seq_len, num_heads, head_dim)
    
    torch.cuda.empty_cache()
    torch.cuda.reset_peak_memory_stats()
    
    # Simulate standard attention
    scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(head_dim)
    attn = torch.softmax(scores, dim=-1)
    output = torch.matmul(attn, v)
    
    standard_memory = torch.cuda.max_memory_allocated()
    
    # Flash attention memory usage
    torch.cuda.empty_cache()
    torch.cuda.reset_peak_memory_stats()
    
    flash_output = flash_attention(q, k, v)
    
    flash_memory = torch.cuda.max_memory_allocated()
    
    print(f"Standard Attention Memory: {standard_memory / 1e6:.2f} MB")
    print(f"Flash Attention Memory: {flash_memory / 1e6:.2f} MB")
    print(f"Memory Reduction: {(1 - flash_memory / standard_memory) * 100:.2f}%")

# Usage example
compare_memory_usage(seq_len=4096, num_heads=8, head_dim=64)

Implementing vLLM for Efficient Inference

vLLM (Variable Length Language Model) is a technique that optimizes inference for transformer models by dynamically managing attention and hidden states. It allows for efficient processing of variable-length sequences and supports features like continuous batching.

import torch
import torch.nn as nn

class vLLMAttention(nn.Module):
    def __init__(self, hidden_size, num_heads):
        super().__init__()
        self.hidden_size = hidden_size
        self.num_heads = num_heads
        self.head_dim = hidden_size // num_heads
        
        self.q_proj = nn.Linear(hidden_size, hidden_size)
        self.k_proj = nn.Linear(hidden_size, hidden_size)
        self.v_proj = nn.Linear(hidden_size, hidden_size)
        self.o_proj = nn.Linear(hidden_size, hidden_size)
        
    def forward(self, hidden_states, past_key_values=None, attention_mask=None):
        batch_size, seq_len, _ = hidden_states.shape
        
        q = self.q_proj(hidden_states)
        k = self.k_proj(hidden_states)
        v = self.v_proj(hidden_states)
        
        q = q.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        k = k.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        v = v.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        
        if past_key_values is not None:
            past_k, past_v = past_key_values
            k = torch.cat([past_k, k], dim=2)
            v = torch.cat([past_v, v], dim=2)
        
        attn_output = flash_attention(q, k, v)
        
        attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, seq_len, self.hidden_size)
        attn_output = self.o_proj(attn_output)
        
        return attn_output, (k, v)

# Usage example
hidden_size, num_heads = 768, 12
vlm_attention = vLLMAttention(hidden_size, num_heads)
hidden_states = torch.randn(2, 16, hidden_size)
output, past_kv = vlm_attention(hidden_states)
print(output.shape)  # Expected output: torch.Size([2, 16, 768])
print(past_kv[0].shape)  # Expected output: torch.Size([2, 12, 16, 64])

Continuous Batching in vLLM

Continuous batching is a key feature of vLLM that allows for efficient processing of multiple sequences with different lengths. It dynamically manages the attention computations to minimize wasted computation on padding tokens.

class ContinuousBatchProcessor:
    def __init__(self, model):
        self.model = model
        self.active_sequences = {}
    
    def process_batch(self, new_tokens):
        outputs = []
        for seq_id, token in new_tokens:
            if seq_id not in self.active_sequences:
                self.active_sequences[seq_id] = []
            
            self.active_sequences[seq_id].append(token)
            
            # Process the sequence
            input_ids = torch.tensor(self.active_sequences[seq_id]).unsqueeze(0)
            with torch.no_grad():
                output = self.model(input_ids)
            
            next_token = output[0, -1].argmax().item()
            outputs.append((seq_id, next_token))
            
            # Remove completed sequences
            if next_token == self.model.config.eos_token_id:
                del self.active_sequences[seq_id]
        
        return outputs

# Usage example (assuming 'model' is a pre-trained transformer)
processor = ContinuousBatchProcessor(model)
new_tokens = [(1, 5), (2, 10), (1, 7), (3, 2)]  # (sequence_id, token)
results = processor.process_batch(new_tokens)
print(results)  # Example output: [(1, 15), (2, 8), (1, 3), (3, 12)]

PagedAttention in vLLM

PagedAttention is a memory management technique used in vLLM to efficiently handle attention computations for long sequences. It divides the attention context into fixed-size pages, allowing for better memory utilization and cache efficiency.

import torch
import math

class PagedAttention:
    def __init__(self, hidden_size, num_heads, page_size=16):
        self.hidden_size = hidden_size
        self.num_heads = num_heads
        self.head_dim = hidden_size // num_heads
        self.page_size = page_size
        self.pages = {}
    
    def compute_attention(self, query, key, value):
        scale = 1 / math.sqrt(self.head_dim)
        scores = torch.matmul(query, key.transpose(-2, -1)) * scale
        attn_weights = torch.softmax(scores, dim=-1)
        return torch.matmul(attn_weights, value)
    
    def forward(self, hidden_states, seq_ids):
        batch_size, seq_len, _ = hidden_states.shape
        q = hidden_states.view(batch_size, seq_len, self.num_heads, self.head_dim)
        
        outputs = []
        for i in range(batch_size):
            seq_id = seq_ids[i]
            if seq_id not in self.pages:
                self.pages[seq_id] = []
            
            # Add new page if needed
            if len(self.pages[seq_id]) * self.page_size < seq_len:
                new_page = hidden_states[i, -self.page_size:].clone()
                self.pages[seq_id].append(new_page)
            
            # Compute attention using paged context
            context = torch.cat(self.pages[seq_id], dim=0)
            k = context.view(-1, self.num_heads, self.head_dim)
            v = k.clone()
            
            attn_output = self.compute_attention(q[i], k, v)
            outputs.append(attn_output)
        
        return torch.stack(outputs)

# Usage example
paged_attention = PagedAttention(hidden_size=768, num_heads=12)
hidden_states = torch.randn(2, 32, 768)
seq_ids = [1, 2]
output = paged_attention.forward(hidden_states, seq_ids)
print(output.shape)  # Expected output: torch.Size([2, 32, 12, 64])

Optimizing Transformer Inference with vLLM

vLLM combines various optimization techniques to achieve efficient transformer inference. Let's implement a simplified version of a vLLM-optimized transformer layer.

import torch
import torch.nn as nn

class vLLMTransformerLayer(nn.Module):
    def __init__(self, hidden_size, num_heads, intermediate_size):
        super().__init__()
        self.attention = vLLMAttention(hidden_size, num_heads)
        self.layer_norm1 = nn.LayerNorm(hidden_size)
        self.layer_norm2 = nn.LayerNorm(hidden_size)
        self.mlp = nn.Sequential(
            nn.Linear(hidden_size, intermediate_size),
            nn.GELU(),
            nn.Linear(intermediate_size, hidden_size)
        )
    
    def forward(self, hidden_states, past_key_values=None, attention_mask=None):
        attn_output, past_kv = self.attention(hidden_states, past_key_values, attention_mask)
        hidden_states = self.layer_norm1(hidden_states + attn_output)
        
        mlp_output = self.mlp(hidden_states)
        hidden_states = self.layer_norm2(hidden_states + mlp_output)
        
        return hidden_states, past_kv

# Usage example
layer = vLLMTransformerLayer(hidden_size=768, num_heads=12, intermediate_size=3072)
input_states = torch.randn(2, 16, 768)
output, past_kv = layer(input_states)
print(output.shape)  # Expected output: torch.Size([2, 16, 768])
print(past_kv[0].shape)  # Expected output: torch.Size([2, 12, 16, 64])

Implementing Efficient Token Generation

Efficient token generation is crucial for language models. Let's implement a function that generates tokens using our vLLM-optimized transformer layer.

def generate_tokens(model, input_ids, max_length, temperature=1.0):
    generated = input_ids.clone()
    past_key_values = None
    
    for _ in range(max_length - input_ids.size(1)):
        with torch.no_grad():
            outputs, past_key_values = model(generated[:, -1:], past_key_values)
            logits = outputs[:, -1, :] / temperature
            next_token = torch.multinomial(torch.softmax(logits, dim=-1), num_samples=1)
            generated = torch.cat((generated, next_token), dim=1)
            
            if next_token.item() == model.config.eos_token_id:
                break
    
    return generated

# Usage example (assuming 'model' is a vLLM-optimized transformer)
input_ids = torch.tensor([[1, 2, 3]])  # Start of sequence
generated_ids = generate_tokens(model, input_ids, max_length=50)
print(generated_ids)

Real-life Example: Text Completion

Let's use our vLLM-optimized model for a practical text completion task.

def complete_text(model, tokenizer, prompt, max_length=50):
    input_ids = tokenizer.encode(prompt, return_tensors="pt")
    generated_ids = generate_tokens(model, input_ids, max_length)
    completed_text = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
    return completed_text

# Usage example (assuming 'model' and 'tokenizer' are properly initialized)
prompt = "The quick brown fox"
completed_text = complete_text(model, tokenizer, prompt)
print(f"Prompt: {prompt}")
print(f"Completed: {completed_text}")

# Example output:
# Prompt: The quick brown fox
# Completed: The quick brown fox jumps over the lazy dog. Its agile movements
# caught the attention of nearby birds, who watched curiously from their perches.

Real-life Example: Dialogue Generation

Let's use our vLLM-optimized model to generate a simple dialogue between two characters.

def generate_dialogue(model, tokenizer, context, num_turns=3):
    dialogue = context
    characters = ["Alice", "Bob"]
    
    for i in range(num_turns):
        for character in characters:
            prompt = f"{dialogue}\n{character}:"
            response = complete_text(model, tokenizer, prompt, max_length=30)
            dialogue += f"\n{character}: {response[len(prompt):].strip()}"
    
    return dialogue

# Usage example
context = "Alice and Bob are discussing their favorite books."
dialogue = generate_dialogue(model, tokenizer, context)
print(dialogue)

# Example output:
# Alice and Bob are discussing their favorite books.
# Alice: My favorite book is "Pride and Prejudice" by Jane Austen. What's yours?
# Bob: I'm a big fan of science fiction. My top pick is "Dune" by Frank Herbert.
# Alice: Oh, I've heard great things about "Dune"! What do you like most about it?
# Bob: The world-building is incredible. The political intrigue and ecology are fascinating.
# Alice: That sounds intriguing. Maybe I should give it a try. Do you have other sci-fi recommendations?
# Bob: Absolutely! You might enjoy "The Hitchhiker's Guide to the Galaxy" for a lighter, humorous take on sci-fi.

Comparing Standard Attention and Flash Attention

Let's compare the performance of standard attention and Flash Attention in terms of speed and memory usage.

import time

def compare_attention_methods(seq_len, num_heads, head_dim, num_runs=5):
    batch_size = 1
    hidden_size = num_heads * head_dim
    
    q = torch.randn(batch_size, seq_len, num_heads, head_dim)
    k = torch.randn(batch_size, seq_len, num_heads, head_dim)
    v = torch.randn(batch_size, seq_len, num_heads, head_dim)
    
    # Standard Attention
    def standard_attention():
        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(head_dim)
        attn = torch.softmax(scores, dim=-1)
        return torch.matmul(attn, v)
    
    # Flash Attention
    def flash_attn():
        return flash_attention(q, k, v)
    
    # Measure execution time
    standard_times = []
    flash_times = []
    
    for _ in range(num_runs):
        start = time.time()
        standard_attention()
        standard_times.append(time.time() - start)
        
        start = time.time()
        flash_attn()
        flash_times.append(time.time() - start)
    
    print(f"Sequence Length: {seq_len}")
    print(f"Standard Attention Avg Time: {sum(standard_times) / num_runs:.4f} s")
    print(f"Flash Attention Avg Time: {sum(flash_times) / num_runs:.4f} s")
    print(f"Speedup: {sum(standard_times) / sum(flash_times):.2f}x")

# Run comparison
compare_attention_methods(seq_len=1024, num_heads=8, head_dim=64)
compare_attention_methods(seq_len=4096, num_heads=8, head_dim=64)

Additional Resources

For more in-depth information on the topics covered in this presentation, consider exploring the following resources:

  1. "Attention Is All You Need" (Vaswani et al., 2017) - The original transformer paper ArXiv: https://arxiv.org/abs/1706.03762
  2. "FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness" (Dao et al., 2022) ArXiv: https://arxiv.org/abs/2205.14135
  3. "Efficient Memory Management for Large Language Model Serving with PagedAttention" (Kwon et al., 2023) ArXiv: https://arxiv.org/abs/2309.06180
  4. "vLLM: Easy, Fast, and Cheap LLM Serving with PagedAttention" (Gu et al., 2023) ArXiv: https://arxiv.org/abs/2309.06180

These papers provide detailed explanations of the concepts and techniques discussed in this presentation, offering a deeper understanding of transformer optimizations and efficient language model serving.

More in BERT and language models

  1. Quantization Effects on Multilingual Language Models12 min

    Quantization is a technique used to reduce the computational and memory requirements of large language models (LLMs) while maintaining their performance. In...

  2. Mastering Temperature and Top_p in Python Language Models10 min

    Temperature and top\p are crucial parameters in language models that control the randomness and diversity of generated text. These settings affect how the...

  3. Reinforcement Learning from Human Feedback! Aligning Language Models Using Python:8 min

    Reinforcement Learning from Human Feedback (RLHF) is a technique used to align language models with human preferences. It involves training two key...

  4. Mastering BERT Transforming NLP with Bidirectional Understanding10 min

    BERT (Bidirectional Encoder Representations from Transformers) revolutionized natural language processing by introducing true bidirectional understanding...

  5. Scaling Laws for Neural Language Models in Python7 min

    Neural language models have shown remarkable improvements in recent years, largely due to increases in model size and training data. This slide introduces...

  6. Language Modeling in NLP with Python11 min

    Language modeling is a fundamental task in Natural Language Processing (NLP) that involves predicting the probability of a sequence of words. It forms the...