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

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

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,26 +143,34 @@ 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
         """
-        # 自注意力子层
-        attn_output = self.attention.forward(x, x, x, mask)
+        # 自注意力子层（Q=K=V，自注意力）
+        attn_output = self.attention.forward(x, x, x, mask, use_cache=use_cache)
         x = self.norm1.forward(x + attn_output)
         
         # 前馈网络子层
         ff_output = self.feed_forward.forward(x)
         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
         # 位置编码
@@ -157,10 +191,55 @@ class SimpleTransformer:
         pe[:, 0::2] = np.sin(position * div_term)  # 偶数维度用正弦
         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   # 词汇表大小