class Block(nn.Module):
def forward(self, x):
x = x + self.attn(self.ln_1(x)) # pre-norm:先 norm 再子层,残差直连
x = x + self.mlp(self.ln_2(x))
return x
# CausalSelfAttention.forward 的骨架
q, k, v = self.c_attn(x).split(self.n_embd, dim=2) # 一次 GEMM 算 QKV
k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # 头换到第 1 维
y = F.scaled_dot_product_attention(q, k, v, is_causal=True) # 融合 kernel
y = y.transpose(1, 2).contiguous().view(B, T, C) # 拼回去
COMMENTS
评论存放在 GitHub Discussions, 用 GitHub 账号登录即可发表,支持 Markdown。 想针对正文某句话说?选中那段文字,点浮出的「评论」即可划线评论;觉得哪里写错了,发表时勾上「同时提交 Issue」。 有人回复你时 GitHub 会按你的通知设置发邮件,不用守在这里。