mirror of
https://github.com/MLSysBook/TinyTorch.git
synced 2026-07-26 03:41:17 -05:00
Vijay Janapa Reddi
6d0afe4949
Added comprehensive documentation clarifying that KV caching is designed ONLY for inference (generation), not training. Key Clarifications: - Cache operations use .data (no gradient tracking) - This is correct and intentional for maximum speed - During generation: no gradients computed (model.eval() mode) - During training: cache not used (standard forward pass) - DO NOT use caching during training Why This is Safe: 1. Training: Uses standard forward pass (full gradient flow) 2. Generation: No backward pass (no gradients needed) 3. Cache is inference optimization, not training component 4. .data usage is correct for generation-only use case Documentation Updates: - Added prominent warning in class docstring - Updated update() method docs - Updated get() method docs - Added inline comments explaining .data usage This addresses gradient flow concerns by making it crystal clear that caching is never used when gradients are needed.
260 lines
10 KiB
Python
Generated
260 lines
10 KiB
Python
Generated
# AUTOGENERATED! DO NOT EDIT! File to edit: ../../modules/source/14_kvcaching/kvcaching_dev.py (unless otherwise specified).
|
||
|
||
__all__ = ['KVCache', 'enable_kv_cache']
|
||
|
||
# Cell
|
||
import numpy as np
|
||
import time
|
||
from typing import Tuple, Optional, Dict, List
|
||
|
||
# Import TinyTorch components from previous modules
|
||
from tinytorch.core.tensor import Tensor
|
||
|
||
# Cell
|
||
class KVCache:
|
||
"""
|
||
Efficient key-value cache for autoregressive generation.
|
||
|
||
Stores K,V matrices for each transformer layer to avoid recomputation
|
||
during sequential token generation. This is THE critical optimization
|
||
that makes production language model serving economically viable.
|
||
|
||
⚠️ IMPORTANT: INFERENCE-ONLY (No Gradient Tracking)
|
||
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
|
||
KV caching is designed ONLY for inference (generation), NOT training.
|
||
- During generation: No gradients computed (model.eval() mode)
|
||
- Cache operations use .data (no gradient tracking)
|
||
- This is correct and intentional for maximum speed
|
||
- DO NOT use caching during training (use standard forward pass)
|
||
|
||
Architecture:
|
||
- Pre-allocates cache tensors with maximum sequence length
|
||
- Tracks current sequence position for efficient O(1) updates
|
||
- Provides update() method to append new K,V pairs without copying
|
||
- Provides get() method to retrieve cached values for attention
|
||
- Handles multiple layers and attention heads properly
|
||
|
||
Memory Layout:
|
||
```
|
||
Layer 0: [Key_cache, Value_cache] # Shape: (batch, num_heads, max_seq, head_dim)
|
||
Layer 1: [Key_cache, Value_cache]
|
||
...
|
||
Layer N: [Key_cache, Value_cache]
|
||
```
|
||
|
||
Performance:
|
||
- Update: O(1) - just index assignment
|
||
- Get: O(1) - just slicing (no data copy)
|
||
- Memory: O(num_layers × batch × heads × max_seq × head_dim)
|
||
"""
|
||
|
||
def __init__(self, batch_size: int, max_seq_len: int, num_layers: int,
|
||
num_heads: int, head_dim: int):
|
||
"""
|
||
Initialize KV cache for efficient generation.
|
||
|
||
Args:
|
||
batch_size: Number of sequences to generate simultaneously
|
||
max_seq_len: Maximum sequence length to support
|
||
num_layers: Number of transformer layers
|
||
num_heads: Number of attention heads per layer
|
||
head_dim: Dimension of each attention head
|
||
"""
|
||
self.batch_size = batch_size
|
||
self.max_seq_len = max_seq_len
|
||
self.num_layers = num_layers
|
||
self.num_heads = num_heads
|
||
self.head_dim = head_dim
|
||
|
||
# Current sequence position (how many tokens are cached)
|
||
self.seq_pos = 0
|
||
|
||
# Cache storage: list of (key_cache, value_cache) tuples per layer
|
||
self.caches = []
|
||
|
||
for layer_idx in range(num_layers):
|
||
# Pre-allocate cache tensors with maximum size
|
||
# Shape: (batch_size, num_heads, max_seq_len, head_dim)
|
||
key_cache = Tensor(np.zeros((batch_size, num_heads, max_seq_len, head_dim)))
|
||
value_cache = Tensor(np.zeros((batch_size, num_heads, max_seq_len, head_dim)))
|
||
|
||
self.caches.append((key_cache, value_cache))
|
||
|
||
def update(self, layer_idx: int, key: Tensor, value: Tensor) -> None:
|
||
"""
|
||
Update cache with new key-value pairs for given layer.
|
||
|
||
This is the core caching operation - efficiently append new K,V
|
||
to the cache without recomputation. This operation is O(1) because
|
||
it's just an indexed assignment.
|
||
|
||
IMPORTANT: KV caching is designed for INFERENCE (generation) only,
|
||
not training. During generation, gradients are not computed. If you
|
||
need gradients, don't use caching (use standard forward pass instead).
|
||
|
||
Args:
|
||
layer_idx: Which transformer layer (0 to num_layers-1)
|
||
key: New key tensor, shape (batch_size, num_heads, 1, head_dim)
|
||
value: New value tensor, shape (batch_size, num_heads, 1, head_dim)
|
||
|
||
Raises:
|
||
ValueError: If layer_idx is out of range or sequence is full
|
||
"""
|
||
if layer_idx >= self.num_layers:
|
||
raise ValueError(f"Layer index {layer_idx} >= num_layers {self.num_layers}")
|
||
|
||
if self.seq_pos >= self.max_seq_len:
|
||
raise ValueError(f"Sequence position {self.seq_pos} >= max_seq_len {self.max_seq_len}")
|
||
|
||
# Get cache for this layer
|
||
key_cache, value_cache = self.caches[layer_idx]
|
||
|
||
# Update cache at current position (efficient O(1) write)
|
||
# Note: We use .data here because caching is inference-only (no gradients needed)
|
||
# This avoids gradient tracking overhead during generation
|
||
key_cache.data[:, :, self.seq_pos:self.seq_pos+1, :] = key.data
|
||
value_cache.data[:, :, self.seq_pos:self.seq_pos+1, :] = value.data
|
||
|
||
# Note: seq_pos is advanced externally via advance() after all layers process
|
||
|
||
def get(self, layer_idx: int) -> Tuple[Tensor, Tensor]:
|
||
"""
|
||
Retrieve cached key-value pairs for attention computation.
|
||
|
||
Returns only the valid portion of the cache (up to current seq_pos).
|
||
This is O(1) because we're just slicing NumPy arrays (view, not copy).
|
||
|
||
IMPORTANT: Returns Tensors without gradient tracking since caching
|
||
is inference-only. The returned tensors can be used in attention
|
||
computation but won't propagate gradients backward.
|
||
|
||
Args:
|
||
layer_idx: Which transformer layer to get cache for
|
||
|
||
Returns:
|
||
(cached_keys, cached_values): Tensors shaped for attention
|
||
Keys: (batch_size, num_heads, seq_pos, head_dim)
|
||
Values: (batch_size, num_heads, seq_pos, head_dim)
|
||
|
||
Raises:
|
||
ValueError: If layer_idx is out of range
|
||
"""
|
||
if layer_idx >= self.num_layers:
|
||
raise ValueError(f"Layer index {layer_idx} >= num_layers {self.num_layers}")
|
||
|
||
# Get cache for this layer
|
||
key_cache, value_cache = self.caches[layer_idx]
|
||
|
||
# Return only the valid portion (up to current sequence position)
|
||
# seq_pos tracks where to write next, so we have seq_pos valid tokens
|
||
valid_len = self.seq_pos
|
||
|
||
# Note: Creating new Tensors from .data (no gradient tracking)
|
||
# This is correct for inference-only caching
|
||
cached_keys = Tensor(key_cache.data[:, :, :valid_len, :])
|
||
cached_values = Tensor(value_cache.data[:, :, :valid_len, :])
|
||
|
||
return cached_keys, cached_values
|
||
|
||
def advance(self) -> None:
|
||
"""
|
||
Advance sequence position after processing current token.
|
||
|
||
Call this after all layers have processed the current token and
|
||
updated their caches. This moves the write pointer forward.
|
||
"""
|
||
self.seq_pos += 1
|
||
|
||
def reset(self) -> None:
|
||
"""
|
||
Reset cache for new generation sequence.
|
||
|
||
Call this when starting a new generation (new prompt).
|
||
Resets the sequence position counter and optionally zeros cache data.
|
||
"""
|
||
self.seq_pos = 0
|
||
|
||
# Zero out caches for clean state (helps with debugging)
|
||
for layer_idx in range(self.num_layers):
|
||
key_cache, value_cache = self.caches[layer_idx]
|
||
key_cache.data.fill(0.0)
|
||
value_cache.data.fill(0.0)
|
||
|
||
def get_memory_usage(self) -> Dict[str, float]:
|
||
"""
|
||
Calculate memory usage of the cache system.
|
||
|
||
Returns:
|
||
Dictionary with memory statistics in MB
|
||
"""
|
||
# Calculate size of one cache tensor
|
||
cache_size = self.batch_size * self.num_heads * self.max_seq_len * self.head_dim
|
||
bytes_per_float = 4 # float32
|
||
|
||
# Each layer has key_cache + value_cache
|
||
total_cache_tensors = self.num_layers * 2
|
||
total_elements = cache_size * total_cache_tensors
|
||
total_bytes = total_elements * bytes_per_float
|
||
total_mb = total_bytes / (1024 * 1024)
|
||
|
||
return {
|
||
'total_mb': total_mb,
|
||
'per_layer_mb': total_mb / self.num_layers,
|
||
'cache_tensors': total_cache_tensors,
|
||
'total_elements': total_elements
|
||
}
|
||
|
||
# Cell
|
||
def enable_kv_cache(batch_size: int, max_seq_len: int, num_layers: int,
|
||
num_heads: int, head_dim: int) -> KVCache:
|
||
"""
|
||
Create and return a KVCache instance for model generation.
|
||
|
||
This function creates a properly sized cache for the model architecture.
|
||
Call this before starting generation, then pass the cache to your
|
||
generation loop.
|
||
|
||
Args:
|
||
batch_size: Number of sequences to generate simultaneously
|
||
max_seq_len: Maximum sequence length to support
|
||
num_layers: Number of transformer layers in model
|
||
num_heads: Number of attention heads per layer
|
||
head_dim: Dimension per attention head (usually embed_dim // num_heads)
|
||
|
||
Returns:
|
||
KVCache instance ready for use
|
||
|
||
Example:
|
||
```python
|
||
# Enable caching for generation
|
||
cache = enable_kv_cache(
|
||
batch_size=1,
|
||
max_seq_len=100,
|
||
num_layers=4,
|
||
num_heads=4,
|
||
head_dim=32
|
||
)
|
||
|
||
# Use in generation loop (pseudocode)
|
||
for step in range(max_new_tokens):
|
||
# Only process new token with cache
|
||
logits = model.forward_cached(new_token, cache)
|
||
next_token = sample(logits)
|
||
```
|
||
"""
|
||
cache = KVCache(batch_size, max_seq_len, num_layers, num_heads, head_dim)
|
||
|
||
print(f"⚡ KV Cache enabled:")
|
||
print(f" Batch size: {batch_size}")
|
||
print(f" Max sequence: {max_seq_len}")
|
||
print(f" Layers: {num_layers}")
|
||
print(f" Heads: {num_heads}")
|
||
print(f" Head dim: {head_dim}")
|
||
|
||
mem_info = cache.get_memory_usage()
|
||
print(f" Memory: {mem_info['total_mb']:.2f} MB")
|
||
print()
|
||
|
||
return cache
|
||
|