(x)
| 355 | |
| 356 | |
| 357 | def use_flash_2(x): |
| 358 | if not flags.get_use_flash_attention_2(): |
| 359 | return False |
| 360 | if flash_attn is None: |
| 361 | return False |
| 362 | if x.device.type != "cuda": |
| 363 | return False |
| 364 | if x.dtype not in (torch.float16, torch.bfloat16): |
| 365 | return False |
| 366 | return True |
| 367 | |
| 368 | |
| 369 | class SelfAttentionBlock(nn.Module): |