squall/flash-attention
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Import torch before flash_attn_2_cuda

Browse Source
This commit is contained in:
Tri Dao 2023-08-19 21:07:33 -07:00
parent 0e8c46ae08
commit d431f16751
1 changed files with 5 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
flash_attn/flash_attn_interface.py
View File

@ -1,8 +1,12 @@
import flash_attn_2_cuda as flash_attn_cuda
import torch import torch
import torch.nn as nn import torch.nn as nn
from einops import rearrange from einops import rearrange
# isort: off
# We need to import the CUDA kernels after importing torch
import flash_attn_2_cuda as flash_attn_cuda
# isort: on
def _get_block_size(device, head_dim, is_dropout, is_causal): def _get_block_size(device, head_dim, is_dropout, is_causal):
# This should match the block sizes in the CUDA kernel # This should match the block sizes in the CUDA kernel