handle code clean

Signed-off-by: weijinqian_v1 <weijinqian@huawei.com>

Author: weijinqian_v1

tests/ut/test_distributed_tensor_parallel.py

@@ -1,14 +1,16 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 # Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
-import pytest
-import torch
 import importlib
 from unittest.mock import MagicMock, patch
+
+import pytest
+import torch
+
 from vllm_ascend.distributed.tensor_parallel import (
     _gather_along_first_dim, _gather_along_last_dim,
     _reduce_scatter_along_first_dim, _reduce_scatter_along_last_dim,
-    all_to_all_sp2hp, all_to_all_hp2sp)
+    all_to_all_hp2sp, all_to_all_sp2hp)
 
 
 # 测试用的固定数据

tests/ut/test_moe_util.py

@@ -1,12 +1,17 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 # Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
-import torch
-import pytest
 import math
+
+import pytest
+import torch
+
+from vllm_ascend.ops.moe_dispatcher.moe_utils import (
+    get_capacity, group_limited_topk, permute, sort_chunks_by_idxs,
+    topk_softmax_with_capacity, unpermute)
+
 import vllm_ascend.patch.worker.patch_common.patch_utils # type: ignore[import]  # isort: skip  # noqa
 
-from vllm_ascend.ops.moe_dispatcher.moe_utils import permute, get_capacity, topk_softmax_with_capacity, group_limited_topk, unpermute, sort_chunks_by_idxs
 
 
 class TestMoeUtils:

tests/ut/test_token_dispatcher.py

@@ -2,14 +2,17 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 # Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
 
-import torch
 import pytest
-import vllm_ascend.patch.worker.patch_common.patch_utils # type: ignore[import]  # isort: skip  # noqa
-
+import torch
 from pytest_mock import MockerFixture
+
+from vllm_ascend.ops.moe_dispatcher.token_dispatcher import (
+    MoEAlltoAllSeqOverLapDispatcher, MoeDispatcherConfig)
 from vllm_ascend.utils import adapt_patch  # noqa E402
 
-from vllm_ascend.ops.moe_dispatcher.token_dispatcher import MoeDispatcherConfig, MoEAlltoAllSeqOverLapDispatcher
+import vllm_ascend.patch.worker.patch_common.patch_utils # type: ignore[import]  # isort: skip  # noqa
+
+
 
 adapt_patch(True)
 

vllm_ascend/models/__init__.py

@@ -8,10 +8,10 @@ def register_model():
     from .deepseek_mtp import CustomDeepSeekMTP  # noqa: F401
     from .deepseek_v2 import CustomDeepseekV2ForCausalLM  # noqa: F401
     from .deepseek_v2 import CustomDeepseekV3ForCausalLM  # noqa: F401
+    from .moe_block import AscendSparseMoeBlock  # noqa: F401
     from .qwen2_5_vl import \
         AscendQwen2_5_VLForConditionalGeneration  # noqa: F401
     from .qwen2_vl import AscendQwen2VLForConditionalGeneration  # noqa: F401
-    from .moe_block import AscendSparseMoeBlock # noqa: F401
     from .qwen3 import CustomQwen3ForCausalLM  # noqa: F401
 
     ModelRegistry.register_model(

vllm_ascend/models/deepseek_dbo.py

@@ -34,8 +34,7 @@
 from transformers import PretrainedConfig
 from vllm.attention import AttentionMetadata
 from vllm.config import CacheConfig, ModelConfig, VllmConfig
-from vllm.distributed import (get_pp_group,
-                              get_tensor_model_parallel_rank,
+from vllm.distributed import (get_pp_group, get_tensor_model_parallel_rank,
                               get_tensor_model_parallel_world_size,
                               get_tp_group, tensor_model_parallel_all_reduce)
 from vllm.distributed.parallel_state import get_dp_group, get_ep_group
@@ -55,8 +54,9 @@
 from vllm.sequence import IntermediateTensors
 
 import vllm_ascend.envs as envs_ascend
-from vllm_ascend.distributed.tensor_parallel import gather_from_sequence_parallel_region
 from vllm_ascend.ascend_forward_context import FusedMoEState
+from vllm_ascend.distributed.tensor_parallel import \
+    gather_from_sequence_parallel_region
 from vllm_ascend.models.deepseek_v2 import (CustomDeepseekV2DecoderLayer,
                                             CustomDeepseekV2MLP,
                                             CustomDeepseekV2MoE)
@@ -69,9 +69,9 @@
 from vllm_ascend.multistream.metadata import (MultiStreamConfig,
                                               MultiStreamStepMetadata,
                                               make_multistream_metadata_ds)
+from vllm_ascend.ops.fused_moe import select_experts
 from vllm_ascend.quantization.w8a8_dynamic import (
     AscendW8A8DynamicLinearMethod, apply_mlp)
-from vllm_ascend.ops.fused_moe import select_experts
 from vllm_ascend.utils import dispose_tensor
 
 VLLM_ASCEND_ENABLE_DBO: bool = envs_ascend.VLLM_ASCEND_ENABLE_DBO

vllm_ascend/models/moe_block.py

@@ -19,22 +19,18 @@
 
 import torch
 import vllm.model_executor.models.qwen3_moe as qwen3
-
 from torch import nn
+from transformers import PretrainedConfig
 from vllm.attention import AttentionMetadata
-from vllm.distributed import (get_tensor_model_parallel_world_size,
-                              get_tp_group)
-from vllm.distributed.parallel_state import get_dp_group
+from vllm.distributed import get_tensor_model_parallel_world_size, get_tp_group
+from vllm.distributed.parallel_state import get_dp_group, get_ep_group
 from vllm.forward_context import get_forward_context
 from vllm.model_executor.layers.linear import ReplicatedLinear
+from vllm.model_executor.layers.quantization import QuantizationConfig
 
 from vllm_ascend.ascend_config import get_ascend_config
-from vllm.distributed.parallel_state import get_ep_group
 from vllm_ascend.ops.fused_moe import AscendFusedMoE
 
-from transformers import PretrainedConfig
-from vllm.model_executor.layers.quantization import QuantizationConfig
-
 
 class AscendSparseMoeBlock(nn.Module):
 

vllm_ascend/models/qwen3_dbo.py

@@ -21,42 +21,44 @@
 # limitations under the License.
 # # Adapted from
 # """Inference-only Qwen3 model."""
-from typing import Optional, Union, List
 from types import SimpleNamespace
+from typing import List, Optional, Union
 
 import torch
 import torch_npu
 from torch import nn
 from transformers import PretrainedConfig
-
-from vllm.model_executor.models.qwen3_moe import Qwen3MoeDecoderLayer, Qwen3MoeModel
-from vllm.config import CacheConfig, VllmConfig
-from vllm.model_executor.layers.quantization import QuantizationConfig
 from vllm.attention import AttentionMetadata
+from vllm.compilation.decorators import support_torch_compile
+from vllm.config import CacheConfig, VllmConfig
+from vllm.distributed import (get_pp_group,
+                              get_tensor_model_parallel_world_size,
+                              get_tp_group)
 from vllm.forward_context import get_forward_context, set_forward_context
-from vllm.distributed import get_tensor_model_parallel_world_size, get_tp_group, \
-    get_pp_group
-from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
+from vllm.model_executor.layers.quantization import QuantizationConfig
+from vllm.model_executor.layers.vocab_parallel_embedding import (
+    ParallelLMHead, VocabParallelEmbedding)
+from vllm.model_executor.models.qwen3_moe import (Qwen3MoeDecoderLayer,
+                                                  Qwen3MoeForCausalLM,
+                                                  Qwen3MoeModel)
 from vllm.model_executor.models.utils import (
     make_empty_intermediate_tensors_factory, make_layers, maybe_prefix)
-from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.sequence import IntermediateTensors
-from vllm.model_executor.models.qwen3_moe import Qwen3MoeForCausalLM
-from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
-from vllm.model_executor.layers.logits_processor import LogitsProcessor
-from vllm.compilation.decorators import support_torch_compile
 
+import vllm_ascend.envs as envs_ascend
+from vllm_ascend.distributed.tensor_parallel import \
+    gather_from_sequence_parallel_region
+from vllm_ascend.multistream.base import MSEventKey
 from vllm_ascend.multistream.context import (
     advance_step_multistream_layer_context, get_multistream_layer_context)
-from vllm_ascend.multistream.base import MSEventKey
 from vllm_ascend.multistream.layers import (MultiStreamPostTransformerLayer,
                                             MultiStreamPreTransformerLayer)
 from vllm_ascend.multistream.metadata import (MultiStreamConfig,
                                               MultiStreamStepMetadata,
                                               make_multistream_metadata_ds)
-from vllm_ascend.ops.fused_moe import select_experts, apply_mlp
-from vllm_ascend.distributed.tensor_parallel import gather_from_sequence_parallel_region
-import vllm_ascend.envs as envs_ascend
+from vllm_ascend.ops.fused_moe import apply_mlp, select_experts
 
 VLLM_ASCEND_ENABLE_DBO: bool = envs_ascend.VLLM_ASCEND_ENABLE_DBO
 

vllm_ascend/multistream/ms_split.py

@@ -4,7 +4,8 @@
 import numpy as np
 import torch
 
-from vllm_ascend.attention.attention_v1 import AscendAttentionState, AscendMetadata
+from vllm_ascend.attention.attention_v1 import (AscendAttentionState,
+                                                AscendMetadata)
 
 from .base import MSAttentionMetadataSplitConfig
 

vllm_ascend/ops/fused_moe.py

@@ -40,11 +40,11 @@
 from vllm_ascend.ascend_config import get_ascend_config
 from vllm_ascend.ascend_forward_context import FusedMoEState
 from vllm_ascend.ops.expert_load_balancer import ExpertLoadBalancer
+from vllm_ascend.ops.moe_dispatcher.token_dispatcher import (
+    MoEAlltoAllSeqOverLapDispatcher, MoeDispatcherConfig)
 from vllm_ascend.utils import (AscendSocVersion, dispose_tensor,
                                get_ascend_soc_version, npu_stream_switch,
                                npu_wait_tensor)
-from vllm_ascend.ops.moe_dispatcher.token_dispatcher import (
-    MoEAlltoAllSeqOverLapDispatcher, MoeDispatcherConfig)
 
 VLLM_ASCEND_MOE_ALL2ALL_BUFFER: bool = envs_ascend.VLLM_ASCEND_MOE_ALL2ALL_BUFFER
 

vllm_ascend/ops/moe_dispatcher/token_dispatcher.py

@@ -24,16 +24,16 @@
 
 import torch
 import torch_npu
-
 from vllm.distributed.parallel_state import get_ep_group
+
 from vllm_ascend.distributed.tensor_parallel import (
     all_gather_last_dim_from_tensor_parallel_region, all_to_all_hp2sp,
     all_to_all_sp2hp, gather_from_sequence_parallel_region,
     reduce_scatter_last_dim_to_tensor_parallel_region)
 from vllm_ascend.ops.comm_utils import async_all_to_all
 from vllm_ascend.ops.moe_dispatcher.moe_utils import (
-    get_capacity, permute, topk_softmax_with_capacity,
-    unpermute)
+    get_capacity, permute, topk_softmax_with_capacity, unpermute)
+
 """ We use the following notation throughout this file:
      H: hidden size
      B: micro batch size