添加KV Cache支持：实现prefill/decode阶段分离

- MultiHeadAttention类添加KV cache机制
- TransformerBlock支持use_cache参数
- SimpleTransformer新增prefill和decode方法
- 添加InferenceEngine推理引擎类
- 更新example.py演示推理过程
- 更新README.md文档

README.md

@@ -1,11 +1,11 @@
 # Simple Transformer
 
-A minimal Transformer implementation using NumPy.
+A minimal Transformer implementation using NumPy with **KV Cache** support for efficient inference.
 
 ## Files
 
-- `transformer.py` - Core Transformer model implementation
+- `transformer.py` - Core Transformer model implementation with KV Cache
-- `example.py` - Usage example
+- `example.py` - Usage examples (training & inference)
 
 ## Usage
 
@@ -21,6 +21,21 @@ The implementation includes:
 - **Feed-Forward Network**: Two-layer fully connected network with ReLU activation
 - **Layer Normalization**: Applied after each sub-layer
 - **Positional Encoding**: Sinusoidal position embeddings
+- **KV Cache**: Efficient inference with prefill/decode separation
+
+## Inference: Prefill vs Decode
+
+### Prefill Stage
+- Process the complete input prompt in parallel
+- Initialize KV cache for all layers
+- Generate the first token
+- Computation: O(n²)
+
+### Decode Stage
+- Process one token at a time
+- Reuse cached Key and Value tensors
+- Only compute new Query
+- Computation: O(n) per step
 
 ## Model Parameters
 
@@ -33,3 +48,8 @@ Default configuration:
 - Max sequence length: 100
 
 Total parameters: ~19.4M
+
+## Classes
+
+- `SimpleTransformer` - Transformer model with prefill/decode methods
+- `InferenceEngine` - Manages inference process with KV cache

example.py

@@ -1,66 +1,160 @@
 import numpy as np
-from transformer import SimpleTransformer, create_padding_mask, create_look_ahead_mask
+from transformer import SimpleTransformer, InferenceEngine, create_padding_mask, create_look_ahead_mask
 
 
-def main():
+def demo_training_forward():
-    """主函数：演示Transformer模型的使用"""
+    """演示训练时的前向传播（无KV cache）"""
+    print("=" * 60)
+    print("训练前向传播演示（无KV cache）")
+    print("=" * 60)
     
     # 模型超参数配置
-    vocab_size = 1000   # 词汇表大小
+    vocab_size = 1000
-    d_model = 512       # 模型维度（嵌入维度）
+    d_model = 512
-    num_heads = 8       # 多头注意力的头数
+    num_heads = 8
-    num_layers = 6      # Transformer编码器层数
+    num_layers = 6
-    d_ff = 2048         # 前馈网络隐藏层维度
+    d_ff = 2048
-    max_seq_len = 100   # 最大序列长度
+    max_seq_len = 100
     
-    # 创建Transformer模型实例
+    # 创建模型
     model = SimpleTransformer(vocab_size, d_model, num_heads, num_layers, d_ff, max_seq_len)
     
-    # 创建随机输入数据（模拟token ids）
+    # 创建随机输入
-    batch_size = 2  # 批次大小
+    batch_size = 2
-    seq_len = 10    # 序列长度
+    seq_len = 10
-    x = np.random.randint(0, vocab_size, (batch_size, seq_len))
+    x = np.random.randint(1, vocab_size, (batch_size, seq_len))
     
-    # 打印模型配置信息
-    print("=== 简单Transformer示例 ===")
-    print(f"词汇表大小: {vocab_size}")
-    print(f"模型维度: {d_model}")
-    print(f"注意力头数: {num_heads}")
-    print(f"编码器层数: {num_layers}")
-    print(f"前馈网络维度: {d_ff}")
-    print(f"最大序列长度: {max_seq_len}")
-    print()
-    
-    # 打印输入信息
     print(f"输入形状: {x.shape}")
