| 172 | # Attention kernels |
| 173 | # ---------------------------- |
| 174 | def flash_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, num_heads: int, compatibility_mode=False, attention_mask=None, return_KV=False): |
| 175 | if attention_mask is not None: |
| 176 | seqlen = q.shape[1] |
| 177 | seqlen_kv = k.shape[1] |
| 178 | q = rearrange(q, "b s (n d) -> (b s) n d", n=num_heads) |
| 179 | k = rearrange(k, "b s (n d) -> (b s) n d", n=num_heads) |
| 180 | v = rearrange(v, "b s (n d) -> (b s) n d", n=num_heads) |
| 181 | cu_seqlens_q = torch.tensor([0, seqlen], device=q.device, dtype=torch.int32) |
| 182 | cu_seqlens_k = torch.tensor([0, seqlen_kv], device=q.device, dtype=torch.int32) |
| 183 | head_mask_type = torch.tensor([1]*num_heads, device=q.device, dtype=torch.int32) |
| 184 | streaming_info = None |
| 185 | base_blockmask = attention_mask |
| 186 | max_seqlen_q_ = seqlen |
| 187 | max_seqlen_k_ = seqlen_kv |
| 188 | p_dropout = 0.0 |
| 189 | x = block_sparse_attn_func( |
| 190 | q, k, v, |
| 191 | cu_seqlens_q, cu_seqlens_k, |
| 192 | head_mask_type, |
| 193 | streaming_info, |
| 194 | base_blockmask, |
| 195 | max_seqlen_q_, max_seqlen_k_, |
| 196 | p_dropout, |
| 197 | deterministic=False, |
| 198 | softmax_scale=None, |
| 199 | is_causal=False, |
| 200 | exact_streaming=False, |
| 201 | return_attn_probs=False, |
| 202 | ).unsqueeze(0) |
| 203 | x = rearrange(x, "b s n d -> b s (n d)", n=num_heads) |
| 204 | elif compatibility_mode: |
| 205 | q = rearrange(q, "b s (n d) -> b n s d", n=num_heads) |
| 206 | k = rearrange(k, "b s (n d) -> b n s d", n=num_heads) |
| 207 | v = rearrange(v, "b s (n d) -> b n s d", n=num_heads) |
| 208 | x = F.scaled_dot_product_attention(q, k, v) |
| 209 | x = rearrange(x, "b n s d -> b s (n d)", n=num_heads) |
| 210 | elif FLASH_ATTN_3_AVAILABLE: |
| 211 | q = rearrange(q, "b s (n d) -> b s n d", n=num_heads) |
| 212 | k = rearrange(k, "b s (n d) -> b s n d", n=num_heads) |
| 213 | v = rearrange(v, "b s (n d) -> b s n d", n=num_heads) |
| 214 | x = flash_attn_interface.flash_attn_func(q, k, v) |
| 215 | if isinstance(x, tuple): |
| 216 | x = x[0] |
| 217 | x = rearrange(x, "b s n d -> b s (n d)", n=num_heads) |
| 218 | elif FLASH_ATTN_2_AVAILABLE: |
| 219 | q = rearrange(q, "b s (n d) -> b s n d", n=num_heads) |
| 220 | k = rearrange(k, "b s (n d) -> b s n d", n=num_heads) |
| 221 | v = rearrange(v, "b s (n d) -> b s n d", n=num_heads) |
| 222 | x = flash_attn.flash_attn_func(q, k, v) |
| 223 | x = rearrange(x, "b s n d -> b s (n d)", n=num_heads) |
| 224 | elif SAGE_ATTN_AVAILABLE: |
| 225 | q = rearrange(q, "b s (n d) -> b n s d", n=num_heads) |
| 226 | k = rearrange(k, "b s (n d) -> b n s d", n=num_heads) |
| 227 | v = rearrange(v, "b s (n d) -> b n s d", n=num_heads) |
| 228 | x = sageattn(q, k, v) |
| 229 | x = rearrange(x, "b n s d -> b s (n d)", n=num_heads) |
| 230 | else: |
| 231 | q = rearrange(q, "b s (n d) -> b n s d", n=num_heads) |