| 283 | """ |
| 284 | |
| 285 | def __init__( |
| 286 | self, |
| 287 | hidden_size: int = 2048, |
| 288 | head_dim: int = 256, |
| 289 | num_heads: int = 4, |
| 290 | expand_v: float = 2, |
| 291 | mode: str = 'chunk', |
| 292 | use_output_gate: bool = True, |
| 293 | use_short_conv: bool = True, |
| 294 | conv_size: int = 4, |
| 295 | conv_bias: bool = False, |
| 296 | layer_idx: int = None, |
| 297 | norm_eps: float = 1e-5, |
| 298 | num_memories: int = 8, |
| 299 | topk: int = 2, |
| 300 | capacity: float = 1.0, |
| 301 | shared_mem: bool = False, |
| 302 | single_kv_proj: bool = False, |
| 303 | **kwargs |
| 304 | ) -> MomAttention: |
| 305 | super().__init__() |
| 306 | self.num_memories = num_memories |
| 307 | self.topk = topk |
| 308 | self.capacity = capacity |
| 309 | self.shared_mem = shared_mem |
| 310 | self.single_kv_proj = single_kv_proj |
| 311 | |
| 312 | self.mode = mode |
| 313 | |
| 314 | self.hidden_size = hidden_size |
| 315 | self.expand_v = expand_v |
| 316 | |
| 317 | self.use_output_gate = use_output_gate |
| 318 | self.use_short_conv = use_short_conv |
| 319 | self.conv_size = conv_size |
| 320 | self.conv_bias = conv_bias |
| 321 | |
| 322 | self.head_dim = head_dim |
| 323 | self.num_heads = num_heads |
| 324 | |
| 325 | self.key_dim = int(self.num_heads * self.head_dim) |
| 326 | self.value_dim = int(self.key_dim * self.expand_v) |
| 327 | self.head_qk_dim = head_dim |
| 328 | self.head_v_dim = int(head_dim * self.expand_v) |
| 329 | self.layer_idx = layer_idx |
| 330 | self.silu = nn.SiLU() |
| 331 | |
| 332 | assert mode in ['chunk', 'fused_recurrent'], f"Not suppoerted mode `{mode}`." |
| 333 | |
| 334 | self.q_proj = nn.Linear(hidden_size, self.key_dim, bias=False) |
| 335 | self.gate = nn.Linear(self.hidden_size, self.num_memories, bias=False) |
| 336 | if self.single_kv_proj: |
| 337 | self.shared_k = nn.Linear(hidden_size, self.key_dim, bias=False) |
| 338 | self.shared_v = nn.Linear(hidden_size, self.value_dim, bias=False) |
| 339 | self.shared_b = nn.Linear(hidden_size, self.num_heads, bias=False) |
| 340 | self.shared_a = nn.Linear(hidden_size, self.num_heads, bias=False) |
| 341 | else: |
| 342 | self.k_proj = nn.ModuleList([ |