-    print(f"输入样本: {x[0]}")
+    print(f"输入示例: {x[0]}")
-    print()
     
-    # 执行前向传播
+    # 训练模式前向传播（不使用KV cache）
     output = model.forward(x)
     
-    # 打印输出信息
     print(f"输出形状: {output.shape}")
-    print(f"输出样本（前5个值）: {output[0, 0, :5]}")
+    print(f"输出示例（前5个值）: {output[0, 0, :5]}")
     print()
 
-    # 打印模型参数统计
+
-    total_params = model.count_parameters()
+def demo_prefill_decode():
-    print(f"总参数量: {total_params:,}")
+    """演示推理时的prefill和decode阶段"""
+    print("=" * 60)
+    print("推理演示（Prefill + Decode）")
+    print("=" * 60)
+    
+    # 模型超参数配置
+    vocab_size = 1000
+    d_model = 128  # 使用较小的模型以便快速演示
+    num_heads = 4
+    num_layers = 2
+    d_ff = 512
+    max_seq_len = 100
+    
+    # 创建模型和推理引擎
+    model = SimpleTransformer(vocab_size, d_model, num_heads, num_layers, d_ff, max_seq_len)
+    engine = InferenceEngine(model)
+    
+    # 模拟输入序列（prompt）
+    batch_size = 1
+    prompt_len = 5
+    prompt = np.random.randint(1, vocab_size, (batch_size, prompt_len))
+    
+    print(f"输入prompt: {prompt[0]}")
+    print(f"Prompt长度: {prompt_len}")
     print()
     
-    # 演示掩码创建
+    # Prefill阶段
-    print("=== 注意力掩码示例 ===")
+    print("--- Prefill阶段 ---")
+    print("处理完整prompt，初始化KV cache")
+    logits = engine.prefill(prompt)
+    print(f"Prefill输出logits形状: {logits.shape}")
+    print(f"当前位置: {engine.current_position}")
     
-    # 填充掩码（用于忽略padding位置）
+    # 采样第一个token
-    padding_mask = create_padding_mask(x)
+    first_token = engine._sample(logits, temperature=0.8)
-    print(f"填充掩码形状: {padding_mask.shape}")
+    print(f"采样的第一个token: {first_token[0, 0]}")
+    print()
     
-    # 前瞻掩码（用于解码器，防止看到未来信息）
+    # Decode阶段
-    look_ahead_mask = create_look_ahead_mask(seq_len)
+    print("--- Decode阶段 ---")
-    print(f"前瞻掩码形状: {look_ahead_mask.shape}")
+    print("逐个生成新token（使用KV cache）")
-    print(f"前瞻掩码示例:\n{look_ahead_mask[:5, :5]}")
+    
+    generated_tokens = [first_token[0, 0]]
+    current_token = first_token
+    
+    # 生成5个token作为演示
+    for i in range(5):
+        logits = engine.decode_step(current_token)
+        current_token = engine._sample(logits, temperature=0.8)
+        generated_tokens.append(current_token[0, 0])
+        print(f"Step {i+1}: 生成token {current_token[0, 0]}, 位置 {engine.current_position}")
+    
+    print()
+    print(f"完整生成序列: {generated_tokens}")
+    print()
+
+
+def demo_auto_regressive():
+    """演示自回归生成"""
+    print("=" * 60)
+    print("自回归生成演示")
+    print("=" * 60)
+    
+    # 模型超参数配置
+    vocab_size = 1000
+    d_model = 128
+    num_heads = 4
+    num_layers = 2
+    d_ff = 512
+    max_seq_len = 100
+    
+    # 创建模型和推理引擎
+    model = SimpleTransformer(vocab_size, d_model, num_heads, num_layers, d_ff, max_seq_len)
+    engine = InferenceEngine(model)
+    
+    # 模拟输入序列
+    batch_size = 1
+    prompt_len = 3
+    prompt = np.random.randint(1, vocab_size, (batch_size, prompt_len))
+    
+    print(f"输入prompt: {prompt[0]}")
+    
+    # 使用generate方法进行自回归生成
+    output = engine.generate(prompt, max_new_tokens=10, temperature=0.8)
+    
+    print(f"生成的完整序列: {output[0]}")
+    print(f"序列长度: {len(output[0])} (原始{prompt_len} + 生成{len(output[0]) - prompt_len})")
+    print()
+
+
+def demo_kv_cache_comparison():
+    """对比有无KV cache的计算差异"""
+    print("=" * 60)
+    print("KV Cache计算对比")
+    print("=" * 60)
+    
+    print("无KV cache（训练模式）:")
+    print("  - 每次forward处理完整序列")
+    print("  - 计算量: O(n^2) 每次")
+    print("  - 适用于训练")
+    print()
+    
+    print("有KV cache（推理模式）:")
+    print("  - Prefill: 处理完整prompt，缓存KV")
+    print("  - Decode: 只处理新token，使用缓存")
+    print("  - 计算量: Prefill O(n^2), Decode O(n) 每步")
+    print("  - 适用于推理")
+    print()
 
 
 if __name__ == "__main__":
