[CI] Fix

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

Author: zzhx1

vllm_ascend/ascend_config.py

@@ -50,11 +50,13 @@ def __init__(self, vllm_config):
         self.lmhead_tensor_parallel_size = additional_config.get(
             "lmhead_tensor_parallel_size", None)
         if self.lmhead_tensor_parallel_size is not None:
-            logger.info(f"Enable lmhead_tensor_parallel_size={self.lmhead_tensor_parallel_size} in pure DP scenario")
-            assert(
+            logger.info(
+                f"Enable lmhead_tensor_parallel_size={self.lmhead_tensor_parallel_size} in pure DP scenario"
+            )
+            assert (
                 vllm_config.parallel_config.tensor_parallel_size == 1
-            ),"lmhead_tensor_parallel_size is only supported in the pure DP scenario"
-            assert(
+            ), "lmhead_tensor_parallel_size is only supported in the pure DP scenario"
+            assert (
                 self.torchair_graph_config.enabled == True
             ), "lmhead_tensor_parallel_size is only supported in graph mode"
 

vllm_ascend/distributed/parallel_state.py

@@ -9,15 +9,18 @@
 _MC2: Optional[GroupCoordinator] = None
 _LMTP: Optional[GroupCoordinator] = None
 
+
 def get_mc2_group() -> GroupCoordinator:
     assert _MC2 is not None, ("mc2 group is not initialized")
     return _MC2
 
+
 def get_lmheadtp_group() -> GroupCoordinator:
     assert _LMTP is not None, (
         "lm head tensor parallel group is not initialized")
     return _LMTP
 
+
 def model_parallel_initialized():
     return (_MC2 is not None)
 
@@ -43,22 +46,23 @@ def init_ascend_model_parallel(parallel_config: ParallelConfig, ):
                                      get_world_group().local_rank,
                                      backend,
                                      group_name="mc2")
-    
+
     lmhead_tensor_parallel_size = parallel_config.lmhead_tensor_parallel_size
     if lmhead_tensor_parallel_size is not None:
         group_ranks = []
         global _LMTP
         num_lmhead_tensor_parallel_groups: int = (world_size //
-                                                lmhead_tensor_parallel_size)
+                                                  lmhead_tensor_parallel_size)
         for i in range(num_lmhead_tensor_parallel_groups):
             ranks = list(
                 range(i * lmhead_tensor_parallel_size,
-                    (i + 1) * lmhead_tensor_parallel_size))
+                      (i + 1) * lmhead_tensor_parallel_size))
             group_ranks.append(ranks)
         _LMTP = init_model_parallel_group(group_ranks,
-                                        get_world_group().local_rank,
-                                        backend,
-                                        group_name="lmheadtp")
+                                          get_world_group().local_rank,
+                                          backend,
+                                          group_name="lmheadtp")
+
 
 def destroy_ascend_model_parallel():
     global _MC2

vllm_ascend/models/deepseek_v2.py

@@ -49,12 +49,11 @@
                                                ReplicatedLinear,
                                                RowParallelLinear,
                                                UnquantizedLinearMethod)
-from vllm.model_executor.layers.logits_processor import LogitsProcessor
 from vllm.model_executor.layers.quantization import QuantizationConfig
 from vllm.model_executor.layers.rotary_embedding import get_rope
 from vllm.model_executor.layers.sampler import get_sampler
-from vllm.model_executor.layers.vocab_parallel_embedding import (
-    ParallelLMHead, VocabParallelEmbedding)
+from vllm.model_executor.layers.vocab_parallel_embedding import \
+    VocabParallelEmbedding
 from vllm.model_executor.model_loader.weight_utils import (
     default_weight_loader, maybe_remap_kv_scale_name)
 from vllm.model_executor.models.deepseek_v2 import \
@@ -68,14 +67,15 @@
     PPMissingLayer, is_pp_missing_parameter,
     make_empty_intermediate_tensors_factory, make_layers, maybe_prefix)
 from vllm.sequence import IntermediateTensors
-from vllm.model_executor.sampling_metadata import SamplingMetadata
 
 from vllm_ascend.ascend_config import get_ascend_config
 from vllm_ascend.ops.fused_moe import AscendFusedMoE
+from vllm_ascend.ops.vocab_parallel_embedding import (CustomLogitsProcessor,
+                                                      CustomParallelLMHead)
 from vllm_ascend.quantization.quant_config import AscendLinearMethod
 from vllm_ascend.quantization.w8a8_dynamic import AscendW8A8DynamicLinearMethod
 from vllm_ascend.utils import dispose_tensor, npu_prefetch
-from vllm_ascend.ops.vocab_parallel_embedding import CustomParallelLMHead, CustomLogitsProcessor
+
 
 class CustomDeepseekV2SiluAndMul(SiluAndMul):
 
@@ -872,7 +872,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
                                                 config.hidden_size,
                                                 quant_config=quant_config,
                                                 prefix=maybe_prefix(
-                                                prefix, "lm_head"))
+                                                    prefix, "lm_head"))
         else:
             self.lm_head = PPMissingLayer()
         self.logits_processor = CustomLogitsProcessor(config.vocab_size)

vllm_ascend/ops/vocab_parallel_embedding.py

@@ -19,40 +19,24 @@
 
 import torch
 from torch.nn import Module
-import torch.distributed as dist
-from torch.nn.parameter import Parameter, UninitializedParameter
-
-from vllm.distributed import (
-    divide,
-    get_tensor_model_parallel_rank,
-    get_tensor_model_parallel_world_size,
-    tensor_model_parallel_all_reduce
-)
-from vllm.model_executor.layers.vocab_parallel_embedding import (
-    VocabParallelEmbedding,
-    DEFAULT_VOCAB_PADDING_SIZE,
-    pad_vocab_size,
-    UnquantizedEmbeddingMethod,
-    ParallelLMHead
-)
+from torch.nn.parameter import Parameter
+from vllm.distributed import (divide, get_tensor_model_parallel_rank,
+                              get_tensor_model_parallel_world_size,
+                              tensor_model_parallel_all_reduce)
 from vllm.model_executor.layers.logits_processor import (
-    LogitsProcessor,
-    _apply_logits_processors,
-    _prune_hidden_states
-)
-from vllm.model_executor.parameter import BasevLLMParameter
-from vllm.model_executor.utils import set_weight_attrs
-from vllm.model_executor.sampling_metadata import SamplingMetadata
+    LogitsProcessor, _apply_logits_processors, _prune_hidden_states)
 from vllm.model_executor.layers.quantization.base_config import (
-    QuantizationConfig,
-    QuantizeMethodBase,
-    method_has_implemented_embedding
-)
+    QuantizationConfig, QuantizeMethodBase, method_has_implemented_embedding)
+from vllm.model_executor.layers.vocab_parallel_embedding import (
+    DEFAULT_VOCAB_PADDING_SIZE, ParallelLMHead, UnquantizedEmbeddingMethod,
+    VocabParallelEmbedding, pad_vocab_size)
+from vllm.model_executor.sampling_metadata import SamplingMetadata
+from vllm.model_executor.utils import set_weight_attrs
 
 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,
         org_vocab_end_index: int, num_org_vocab_padding: int,
@@ -105,7 +89,6 @@ def vocab_parallel_embedding_forward(self, input_):
 
 
 class CustomParallelLMHead(ParallelLMHead):
-    
     """Costom Parallelized LM head, added the feature of lmheadTP in pure dp scenario
     
     Output logits weight matrices used in the Sampler. The weight and bias
@@ -120,6 +103,7 @@ class CustomParallelLMHead(ParallelLMHead):
         org_num_embeddings: original vocabulary size (without LoRA).
         padding_size: padding size for the vocabulary.
     """
+
     def __init__(self,
                  num_embeddings: int,
                  embedding_dim: int,
@@ -128,7 +112,7 @@ def __init__(self,
                  org_num_embeddings: Optional[int] = None,
                  padding_size: int = DEFAULT_VOCAB_PADDING_SIZE,
                  quant_config: Optional[QuantizationConfig] = None,
-                 prefix: str = ""): 
+                 prefix: str = ""):
         Module.__init__(self)
 
         if _enable_lmhead_tp():
