diff --git a/lmdeploy/turbomind/deploy/module.py b/lmdeploy/turbomind/deploy/module.py index a349e85b5..f2a6e4078 100644 --- a/lmdeploy/turbomind/deploy/module.py +++ b/lmdeploy/turbomind/deploy/module.py @@ -229,6 +229,8 @@ def _export(self, idx: int, qkvo, kind: str, pack_fn, **kwargs): if is_lora_a: qkv, o = map(transpose, qkvo) else: + if self.model.repeat_kv: + qkvo = self._repeat_kv(qkvo, kind) qkv, o = self._reorder_and_merge(qkvo) self.model.save_split(pack_fn(qkv), self._attn.format(idx, 'w_qkv', kind),