CausalSelfAttention第29-76行—— 最重要class CausalSelfAttention(nn.Module):def __init__(self, config):super().__init__()assert config.n_embd % config.n_head 0# QKV 投影一个线性层同时计算 Q, K, Vself.c_attn nn.Linear(config.n_embd, 3 * config.n_embd, biasconfig.bias)# 输出投影self.c_proj nn.Linear(config.n_embd, config.n_embd, biasconfig.bias)# 正则化self.attn_dropout nn.Dropout(config.dropout)self.resid_dropout nn.Dropout(config.dropout)self.n_head config.n_headself.n_embd config.n_embdself.dropout config.dropout# Flash Attention 检查PyTorch 2.0self.flash hasattr(torch.nn.functional, scaled_dot_product_attention)if not self.flash:# 因果掩码确保只能看到左边的tokenself.register_buffer(bias, torch.tril(torch.ones(config.block_size, config.block_size)).view(1, 1, config.block_size, config.block_size))def forward(self, x):B, T, C x.size() # Batch, Time, Channel# 1. 投影得到 Q, K, Vq, k, v self.c_attn(x).split(self.n_embd, dim2)# 2. 多头重塑[B, T, C] → [B, heads, T, head_dim]k k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)q q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)v v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)# 3. 计算注意力if self.flash:# 高效实现PyTorch 2.0y torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_maskNone,dropout_pself.dropout if self.training else 0,is_causalTrue)else:# 手动实现att (q k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))att att.masked_fill(self.bias[:,:,:T,:T] 0, float(-inf))att F.softmax(att, dim-1)att self.attn_dropout(att)y att v# 4. 合并多头输出y y.transpose(1, 2).contiguous().view(B, T, C)# 5. 输出投影y self.resid_dropout(self.c_proj(y))return y数据流详解输入: [B, T, C] (Batch , Sequence, Embedding Dim)│▼┌─────────────────────────────────────┐│ c_attn: Linear(C → 3*C) ││ 一次投影同时得到 Q, K, V │└─────────────────────────────────────┘│▼┌─────────────────────────────────────┐│ 重塑为多头 ││ [B, T, 3*C] → [B, heads, T, hs] ││ 其中 hs C / heads │└─────────────────────────────────────┘│▼┌─────────────────────────────────────┐│ Attention(Q, K, V) ││ softmax(QK^T / √hs) V │└─────────────────────────────────────┘│▼┌─────────────────────────────────────┐│ 合并多头 输出投影 ││ [B, heads, T, hs] → [B, T, C] │└─────────────────────────────────────┘为什么叫Causal因果因为掩码矩阵 self.bias 是下三角矩阵[[1, 0, 0, 0],[1, 1, 0, 0],[1, 1, 1, 0],[1, 1, 1, 1]]这确保位置 i 只能看到位置 0~i看不到未来。 
长河君 · 2026/8/25 6:48:13 · 阅读 1.2k nanoGPT 逐行讲解——CausalSelfAttention
CausalSelfAttention第29-76行—— 最重要class CausalSelfAttention(nn.Module):def __init__(self, config):super().__init__()assert config.n_embd % config.n_head 0# QKV 投影一个线性层同时计算 Q, K, Vself.c_attn nn.Linear(config.n_embd, 3 * config.n_embd, biasconfig.bias)# 输出投影self.c_proj nn.Linear(config.n_embd, config.n_embd, biasconfig.bias)# 正则化self.attn_dropout nn.Dropout(config.dropout)self.resid_dropout nn.Dropout(config.dropout)self.n_head config.n_headself.n_embd config.n_embdself.dropout config.dropout# Flash Attention 检查PyTorch 2.0self.flash hasattr(torch.nn.functional, scaled_dot_product_attention)if not self.flash:# 因果掩码确保只能看到左边的tokenself.register_buffer(bias, torch.tril(torch.ones(config.block_size, config.block_size)).view(1, 1, config.block_size, config.block_size))def forward(self, x):B, T, C x.size() # Batch, Time, Channel# 1. 投影得到 Q, K, Vq, k, v self.c_attn(x).split(self.n_embd, dim2)# 2. 多头重塑[B, T, C] → [B, heads, T, head_dim]k k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)q q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)v v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)# 3. 计算注意力if self.flash:# 高效实现PyTorch 2.0y torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_maskNone,dropout_pself.dropout if self.training else 0,is_causalTrue)else:# 手动实现att (q k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))att att.masked_fill(self.bias[:,:,:T,:T] 0, float(-inf))att F.softmax(att, dim-1)att self.attn_dropout(att)y att v# 4. 合并多头输出y y.transpose(1, 2).contiguous().view(B, T, C)# 5. 输出投影y self.resid_dropout(self.c_proj(y))return y数据流详解输入: [B, T, C] (Batch , Sequence, Embedding Dim)│▼┌─────────────────────────────────────┐│ c_attn: Linear(C → 3*C) ││ 一次投影同时得到 Q, K, V │└─────────────────────────────────────┘│▼┌─────────────────────────────────────┐│ 重塑为多头 ││ [B, T, 3*C] → [B, heads, T, hs] ││ 其中 hs C / heads │└─────────────────────────────────────┘│▼┌─────────────────────────────────────┐│ Attention(Q, K, V) ││ softmax(QK^T / √hs) V │└─────────────────────────────────────┘│▼┌─────────────────────────────────────┐│ 合并多头 输出投影 ││ [B, heads, T, hs] → [B, T, C] │└─────────────────────────────────────┘为什么叫Causal因果因为掩码矩阵 self.bias 是下三角矩阵[[1, 0, 0, 0],[1, 1, 0, 0],[1, 1, 1, 0],[1, 1, 1, 1]]这确保位置 i 只能看到位置 0~i看不到未来。