Fix isort and yapf

root · CalvinXKY · clrs97 · commit 3b372e4a72c0 · 2025-09-03T14:55:33.000+08:00
Co-authored-by: CalvinXKY &lt;kyxiezju@163.com&gt;
Signed-off-by: clrs97 &lt;524936896@qq.com&gt;
diff --git a/vllm_ascend/ascend_config.py b/vllm_ascend/ascend_config.py
@@ -63,9 +63,12 @@ def __init__(self, vllm_config):
                 )
 
         self.enable_mla_sp = additional_config.get("enable_mla_sp", False)
-        self.o_shard_parallel_size = int(additional_config.get("o_shard_parallel_size", -1))
+        self.o_shard_parallel_size = int(
+            additional_config.get("o_shard_parallel_size", -1))
         self.enable_o_shard = self.o_shard_parallel_size > 0
-        self.o_shard_full_layers = int(additional_config.get("o_shard_full_layers", 0))
+        self.o_shard_full_layers = int(
+            additional_config.get("o_shard_full_layers", 0))
+
 
 class TorchairGraphConfig:
     """
diff --git a/vllm_ascend/attention/mla_v1.py b/vllm_ascend/attention/mla_v1.py
@@ -2,13 +2,15 @@
 from typing import TYPE_CHECKING, NamedTuple, Optional, Tuple, Type, TypeVar
 
 import torch
+import torch.distributed as dist
 import torch_npu
 from torch import nn
 from vllm.attention.backends.abstract import (AttentionBackend,
                                               AttentionMetadata,
                                               MLAAttentionImpl)
 from vllm.config import VllmConfig, get_current_vllm_config
 from vllm.distributed import get_tensor_model_parallel_world_size, get_tp_group
+from vllm.distributed.parallel_state import get_dp_group
 from vllm.model_executor.layers.linear import (LinearBase,
                                                UnquantizedLinearMethod)
 from vllm.utils import cdiv, round_down
@@ -17,19 +19,16 @@
 from vllm_ascend.attention.attention_v1 import AscendAttentionState
 from vllm_ascend.attention.utils import (AscendCommonAttentionMetadata,
                                          split_decodes_and_prefills)
+from vllm_ascend.distributed.parallel_state import get_mla_sp_world_group
+from vllm_ascend.mla_sp_context import get_sp_context
 from vllm_ascend.multistream.base import MSAttentionMetadataSplitConfig
 from vllm_ascend.multistream.context import get_multistream_comm_context
 from vllm_ascend.multistream.ms_split import model_input_split_v1_mla_attn
 from vllm_ascend.ops.attention import vanilla_chunked_prefill_mla
+from vllm_ascend.ops.shard import RowShardLinear
 from vllm_ascend.utils import npu_prefetch
 from vllm_ascend.worker.npu_input_batch import InputBatch
 
-import torch.distributed as dist
-from vllm.distributed.parallel_state import get_dp_group
-from vllm_ascend.distributed.parallel_state import get_mla_sp_world_group
-from vllm_ascend.mla_sp_context import get_sp_context
-from vllm_ascend.ops.shard import RowShardLinear
-
 if TYPE_CHECKING:
     from vllm.v1.core.sched.output import SchedulerOutput
 
@@ -968,7 +967,7 @@ def _mla_preprocess(self, hidden_states, kv_cache, attn_metadata,
                 prefill_q_nope, prefill_q_pe, prefill_k_nope, prefill_k_pe,
                 prefill_value)
         return decode_preprocess_res, prefill_preprocess_res
-    
+
     def _forward_prefill_sp(
         self,
         hidden_states: torch.Tensor,
@@ -982,10 +981,13 @@ def _forward_prefill_sp(
                      enabled=self.enable_prefetch)
         # Split inputs from local DP to each device.
         dp_sp_hidden_states = hidden_states
-        rank_sp_hidden_states = dp_sp_hidden_states[sp_context.my_rank_sp_start_token_within_dp:sp_context.my_rank_sp_end_token_within_dp]
+        rank_sp_hidden_states = dp_sp_hidden_states[
+            sp_context.my_rank_sp_start_token_within_dp:sp_context.
+            my_rank_sp_end_token_within_dp]
         sp_tokens = rank_sp_hidden_states.shape[0]
         if sp_tokens == 0:
-            rank_sp_hidden_states = nn.functional.pad(rank_sp_hidden_states, (0, 0, 0, 1))
+            rank_sp_hidden_states = nn.functional.pad(rank_sp_hidden_states,
+                                                      (0, 0, 0, 1))
             sp_tokens = 1
         # MLA prefill:
         # 1. Perform q_a_proj and q_a_layernorm to obtain q_c
@@ -997,19 +999,23 @@ def _forward_prefill_sp(
         # Rearrange down_proj outputs across DP.
         sp_output = torch.cat([sp_hidden_states_or_q_c, sp_kv_no_split], dim=1)
         if sp_tokens < sp_context.num_tokens_per_rank:
-            sp_output = nn.functional.pad(sp_output, (0, 0, 0, sp_context.num_tokens_per_rank - sp_tokens))
+            sp_output = nn.functional.pad(
+                sp_output,
+                (0, 0, 0, sp_context.num_tokens_per_rank - sp_tokens))
         global_sp_output = get_mla_sp_world_group().all_gather(sp_output, 0)
         my_dp = sp_context.my_dp
-        dp_output = global_sp_output[sp_context.start_token_of_dp[my_dp]:sp_context.end_token_of_dp[my_dp]]
-        prefill_q_c, prefill_kv_no_split = dp_output.split([self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim], dim=-1)
+        dp_output = global_sp_output[sp_context.start_token_of_dp[my_dp]:
+                                     sp_context.end_token_of_dp[my_dp]]
+        prefill_q_c, prefill_kv_no_split = dp_output.split(
+            [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim],
+            dim=-1)
 
         if attn_metadata is None:
             # Dummy run, just construct the attention outputs.
             output_prefill = torch.empty(
                 [prefill_q_c.shape[0], self.num_heads * self.v_head_dim],
                 dtype=hidden_states.dtype,
-                device=hidden_states.device
-            )
+                device=hidden_states.device)
         else:
             # Preprocess prefill tokens, write kv cache and get:
             # prefill_q_nope, prefill_q_pe, prefill_k_nope, prefill_k_pe, prefill_value
@@ -1025,7 +1031,7 @@ def _forward_prefill_sp(
             prefill_k_pe, prefill_k_c_normed = self.exec_kv_prefill(
                 prefill_kv_no_split, cos, sin, kv_cache, prefill_slots)
             prefill_k_pe = prefill_k_pe.view(prefill_q_c.shape[0],
-                                                self.num_kv_heads, -1)
+                                             self.num_kv_heads, -1)
             prefill_k_nope, prefill_value = self.kv_b_proj(
                 prefill_k_c_normed)[0].view(
                     -1, self.num_heads,
@@ -1034,10 +1040,11 @@ def _forward_prefill_sp(
             prefill_k_pe = prefill_k_pe.expand(
                 (*prefill_k_nope.shape[:-1], -1))
             # Attention outputs.
-            output_prefill = self._forward_prefill(
-                    prefill_q_nope, prefill_q_pe,
-                    prefill_k_nope, prefill_k_pe,
-                    prefill_value, kv_cache, attn_metadata)
+            output_prefill = self._forward_prefill(prefill_q_nope,
+                                                   prefill_q_pe,
+                                                   prefill_k_nope,
+                                                   prefill_k_pe, prefill_value,
+                                                   kv_cache, attn_metadata)
 
         # Rearrange attention outputs across DP to run SP.
         sp_world_group = get_mla_sp_world_group()
@@ -1050,33 +1057,38 @@ def _forward_prefill_sp(
         if get_dp_group().world_size == 1:
             padded_len = sp_context.num_tokens_per_rank * sp_world_group.world_size
             if sp_context.num_global_tokens < padded_len:
-                sp_send = nn.functional.pad(sp_send, (0, 0, 0, padded_len - sp_context.num_global_tokens))
-            sp_output = torch.empty(
-                [sp_context.num_tokens_per_rank * tp_size, self.num_heads * self.v_head_dim],
-                dtype=sp_send.dtype,
-                device=sp_send.device
-            )
+                sp_send = nn.functional.pad(
+                    sp_send,
+                    (0, 0, 0, padded_len - sp_context.num_global_tokens))
+            sp_output = torch.empty([
+                sp_context.num_tokens_per_rank * tp_size,
+                self.num_heads * self.v_head_dim
+            ],
+                                    dtype=sp_send.dtype,
+                                    device=sp_send.device)
             dist.all_to_all_single(
                 output=sp_output,
                 input=sp_send,
                 group=sp_world_group.device_group,
             )
-            sp_output = sp_output.reshape(sp_context.num_tokens_per_rank, tp_size * self.num_heads * self.v_head_dim)
+            sp_output = sp_output.reshape(
+                sp_context.num_tokens_per_rank,
+                tp_size * self.num_heads * self.v_head_dim)
             sp_output = sp_output[:num_sp_tokens]
         else:
             sp_output = torch.empty(
                 [num_sp_tokens * tp_size, self.num_heads * self.v_head_dim],
                 dtype=sp_send.dtype,
-                device=sp_send.device
-            )
+                device=sp_send.device)
             dist.all_to_all_single(
                 output=sp_output,
                 input=sp_send,
                 output_split_sizes=sp_context.output_split_sizes,
                 input_split_sizes=sp_context.input_split_sizes,
                 group=sp_world_group.device_group,
             )
-            sp_output = sp_output.reshape(num_sp_tokens, tp_size * self.num_heads * self.v_head_dim)
+            sp_output = sp_output.reshape(
+                num_sp_tokens, tp_size * self.num_heads * self.v_head_dim)
         sp_tokens = sp_output.shape[0]
         if sp_tokens == 0:
             sp_output = nn.functional.pad(sp_output, (0, 0, 0, 1))
@@ -1085,7 +1097,9 @@ def _forward_prefill_sp(
         o_output = self.o_proj(sp_output)[0]
         del sp_output
         if sp_tokens < sp_context.num_tokens_per_rank:
-            o_output = nn.functional.pad(o_output, (0, 0, 0, sp_context.num_tokens_per_rank - sp_tokens))
+            o_output = nn.functional.pad(
+                o_output,
+                (0, 0, 0, sp_context.num_tokens_per_rank - sp_tokens))
         dp_output = get_tp_group().all_gather(o_output, 0)
         dp_output = dp_output[:sp_context.num_my_dp_sp_tokens]
         return dp_output
@@ -1101,7 +1115,8 @@ def forward(
         assert output is not None, "Output tensor must be provided."
         if get_sp_context() is not None:
             # SP across DP
-            output[...] = self._forward_prefill_sp(hidden_states, kv_cache, attn_metadata)
+            output[...] = self._forward_prefill_sp(hidden_states, kv_cache,
+                                                   attn_metadata)
             return output
         if attn_metadata is None:
             # Profiling run.
diff --git a/vllm_ascend/distributed/parallel_state.py b/vllm_ascend/distributed/parallel_state.py
@@ -93,6 +93,7 @@ def init_ascend_model_parallel(parallel_config: ParallelConfig, ):
 
     init_ascend_mla_sp_model_parallel()
 
+
 def get_mlp_tensor_model_parallel_world_size():
     """Return world size for the tensor model parallel group."""
     return get_mlp_tp_group().world_size
@@ -102,19 +103,24 @@ def get_mlp_tensor_model_parallel_rank():
     """Return world size for the tensor model parallel group."""
     return get_mlp_tp_group().rank_in_group
 
+
 # vllm-ascend will maintain its own MLA SP world GroupCoordinator and o_proj sharding GroupCoordinator for
 # customize parallel solution
 _MLA_SP_WORLD: Optional[GroupCoordinator] = None
 _O_SHARD: Optional[GroupCoordinator] = None
 
+
 def get_mla_sp_world_group() -> GroupCoordinator:
-    assert _MLA_SP_WORLD is not None, ("MLA sequence parallel world group is not initialized")
+    assert _MLA_SP_WORLD is not None, (
+        "MLA sequence parallel world group is not initialized")
     return _MLA_SP_WORLD
 
+
 def get_o_shard_group() -> GroupCoordinator:
     assert _O_SHARD is not None, ("o_proj sharding group is not initialized")
     return _O_SHARD
 
+
 def init_ascend_mla_sp_model_parallel():
     from vllm_ascend.ascend_config import get_ascend_config
     ascend_config = get_ascend_config()
@@ -138,13 +144,16 @@ def init_ascend_mla_sp_model_parallel():
         num_o_shard_parallel_groups = world_size // o_shard_parallel_size
         group_ranks = []
         for i in range(num_o_shard_parallel_groups):
-            ranks = list(range(i * o_shard_parallel_size, (i + 1) * o_shard_parallel_size))
+            ranks = list(
+                range(i * o_shard_parallel_size,
+                      (i + 1) * o_shard_parallel_size))
             group_ranks.append(ranks)
         _O_SHARD = init_model_parallel_group(group_ranks,
                                              get_world_group().local_rank,
                                              backend,
                                              group_name="o_shard")
 
+
 def destroy_ascend_model_parallel():
     global _MC2
     if _MC2:
diff --git a/vllm_ascend/mla_sp_context.py b/vllm_ascend/mla_sp_context.py
@@ -1,14 +1,15 @@
+from dataclasses import dataclass
 from typing import Optional, Union
 
 import torch
 from torch import nn
-
-from dataclasses import dataclass
 from vllm.attention import AttentionMetadata
-from vllm.distributed.parallel_state import (get_dp_group, get_tp_group)
+from vllm.distributed.parallel_state import get_dp_group, get_tp_group
 from vllm.forward_context import get_forward_context
+
 from vllm_ascend.distributed.parallel_state import get_mla_sp_world_group
 
+
 @dataclass
 class SPContext:
     num_global_tokens: int
@@ -29,14 +30,18 @@ class SPContext:
     input_split_sizes: list[int]
     output_split_sizes: list[int]
 
+
 _sp_context: Optional[SPContext] = None
 
+
 def get_sp_context() -> Optional[SPContext]:
     return _sp_context
 
+
 def set_sp_context(
     input_ids: torch.Tensor,
-    attn_metadata: Optional[Union["AttentionMetadata", dict[str, "AttentionMetadata"]]] = None,
+    attn_metadata: Optional[Union["AttentionMetadata",
+                                  dict[str, "AttentionMetadata"]]] = None,
 ):
     global _sp_context
     _sp_context = None
@@ -66,8 +71,11 @@ def set_sp_context(
         assert num_input_tokens == 1, "Length of dummy run must be 1."
 
     sp_metadata = torch.cat([
-        torch.tensor([sp_enabled, num_input_tokens], device=input_ids.device, dtype=torch.int32),
-        nn.functional.pad(input_ids, (0, max_num_tokens_across_dp - num_input_tokens)),
+        torch.tensor([sp_enabled, num_input_tokens],
+                     device=input_ids.device,
+                     dtype=torch.int32),
+        nn.functional.pad(input_ids,
+                          (0, max_num_tokens_across_dp - num_input_tokens)),
     ]).unsqueeze(0)
     sp_metadata_across_dp = dp_group.all_gather(sp_metadata, 0)
     for i in range(dp_group.world_size):
@@ -86,13 +94,17 @@ def set_sp_context(
         num_global_tokens += num_tokens
         end_token_of_dp.append(num_global_tokens)
 
-    num_tokens_per_rank = calc_div_ceil(num_global_tokens, sp_world_group.world_size)
+    num_tokens_per_rank = calc_div_ceil(num_global_tokens,
+                                        sp_world_group.world_size)
     num_tokens_per_dp = num_tokens_per_rank * tp_group.world_size
-    global_tokens = torch.empty(num_global_tokens, dtype=input_ids.dtype, device=input_ids.device)
+    global_tokens = torch.empty(num_global_tokens,
+                                dtype=input_ids.dtype,
+                                device=input_ids.device)
     for i in range(dp_group.world_size):
         row = sp_metadata_across_dp[i]
         num_tokens = row[1]
-        global_tokens[start_token_of_dp[i]:end_token_of_dp[i]] = row[2:num_tokens+2]
+        global_tokens[start_token_of_dp[i]:end_token_of_dp[i]] = row[
+            2:num_tokens + 2]
 
     dp_sp_start_token = []
     dp_sp_end_token = []
@@ -101,17 +113,22 @@ def set_sp_context(
     for i in range(dp_group.world_size):
         start_token = i * num_tokens_per_dp
         dp_sp_start_token.append(start_token)
-        dp_sp_end_token.append(min(start_token + num_tokens_per_dp, num_global_tokens))
+        dp_sp_end_token.append(
+            min(start_token + num_tokens_per_dp, num_global_tokens))
     for i in range(sp_world_group.world_size):
         start_token = i * num_tokens_per_rank
         rank_sp_start_token.append(start_token)
-        rank_sp_end_token.append(min(start_token + num_tokens_per_rank, num_global_tokens))
+        rank_sp_end_token.append(
+            min(start_token + num_tokens_per_rank, num_global_tokens))
 
     my_dp = dp_group.rank_in_group
     my_rank = sp_world_group.rank_in_group
     my_rank_sp_start_token_within_dp = tp_group.rank_in_group * num_tokens_per_rank
-    my_rank_sp_end_token_within_dp = min(my_rank_sp_start_token_within_dp + num_tokens_per_rank, max(0, dp_sp_end_token[my_dp] - dp_sp_start_token[my_dp]))
-    num_my_dp_sp_tokens = max(0, dp_sp_end_token[my_dp] - dp_sp_start_token[my_dp])
+    my_rank_sp_end_token_within_dp = min(
+        my_rank_sp_start_token_within_dp + num_tokens_per_rank,
+        max(0, dp_sp_end_token[my_dp] - dp_sp_start_token[my_dp]))
+    num_my_dp_sp_tokens = max(
+        0, dp_sp_end_token[my_dp] - dp_sp_start_token[my_dp])
 
     tp_size = tp_group.world_size
     input_split_sizes = []
@@ -138,7 +155,8 @@ def set_sp_context(
 
     forward_context.with_prefill = True
     forward_context.max_tokens_across_dp = num_tokens_per_dp
-    forward_context.padded_num_tokens = calc_div_ceil(num_tokens_per_dp, tp_size) * tp_size
+    forward_context.padded_num_tokens = calc_div_ceil(num_tokens_per_dp,
+                                                      tp_size) * tp_size
     from vllm_ascend.ascend_forward_context import FusedMoEState
     if forward_context.fused_moe_state == FusedMoEState.NaiveMulticast:
         forward_context.fused_moe_state = FusedMoEState.AllGather
@@ -170,5 +188,6 @@ def set_sp_context(
         output_split_sizes=output_split_sizes,
     )
 
+
 def calc_div_ceil(up: int, down: int) -> int:
     return (up + down - 1) // down
diff --git a/vllm_ascend/models/deepseek_v2.py b/vllm_ascend/models/deepseek_v2.py
diff --git a/vllm_ascend/ops/shard.py b/vllm_ascend/ops/shard.py
