fix bug

Signed-off-by: zzhx1 <zzh_201018@outlook.com>

Author: zzhx1

vllm_ascend/ascend_config.py

@@ -57,9 +57,6 @@ def __init__(self, vllm_config):
             assert(
                 self.torchair_graph_config.enabled == True
             ), "lmhead_tensor_parallel_size is only supported in graph mode"
-            assert(
-                vllm_config.kv_transfer_config is not None and vllm_config.kv_transfer_config.is_kv_consumer
-            ),"lmhead_tensor_parallel_size is only supported in pd scenario and can only be used in D node."
 
 
 class TorchairGraphConfig:

vllm_ascend/ops/vocab_parallel_embedding.py

@@ -41,7 +41,7 @@
     _prune_hidden_states
 )
 from vllm.model_executor.parameter import BasevLLMParameter
-from vllm.model_executor.utils import set_weight_attrs, _enable_lmhead_tp
+from vllm.model_executor.utils import set_weight_attrs
 from vllm.model_executor.sampling_metadata import SamplingMetadata
 from vllm.model_executor.layers.quantization.base_config import (
     QuantizationConfig,
@@ -51,7 +51,7 @@
 
 from vllm_ascend.distributed.parallel_state import get_lmheadtp_group
 from vllm_ascend.ascend_config import get_ascend_config
-
+from vllm_ascend.utils import _enable_lmhead_tp
 
 def get_masked_input_and_mask(
         input_: torch.Tensor, org_vocab_start_index: int,

vllm_ascend/worker/model_runner_v1.py

@@ -57,7 +57,7 @@
 from vllm.sampling_params import SamplingType
 from vllm.sequence import IntermediateTensors
 from vllm.utils import (STR_DTYPE_TO_TORCH_DTYPE, DeviceMemoryProfiler,
-                        LazyLoader, cdiv, _enable_lmhead_tp)
+                        LazyLoader, cdiv)
 from vllm.v1.kv_cache_interface import (FullAttentionSpec, KVCacheConfig,
                                         KVCacheSpec)
 from vllm.v1.outputs import (EMPTY_MODEL_RUNNER_OUTPUT, LogprobsTensors,
@@ -87,7 +87,7 @@
 from vllm_ascend.utils import (ACL_FORMAT_FRACTAL_ND, ACL_FORMAT_FRACTAL_NZ,
                                ProfileExecuteDuration, is_310p,
                                maybe_converting_weight_acl_format,
-                               vllm_version_is)
+                               vllm_version_is, _enable_lmhead_tp)
 from vllm_ascend.worker.eagle_proposer_v1 import EagleProposer
 from vllm_ascend.worker.mtp_proposer_v1 import MtpProposer
 from vllm_ascend.worker.npu_input_batch import CachedRequestState, InputBatch