[CI][RFC] Replace black-jupyter with ruff-format (#37210)

Co-authored-by: Alison Shao <a.shao@wustl.edu>

test/registered/unit/layers/quantization/test_deepgemm_ue8m0_requant.py

@@ -39,9 +39,10 @@ class TestDeepGemmUE8M0Requant(CustomTestCase):
     def test_helper_requants_supported_deepgemm_bf16_once(self):
         weight, weight_scale = _make_params()
 
-        with self._enabled_deepgemm_ue8m0(), patch.object(
-            fp8_utils, "requant_weight_ue8m0_inplace"
-        ) as requant:
+        with (
+            self._enabled_deepgemm_ue8m0(),
+            patch.object(fp8_utils, "requant_weight_ue8m0_inplace") as requant,
+        ):
             fired = fp8_utils.requant_block_scale_ue8m0_for_deepgemm(
                 weight,
                 weight_scale,
@@ -67,9 +68,10 @@ class TestDeepGemmUE8M0Requant(CustomTestCase):
     def test_helper_skips_non_bf16_output(self):
         weight, weight_scale = _make_params()
 
-        with self._enabled_deepgemm_ue8m0(), patch.object(
-            fp8_utils, "requant_weight_ue8m0_inplace"
-        ) as requant:
+        with (
+            self._enabled_deepgemm_ue8m0(),
+            patch.object(fp8_utils, "requant_weight_ue8m0_inplace") as requant,
+        ):
             fired = fp8_utils.requant_block_scale_ue8m0_for_deepgemm(
                 weight,
                 weight_scale,
@@ -86,9 +88,10 @@ class TestDeepGemmUE8M0Requant(CustomTestCase):
     def test_helper_skips_shape_deepgemm_will_not_run(self):
         weight, weight_scale = _make_params(n=96, k=128)
 
-        with self._enabled_deepgemm_ue8m0(), patch.object(
-            fp8_utils, "requant_weight_ue8m0_inplace"
-        ) as requant:
+        with (
+            self._enabled_deepgemm_ue8m0(),
+            patch.object(fp8_utils, "requant_weight_ue8m0_inplace") as requant,
+        ):
             fired = fp8_utils.requant_block_scale_ue8m0_for_deepgemm(
                 weight,
                 weight_scale,
@@ -105,9 +108,10 @@ class TestDeepGemmUE8M0Requant(CustomTestCase):
     def test_helper_skips_non_deepgemm_runner(self):
         weight, weight_scale = _make_params()
 
-        with self._enabled_deepgemm_ue8m0(), patch.object(
-            fp8_utils, "requant_weight_ue8m0_inplace"
-        ) as requant:
+        with (
+            self._enabled_deepgemm_ue8m0(),
+            patch.object(fp8_utils, "requant_weight_ue8m0_inplace") as requant,
+        ):
             fired = fp8_utils.requant_block_scale_ue8m0_for_deepgemm(
                 weight,
                 weight_scale,
@@ -125,9 +129,10 @@ class TestDeepGemmUE8M0Requant(CustomTestCase):
         weight, weight_scale = _make_params()
         unsupported_block_size = [128, 256]
 
-        with self._enabled_deepgemm_ue8m0(), patch.object(
-            fp8_utils, "requant_weight_ue8m0_inplace"
-        ) as requant:
+        with (
+            self._enabled_deepgemm_ue8m0(),
+            patch.object(fp8_utils, "requant_weight_ue8m0_inplace") as requant,
+        ):
             fired = fp8_utils.requant_block_scale_ue8m0_for_deepgemm(
                 weight,
                 weight_scale,
@@ -154,9 +159,10 @@ class TestDeepGemmUE8M0Requant(CustomTestCase):
         layer.weight, layer.weight_scale = _make_params()
         layer.orig_dtype = torch.bfloat16
 
-        with self._enabled_deepgemm_ue8m0(), patch.object(
-            fp8_utils, "requant_weight_ue8m0_inplace"
-        ) as requant:
+        with (
+            self._enabled_deepgemm_ue8m0(),
+            patch.object(fp8_utils, "requant_weight_ue8m0_inplace") as requant,
+        ):
             scheme.process_weights_after_loading(layer)
             scheme.process_weights_after_loading(layer)
 
@@ -179,18 +185,22 @@ class TestDeepGemmUE8M0Requant(CustomTestCase):
             weight_scale.format_ue8m0 = True
             return True
 
-        with patch.multiple(
-            fp8_quant,
-            _is_cpu=False,
-            _is_fp8_fnuz=False,
-            _use_aiter=False,
-        ), patch.object(
-            method, "is_deepgemm_moe_runner_backend_enabled", return_value=True
-        ), patch.object(
-            fp8_quant,
-            "requant_block_scale_ue8m0_for_deepgemm",
-            side_effect=_mark_ue8m0,
-        ) as requant:
+        with (
+            patch.multiple(
+                fp8_quant,
+                _is_cpu=False,
+                _is_fp8_fnuz=False,
+                _use_aiter=False,
+            ),
+            patch.object(
+                method, "is_deepgemm_moe_runner_backend_enabled", return_value=True
+            ),
+            patch.object(
+                fp8_quant,
+                "requant_block_scale_ue8m0_for_deepgemm",
+                side_effect=_mark_ue8m0,
+            ) as requant,
+        ):
             method.process_weights_after_loading_block_quant(layer)
 
         self.assertEqual(

test/registered/unit/layers/quantization/test_flashinfer_trtllm_fp8_fallback.py

@@ -54,17 +54,16 @@ class TestFlashinferTrtllmFp8Fallback(CustomTestCase):
         trtllm_spy = MagicMock(return_value=torch.zeros((M, N), dtype=dtype))
         quant_spy = MagicMock(return_value=(MagicMock(), MagicMock()))
 
-        with patch.object(
-            fp8_utils,
-            "_get_flashinfer_groupwise_backend",
-            return_value="trtllm",
-            create=True,
-        ), patch.object(
-            fp8_utils, "gemm_fp8_nt_groupwise", trtllm_spy, create=True
-        ), patch.object(
-            fp8_utils, "triton_w8a8_block_fp8_linear", triton_spy
-        ), patch.object(
-            fp8_utils, "sglang_per_token_group_quant_fp8", quant_spy
+        with (
+            patch.object(
+                fp8_utils,
+                "_get_flashinfer_groupwise_backend",
+                return_value="trtllm",
+                create=True,
+            ),
+            patch.object(fp8_utils, "gemm_fp8_nt_groupwise", trtllm_spy, create=True),
+            patch.object(fp8_utils, "triton_w8a8_block_fp8_linear", triton_spy),
+            patch.object(fp8_utils, "sglang_per_token_group_quant_fp8", quant_spy),
         ):
             fp8_utils.flashinfer_gemm_w8a8_block_fp8_linear_with_fallback(
                 input_2d, weight, BLOCK_SIZE, weight_scale

test/registered/unit/layers/quantization/test_mxfp4_flashinfer_activation_prep.py

@@ -25,16 +25,20 @@ class TestMxfp4FlashinferActivationPrep(CustomTestCase):
         x_quant = torch.empty(3, 64, dtype=torch.float8_e4m3fn)
         x_scale = torch.arange(6, dtype=torch.uint8).reshape(3, 2)
 
-        with patch(
-            "sglang.srt.layers.moe.route_quant_handoff.take", return_value=None
-        ) as take, patch(
-            "sglang.srt.layers.quantization.mxfp4._is_sm107_supported",
-            return_value=True,
-        ), patch(
-            "sglang.srt.layers.quantization.fp8_utils.flashinfer_mxfp8_quantize",
-            return_value=(x_quant, x_scale),
-            create=True,
-        ) as quantize:
+        with (
+            patch(
+                "sglang.srt.layers.moe.route_quant_handoff.take", return_value=None
+            ) as take,
+            patch(
+                "sglang.srt.layers.quantization.mxfp4._is_sm107_supported",
+                return_value=True,
+            ),
+            patch(
+                "sglang.srt.layers.quantization.fp8_utils.flashinfer_mxfp8_quantize",
+                return_value=(x_quant, x_scale),
+                create=True,
+            ) as quantize,
+        ):
             actual_x, packed_topk, actual_quant, actual_scale = (
                 _prepare_flashinfer_mxfp8_activations(x, 64)
             )
@@ -51,19 +55,24 @@ class TestMxfp4FlashinferActivationPrep(CustomTestCase):
         x_quant = torch.empty(3, 64, dtype=torch.float8_e4m3fn)
         x_scale = torch.arange(6, dtype=torch.uint8).reshape(3, 2)
 
-        with patch(
-            "sglang.srt.layers.moe.route_quant_handoff.take", return_value=None
-        ) as take, patch(
-            "sglang.srt.layers.quantization.mxfp4._is_sm107_supported",
-            return_value=True,
-        ), patch.object(
-            per_token_group_quant_module,
-            "per_token_group_quant",
-            return_value=(x_quant, x_scale),
-        ) as quantize, patch(
-            "sglang.srt.layers.quantization.fp8_utils.flashinfer_mxfp8_quantize",
-            create=True,
-        ) as flashinfer_quantize:
+        with (
+            patch(
+                "sglang.srt.layers.moe.route_quant_handoff.take", return_value=None
+            ) as take,
+            patch(
+                "sglang.srt.layers.quantization.mxfp4._is_sm107_supported",
+                return_value=True,
+            ),
+            patch.object(
+                per_token_group_quant_module,
+                "per_token_group_quant",
+                return_value=(x_quant, x_scale),
+            ) as quantize,
+            patch(
+                "sglang.srt.layers.quantization.fp8_utils.flashinfer_mxfp8_quantize",
+                create=True,
+            ) as flashinfer_quantize,
+        ):
             actual_x, packed_topk, actual_quant, actual_scale = (
                 _prepare_flashinfer_mxfp8_activations(x, 64)
             )
@@ -81,19 +90,22 @@ class TestMxfp4FlashinferActivationPrep(CustomTestCase):
         x_quant = torch.empty(3, 64, dtype=torch.float8_e4m3fn)
         x_scale = torch.arange(6, dtype=torch.uint8).reshape(3, 2)
 
-        with patch(
-            "sglang.srt.layers.moe.route_quant_handoff.take", return_value=None
-        ), patch(
-            "sglang.srt.layers.quantization.mxfp4._is_sm107_supported",
-            return_value=False,
-        ), patch.object(
-            per_token_group_quant_module,
-            "per_token_group_quant",
-            return_value=(x_quant, x_scale),
-        ) as quantize, patch(
-            "sglang.srt.layers.quantization.fp8_utils.flashinfer_mxfp8_quantize",
-            create=True,
-        ) as flashinfer_quantize:
+        with (
+            patch("sglang.srt.layers.moe.route_quant_handoff.take", return_value=None),
+            patch(
+                "sglang.srt.layers.quantization.mxfp4._is_sm107_supported",
+                return_value=False,
+            ),
+            patch.object(
+                per_token_group_quant_module,
+                "per_token_group_quant",
+                return_value=(x_quant, x_scale),
+            ) as quantize,
+            patch(
+                "sglang.srt.layers.quantization.fp8_utils.flashinfer_mxfp8_quantize",
+                create=True,
+            ) as flashinfer_quantize,
+        ):
             actual_x, packed_topk, actual_quant, actual_scale = (
                 _prepare_flashinfer_mxfp8_activations(x, 64)
             )
@@ -111,11 +123,14 @@ class TestMxfp4FlashinferActivationPrep(CustomTestCase):
         x_quant = torch.empty(3, 128, dtype=torch.float8_e4m3fn)
         x_scale = torch.arange(12, dtype=torch.uint8)
 
-        with patch(
-            "sglang.srt.layers.quantization.fp8_utils.flashinfer_mxfp8_quantize",
-            return_value=(x_quant, x_scale),
-            create=True,
-        ) as quantize, patch("sglang.srt.layers.moe.route_quant_handoff.take") as take:
+        with (
+            patch(
+                "sglang.srt.layers.quantization.fp8_utils.flashinfer_mxfp8_quantize",
+                return_value=(x_quant, x_scale),
+                create=True,
+            ) as quantize,
+            patch("sglang.srt.layers.moe.route_quant_handoff.take") as take,
+        ):
             actual_x, packed_topk, actual_quant, actual_scale = (
                 _prepare_flashinfer_mxfp8_activations(x, 128)
             )
@@ -133,13 +148,16 @@ class TestMxfp4FlashinferActivationPrep(CustomTestCase):
         x_quant = torch.empty(2, 64, dtype=torch.float8_e4m3fn)
         x_scale = torch.arange(4, dtype=torch.uint8).reshape(2, 2)
 
-        with patch(
-            "sglang.srt.layers.moe.route_quant_handoff.take",
-            return_value=(packed_topk, x_quant, x_scale),
-        ), patch(
-            "sglang.srt.layers.quantization.fp8_utils.flashinfer_mxfp8_quantize",
-            create=True,
-        ) as quantize:
+        with (
+            patch(
+                "sglang.srt.layers.moe.route_quant_handoff.take",
+                return_value=(packed_topk, x_quant, x_scale),
+            ),
+            patch(
+                "sglang.srt.layers.quantization.fp8_utils.flashinfer_mxfp8_quantize",
+                create=True,
+            ) as quantize,
+        ):
             actual_x, actual_packed, actual_quant, actual_scale = (
                 _prepare_flashinfer_mxfp8_activations(x, 64)
             )

test/registered/unit/layers/quantization/test_mxfp4_sm90_cutlass.py

@@ -585,9 +585,9 @@ def test_humming_range_ignores_prerounded_hidden_tail():
         "w13 Humming residual changed with the never-written hidden tail; "
         "the _UE8M0_ONE fill leaked into the per-expert E8M0 range"
     )
-    assert torch.equal(
-        residuals[0][1], residuals[1][1]
-    ), "w2 Humming residual changed with the never-written hidden tail"
+    assert torch.equal(residuals[0][1], residuals[1][1]), (
+        "w2 Humming residual changed with the never-written hidden tail"
+    )
 
 
 @pytest.mark.skipif(

test/registered/unit/layers/quantization/test_nvfp4_moe_backends.py

@@ -43,17 +43,17 @@ class TestNvFp4MoeBackends(CustomTestCase):
         quant_config = ModelOptFp4Config(
             is_checkpoint_nvfp4_serialized=True, group_size=16
         )
-        with get_context().override_server_args(
-            model_path="dummy"
-        ), get_flags().moe.override(
-            runner_backend=MoeRunnerBackend(backend)
-        ), get_parallel().override(
-            moe_ep_size=1,
-            moe_ep_rank=0,
-            moe_tp_size=1,
-            moe_tp_rank=0,
-            tp_size=1,
-            tp_rank=0,
+        with (
+            get_context().override_server_args(model_path="dummy"),
+            get_flags().moe.override(runner_backend=MoeRunnerBackend(backend)),
+            get_parallel().override(
+                moe_ep_size=1,
+                moe_ep_rank=0,
+                moe_tp_size=1,
+                moe_tp_rank=0,
+                tp_size=1,
+                tp_rank=0,
+            ),
         ):
             layer = FusedMoE(
                 num_experts=E,