From 4c8ff9154e76c68e7114292bd527c22f45fbf586 Mon Sep 17 00:00:00 2001 From: Katherine Crowson Date: Mon, 25 Sep 2023 10:47:34 -0700 Subject: [PATCH] Fix NameError and typo in ApplyRotaryEmbQKV_ (#569) --- flash_attn/layers/rotary.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/flash_attn/layers/rotary.py b/flash_attn/layers/rotary.py index bd05258..4ec049e 100644 --- a/flash_attn/layers/rotary.py +++ b/flash_attn/layers/rotary.py @@ -193,16 +193,16 @@ class ApplyRotaryEmbQKV_(torch.autograd.Function): sin_k = sin if sin_k is None else sin_k dq, dk = dqkv[:, :, 0], dqkv[:, :, 1] 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( dk, cos_k, sin_k, seqlen_offsets, - interleaved=interleaved, + interleaved=ctx.interleaved, inplace=True, - conjudate=True, + conjugate=True, ) return dqkv, None, None, None, None, None, None