[CI] Fix lint brought by #27527 (#28988)

Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>

test/manual/test_create_custom_4d_mask.py

@@ -43,6 +43,7 @@ _DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 # Standalone reference implementation (original loop-based code, pre-a475156d)
 # ---------------------------------------------------------------------------
 
+
 def _create_custom_4d_mask_reference(
     sequence_length, dtype, device, batch_size, token_type_ids
 ):
@@ -79,6 +80,7 @@ def _create_custom_4d_mask_reference(
 # python/sglang/srt/models/deepseek_ocr.py)
 # ---------------------------------------------------------------------------
 
+
 def _create_custom_4d_mask_new(
     sequence_length, dtype, device, batch_size, token_type_ids
 ):
@@ -117,6 +119,7 @@ def _create_custom_4d_mask_new(
 # Helpers
 # ---------------------------------------------------------------------------
 
+
 def _make_token_type_ids(batch_size, seq_len, image_fraction, device):
     """First `image_fraction` tokens per sequence are image (0), rest are text (1).
 
@@ -172,6 +175,7 @@ def _bench(fn, run_device, n=50, **kwargs):
 # Accuracy tests
 # ---------------------------------------------------------------------------
 
+
 class TestAccuracy(unittest.TestCase):
     """Verify new implementation produces identical masks to the reference."""
 
@@ -280,11 +284,17 @@ class TestAccuracy(unittest.TestCase):
         for q in range(S):
             for k in range(S):
                 if k <= q:
-                    self.assertEqual(mask2d[q, k].item(), 0.0,
+                    self.assertEqual(
-                                     f"text[{q}] should attend to text[{k}]")
+                        mask2d[q, k].item(),
+                        0.0,
+                        f"text[{q}] should attend to text[{k}]",
+                    )
                 else:
-                    self.assertEqual(mask2d[q, k].item(), min_val,
+                    self.assertEqual(
-                                     f"text[{q}] should NOT attend to text[{k}]")
+                        mask2d[q, k].item(),
+                        min_val,
+                        f"text[{q}] should NOT attend to text[{k}]",
+                    )
 
     def test_image_full_attention(self):
         """Image tokens must attend to all other image tokens (bidirectional)."""
@@ -296,8 +306,9 @@ class TestAccuracy(unittest.TestCase):
         mask2d = out.cpu()[0, 0]
         for q in range(n_img):
             for k in range(n_img):
-                self.assertEqual(mask2d[q, k].item(), 0.0,
+                self.assertEqual(
-                                 f"image[{q}] should attend to image[{k}]")
+                    mask2d[q, k].item(), 0.0, f"image[{q}] should attend to image[{k}]"
+                )
 
     def test_text_attends_to_image(self):
         """Every text token must attend to every image token."""
@@ -309,8 +320,9 @@ class TestAccuracy(unittest.TestCase):
         mask2d = out.cpu()[0, 0]
         for q in range(n_img, S):
             for k in range(n_img):
-                self.assertEqual(mask2d[q, k].item(), 0.0,
+                self.assertEqual(
-                                 f"text[{q}] should attend to image[{k}]")
+                    mask2d[q, k].item(), 0.0, f"text[{q}] should attend to image[{k}]"
+                )
 
     # --- dtype coverage ---
 
@@ -352,8 +364,9 @@ class TestPerformance(unittest.TestCase):
         cls.dtype = torch.float32
 
     def _run_case(self, batch_size, seq_len, image_fraction):
-        ids = _make_token_type_ids(batch_size, seq_len, image_fraction,
+        ids = _make_token_type_ids(
-                                   device=self.device)
+            batch_size, seq_len, image_fraction, device=self.device
+        )
         kwargs = dict(
             sequence_length=seq_len,
             dtype=self.dtype,
@@ -361,10 +374,15 @@ class TestPerformance(unittest.TestCase):
             batch_size=batch_size,
             token_type_ids=ids,
         )
-        t_ref = _bench(_create_custom_4d_mask_reference, run_device=self.device,
+        t_ref = _bench(
-                       n=BENCH_ITERS, **kwargs)
+            _create_custom_4d_mask_reference,
-        t_new = _bench(_create_custom_4d_mask_new, run_device=self.device,
+            run_device=self.device,
-                       n=BENCH_ITERS, **kwargs)
+            n=BENCH_ITERS,
+            **kwargs,
+        )
+        t_new = _bench(
+            _create_custom_4d_mask_new, run_device=self.device, n=BENCH_ITERS, **kwargs
+        )
         speedup = t_ref / t_new
         dev_tag = "CUDA" if "cuda" in str(self.device) else "CPU"
         print(
@@ -406,6 +424,7 @@ class TestPerformance(unittest.TestCase):
 # PyTorch profiler (optional – triggered by --profile or PROFILE_TRACES=1)
 # ---------------------------------------------------------------------------
 
+
 def run_profiler_traces(output_dir: str = "./pt_traces", device: str = _DEVICE):
     """
     Capture Chrome-trace JSON files for both implementations.