Update raven_modeling_minimal.py

raven_modeling_minimal.py

@@ -22,7 +22,7 @@ import torch.nn.functional as F
 from transformers import GenerationConfig
 from transformers.models.llama.modeling_llama import LlamaRotaryEmbedding, apply_rotary_pos_emb
 
-torch.backends.cuda.enable_math_sdp(False)
+# torch.backends.cuda.enable_math_sdp(False)
 
 
 class RavenPreTrainedModel(PreTrainedModel):
@@ -419,16 +419,11 @@ class CausalSelfAttention(torch.nn.Module):
             if q.shape[2] < k.shape[2]:
                 if q.shape[2] > 1:
                     bias = attn_bias.causal_lower_right(q.shape[2], k.shape[2])
-                    y = torch.nn.functional.scaled_dot_product_attention(q, k, v, bias, dropout_p=0.0)
+                    y = torch.nn.functional.scaled_dot_product_attention(q, k, v, bias, dropout_p=0.0, enable_gqa=True)
                 else:
-                    y = torch.nn.functional.scaled_dot_product_attention(q, k, v, dropout_p=0.0, is_causal=False)
+                    y = torch.nn.functional.scaled_dot_product_attention(q, k, v, dropout_p=0.0, is_causal=False, enable_gqa=True)
             else:
-                if self.n_head != self.n_kv_heads: # GQA case
-                    repeat_factor = self.n_head // self.n_kv_heads
-                    k = k.repeat_interleave(repeat_factor, dim=1)
-                    v = v.repeat_interleave(repeat_factor, dim=1)
-
-                y = torch.nn.functional.scaled_dot_product_attention(q, k, v, dropout_p=0.0, is_causal=True)
+                y = torch.nn.functional.scaled_dot_product_attention(q, k, v, dropout_p=0.0, is_causal=True, enable_gqa=True)
         y = y.transpose(1, 2).reshape(B, S, E).contiguous()  # reshape is a view if possible (it mostly is)
         return self.proj(y)
 
@@ -651,7 +646,7 @@ class RavenForCausalLM(RavenPreTrainedModel, GenerationMixin):
         if position_ids is None and cache_position is None:
             position_ids = torch.arange(input_ids.shape[1], device=self.device).unsqueeze(0)
         elif cache_position is not None:
-            position_ids = torch.tensor([cache_position], device=self.device).unsqueeze(0)
+            position_ids = cache_position.unsqueeze(0)
 
         if input_embeds is None:
             input_embeds = self.transformer.wte(input_ids)  # type: ignore # types broken in 2.6+