-    main()
+    # 演示训练前向传播
+    demo_training_forward()
+    
+    # 演示prefill和decode
+    demo_prefill_decode()
+    
+    # 演示自回归生成
+    demo_auto_regressive()
+    
+    # 对比说明
+    demo_kv_cache_comparison()

transformer.py

@@ -44,7 +44,7 @@ class LayerNorm:
 
 
 class MultiHeadAttention:
-    """多头注意力机制"""
+    """多头注意力机制（支持KV cache）"""
     def __init__(self, d_model, num_heads):
         self.d_model = d_model  # 模型维度
         self.num_heads = num_heads  # 注意力头数
@@ -56,13 +56,23 @@ class MultiHeadAttention:
         self.W_v = Linear(d_model, d_model)  # Value变换
         self.W_o = Linear(d_model, d_model)  # 输出变换
         
-    def forward(self, Q, K, V, mask=None):
+        # KV cache：用于推理时缓存历史Key和Value
+        self.k_cache = None  # 缓存的Key [batch_size, num_heads, seq_len, d_k]
+        self.v_cache = None  # 缓存的Value [batch_size, num_heads, seq_len, d_k]
+        
+    def clear_cache(self):
+        """清除KV cache"""
+        self.k_cache = None
+        self.v_cache = None
+        
+    def forward(self, Q, K, V, mask=None, use_cache=False):
         """
         前向传播
         Q: Query张量 [batch_size, seq_len, d_model]
         K: Key张量 [batch_size, seq_len, d_model]
         V: Value张量 [batch_size, seq_len, d_model]
         mask: 注意力掩码
+        use_cache: 是否使用KV cache（推理时设为True）
         """
         batch_size = Q.shape[0]
         seq_len = Q.shape[1]
@@ -72,7 +82,23 @@ class MultiHeadAttention:
         K = self.W_k.forward(K).reshape(batch_size, seq_len, self.num_heads, self.d_k).transpose(0, 2, 1, 3)
         V = self.W_v.forward(V).reshape(batch_size, seq_len, self.num_heads, self.d_k).transpose(0, 2, 1, 3)
         
+        # KV cache逻辑
+        if use_cache:
+            if self.k_cache is None:
+                # Prefill阶段：首次计算，缓存完整的K和V
+                self.k_cache = K
+                self.v_cache = V
+            else:
+                # Decode阶段：拼接历史缓存和新的K、V
+                self.k_cache = np.concatenate([self.k_cache, K], axis=2)
+                self.v_cache = np.concatenate([self.v_cache, V], axis=2)
+                # 使用完整的K和V进行注意力计算
+                K = self.k_cache
+                V = self.v_cache
+        
         # 计算缩放点积注意力分数
+        # Q: [batch, heads, q_len, d_k]
+        # K: [batch, heads, kv_len, d_k]（decode时kv_len > q_len）
         attn_scores = Q @ K.transpose(0, 1, 3, 2) / np.sqrt(self.d_k)
         
         # 应用掩码（如果提供）
