fix: 修复 MLA 多个 bug 并缩小测试模型参数

- MLA kv_b_proj 输出维度和 q_rope 切分偏移修复
- 打通 MLA 配置从 ModelConfig 到 DecoderBlock 的传递路径
- rope_theta 配置不再被忽略，MLA 使用 qk_rope_head_dim
- tie_weight 使用 is True 避免 None 隐式生效
- norm_eps/rope base 类型标注修正
- 测试模型参数缩小 (dim=8, head_dim=4)
- 新增 6 种架构配置 × 2 场景的前向传播测试

astrai/model/components/attention.py

@@ -143,7 +143,7 @@ class MLA(nn.Module):
 
         self.kv_b_proj = Linear(
             kv_lora_rank,
-            n_kv_heads * (self.head_dim + qk_rope_head_dim + self.head_dim),
+            n_kv_heads * (2 * self.head_dim),
         )
 
         self.o_proj = Linear(dim, dim, bias=False)
@@ -176,7 +176,7 @@ class MLA(nn.Module):
 
         q_nope, q_rope = (
             q[..., : self.qk_nope_head_dim],
-            q[..., self.qk_rope_head_dim :],
+            q[..., self.qk_nope_head_dim :],
         )
         q_rope = apply_rotary_emb(q_rope, rotary_emb)
         k_rope = apply_rotary_emb(k_rope, rotary_emb)

astrai/model/components/decoder_block.py

@@ -16,13 +16,13 @@ class DecoderBlock(nn.Module):
         n_heads: int,
         dim_ffn: int,
         n_kv_heads: int,
-        norm_eps: int,
+        norm_eps: float,
         use_qk_norm: bool,
         use_gated_attention: bool,
         layer_id: int,
         attn_type: str = "gqa",
         ffn_type: str = "mlp",
-        **moe_kwargs,
+        **kwargs,
     ):
         super().__init__()
         self.attention = AttnFactory.create(
@@ -34,10 +34,11 @@ class DecoderBlock(nn.Module):
             norm_eps=norm_eps,
             use_gated_attention=use_gated_attention,
             layer_id=layer_id,
+            **kwargs,
         )
         self.input_norm = RMSNorm(dim, norm_eps)
         self.post_attention_norm = RMSNorm(dim, norm_eps)
-        self.mlp = FFNFactory.create(ffn_type, dim, dim_ffn, **moe_kwargs)
+        self.mlp = FFNFactory.create(ffn_type, dim, dim_ffn, **kwargs)
 
     def forward(
         self,

astrai/model/components/rope.py

@@ -30,7 +30,7 @@ def apply_rotary_emb(x: torch.Tensor, freqs_cis: Tensor) -> Tensor:
 
 
 class RotaryEmbedding(nn.Module):
-    def __init__(self, dim: int, max_len: int, base: int = 10000):
+    def __init__(self, dim: int, max_len: int, base: float = 10000):
         super().__init__()
         self.dim = dim
         self.max_len = max_len