fix

Author: wangzaijun

lightllm/models/bloom/layer_infer/transformer_layer_infer.py

@@ -3,7 +3,7 @@
 import torch.functional as F
 import torch.distributed as dist
 import numpy as np
-
+from typing import Tuple
 from lightllm.common.basemodel import TransformerLayerInferTpl
 from lightllm.models.bloom.layer_weights.transformer_layer_weight import BloomTransformerLayerWeight
 from lightllm.models.bloom.triton_kernel.context_flashattention_nopad import context_attention_fwd
@@ -43,7 +43,9 @@ def _ffn_norm(self, input, infer_state: InferStateInfo, layer_weight: BloomTrans
             eps=self.eps_,
         )
 
-    def _get_qkv(self, input, infer_state: InferStateInfo, layer_weight: BloomTransformerLayerWeight) -> torch.Tensor:
+    def _get_qkv(
+        self, input, infer_state: InferStateInfo, layer_weight: BloomTransformerLayerWeight
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
         q = layer_weight.q_proj.mm(input.view(-1, self.embed_dim_))
         cache_kv = self._pre_cache_kv(infer_state=infer_state, layer_weight=layer_weight)
         cache_kv = layer_weight.kv_proj.mm(