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

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

benchmark/asr/bench_sglang.py

@@ -343,7 +343,7 @@ def run_evaluation(args):
         print("\n" + "=" * 20 + " Sample Predictions " + "=" * 20)
         num_to_show = min(args.print_n, len(results))
         for i in range(num_to_show):
-            print(f"Sample {i+1}:")
+            print(f"Sample {i + 1}:")
             print(f"  REF: {references[i]}")
             print(f"  PRED: {predictions[i]}")
             print("-" * 40)

benchmark/bench_linear_attention/bench_kda_flashinfer_mtp.py

@@ -243,9 +243,7 @@ def run(task, fi, tri, device, dtype, args):
                     )  # noqa: E731
                 else:
                     inp = make_decode_inputs(B, H, HV, K, V, pool, device, dtype)
-                    corr = lambda kern: call_decode(
-                        kern, inp, inp["ssm"].clone()
-                    )  # noqa: E731
+                    corr = lambda kern: call_decode(kern, inp, inp["ssm"].clone())  # noqa: E731
                     ssm_t = inp["ssm"].clone()
                     timed = lambda kern: call_decode(kern, inp, ssm_t)  # noqa: E731
 

benchmark/gsm8k/bench_sglang.py

@@ -53,9 +53,9 @@ def main(args):
     if args.enable_thinking:
         from transformers import AutoTokenizer
 
-        assert (
-            args.tokenizer_path is not None
-        ), "--tokenizer-path is required when --enable-thinking is set"
+        assert args.tokenizer_path is not None, (
+            "--tokenizer-path is required when --enable-thinking is set"
+        )
         tokenizer = AutoTokenizer.from_pretrained(
             args.tokenizer_path, trust_remote_code=True
         )

benchmark/hf3fs/bench_client.py

@@ -14,11 +14,11 @@ def print_stats(x: List[int]):
     x = sorted(x)
     lenx = len(x)
     print(
-        f"mean = {sum(x)/len(x):.2f}, "
+        f"mean = {sum(x) / len(x):.2f}, "
         f"min = {min(x):.2f}, "
-        f"p25 = {x[int(lenx*0.25)]:.2f}, "
-        f"p50 = {x[int(lenx*0.5)]:.2f}, "
-        f"p75 = {x[int(lenx*0.75)]:.2f}, "
+        f"p25 = {x[int(lenx * 0.25)]:.2f}, "
+        f"p50 = {x[int(lenx * 0.5)]:.2f}, "
+        f"p75 = {x[int(lenx * 0.75)]:.2f}, "
         f"max = {max(x):.2f}"
     )
 

benchmark/hf3fs/bench_storage.py

@@ -18,11 +18,11 @@ def print_stats(x: List[int]):
     x = sorted(x)
     lenx = len(x)
     print(
-        f"mean = {sum(x)/len(x):.2f}, "
+        f"mean = {sum(x) / len(x):.2f}, "
         f"min = {min(x):.2f}, "
-        f"p25 = {x[int(lenx*0.25)]:.2f}, "
-        f"p50 = {x[int(lenx*0.5)]:.2f}, "
-        f"p75 = {x[int(lenx*0.75)]:.2f}, "
+        f"p25 = {x[int(lenx * 0.25)]:.2f}, "
+        f"p50 = {x[int(lenx * 0.5)]:.2f}, "
+        f"p75 = {x[int(lenx * 0.75)]:.2f}, "
         f"max = {max(x):.2f}"
     )
 

benchmark/hf3fs/bench_zerocopy.py

@@ -109,7 +109,7 @@ elif hicache_mem_layout == "layer_first":
     for operation in operations:
         cache_controller.generic_page_backup(operation, batch_size=128)
 tok = time.monotonic()
