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

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

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():
-    """主函数：演示Transformer模型的使用"""
+def demo_training_forward():
+    """演示训练时的前向传播（无KV cache）"""
+    print("=" * 60)
+    print("训练前向传播演示（无KV cache）")
+    print("=" * 60)
     
     # 模型超参数配置
-    vocab_size = 1000   # 词汇表大小
-    d_model = 512       # 模型维度（嵌入维度）
-    num_heads = 8       # 多头注意力的头数
-    num_layers = 6      # Transformer编码器层数
-    d_ff = 2048         # 前馈网络隐藏层维度
-    max_seq_len = 100   # 最大序列长度
+    vocab_size = 1000
+    d_model = 512
+    num_heads = 8
+    num_layers = 6
+    d_ff = 2048
+    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  # 批次大小
-    seq_len = 10    # 序列长度
-    x = np.random.randint(0, vocab_size, (batch_size, seq_len))
+    # 创建随机输入
+    batch_size = 2
+    seq_len = 10
+    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()
+    print(f"输入示例: {x[0]}")
     
-    # 执行前向传播
+    # 训练模式前向传播（不使用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()
+
+
+def demo_prefill_decode():
+    """演示推理时的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()
     
-    # 打印模型参数统计
-    total_params = model.count_parameters()
-    print(f"总参数量: {total_params:,}")
+    # Prefill阶段
+    print("--- Prefill阶段 ---")
+    print("处理完整prompt，初始化KV cache")
+    logits = engine.prefill(prompt)
+    print(f"Prefill输出logits形状: {logits.shape}")
+    print(f"当前位置: {engine.current_position}")
+    
+    # 采样第一个token
+    first_token = engine._sample(logits, temperature=0.8)
+    print(f"采样的第一个token: {first_token[0, 0]}")
     print()
     
-    # 演示掩码创建
-    print("=== 注意力掩码示例 ===")
+    # Decode阶段
+    print("--- Decode阶段 ---")
+    print("逐个生成新token（使用KV cache）")
     
-    # 填充掩码（用于忽略padding位置）
-    padding_mask = create_padding_mask(x)
-    print(f"填充掩码形状: {padding_mask.shape}")
+    generated_tokens = [first_token[0, 0]]
+    current_token = first_token
     
-    # 前瞻掩码（用于解码器，防止看到未来信息）
-    look_ahead_mask = create_look_ahead_mask(seq_len)
-    print(f"前瞻掩码形状: {look_ahead_mask.shape}")
-    print(f"前瞻掩码示例:\n{look_ahead_mask[:5, :5]}")
+    # 生成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()