self.input_layernorm = RMSNorm(...)

self.self_attn = Attention(..., rope=True, num_key_value_heads=...)

self.post_attention_layernorm = RMSNorm(...)

self.mlp = SwiGLU(...) or MoE(...)