4.1 从矩阵算子拼成 Decoder Block

x → RMSNorm → QKV Linear → RoPE(Q,K) → Attention(KV 写入/读取)
  → Output Linear → Residual → RMSNorm → Gate/Up Linear
  → SiLU(gate) * up → Down Linear → Residual

import torch

def rmsnorm(x, weight, eps=1e-6):
    xf = x.float()
    inv = torch.rsqrt(xf.square().mean(dim=-1, keepdim=True) + eps)
    return (xf * inv).to(x.dtype) * weight

def apply_rope(x, cos, sin):
    # x: [..., head_dim]；约定相邻两维组成一对
    even, odd = x[..., 0::2], x[..., 1::2]
    y = torch.empty_like(x)
    y[..., 0::2] = even * cos - odd * sin
    y[..., 1::2] = even * sin + odd * cos
    return y