minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source 
Co-authored-by: Alison Shao <a.shao@wustl.edu>
This commit is contained in:
Alex Nails
2026-09-02 19:46:08 -07:00
committed by GitHub
co-authored by Alison Shao
parent 2641e427be
commit 28262c20df
1411 changed files with 7766 additions and 8176 deletions
Download Patch File Download Diff File Expand all files Collapse all files
benchmark/kernels/bench_fused_rmsnorm_fp8_quant.py
+3 -3
View File
@@ -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
+3 -3
View File
@@ -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
+21 -15
View File
@@ -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
+1 -2
View File
@@ -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
+1 -1
View File
@@ -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
+1 -1
View File
@@ -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
+1 -2
View File
@@ -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
+1 -3
View File
@@ -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
+1 -3
View File
@@ -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
+8 -8
View File
@@ -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"