fix

Author: hiworldwzj

lightllm/models/llama/layer_infer/post_layer_infer.py

@@ -12,7 +12,6 @@
 from lightllm.common.basemodel import PostLayerInferTpl
 from lightllm.utils.infer_utils import mark_cost_time
 from lightllm.distributed.communication_op import all_gather
-from lightllm.utils.envs_utils import get_env_start_args
 
 
 class LlamaPostLayerInfer(PostLayerInferTpl):

lightllm/models/qwen3_moe/layer_infer/transformer_layer_infer.py

@@ -120,6 +120,11 @@ def _tpsp_get_qkv(
             infer_state.position_cos,
             infer_state.position_sin,
         )
+
+        if infer_state.need_dp_prefill_balance:
+            q = infer_state._all_to_all_unbalance_get(data=q)
+            cache_kv = infer_state._all_to_all_unbalance_get(data=cache_kv)
+        
         return q, cache_kv
 
     def _moe_ffn(

lightllm/models/stablelm/layer_infer/transformer_layer_infer.py

@@ -3,7 +3,7 @@
 import torch.distributed as dist
 import numpy as np
 from functools import partial
-
+from typing import Tuple
 from lightllm.models.llama.triton_kernel.rotary_emb import rotary_emb_fwd
 from lightllm.models.bloom.triton_kernel.layernorm import layernorm_forward
 from lightllm.models.stablelm.layer_weights.transformer_layer_weight import StablelmTransformerLayerWeight
@@ -38,6 +38,10 @@ def _get_qkv(
         )
         return q, cache_kv
 
+    def _tpsp_get_qkv(self, input, infer_state, layer_weight) -> Tuple[torch.Tensor, torch.Tensor]:
+        # TODO
+        raise Exception("not impl")
+
     def _get_o(
         self, input, infer_state: LlamaInferStateInfo, layer_weight: StablelmTransformerLayerWeight
     ) -> torch.Tensor:
@@ -46,6 +50,11 @@ def _get_o(
         )
         return o_tensor
 
+    def _tpsp_get_o(self, input, infer_state, layer_weight) -> Tuple[torch.Tensor, torch.Tensor]:
+        # TODO
+        raise Exception("not impl")
+
+
     def _att_norm(
         self, input, infer_state: LlamaInferStateInfo, layer_weight: StablelmTransformerLayerWeight
     ) -> torch.Tensor: