update

lmdeploy/turbomind/deploy/module.py

@@ -191,7 +191,7 @@ def __init__(self, model: BaseOutputModel):
         self.attn_bias = model.model_config.attn_bias
 
     def _reorder_and_merge(self, qkvo):
-        q, k, v, o = map(transpose, qkvo)
+        q, k, v, o = qkvo
         # reorder output dim for tm's rotary embedding layout
         if self.model.permute_qk:
             q = permute_v2(q, self.head_dim)
@@ -210,8 +210,9 @@ def _repeat_kv(self, qkvo, kind: str):
 
         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)
+            x = x.t().reshape(-1, head_dim, dim)
+            x = x.repeat(1, self.model.repeat_kv, 1)
+            x = x.reshape(-1, dim).t()
             return x
 
         k, v = map(_repeat, (k, v))
@@ -229,6 +230,7 @@ def _export(self, idx: int, qkvo, kind: str, pack_fn, **kwargs):
         if is_lora_a:
             qkv, o = map(transpose, qkvo)
         else:
+            qkvo = tuple(map(transpose, qkvo))
             if self.model.repeat_kv:
                 qkvo = self._repeat_kv(qkvo, kind)
             qkv, o = self._reorder_and_merge(qkvo)