File size: 1,168 Bytes
fb6b7f8 be8ee5f fb6b7f8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 | import torch.nn as nn
from .SwiGLU import SwiGLU
from .GroupedQueryAttention import GroupedQueryAttention
class DecoderBlock(nn.Module):
def __init__(self, d_model: int, num_heads: int, num_kv_heads: int, ff_dim: int, dropout: float):
super().__init__()
self.gqa = GroupedQueryAttention(d_model, num_heads, num_kv_heads, dropout)
self.ffn = SwiGLU(d_model, ff_dim)
self.norm1 = nn.RMSNorm(d_model, eps=1e-6)
self.norm2 = nn.RMSNorm(d_model, eps=1e-6)
self.drop = nn.Dropout(dropout)
def forward(self, x, cos, sin, attn_mask=None):
x = x + self.drop(self.gqa(self.norm1(x), cos, sin, attn_mask))
x = x + self.drop(self.ffn(self.norm2(x)))
return x
def prefill(self, x, cos, sin):
attn, kv = self.gqa.prefill(self.norm1(x), cos, sin)
x = x + self.drop(attn)
x = x + self.drop(self.ffn(self.norm2(x)))
return x, list(kv)
def forward_with_cache(self, x, kv, cache_len: int, cos, sin):
x = x + self.drop(self.gqa.forward_with_cache(self.norm1(x), kv, cache_len, cos, sin))
x = x + self.drop(self.ffn(self.norm2(x)))
return x
|