| 196 | self.backend = backend |
| 197 | |
| 198 | def forward( |
| 199 | self, |
| 200 | x, |
| 201 | context=None, |
| 202 | mask=None, |
| 203 | additional_tokens=None, |
| 204 | n_times_crossframe_attn_in_self=0, |
| 205 | ): |
| 206 | h = self.heads |
| 207 | |
| 208 | if additional_tokens is not None: |
| 209 | # get the number of masked tokens at the beginning of the output sequence |
| 210 | n_tokens_to_mask = additional_tokens.shape[1] |
| 211 | # add additional token |
| 212 | x = torch.cat([additional_tokens, x], dim=1) |
| 213 | |
| 214 | q = self.to_q(x) |
| 215 | context = default(context, x) |
| 216 | k = self.to_k(context) |
| 217 | v = self.to_v(context) |
| 218 | |
| 219 | if n_times_crossframe_attn_in_self: |
| 220 | # reprogramming cross-frame attention as in https://arxiv.org/abs/2303.13439 |
| 221 | assert x.shape[0] % n_times_crossframe_attn_in_self == 0 |
| 222 | n_cp = x.shape[0] // n_times_crossframe_attn_in_self |
| 223 | k = repeat(k[::n_times_crossframe_attn_in_self], "b ... -> (b n) ...", n=n_cp) |
| 224 | v = repeat(v[::n_times_crossframe_attn_in_self], "b ... -> (b n) ...", n=n_cp) |
| 225 | |
| 226 | q, k, v = map(lambda t: rearrange(t, "b n (h d) -> b h n d", h=h), (q, k, v)) |
| 227 | |
| 228 | ## old |
| 229 | """ |
| 230 | sim = einsum('b i d, b j d -> b i j', q, k) * self.scale |
| 231 | del q, k |
| 232 | |
| 233 | if exists(mask): |
| 234 | mask = rearrange(mask, 'b ... -> b (...)') |
| 235 | max_neg_value = -torch.finfo(sim.dtype).max |
| 236 | mask = repeat(mask, 'b j -> (b h) () j', h=h) |
| 237 | sim.masked_fill_(~mask, max_neg_value) |
| 238 | |
| 239 | # attention, what we cannot get enough of |
| 240 | sim = sim.softmax(dim=-1) |
| 241 | |
| 242 | out = einsum('b i j, b j d -> b i d', sim, v) |
| 243 | """ |
| 244 | ## new |
| 245 | with sdp_kernel(**BACKEND_MAP[self.backend]): |
| 246 | # print("dispatching into backend", self.backend, "q/k/v shape: ", q.shape, k.shape, v.shape) |
| 247 | out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask) # scale is dim_head ** -0.5 per default |
| 248 | |
| 249 | del q, k, v |
| 250 | out = rearrange(out, "b h n d -> b n (h d)", h=h) |
| 251 | |
| 252 | if additional_tokens is not None: |
| 253 | # remove additional token |
| 254 | out = out[:, n_tokens_to_mask:] |
| 255 | return self.to_out(out) |