Fix GQA in oracle-top-p and magicpig

Author: apd10

sparse_attention_hub/sparse_attention/research_attention/maskers/fixed/implementations/oracle_top_p.py

@@ -10,6 +10,10 @@
     MaskerConfig,
     MaskerRegistry,
 )
+from sparse_attention_hub.sparse_attention.utils.kv_utils import (
+    _get_num_key_value_groups,
+    repeat_kv,
+)
 from sparse_attention_hub.sparse_attention.utils.mask import Mask
 
 from ..base import TopPMasker, TopPMaskerConfig
@@ -71,6 +75,8 @@ def _compute_attention_scores(
         self, keys: torch.Tensor, queries: torch.Tensor
     ) -> torch.Tensor:
         """Compute exp(attention scores) between queries and keys."""
+        ngroups = _get_num_key_value_groups(queries, keys)
+        keys = repeat_kv(keys, ngroups)
         raw_attention_scores = queries @ keys.transpose(-2, -1)
         _max_attention_score = raw_attention_scores.max(dim=-1, keepdim=True)[0]
         adjusted = torch.exp(raw_attention_scores - _max_attention_score)

sparse_attention_hub/sparse_attention/research_attention/maskers/sampling/implementations/magic_pig.py

@@ -15,6 +15,11 @@
     MaskerConfig,
     MaskerRegistry,
 )
+from sparse_attention_hub.sparse_attention.utils.kv_utils import (
+    _get_num_key_value_groups,
+    repeat_kv,
+)
+
 from sparse_attention_hub.sparse_attention.utils.mask import Mask
 
 from ..base import SamplingMasker, SamplingMaskerConfig
@@ -308,6 +313,9 @@ def add_mask(
         seq_len_queries: int = queries.shape[2]
         seq_len_keys: int = keys.shape[2]
 
+        ngroups = _get_num_key_value_groups(queries, keys)
+        keys = repeat_kv(keys, ngroups)
+
         probabilities: torch.Tensor = self._compute_probabilities(keys, queries)
         matches: torch.Tensor = self._compute_lsh_matches(keys, queries)
         dense_mask: torch.Tensor = matches * probabilities