repeat kv for GQA (#462)

Author: lchu6

mamba_ssm/modules/mha.py

@@ -173,6 +173,8 @@ def _update_kvcache_attention(self, q, kv, inference_params):
             # TODO: this only uses seqlen_offset and not lengths_per_sample.
             kv = self._update_kv_cache(kv, inference_params)
             k, v = kv.unbind(dim=-3)
+            k = torch.repeat_interleave(k, dim=2, repeats=self.num_heads // self.num_heads_kv)
+            v = torch.repeat_interleave(v, dim=2, repeats=self.num_heads // self.num_heads_kv)
             return F.scaled_dot_product_attention(
                 q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), is_causal=self.causal, scale=self.softmax_scale
             ).transpose(1, 2)
@@ -275,6 +277,8 @@ def forward(self, x, inference_params=None):
                 )
             if inference_params is None:
                 k, v = kv.unbind(dim=-3)
+                k = torch.repeat_interleave(k, dim=2, repeats=self.num_heads // self.num_heads_kv)
+                v = torch.repeat_interleave(v, dim=2, repeats=self.num_heads // self.num_heads_kv)
                 context = F.scaled_dot_product_attention(
                     q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), is_causal=self.causal, scale=self.softmax_scale
                 ).transpose(1, 2)