feat: no flash attention during inference

Signed-off-by: jupyterjazz <[email protected]>

rotary.py

@@ -6,11 +6,13 @@ from typing import Optional, Tuple, Union
 
 import torch
 from einops import rearrange, repeat
-try:
-    from flash_attn.ops.triton.rotary import apply_rotary
-except ImportError:
-    def apply_rotary(*args, **kwargs):
-        raise RuntimeError('RoPE requires flash-attention to be installed')
+
+if torch.cuda.is_available():
+    try:
+        from flash_attn.ops.triton.rotary import apply_rotary
+    except ImportError:
+        def apply_rotary(*args, **kwargs):
+            raise RuntimeError('RoPE requires flash-attention to be installed')
 
 
 def rotate_half(x, interleaved=False):
@@ -60,6 +62,7 @@ class ApplyRotaryEmb(torch.autograd.Function):
             interleaved=interleaved,
             inplace=inplace,
         )
+
         if isinstance(seqlen_offsets, int):
             ctx.save_for_backward(cos, sin, cu_seqlens)  # Can't save int with save_for_backward
             ctx.seqlen_offsets = seqlen_offsets
@@ -82,6 +85,7 @@ class ApplyRotaryEmb(torch.autograd.Function):
         # "[CUDA]: invalid device context", and cloning makes it work. Idk why. Triton 2.1.0 works.
         if not ctx.interleaved and not ctx.inplace:
             do = do.clone()
+
         dx = apply_rotary(
             do,
             cos,
@@ -150,21 +154,37 @@ class ApplyRotaryEmbQKV_(torch.autograd.Function):
         # batch, seqlen, three, nheads, headdim = qkv.shape
         assert qkv.shape[-3] == 3
         if cos_k is None and sin_k is None and qkv.is_contiguous():
-            # Call 1 kernel instead of 2 kernels
-            # We need qkv to be contiguous so that when we reshape to combine (3, nheads)
-            # dimensions, we get the same tensor
-            qk = rearrange(qkv[..., :2, :, :], "... t h d -> ... (t h) d")
-            # qk = qkv[:, :, :2].reshape(batch, seqlen, -1, headdim)
-            apply_rotary(
-                qk,
-                cos,
-                sin,
-                seqlen_offsets=seqlen_offsets,
-                interleaved=interleaved,
-                inplace=True,
-                cu_seqlens=cu_seqlens,
-                max_seqlen=max_seqlen,
-            )
+
+            if torch.cuda.is_available():
+                # Call 1 kernel instead of 2 kernels
+                # We need qkv to be contiguous so that when we reshape to combine (3, nheads)
+                # dimensions, we get the same tensor
+                qk = rearrange(qkv[..., :2, :, :], "... t h d -> ... (t h) d")
+                # qk = qkv[:, :, :2].reshape(batch, seqlen, -1, headdim)
+                apply_rotary(
+                    qk,
+                    cos,
+                    sin,
+                    seqlen_offsets=seqlen_offsets,
+                    interleaved=interleaved,
+                    inplace=True,
+                    cu_seqlens=cu_seqlens,
+                    max_seqlen=max_seqlen,
+                )
+            else:
+                q_rot = apply_rotary_emb_torch(
+                    qkv[:, :, 0],
+                    cos,
+                    sin,
+                    interleaved=interleaved,
+                )
+                k_rot = apply_rotary_emb_torch(
+                    qkv[:, :, 1],
+                    cos,
+                    sin,
+                    interleaved=interleaved,
+                )
+                qkv = torch.stack((q_rot, k_rot, qkv[:, :, 2]), dim=2)
         else:
             cos_k = cos if cos_k is None else cos_k
             sin_k = sin if sin_k is None else sin_k
@@ -228,7 +248,6 @@ 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,