ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Update flex_attention.py

Browse source
This commit is contained in:
Daniel Han 2024-11-06 15:56:52 -08:00
commit 16dcb3ee93
1 changed files with 3 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

8
unsloth/kernels/flex_attention.py
View file

@ -42,15 +42,13 @@ if not HAS_FLEX_ATTENTION:
# Below fails on compiled_autograd, so disable it
try:
disable_compiled_autograd = torch._dynamo.compiled_autograd.disable
disable_compile = torch._dynamo.compiled_autograd.disable
except:
disable_compiled_autograd = lambda f: f
disable_compile = lambda f: f
pass
# Logit softcapping
@disable_compiled_autograd(
torch.compile(fullgraph = True, dynamic = True, options = torch_compile_options)
)
@disable_compile(torch.compile(fullgraph = True, dynamic = True, options = torch_compile_options))
def slow_attention_softcapping(Q, K, V, causal_mask, self, bsz, q_len):
n_heads = self.num_heads
head_dim = self.head_dim