ViperEkura
0ba8c70ce1
- 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 场景的前向传播测试
60 lines
1.6 KiB
Python
60 lines
1.6 KiB
Python
from typing import Optional
|
|
|
|
import torch.nn as nn
|
|
from torch import Tensor
|
|
|
|
from astrai.inference.core.cache import KvcacheView
|
|
from astrai.model.components.attention import AttnFactory
|
|
from astrai.model.components.mlp import FFNFactory
|
|
from astrai.model.components.norm import RMSNorm
|
|
|
|
|
|
class DecoderBlock(nn.Module):
|
|
def __init__(
|
|
self,
|
|
dim: int,
|
|
n_heads: int,
|
|
dim_ffn: int,
|
|
n_kv_heads: int,
|
|
norm_eps: float,
|
|
use_qk_norm: bool,
|
|
use_gated_attention: bool,
|
|
layer_id: int,
|
|
attn_type: str = "gqa",
|
|
ffn_type: str = "mlp",
|
|
**kwargs,
|
|
):
|
|
super().__init__()
|
|
self.attention = AttnFactory.create(
|
|
attn_type,
|
|
dim=dim,
|
|
n_heads=n_heads,
|
|
n_kv_heads=n_kv_heads,
|
|
use_qk_norm=use_qk_norm,
|
|
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, **kwargs)
|
|
|
|
def forward(
|
|
self,
|
|
x: Tensor,
|
|
rotary_emb: Tensor,
|
|
attention_mask: Optional[Tensor] = None,
|
|
paged_cache: Optional[KvcacheView] = None,
|
|
) -> Tensor:
|
|
attn_output = self.attention(
|
|
self.input_norm(x),
|
|
rotary_emb,
|
|
attention_mask,
|
|
paged_cache,
|
|
)
|
|
x = attn_output + x
|
|
x = self.mlp(self.post_attention_norm(x)) + x
|
|
|
|
return x
|