replicate kv for some models when tp is divisble by kv_head_num

lmdeploy/turbomind/deploy/module.py

@@ -202,6 +202,26 @@ def _reorder_and_merge(self, qkvo):
             o = torch.zeros_like(q)
         return qkv, o
 
+    def _repeat_kv(self, qkvo, kind: str):
+        """replicate kv."""
+        q, k, v, o = qkvo
+        head_dim = self.model.model_config.size_per_head
+        hidden_dim = self.model.model_config.hidden_units
+
+        def _repeat(x):
+            dim = hidden_dim if kind != 'bias' else 1
+            x = x.view(-1, head_dim, dim).repeat(1, self.model.repeat_kv, 1)
+            x = x.reshape(-1, dim)
+            return x
+
+        k, v = map(_repeat, (k, v))
+        if kind == 'bias':
+            if o is None:
+                o = torch.zeros(hidden_dim, dtype=q.dtype, device=q.device)
+            q, k, v, o = map(torch.squeeze, (q, k, v, o))
+
+        return (q, k, v, o)
+
     def _export(self, idx: int, qkvo, kind: str, pack_fn, **kwargs):
         if all(x is None for x in qkvo):
             return

lmdeploy/turbomind/deploy/target_model/base.py

@@ -78,6 +78,17 @@ def __init__(self,
             self.model_config.expert_inter_size = _pad_inter_size(
                 self.model_config.expert_inter_size,
                 self.model_config.group_size, self.tensor_para_size)
+
+        # head_num is divisble by tp but kv_head_num is not
+        # and tp is divisble by kv_head_num
+        assert self.model_config.head_num % self.tensor_para_size == 0
+        self.repeat_kv = 0
+        if (self.tensor_para_size > self.model_config.kv_head_num and
+                self.tensor_para_size % self.model_config.kv_head_num == 0):
+            self.repeat_kv = (self.tensor_para_size //
+                              self.model_config.kv_head_num)
+            self.model_config.kv_head_num = self.tensor_para_size
+
         self.model_config.verify()
         assert self.model_config.kv_head_num % self.tensor_para_size == 0