| import torch | |
| import torch.nn as nn | |
| from .rmsnorm import RMSNorm | |
| from .attention import GroupedQueryAttention | |
| from .mlp import SwiGLU | |
| class DecoderBlock(nn.Module): | |
| """One transformer block: attention sub-layer, then MLP sub-layer. | |
| Each sub-layer is pre-normed and wrapped in a residual (add-back) connection. | |
| """ | |
| def __init__(self, dim, n_heads, n_kv_heads, head_dim, hidden_dim, eps=1e-6): | |
| super().__init__() | |
| # a norm before each sub-layer | |
| self.attn_norm = RMSNorm(dim, eps) | |
| self.mlp_norm = RMSNorm(dim, eps) | |
| self.attn = GroupedQueryAttention(dim, n_heads, n_kv_heads, head_dim) | |
| self.mlp = SwiGLU(dim, hidden_dim) | |
| def forward(self, x, cos, sin, block_mask): | |
| # attention sub-layer: normalize, attend, add the result back | |
| x = x + self.attn(self.attn_norm(x), cos, sin, block_mask) | |
| # mlp sub-layer: normalize, process per-token, add the result back | |
| x = x + self.mlp(self.mlp_norm(x)) | |
| return x |