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

Remove redundant shape asserts in rotary embeddings

Browse Source
This commit is contained in:
Alexander Ploshkin 2022-12-15 18:13:21 +04:00
parent 04c4c6106e
commit 96656b9323
1 changed files with 0 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
flash_attn/layers/rotary.py
View File

@ -43,8 +43,6 @@ class ApplyRotaryEmb(torch.autograd.Function):
rotary_dim *= 2
assert rotary_dim <= headdim
assert seqlen <= rotary_seqlen
assert cos.shape == (rotary_seqlen, rotary_dim // 2)
assert sin.shape == (rotary_seqlen, rotary_dim // 2)
x1, x2 = x[..., :rotary_dim].chunk(2, dim=-1)
out = torch.empty_like(x) if not inplace else x
o1, o2 = out[..., :rotary_dim].chunk(2, dim=-1) if not inplace else (x1, x2)