| 166 | return self.down_proj(up * F.silu(gate)) |
| 167 | |
| 168 | class GQA(nn.Module): |
| 169 | def __init__(self, |
| 170 | dim: int, |
| 171 | n_head: int, |
| 172 | shape_rotator: ShapeRotator, |
| 173 | kv_heads: Optional[int] = None, |
| 174 | eps: float = 1e-5, |
| 175 | causal: bool = True,): |
| 176 | super().__init__() |
| 177 | self.n_heads = n_head |
| 178 | self.kv_heads = default(kv_heads, n_head) |
| 179 | self.head_dim = dim // n_head |
| 180 | self.causal = causal |
| 181 | |
| 182 | self.proj_qkv = Linear(dim, self.head_dim*(n_head+2*self.kv_heads)) |
| 183 | |
| 184 | self.norm_q = Norm(self.head_dim*n_head, eps=eps) |
| 185 | self.norm_k = Norm(self.head_dim*self.kv_heads, eps=eps) |
| 186 | |
| 187 | self.attn_out = Linear(dim, dim) |
| 188 | |
| 189 | self.shape_rotator = shape_rotator |
| 190 | |
| 191 | def _sdpa(self, q: Tensor, k: Tensor, v: Tensor) -> Tensor: |
| 192 | k = k.repeat_interleave(self.n_heads // self.kv_heads, dim=2) |
| 193 | v = v.repeat_interleave(self.n_heads // self.kv_heads, dim=2) |
| 194 | x = F.scaled_dot_product_attention( |
| 195 | q.transpose(1, 2), |
| 196 | k.transpose(1, 2), |
| 197 | v.transpose(1, 2), |
| 198 | is_causal=False if (q.size(1) != k.size(1)) else self.causal, |
| 199 | ) |
| 200 | x = x.transpose(1, 2).contiguous() |
| 201 | return x |
| 202 | |
| 203 | def _attend(self, q: Tensor, k: Tensor, v: Tensor, kv_cache: Optional[Tensor] = None,): |
| 204 | cache_len = get_cache_len(kv_cache) |
| 205 | q, k = self.shape_rotator.rotate(q, k, offset=cache_len) |
| 206 | if exists(kv_cache): |
| 207 | k = T.cat([kv_cache[:, :cache_len, 0], k], dim=1) |
| 208 | v = T.cat([kv_cache[:, :cache_len, 1], v], dim=1) |
| 209 | kv_cache[:, :k.size(1), 0] = k |
| 210 | kv_cache[:, :v.size(1), 1] = v |
| 211 | x = self._sdpa(q, k, v) |
| 212 | return self.attn_out(rearrange(x, 'b s h d -> b s (h d)')) |
| 213 | |
| 214 | def _project(self, x): |
| 215 | full_q, full_k, full_v = self.proj_qkv(x).chunk(3, dim=-1) |
| 216 | normed_full_q = self.norm_q(full_q).to(full_q.dtype) |
| 217 | normed_full_k = self.norm_k(full_k).to(full_k.dtype) |
| 218 | |
| 219 | q = rearrange(normed_full_q, 'b s (h d) -> b s h d', h=self.n_heads) |
| 220 | k = rearrange(normed_full_k, 'b s (h d) -> b s h d', h=self.kv_heads) |
| 221 | v = rearrange(full_v, 'b s (h d) -> b s h d', h=self.kv_heads) |
| 222 | return q, k, v |
| 223 | |
| 224 | def forward(self, |
| 225 | x: Tensor, |