拆分transformer.py为独立模块

按Transformer架构组件拆分：
- utils.py: 激活函数和掩码工具
- layers.py: 基础层（Linear, LayerNorm）
- attention.py: 多头注意力机制
- feedforward.py: 前馈网络
- block.py: Transformer编码器块
- model.py: SimpleTransformer模型
- inference.py: 推理引擎
- 删除旧的transformer.py文件

README.md

@@ -4,7 +4,15 @@ A minimal Transformer implementation using NumPy with **KV Cache** support for e
 
 ## Files
 
-- `transformer.py` - Core Transformer model implementation with KV Cache
+- `transformer/` - Transformer model package (按架构组件拆分)
+  - `__init__.py` - 包初始化，导出所有组件
+  - `utils.py` - 激活函数和掩码工具
+  - `layers.py` - 基础层（Linear, LayerNorm）
+  - `attention.py` - 多头注意力机制
+  - `feedforward.py` - 前馈网络
+  - `block.py` - Transformer编码器块
+  - `model.py` - SimpleTransformer模型
+  - `inference.py` - 推理引擎
 - `example.py` - Usage examples (training & inference)
 
 ## Usage

transformer.py

@@ -1,438 +0,0 @@
-import numpy as np
-
-
-def softmax(x, axis=-1):
-    """Softmax激活函数，用于将 logits 转换为概率分布"""
-    # 减去最大值防止数值溢出
-    e_x = np.exp(x - np.max(x, axis=axis, keepdims=True))
-    return e_x / np.sum(e_x, axis=axis, keepdims=True)
-
-
-def relu(x):
-    """ReLU激活函数"""
-    return np.maximum(0, x)
-
-
-class Linear:
-    """全连接层（线性层）"""
-    def __init__(self, in_features, out_features):
-        # 使用He初始化权重
-        self.weight = np.random.randn(in_features, out_features) * np.sqrt(2.0 / in_features)
-        # 偏置初始化为0
-        self.bias = np.zeros(out_features)
-        
-    def forward(self, x):
-        """前向传播：y = xW + b"""
-        return x @ self.weight + self.bias
-
-
-class LayerNorm:
-    """层归一化（Layer Normalization）"""
-    def __init__(self, d_model, eps=1e-6):
-        # 缩放参数
-        self.gamma = np.ones(d_model)
-        # 偏移参数
-        self.beta = np.zeros(d_model)
-        # 防止除零的小常数
-        self.eps = eps
-        
-    def forward(self, x):
-        """前向传播：对最后一个维度进行归一化"""
-        mean = np.mean(x, axis=-1, keepdims=True)
-        std = np.std(x, axis=-1, keepdims=True)
-        return self.gamma * (x - mean) / (std + self.eps) + self.beta
-
-
-class MultiHeadAttention:
-    """多头注意力机制（支持KV cache）"""
-    def __init__(self, d_model, num_heads):
-        self.d_model = d_model  # 模型维度
-        self.num_heads = num_heads  # 注意力头数
-        self.d_k = d_model // num_heads  # 每个头的维度
-        
-        # 定义Q、K、V的线性变换层
-        self.W_q = Linear(d_model, d_model)  # Query变换
-        self.W_k = Linear(d_model, d_model)  # Key变换
-        self.W_v = Linear(d_model, d_model)  # Value变换
-        self.W_o = Linear(d_model, d_model)  # 输出变换
-        
-        # 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]
-        
-        # 线性变换并分割成多头
-        Q = self.W_q.forward(Q).reshape(batch_size, seq_len, self.num_heads, self.d_k).transpose(0, 2, 1, 3)
-        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)
-        
-        # 应用掩码（如果提供）
-        if mask is not None:
-            attn_scores = np.where(mask == 0, -1e9, attn_scores)
-        
-        # Softmax得到注意力权重
-        attn_probs = softmax(attn_scores, axis=-1)
-        
-        # 加权求和
-        attn_output = attn_probs @ V
-        
-        # 拼接多头输出
-        attn_output = attn_output.transpose(0, 2, 1, 3).reshape(batch_size, seq_len, self.d_model)
-        
-        # 最终线性变换
-        output = self.W_o.forward(attn_output)
-        return output
-
-
-class FeedForward:
-    """前馈神经网络（两层全连接网络）"""
-    def __init__(self, d_model, d_ff):
-        # 第一层：扩展维度
-        self.linear1 = Linear(d_model, d_ff)
-        # 第二层：恢复维度
-        self.linear2 = Linear(d_ff, d_model)
-        
-    def forward(self, x):
-        """前向传播：线性 -> ReLU -> 线性"""
-        return self.linear2.forward(relu(self.linear1.forward(x)))
-
-
-class TransformerBlock:
-    """Transformer编码器块"""
-    def __init__(self, d_model, num_heads, d_ff):
-        # 多头自注意力层
-        self.attention = MultiHeadAttention(d_model, num_heads)
-        # 前馈网络层
-        self.feed_forward = FeedForward(d_model, d_ff)
-        # 两个层归一化
-        self.norm1 = LayerNorm(d_model)
-        self.norm2 = LayerNorm(d_model)
-        
-    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, 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编码器模型（支持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
-        # 位置编码
-        self.positional_encoding = self._create_positional_encoding(max_seq_len, d_model)
-        # Transformer编码器层堆叠
-        self.transformer_blocks = [
-            TransformerBlock(d_model, num_heads, d_ff) for _ in range(num_layers)
-        ]
-        
-    def _create_positional_encoding(self, max_seq_len, d_model):
-        """
-        创建正弦余弦位置编码
-        使用不同频率的正弦和余弦函数生成位置信息
-        """
-        pe = np.zeros((max_seq_len, d_model))
-        position = np.arange(0, max_seq_len).reshape(-1, 1).astype(float)
-        div_term = np.exp(np.arange(0, d_model, 2).astype(float) * -(np.log(10000.0) / d_model))
-        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: 注意力掩码
-        """
-        seq_len = x.shape[1]
-        
-        # 词嵌入 + 缩放
-        x = self.embedding[x] * np.sqrt(self.d_model)
-        
-        # 加上位置编码
-        x = x + self.positional_encoding[:seq_len, :]
-        
-        # 通过所有Transformer块
-        for block in self.transformer_blocks:
-            x = block.forward(x, mask)
-            
-        return x
-    
-    def count_parameters(self):
-        """统计模型参数数量"""
-        count = 0
-        count += self.embedding.size  # 嵌入层参数
-        for block in self.transformer_blocks:
-            # 注意力层参数
-            count += block.attention.W_q.weight.size + block.attention.W_q.bias.size
-            count += block.attention.W_k.weight.size + block.attention.W_k.bias.size
-            count += block.attention.W_v.weight.size + block.attention.W_v.bias.size
-            count += block.attention.W_o.weight.size + block.attention.W_o.bias.size
-            # 前馈网络参数
-            count += block.feed_forward.linear1.weight.size + block.feed_forward.linear1.bias.size
-            count += block.feed_forward.linear2.weight.size + block.feed_forward.linear2.bias.size
-            # 层归一化参数
-            count += block.norm1.gamma.size + block.norm1.beta.size
-            count += block.norm2.gamma.size + block.norm2.beta.size
-        return count
-
-
-def create_padding_mask(seq, pad_idx=0):
-    """
-    创建填充掩码
-    用于忽略序列中的填充位置（padding tokens）
-    """
-    return (seq != pad_idx).astype(float)[:, np.newaxis, np.newaxis, :]
-
-
-def create_look_ahead_mask(size):
-    """
-    创建前瞻掩码（下三角掩码）
-    用于解码器中，防止位置i看到i之后的信息
-    """
-    mask = np.triu(np.ones((size, size)), k=1).astype(bool)
-    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   # 词汇表大小
-    d_model = 512       # 模型维度
-    num_heads = 8       # 注意力头数
-    num_layers = 6      # Transformer层数
-    d_ff = 2048         # 前馈网络隐藏层维度
-    max_seq_len = 100   # 最大序列长度
-    
-    # 创建模型
-    model = SimpleTransformer(vocab_size, d_model, num_heads, num_layers, d_ff, max_seq_len)
-    
-    # 创建随机输入
-    batch_size = 2
-    seq_len = 10
-    x = np.random.randint(0, vocab_size, (batch_size, seq_len))
-    
-    # 前向传播
-    output = model.forward(x)
-    print(f"输入形状: {x.shape}")
-    print(f"输出形状: {output.shape}")
-    print(f"模型参数量: {model.count_parameters():,}")

transformer/__init__.py

@@ -0,0 +1,33 @@
+"""
+Transformer模型实现包
+按架构组件拆分：
+- utils: 激活函数和掩码工具
+- layers: 基础层（Linear, LayerNorm）
+- attention: 多头注意力机制
+- feedforward: 前馈网络
+- block: Transformer块
+- model: SimpleTransformer模型
+- inference: 推理引擎
+"""
+
+from .utils import softmax, relu, create_padding_mask, create_look_ahead_mask
+from .layers import Linear, LayerNorm
+from .attention import MultiHeadAttention
+from .feedforward import FeedForward
+from .block import TransformerBlock
+from .model import SimpleTransformer
+from .inference import InferenceEngine
+
+__all__ = [
+    'softmax',
+    'relu',
+    'create_padding_mask',
+    'create_look_ahead_mask',
+    'Linear',
+    'LayerNorm',
+    'MultiHeadAttention',
+    'FeedForward',
+    'TransformerBlock',
+    'SimpleTransformer',
+    'InferenceEngine'
+]

transformer/attention.py

@@ -0,0 +1,84 @@
+"""
+多头注意力机制模块
+包含MultiHeadAttention类，支持KV cache
+"""
+
+import numpy as np
+from .layers import Linear
+from .utils import softmax
+
+
+class MultiHeadAttention:
+    """多头注意力机制（支持KV cache）"""
+    def __init__(self, d_model, num_heads):
+        self.d_model = d_model  # 模型维度
+        self.num_heads = num_heads  # 注意力头数
+        self.d_k = d_model // num_heads  # 每个头的维度
+        
+        # 定义Q、K、V的线性变换层
+        self.W_q = Linear(d_model, d_model)  # Query变换
+        self.W_k = Linear(d_model, d_model)  # Key变换
+        self.W_v = Linear(d_model, d_model)  # Value变换
+        self.W_o = Linear(d_model, d_model)  # 输出变换
+        
+        # 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]
+        
+        # 线性变换并分割成多头
+        Q = self.W_q.forward(Q).reshape(batch_size, seq_len, self.num_heads, self.d_k).transpose(0, 2, 1, 3)
+        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)
+        
+        # 应用掩码（如果提供）
+        if mask is not None:
+            attn_scores = np.where(mask == 0, -1e9, attn_scores)
+        
+        # Softmax得到注意力权重
+        attn_probs = softmax(attn_scores, axis=-1)
+        
+        # 加权求和
+        attn_output = attn_probs @ V
+        
+        # 拼接多头输出
+        attn_output = attn_output.transpose(0, 2, 1, 3).reshape(batch_size, seq_len, self.d_model)
+        
+        # 最终线性变换
+        output = self.W_o.forward(attn_output)
+        return output

transformer/block.py

@@ -0,0 +1,42 @@
+"""
+Transformer块模块
+包含TransformerBlock类
+"""
+
+from .attention import MultiHeadAttention
+from .feedforward import FeedForward
+from .layers import LayerNorm
+
+
+class TransformerBlock:
+    """Transformer编码器块"""
+    def __init__(self, d_model, num_heads, d_ff):
+        # 多头自注意力层
+        self.attention = MultiHeadAttention(d_model, num_heads)
+        # 前馈网络层
+        self.feed_forward = FeedForward(d_model, d_ff)
+        # 两个层归一化
+        self.norm1 = LayerNorm(d_model)
+        self.norm2 = LayerNorm(d_model)
+        
+    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, 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()

transformer/feedforward.py

@@ -0,0 +1,20 @@
+"""
+前馈网络模块
+包含FeedForward类
+"""
+
+from .layers import Linear
+from .utils import relu
+
+
+class FeedForward:
+    """前馈神经网络（两层全连接网络）"""
+    def __init__(self, d_model, d_ff):
+        # 第一层：扩展维度
+        self.linear1 = Linear(d_model, d_ff)
+        # 第二层：恢复维度
+        self.linear2 = Linear(d_ff, d_model)
+        
+    def forward(self, x):
+        """前向传播：线性 -> ReLU -> 线性"""
+        return self.linear2.forward(relu(self.linear1.forward(x)))

transformer/inference.py

@@ -0,0 +1,128 @@
+"""
+推理引擎模块
+包含InferenceEngine类，管理prefill和decode阶段
+"""
+
+import numpy as np
+from .utils import softmax, create_look_ahead_mask
+
+
+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

transformer/layers.py

@@ -0,0 +1,36 @@
+"""
+基础层模块
+包含Linear和LayerNorm层
+"""
+
+import numpy as np
+
+
+class Linear:
+    """全连接层（线性层）"""
+    def __init__(self, in_features, out_features):
+        # 使用He初始化权重
+        self.weight = np.random.randn(in_features, out_features) * np.sqrt(2.0 / in_features)
+        # 偏置初始化为0
+        self.bias = np.zeros(out_features)
+        
+    def forward(self, x):
+        """前向传播：y = xW + b"""
+        return x @ self.weight + self.bias
+
+
+class LayerNorm:
+    """层归一化（Layer Normalization）"""
+    def __init__(self, d_model, eps=1e-6):
+        # 缩放参数
+        self.gamma = np.ones(d_model)
+        # 偏移参数
+        self.beta = np.zeros(d_model)
+        # 防止除零的小常数
+        self.eps = eps
+        
+    def forward(self, x):
+        """前向传播：对最后一个维度进行归一化"""
+        mean = np.mean(x, axis=-1, keepdims=True)
+        std = np.std(x, axis=-1, keepdims=True)
+        return self.gamma * (x - mean) / (std + self.eps) + self.beta

transformer/model.py

@@ -0,0 +1,117 @@
+"""
+Transformer模型模块
+包含SimpleTransformer类
+"""
+
+import numpy as np
+from .block import TransformerBlock
+
+
+class SimpleTransformer:
+    """简化的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
+        # 位置编码
+        self.positional_encoding = self._create_positional_encoding(max_seq_len, d_model)
+        # Transformer编码器层堆叠
+        self.transformer_blocks = [
+            TransformerBlock(d_model, num_heads, d_ff) for _ in range(num_layers)
+        ]
+        
+    def _create_positional_encoding(self, max_seq_len, d_model):
+        """
+        创建正弦余弦位置编码
+        使用不同频率的正弦和余弦函数生成位置信息
+        """
+        pe = np.zeros((max_seq_len, d_model))
+        position = np.arange(0, max_seq_len).reshape(-1, 1).astype(float)
+        div_term = np.exp(np.arange(0, d_model, 2).astype(float) * -(np.log(10000.0) / d_model))
+        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: 注意力掩码
+        """
+        seq_len = x.shape[1]
+        
+        # 词嵌入 + 缩放
+        x = self.embedding[x] * np.sqrt(self.d_model)
+        
+        # 加上位置编码
+        x = x + self.positional_encoding[:seq_len, :]
+        
+        # 通过所有Transformer块
+        for block in self.transformer_blocks:
+            x = block.forward(x, mask)
+            
+        return x
+    
+    def count_parameters(self):
+        """统计模型参数数量"""
+        count = 0
+        count += self.embedding.size  # 嵌入层参数
+        for block in self.transformer_blocks:
+            # 注意力层参数
+            count += block.attention.W_q.weight.size + block.attention.W_q.bias.size
+            count += block.attention.W_k.weight.size + block.attention.W_k.bias.size
+            count += block.attention.W_v.weight.size + block.attention.W_v.bias.size
+            count += block.attention.W_o.weight.size + block.attention.W_o.bias.size
+            # 前馈网络参数
+            count += block.feed_forward.linear1.weight.size + block.feed_forward.linear1.bias.size
+            count += block.feed_forward.linear2.weight.size + block.feed_forward.linear2.bias.size
+            # 层归一化参数
+            count += block.norm1.gamma.size + block.norm1.beta.size
+            count += block.norm2.gamma.size + block.norm2.beta.size
+        return count

transformer/utils.py

@@ -0,0 +1,40 @@
+"""
+工具函数模块
+包含激活函数和掩码创建函数
+"""
+
+import numpy as np
+
+
+def softmax(x, axis=-1):
+    """Softmax激活函数，用于将 logits 转换为概率分布"""
+    # 减去最大值防止数值溢出
+    e_x = np.exp(x - np.max(x, axis=axis, keepdims=True))
+    return e_x / np.sum(e_x, axis=axis, keepdims=True)
+
+
+def relu(x):
+    """ReLU激活函数"""
+    return np.maximum(0, x)
+
+
+def create_padding_mask(seq, pad_idx=0):
+    """
+    创建填充掩码
+    用于忽略序列中的填充位置（padding tokens）
+    seq: 输入序列 [batch_size, seq_len]
+    pad_idx: 填充token的索引
+    返回: 掩码 [batch_size, 1, 1, seq_len]
+    """
+    return (seq != pad_idx).astype(float)[:, np.newaxis, np.newaxis, :]
+
+
+def create_look_ahead_mask(size):
+    """
+    创建前瞻掩码（下三角掩码）
+    用于解码器中，防止位置i看到i之后的信息
+    size: 序列长度
+    返回: 掩码 [size, size]
+    """
+    mask = np.triu(np.ones((size, size)), k=1).astype(bool)
+    return (~mask).astype(float)