-print(f"{tok-tik:.6f} s")
+print(f"{tok - tik:.6f} s")
 
 operations = [
     PrefetchOperation(
@@ -137,4 +137,4 @@ elif hicache_mem_layout == "layer_first":
     for operation in operations:
         cache_controller.generic_page_transfer(operation, batch_size=128)
 tok = time.monotonic()
-print(f"{tok-tik:.6f} s")
+print(f"{tok - tik:.6f} s")

benchmark/hicache/bench_mix.py

@@ -457,7 +457,7 @@ class WorkloadGenerator:
             try:
                 user_data, response = self.response_queue.get(timeout=10)
                 logger.info(
-                    f"{((time.perf_counter()-self.start_time)/self.duration*100):.2f}%"
+                    f"{((time.perf_counter() - self.start_time) / self.duration * 100):.2f}%"
                 )
                 if not response.success:
                     raise ValueError(f"Request failed with error: {response.error}")
@@ -540,10 +540,10 @@ class WorkloadGenerator:
         output_stats = self.user_generator.output_stats
         print(f"round_ratios: {user_stats}")
         print(
-            f"mean_new_tokens_per_round: {[int(a/b) if b > 0 else 0 for a, b in input_stats]}"
+            f"mean_new_tokens_per_round: {[int(a / b) if b > 0 else 0 for a, b in input_stats]}"
         )
         print(
-            f"mean_return_tokens_per_round: {[int(a/b) if b > 0 else 0 for a, b in output_stats]}"
+            f"mean_return_tokens_per_round: {[int(a / b) if b > 0 else 0 for a, b in output_stats]}"
         )
         return performance_data
 

benchmark/hicache/bench_serving.py

@@ -75,9 +75,9 @@ async def async_request_openai_completions(
     pbar: Optional[tqdm] = None,
 ) -> RequestFuncOutput:
     api_url = request_func_input.api_url
-    assert api_url.endswith(
-        "completions"
-    ), "OpenAI Completions API URL must end with 'completions'."
+    assert api_url.endswith("completions"), (
+        "OpenAI Completions API URL must end with 'completions'."
+    )
 
     async with aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session:
         payload = {

benchmark/hicache/nextqa.py

@@ -120,7 +120,7 @@ class NExTQALoader(VideoLoader):
         video = Video(video_path, num_frames)
         prompt = entry["question"] + "?"
         if self.task == "MC":  # add choices
-            prompt += f' a0: {entry["a0"]}, a1: {entry["a1"]}, a2: {entry["a2"]}, a3: {entry["a3"]}'
+            prompt += f" a0: {entry['a0']}, a1: {entry['a1']}, a2: {entry['a2']}, a3: {entry['a3']}"
         return VideoPrompt(video_path, num_frames, prompt)
 
     def __iter__(self):

benchmark/kernels/bench_fused_rmsnorm_fp8_quant.py

@@ -149,9 +149,9 @@ def _check_correctness():
             cos = torch.nn.functional.cosine_similarity(
                 (q.float() * scale).flatten(), ref_deq.flatten(), dim=0
             ).item()
-            assert (
-                cos > 0.99
-            ), f"{provider} h={hidden_size} residual={add_residual} cos={cos:.4f}"
+            assert cos > 0.99, (
+                f"{provider} h={hidden_size} residual={add_residual} cos={cos:.4f}"
+            )
     print("correctness check passed (all fused providers vs unfused within FP8)")
 
 

benchmark/kernels/deepep/deepep_utils.py

@@ -191,9 +191,9 @@ def bench_kineto(
     kernel_names = (kernel_names,) if isinstance(kernel_names, str) else kernel_names
     assert all([isinstance(name, str) for name in kernel_names])
     for name in kernel_names:
-        assert (
-            sum([name in line for line in prof_lines]) == 1
-        ), f"Errors of the kernel {name} in the profiling table"
+        assert sum([name in line for line in prof_lines]) == 1, (
+            f"Errors of the kernel {name} in the profiling table"
+        )
 
     # Save chrome traces
     if trace_path is not None:

benchmark/kernels/deepep/tuning_deepep.py

@@ -155,7 +155,7 @@ def test_main(
                 for with_topk in (False, True):
                     if local_rank == 0:
                         print(
-                            f'[testing] Running with {"FP8" if isinstance(current_x, tuple) else "BF16"}, {"with" if with_topk else "without"} top-k (async={async_mode}, previous={previous_mode}) ...',
+                            f"[testing] Running with {'FP8' if isinstance(current_x, tuple) else 'BF16'}, {'with' if with_topk else 'without'} top-k (async={async_mode}, previous={previous_mode}) ...",
                             flush=True,
                             end="",
                         )
@@ -198,9 +198,9 @@ def test_main(
 
                     # Checks
                     recv_gbl_rank_prefix_sum = handle[-4]
-                    assert gbl_num_tokens_per_rank[rank].item() == recv_x.size(
-                        0
-                    ), f"{gbl_num_tokens_per_rank[rank].item()} != {recv_x.size(0)}"
+                    assert gbl_num_tokens_per_rank[rank].item() == recv_x.size(0), (
+                        f"{gbl_num_tokens_per_rank[rank].item()} != {recv_x.size(0)}"
+                    )
                     assert (
                         gbl_num_tokens_per_expert.view(num_ranks, -1)[rank].tolist()
                         == recv_num_tokens_per_expert_list
@@ -325,11 +325,14 @@ def test_main(
                 tune_args = {"x": current_x, "handle": handle, "config": config}
                 t = bench(lambda: buffer.dispatch(**tune_args))[0]
                 if t < best_time:
-                    best_time, best_results = t, (
-                        num_sms,
-                        nvl_chunk_size,
-                        rdma_chunk_size,
-                        config_kwargs,
+                    best_time, best_results = (
+                        t,
+                        (
+                            num_sms,
+                            nvl_chunk_size,
+                            rdma_chunk_size,
+                            config_kwargs,
+                        ),
                     )
                 if local_rank == 0:
                     print(
@@ -338,7 +341,7 @@ def test_main(
                     )
         if local_rank == 0:
             print(
-                f'[tuning] Best dispatch ({"FP8" if isinstance(current_x, tuple) else "BF16"}): SMs {best_results[0]}, NVL chunk {best_results[1]}, RDMA chunk {best_results[2]}: {rdma_send_bytes / 1e9 / best_time:.2f} GB/s (RDMA), {nvl_recv_bytes / 1e9 / best_time:.2f} GB/s (NVL)',
+                f"[tuning] Best dispatch ({'FP8' if isinstance(current_x, tuple) else 'BF16'}): SMs {best_results[0]}, NVL chunk {best_results[1]}, RDMA chunk {best_results[2]}: {rdma_send_bytes / 1e9 / best_time:.2f} GB/s (RDMA), {nvl_recv_bytes / 1e9 / best_time:.2f} GB/s (NVL)",
                 flush=True,
             )
             print("", flush=True)
@@ -399,11 +402,14 @@ def test_main(
                     flush=True,
                 )
                 if t < best_time:
-                    best_time, best_results = t, (
-                        num_sms,
-                        nvl_chunk_size,
-                        rdma_chunk_size,
-                        config_kwargs,
+                    best_time, best_results = (
+                        t,
+                        (
+                            num_sms,
+                            nvl_chunk_size,
+                            rdma_chunk_size,
+                            config_kwargs,
+                        ),
                     )
 
     if local_rank == 0:

benchmark/kernels/deepseek/benchmark_deepgemm_fp8_gemm.py

@@ -59,7 +59,6 @@ def tl_gemm(
             bx,
             by,
         ):
-
             A_shared = T.alloc_shared(A_shared_shape, in_dtype)
             B_shared = T.alloc_shared(B_shared_shape, in_dtype)
             C_shared = T.alloc_shared(C_shared_shape, out_dtype)
@@ -350,7 +349,7 @@ def get_benchmark(tp_size):
         tflops = flops / (ms * 1e-3) / 1e12
 
         # Print shape-specific results with TFLOPS
-        print(f"Time: {ms*1000:.2f} ms, TFLOPS: {tflops:.2f}")
+        print(f"Time: {ms * 1000:.2f} ms, TFLOPS: {tflops:.2f}")
         return ms * 1000, max_ms * 1000, min_ms * 1000  # convert to ms
 
     return benchmark

benchmark/kernels/deepseek/benchmark_deepgemm_fp8_gemm_blackwell.py

@@ -224,7 +224,7 @@ def _benchmark(m, n, k, tp_size, provider):
     tflops = flops / (ms * 1e-3) / 1e12
 
     # Print shape-specific results with TFLOPS
-    print(f"Time: {ms*1000:.2f} us, TFLOPS: {tflops:.2f}")
+    print(f"Time: {ms * 1000:.2f} us, TFLOPS: {tflops:.2f}")
     return ms, max_ms, min_ms
 
 

benchmark/kernels/deepseek/benchmark_deepgemm_fp8_group_gemm.py

@@ -435,7 +435,7 @@ def get_benchmark(tp_size):
         flops = 2 * m * n * k  # multiply-adds
         tflops = flops / (ms * 1e-3) / 1e12
 
-        print(f"Time: {ms*1000:.2f} ms, TFLOPS: {tflops:.2f}")
+        print(f"Time: {ms * 1000:.2f} ms, TFLOPS: {tflops:.2f}")
         return ms * 1000, max_ms * 1000, min_ms * 1000  # convert to ms
 
     return benchmark

benchmark/kernels/deepseek/benchmark_q8kv8_kv_gather.py

@@ -243,8 +243,7 @@ def main():
             else:
                 speedup = f"{legacy_us / us:.2f}x"
             print(
-                f"{s_q:>5} {kv_len:>7} {name:>10} {us:>9.1f} "
-                f"{tbps:>9.3f} {speedup:>8}"
+                f"{s_q:>5} {kv_len:>7} {name:>10} {us:>9.1f} {tbps:>9.3f} {speedup:>8}"
             )
         print()
 

benchmark/kernels/elementwise/benchmark_concat_mla.py

@@ -143,9 +143,7 @@ output_exp = execute_and_get_output(fn_cuda, data)
 if not torch.all(output_ref == output_exp):
     abs_delta = torch.abs(output_ref - output_exp)
     raise AssertionError(
-        f"{output_ref=} {output_exp=} "
-        f"{abs_delta=} "
-        f"{torch.argwhere(abs_delta != 0.0)=} "
+        f"{output_ref=} {output_exp=} {abs_delta=} {torch.argwhere(abs_delta != 0.0)=} "
     )
 
 

benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton_sep.py

@@ -535,7 +535,6 @@ class BestConfigTrace:
 
 
 class BenchmarkWorker:
-
     def __init__(self, seed: int, server_args: ServerArgs) -> None:
         torch.set_default_device("cuda")
         torch.cuda.manual_seed_all(0)
@@ -729,8 +728,7 @@ class BenchmarkWorker:
                 down_use_tma_map[block_m] = time_cost_all[2] > time_cost_all[3]
 
             print(
-                f"Round 1 done. Down TMA decisions per BLOCK_SIZE_M: "
-                f"{down_use_tma_map}"
+                f"Round 1 done. Down TMA decisions per BLOCK_SIZE_M: {down_use_tma_map}"
             )
 
             # === Round 2: Up with c_sorted from round 1 ===

benchmark/kernels/lora_csgmv/tune_lora_csgmv.py

@@ -470,9 +470,9 @@ def _tune_shrink(
     device: torch.device,
 ) -> tuple:
     """Tune shrink kernel for one layer type. Returns (best_configs, results)."""
-    print(f"\n{'='*80}")
+    print(f"\n{'=' * 80}")
     print(f"Tuning SHRINK — {label} (K={K}, N={N}, slices={num_slices})")
-    print(f"{'='*80}")
+    print(f"{'=' * 80}")
 
     search = get_shrink_search_space()
     print(f"Search space: {len(search)} configs")
@@ -508,7 +508,7 @@ def _tune_shrink(
                 best_config = config
             if (i + 1) % 20 == 0:
                 print(
-                    f"  chunk={chunk_size}: {i+1}/{len(search)} tested, best={best_time:.3f}ms"
+                    f"  chunk={chunk_size}: {i + 1}/{len(search)} tested, best={best_time:.3f}ms"
                 )
 
         best_configs[chunk_size] = sort_config(best_config)
@@ -533,9 +533,9 @@ def _tune_expand(
     device: torch.device,
 ) -> tuple:
     """Tune expand kernel for one layer type. Returns (best_configs, results)."""
-    print(f"\n{'='*80}")
+    print(f"\n{'=' * 80}")
     print(f"Tuning EXPAND — {label} (output_dim={output_dim}, slices={num_slices})")
-    print(f"{'='*80}")
+    print(f"{'=' * 80}")
 
     search = get_expand_search_space()
     print(f"Search space: {len(search)} configs")
@@ -584,7 +584,7 @@ def _tune_expand(
                 best_config = config
             if (i + 1) % 50 == 0:
                 print(
-                    f"  chunk={chunk_size}: {i+1}/{len(search)} tested, best={best_time:.3f}ms"
+                    f"  chunk={chunk_size}: {i + 1}/{len(search)} tested, best={best_time:.3f}ms"
                 )
 
         best_configs[chunk_size] = sort_config(best_config)
@@ -673,9 +673,9 @@ def main(args: argparse.Namespace):
             )
 
     # --- Summary ---
-    print(f"\n{'='*80}")
+    print(f"\n{'=' * 80}")
     print(f"SUMMARY")
-    print(f"{'='*80}")
+    print(f"{'=' * 80}")
     print(
         f"\n{'layer':<10} {'kernel':<8} {'K/dim':>6} {'chunk':>6}"
         f" {'baseline':>10} {'tuned':>10} {'speedup':>8}  config"

benchmark/lean_kernel_sweep.py

@@ -137,19 +137,21 @@ def main():
                 # b32 x 128K on the 8-KV-head config exceeds the microbench's single
                 # contiguous KV tensor (faults the GPU); real serving uses a paged pool.
                 if H_KV == 8 and B == 32 and S == 131072:
-                    print(f"{B:>5} {S//1024:>5}K {'skipped (contiguous-KV limit)':>30}")
+                    print(
+                        f"{B:>5} {S // 1024:>5}K {'skipped (contiguous-KV limit)':>30}"
+                    )
                     rows.append(f"{name},{H_Q},{H_KV},{B},{S},,,,skip,")
                     continue
                 try:
                     std, lean, cos, gate = run(H_Q, H_KV, B, S)
                 except torch.cuda.OutOfMemoryError:
                     torch.cuda.empty_cache()
-                    print(f"{B:>5} {S//1024:>5}K {'OOM':>9}")
+                    print(f"{B:>5} {S // 1024:>5}K {'OOM':>9}")
                     rows.append(f"{name},{H_Q},{H_KV},{B},{S},,,,OOM,")
                     continue
                 sp = std / lean
                 print(
-                    f"{B:>5} {S//1024:>5}K {std:>9.3f} {lean:>9.3f} {sp:>7.2f}x {('ON' if gate else 'OFF'):>5} {cos:>7.4f}"
+                    f"{B:>5} {S // 1024:>5}K {std:>9.3f} {lean:>9.3f} {sp:>7.2f}x {('ON' if gate else 'OFF'):>5} {cos:>7.4f}"
                 )
                 rows.append(
                     f"{name},{H_Q},{H_KV},{B},{S},{std:.4f},{lean:.4f},{sp:.4f},{cos:.4f},{int(gate)}"

benchmark/mmlu/bench_sglang.py

@@ -163,7 +163,7 @@ def main(args):
     pt = 0
     for subject, num_qs in zip(subjects[: args.nsub], num_questions):
         print(
-            f"subject: {subject}, #q:{num_qs}, acc: {np.mean(cors[pt: pt + num_qs]):.3f}"
+            f"subject: {subject}, #q:{num_qs}, acc: {np.mean(cors[pt : pt + num_qs]):.3f}"
         )
         pt += num_qs
     assert pt == len(cors)

benchmark/ocr/bench_sglang.py

@@ -502,7 +502,7 @@ async def process_sample(
                         }
                     )
                 print(
-                    f"[INPUT ]   [{i+1}] type={ttype!r:12s}  expected: {expected[:120]}",
+                    f"[INPUT ]   [{i + 1}] type={ttype!r:12s}  expected: {expected[:120]}",
                     flush=True,
                 )
             # Print OCR output (truncate long outputs)

benchmark/ocr/generate_report.py

@@ -159,12 +159,12 @@ def _render_sample(sample: dict, failures_only: bool) -> str:
     if failures_only and passed == total and not error:
         return ""
 
-    pct = f"{100*passed//total}%" if total else "—"
+    pct = f"{100 * passed // total}%" if total else "—"
     header_cls = "fail" if (error or passed < total) else "pass"
 
     parts = [f'<div class="sample">']
     parts.append(
-        f'<details {"open" if (error or passed < total) else ""}>'
+        f"<details {'open' if (error or passed < total) else ''}>"
         f'<summary class="sample-header {header_cls}">'
         f"<span>📄 {html.escape(pdf)}  &nbsp;·&nbsp; page {page}</span>"
         f"<span>"
@@ -218,7 +218,7 @@ def _render_sample(sample: dict, failures_only: bool) -> str:
                         f'<div class="rendered">{_latex_to_display(latex)}</div>'
                     )
             elif ttype in ("present", "absent", "text_presence", "text_absence"):
-                parts.append(f'<pre>{html.escape(ti.get("text", ""))}</pre>')
+                parts.append(f"<pre>{html.escape(ti.get('text', ''))}</pre>")
             elif ttype in ("order", "natural_reading_order"):
                 before = ti.get("before", "")
                 after = ti.get("after", "")
@@ -236,7 +236,7 @@ def _render_sample(sample: dict, failures_only: bool) -> str:
                         parts.append(f'<div class="rendered">{m}</div>')
                     if len(matches) > 6:
                         parts.append(
-                            f'<p style="color:#888;font-size:0.8em">… and {len(matches)-6} more</p>'
+                            f'<p style="color:#888;font-size:0.8em">… and {len(matches) - 6} more</p>'
                         )
                 else:
                     parts.append(

benchmark/prefill_only/util.py

@@ -383,14 +383,14 @@ async def send_warmup_requests(
                 http_url, data=request_json, headers=headers
             ) as resp:
                 if resp.status == 200:
-                    print(f"Warmup request {i+1}/{num_warmup} completed successfully")
+                    print(f"Warmup request {i + 1}/{num_warmup} completed successfully")
                 else:
                     print(
-                        f"Warmup request {i+1}/{num_warmup} failed with status {resp.status}"
+                        f"Warmup request {i + 1}/{num_warmup} failed with status {resp.status}"
                     )
 
         except Exception as e:
-            print(f"Warmup request {i+1}/{num_warmup} failed with error: {e}")
+            print(f"Warmup request {i + 1}/{num_warmup} failed with error: {e}")
 
     print("HTTP warmup requests completed")
 
@@ -745,7 +745,6 @@ async def run_generic_benchmark(
     async with aiohttp.ClientSession(
         timeout=aiohttp.ClientTimeout(total=300)
     ) as session:
-
         # Send START_PROFILE if profiling is enabled
         if config.profile:
             await send_profile_request("START_PROFILE", http_url, session=session)

benchmark/scheduler/bench_token_storage.py

@@ -254,7 +254,7 @@ def print_breakdown(title: str, results: dict[str, dict[str, float]]) -> None:
 
 
 def microbench_torch_tensor_paths(
-    sizes: tuple[int, ...] = (1_000, 10_000, 100_000)
+    sizes: tuple[int, ...] = (1_000, 10_000, 100_000),
 ) -> None:
     """Compare three CPU-buffer -> pinned cuda tensor paths.
 
@@ -293,9 +293,11 @@ def microbench_torch_tensor_paths(
         ),
         (
             "(C) from_numpy(frombuf(array('q'))).pin() -> cuda",
-            lambda x: torch.from_numpy(np.frombuffer(x, dtype=np.int64))
-            .pin_memory()
-            .to("cuda", non_blocking=True),
+            lambda x: (
+                torch.from_numpy(np.frombuffer(x, dtype=np.int64))
+                .pin_memory()
+                .to("cuda", non_blocking=True)
+            ),
         ),
     ]:
         cells = []