@@ -137,7 +121,7 @@ def __init__(self,
         else:
             tp_rank = get_tensor_model_parallel_rank()
             self.tp_size = get_tensor_model_parallel_world_size()
-        
+
         self.num_embeddings = num_embeddings
         self.padding_size = padding_size
         self.org_vocab_size = org_num_embeddings or num_embeddings
@@ -197,7 +181,7 @@ def __init__(self,
                                          self.num_embeddings_padded,
                                          params_dtype=params_dtype,
                                          weight_loader=self.weight_loader)
-        
+
         self.quant_config = quant_config
         if bias:
             self.bias = Parameter(
@@ -209,25 +193,24 @@ def __init__(self,
             })
         else:
             self.register_parameter("bias", None)
-        
+
+
 class CustomLogitsProcessor(LogitsProcessor):
     """Custom logits processor extending base LogitsProcessor functionality.
     Added the feature of lmheadTP in pure dp scenario
     """
-    
+
     def __init__(self,
                  vocab_size: int,
                  org_vocab_size: Optional[int] = None,
                  scale: float = 1.0,
                  logits_as_input: bool = False,
                  soft_cap: Optional[float] = None) -> None:
-        super().__init__(
-            vocab_size=vocab_size,
-            org_vocab_size=org_vocab_size,
-            scale=scale,
-            logits_as_input=logits_as_input,
-            soft_cap=soft_cap
-        )
+        super().__init__(vocab_size=vocab_size,
+                         org_vocab_size=org_vocab_size,
+                         scale=scale,
+                         logits_as_input=logits_as_input,
+                         soft_cap=soft_cap)
 
     def forward(
         self,
@@ -258,15 +241,15 @@ def forward(
             if sampling_metadata is not None and \
                 sampling_metadata.seq_groups is not None:
                 logits = _apply_logits_processors(logits, sampling_metadata)
-        
+
         return logits
 
     def _get_logits(
-            self,
-            hidden_states: torch.Tensor,
-            lm_head: CustomParallelLMHead,
-            embedding_bias: Optional[torch.Tensor],
-        ) -> Optional[torch.Tensor]:
+        self,
+        hidden_states: torch.Tensor,
+        lm_head: CustomParallelLMHead,
+        embedding_bias: Optional[torch.Tensor],
+    ) -> Optional[torch.Tensor]:
         """
         Compute logits for next token prediction using parallel processing.
         
@@ -281,16 +264,15 @@ def _get_logits(
 
         if _enable_lmhead_tp():
             # Gather hidden states from all devices in tensor parallel group
-            gathered_hidden_states = get_lmheadtp_group().all_gather(hidden_states, dim=0)
+            gathered_hidden_states = get_lmheadtp_group().all_gather(
+                hidden_states, dim=0)
         else:
             gathered_hidden_states = hidden_states
 
         # Compute logits using quantized matrix multiplication
-        local_logits = lm_head.quant_method.apply(
-            lm_head,
-            gathered_hidden_states,
-            bias=embedding_bias
-        )
+        local_logits = lm_head.quant_method.apply(lm_head,
+                                                  gathered_hidden_states,
+                                                  bias=embedding_bias)
 
         if _enable_lmhead_tp():
             logits = get_lmheadtp_group().all_to_all(local_logits)
@@ -301,6 +283,5 @@ def _get_logits(
         # Remove paddings in vocab (if any)
         if logits is not None:
             logits = logits[..., :self.org_vocab_size]
-            
+
         return logits
-    

vllm_ascend/platform.py

@@ -134,8 +134,7 @@ def check_and_update_config(cls, vllm_config: VllmConfig) -> None:
         if parallel_config:
             # assign lmhead tensor parallel size
             parallel_config.lmhead_tensor_parallel_size = (
-                ascend_config.lmhead_tensor_parallel_size
-            )
+                ascend_config.lmhead_tensor_parallel_size)
 
         if model_config is None:
             logger.warning("Model config is missing. This may indicate "

vllm_ascend/worker/model_runner_v1.py

@@ -85,14 +85,13 @@
 from vllm_ascend.torchair.utils import (check_torchair_cache_exist,
                                         write_kv_cache_bytes_to_file)
 from vllm_ascend.utils import (ACL_FORMAT_FRACTAL_ND, ACL_FORMAT_FRACTAL_NZ,
-                               ProfileExecuteDuration, is_310p,
-                               maybe_converting_weight_acl_format,
-                               vllm_version_is, _enable_lmhead_tp)
+                               ProfileExecuteDuration, _enable_lmhead_tp,
+                               is_310p, maybe_converting_weight_acl_format,
+                               vllm_version_is)
 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
 
-
 if not vllm_version_is("0.10.0"):
     from vllm.tasks import GenerationTask, SupportedTask
     from vllm.v1.worker.kv_connector_model_runner_mixin import \
@@ -1334,8 +1333,8 @@ def _process_reqs(
         aux_hidden_states = None
         if self.use_aux_hidden_state_outputs:
             hidden_states, aux_hidden_states = hidden_states
-        
-        if _enable_lmhead_tp(): # 
+
+        if _enable_lmhead_tp():  #
             if not with_prefill:
                 max_num_reqs_across_dp = padded_num_tokens_across_dp
             else:
@@ -1998,7 +1997,7 @@ def _dummy_run(
                     if self.use_spec_decode and isinstance(
                             self.drafter, EagleProposer):
                         self.drafter.dummy_run(num_tokens)
-            
+
             if _enable_lmhead_tp() and not self.in_profile_run:
                 if not with_prefill:
                     max_num_reqs_across_dp = num_reqs
@@ -2008,7 +2007,7 @@ def _dummy_run(
                                             device=hidden_states.device,
                                             dtype=torch.int32)
                 model.compute_logits(hidden_states[dummy_indices], None)
-            
+
             if self.speculative_config and self.speculative_config.method == "deepseek_mtp":
                 assert isinstance(self.drafter, MtpProposer)
                 self.drafter.dummy_run(