@@ -117,14 +143,17 @@ class TransformerBlock:
         self.norm1 = LayerNorm(d_model)
         self.norm2 = LayerNorm(d_model)
         
-    def forward(self, x, mask=None):
+    def forward(self, x, mask=None, use_cache=False):
         """
         前向传播（残差连接 + 层归一化）
         1. 自注意力 -> 残差连接 -> 层归一化
         2. 前馈网络 -> 残差连接 -> 层归一化
+        x: 输入张量 [batch_size, seq_len, d_model]
+        mask: 注意力掩码
+        use_cache: 是否使用KV cache
         """
-        # 自注意力子层
+        # 自注意力子层（Q=K=V，自注意力）
-        attn_output = self.attention.forward(x, x, x, mask)
+        attn_output = self.attention.forward(x, x, x, mask, use_cache=use_cache)
         x = self.norm1.forward(x + attn_output)
         
         # 前馈网络子层
@@ -132,11 +161,16 @@ class TransformerBlock:
         x = self.norm2.forward(x + ff_output)
         return x
     
+    def clear_cache(self):
+        """清除该层的KV cache"""
+        self.attention.clear_cache()
+
 
 class SimpleTransformer:
-    """简化的Transformer编码器模型"""
+    """简化的Transformer编码器模型（支持prefill/decode分离）"""
     def __init__(self, vocab_size, d_model, num_heads, num_layers, d_ff, max_seq_len):
         self.d_model = d_model
+        self.vocab_size = vocab_size
         # 词嵌入层
         self.embedding = np.random.randn(vocab_size, d_model) * 0.01
         # 位置编码
@@ -158,9 +192,54 @@ class SimpleTransformer:
         pe[:, 1::2] = np.cos(position * div_term)  # 奇数维度用余弦
         return pe
     
+    def clear_cache(self):
+        """清除所有层的KV cache"""
+        for block in self.transformer_blocks:
+            block.clear_cache()
+        
+    def prefill(self, x, mask=None):
+        """
+        Prefill阶段：处理完整的输入序列
+        x: 输入序列 [batch_size, seq_len]
+        mask: 注意力掩码
+        返回: 输出张量 [batch_size, seq_len, d_model]
+        """
+        seq_len = x.shape[1]
+        
+        # 词嵌入 + 缩放
+        x = self.embedding[x] * np.sqrt(self.d_model)
+        
+        # 加上位置编码
+        x = x + self.positional_encoding[:seq_len, :]
+        
+        # 通过所有Transformer块，使用KV cache
+        for block in self.transformer_blocks:
+            x = block.forward(x, mask, use_cache=True)
+            
+        return x
+    
+    def decode(self, x, position):
+        """
+        Decode阶段：处理单个新token
+        x: 新token [batch_size, 1]
+        position: 当前token在序列中的位置
+        返回: 输出张量 [batch_size, 1, d_model]
+        """
+        # 词嵌入 + 缩放
+        x = self.embedding[x] * np.sqrt(self.d_model)
+        
+        # 加上位置编码（使用当前位置）
+        x = x + self.positional_encoding[position:position+1, :]
+        
+        # 通过所有Transformer块，使用KV cache
+        for block in self.transformer_blocks:
+            x = block.forward(x, use_cache=True)
+            
+        return x
+    
     def forward(self, x, mask=None):
         """
-        前向传播
+        前向传播（不使用KV cache，用于训练）
         x: 输入序列 [batch_size, seq_len]
         mask: 注意力掩码
         """
@@ -214,6 +293,127 @@ def create_look_ahead_mask(size):
     return (~mask).astype(float)
 
 
