Fix

Author: mseeger

litgpt/attention.py

@@ -272,15 +272,24 @@ def scaled_dot_product_attention(
             # in this case.
             key = k_and_v.keys()
             value = k_and_v.values()
+            is_causal = mask is None
+            enable_gqa = self.config.n_query_groups < self.config.n_head
+            if is_causal and enable_gqa:
+                # Some efficient kernels have not implemented
+                # `enabla_gqa=True`. It is better to extend keys, values in
+                # this case.
+                q_per_kv = self.config.n_head // self.config.n_query_groups
+                key = key.repeat_interleave(q_per_kv, dim=1)
+                value = value.repeat_interleave(q_per_kv, dim=1)
             kwargs = dict(
                 query=query,
                 key=key,
                 value=value,
                 attn_mask=mask,
                 dropout_p=0.0,
                 scale=scale,
-                is_causal=mask is None,
-                enable_gqa=self.config.n_query_groups < self.config.n_head,
+                is_causal=is_causal,
+                enable_gqa=enable_gqa,
             )
             self._filter_sdpa_kernels(**kwargs)
             if self._sdpa_kernels is not None: