fix yapf

Signed-off-by: taoyuxiang <[email protected]>

Author: NicholasTao

tests/ut/models/test_qwen3_moe.py

@@ -18,6 +18,7 @@
 import pytest
 import torch
 from vllm.model_executor.models.qwen3_moe import Qwen3MoeForCausalLM
+
 from vllm_ascend.models.qwen3_moe import (CustomQwen3MoeAttention,
                                           CustomQwen3MoeForCausalLM)
 
@@ -49,7 +50,8 @@ def test_packed_modules_mapping_structure(self):
         assert CustomQwen3MoeForCausalLM.packed_modules_mapping == expected_mapping
 
 
-class TestNormalizeQKVWithFixedInput(unittest.TestCase):
+class TestCustomQwen3MoeAttention(unittest.TestCase):
+
     def setUp(self):
         self.batch = 2
         self.seq_len = 3
@@ -60,20 +62,16 @@ def setUp(self):
 
         total_dim = self.q_size + 2 * self.kv_size
 
-        self.qkv = torch.arange(
-            self.batch * self.seq_len * total_dim,
-            dtype=torch.float32
-        ).reshape(self.batch, self.seq_len, total_dim)
+        self.qkv = torch.arange(self.batch * self.seq_len * total_dim,
+                                dtype=torch.float32).reshape(
+                                    self.batch, self.seq_len, total_dim)
 
     def test_constant_input_normalization(self):
-        ones_qkv = torch.ones(
-            (1, 1, self.q_size + 2 * self.kv_size),
-            dtype=torch.float32
-        )
+        ones_qkv = torch.ones((1, 1, self.q_size + 2 * self.kv_size),
+                              dtype=torch.float32)
 
         q, k, v = CustomQwen3MoeAttention.normalize_qkv(
-            ones_qkv, self.q_size, self.kv_size, self.head_dim, self.rms_eps
-        )
+            ones_qkv, self.q_size, self.kv_size, self.head_dim, self.rms_eps)
 
         norm_val = 1.0 / math.sqrt(1.0 + self.rms_eps)
 

tests/ut/ops/test_rotary_embedding.py

@@ -341,9 +341,7 @@ def test_set_cos_sin_cache(self):
         # mock out register_buffer
         model.register_buffer = MagicMock()
         # call the private method via name mangling
-        model._set_cos_sin_cache(seq_len=8,
-                                 device="cpu",
-                                 dtype=torch.float32)
+        model._set_cos_sin_cache(seq_len=8, device="cpu", dtype=torch.float32)
         # expect three calls: inv_freq, cos, sin
         assert model.register_buffer.call_count == 3
         names = [call.args[0] for call in model.register_buffer.call_args_list]