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

Fix NameError and typo in ApplyRotaryEmbQKV_ (#569)

Browse Source
This commit is contained in:
Katherine Crowson 2023-09-25 10:47:34 -07:00 committed by GitHub
parent 0a1d03c7ea
commit 4c8ff9154e
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 3 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
flash_attn/layers/rotary.py
View File

@ -193,16 +193,16 @@ class ApplyRotaryEmbQKV_(torch.autograd.Function):
sin_k = sin if sin_k is None else sin_k sin_k = sin if sin_k is None else sin_k
dq, dk = dqkv[:, :, 0], dqkv[:, :, 1] dq, dk = dqkv[:, :, 0], dqkv[:, :, 1]
apply_rotary( apply_rotary(
dq, cos, sin, seqlen_offsets, interleaved=interleaved, inplace=True, conjugate=True dq, cos, sin, seqlen_offsets, interleaved=ctx.interleaved, inplace=True, conjugate=True
) )
apply_rotary( apply_rotary(
dk, dk,
cos_k, cos_k,
sin_k, sin_k,
seqlen_offsets, seqlen_offsets,
interleaved=interleaved, interleaved=ctx.interleaved,
inplace=True, inplace=True,
conjudate=True, conjugate=True,
) )
return dqkv, None, None, None, None, None, None return dqkv, None, None, None, None, None, None