+class InferenceEngine:
+    """推理引擎：管理prefill和decode阶段"""
+    def __init__(self, model):
+        """
+        初始化推理引擎
+        model: SimpleTransformer模型实例
+        """
+        self.model = model
+        self.generated_tokens = []  # 已生成的token列表
+        self.current_position = 0   # 当前位置
+        
+    def reset(self):
+        """重置推理状态"""
+        self.model.clear_cache()
+        self.generated_tokens = []
+        self.current_position = 0
+        
+    def prefill(self, input_ids):
+        """
+        Prefill阶段：处理完整的输入序列
+        input_ids: 输入token序列 [batch_size, seq_len]
+        返回: 下一个token的logits [batch_size, vocab_size]
+        """
+        self.reset()
+        batch_size = input_ids.shape[0]
+        seq_len = input_ids.shape[1]
+        
+        # 创建前瞻掩码（防止看到未来信息）
+        mask = create_look_ahead_mask(seq_len)
+        
+        # 执行prefill前向传播（使用KV cache）
+        output = self.model.prefill(input_ids, mask)
+        
+        # 获取最后一个位置的输出（用于预测下一个token）
+        last_output = output[:, -1, :]  # [batch_size, d_model]
+        
+        # 简单的logits计算：使用嵌入矩阵的转置作为输出投影
+        # logits = last_output @ embedding.T
+        logits = last_output @ self.model.embedding.T  # [batch_size, vocab_size]
+        
+        # 更新状态
+        self.current_position = seq_len
+        
+        return logits
+    
+    def decode_step(self, input_token):
+        """
+        Decode阶段：处理单个新token
+        input_token: 新token [batch_size, 1]
+        返回: 下一个token的logits [batch_size, vocab_size]
+        """
+        # 执行decode前向传播（使用KV cache）
+        output = self.model.decode(input_token, self.current_position)
+        
+        # 获取输出（单token，所以直接取[:, 0, :]）
+        last_output = output[:, 0, :]  # [batch_size, d_model]
+        
+        # 计算logits
+        logits = last_output @ self.model.embedding.T  # [batch_size, vocab_size]
+        
+        # 更新位置
+        self.current_position += 1
+        
+        return logits
+    
+    def generate(self, input_ids, max_new_tokens=50, temperature=1.0):
+        """
+        自回归生成文本
+        input_ids: 输入token序列 [batch_size, seq_len]
+        max_new_tokens: 最大生成token数
+        temperature: 温度参数（控制随机性）
+        返回: 完整的生成序列 [batch_size, seq_len + max_new_tokens]
+        """
+        batch_size = input_ids.shape[0]
+        
+        # 保存原始输入
+        generated = input_ids.copy()
+        
+        # Prefill阶段
+        logits = self.prefill(input_ids)
+        
+        # 采样第一个生成的token
+        next_token = self._sample(logits, temperature)
+        generated = np.concatenate([generated, next_token], axis=1)
+        
+        # Decode阶段：逐个生成token
+        for i in range(max_new_tokens - 1):
+            logits = self.decode_step(next_token)
+            next_token = self._sample(logits, temperature)
+            generated = np.concatenate([generated, next_token], axis=1)
+            
+            # 检查是否生成了结束符（这里用0作为结束符）
+            if np.all(next_token == 0):
+                break
+                
+        return generated
+    
+    def _sample(self, logits, temperature=1.0):
+        """
+        从logits中采样
+        logits: [batch_size, vocab_size]
+        temperature: 温度参数
+        返回: 采样的token [batch_size, 1]
+        """
+        # 应用温度缩放
+        logits = logits / temperature
+        
+        # 计算概率分布
+        probs = softmax(logits, axis=-1)
+        
+        # 按概率采样
+        batch_size = logits.shape[0]
+        next_tokens = np.zeros((batch_size, 1), dtype=int)
+        
+        for i in range(batch_size):
+            # 使用numpy的random.choice采样
+            next_tokens[i, 0] = np.random.choice(logits.shape[1], p=probs[i])
+            
+        return next_tokens
+
+
 if __name__ == "__main__":
     # 模型超参数
     vocab_size = 1000   # 词汇表大小