Add ruff to pre-commit (#1201)

<!-- .github/pull_request_template.md -->

## 📌 Description

Add ruff to pre-commit and reformat code to pass ruff rules.

<!-- What does this PR do? Briefly describe the changes and why they’re
needed. -->

## 🔍 Related Issues

<!-- Link any related issues here -->

## 🚀 Pull Request Checklist

Thank you for contributing to FlashInfer! Before we review your pull
request, please make sure the following items are complete.

### ✅ Pre-commit Checks

- [x] I have installed `pre-commit` by running `pip install pre-commit`
(or used your preferred method).
- [x] I have installed the hooks with `pre-commit install`.
- [x] I have run the hooks manually with `pre-commit run --all-files`
and fixed any reported issues.

> If you are unsure about how to set up `pre-commit`, see [the
pre-commit documentation](https://pre-commit.com/).

## 🧪 Tests

- [x] Tests have been added or updated as needed.
- [x] All tests are passing (`unittest`, etc.).

## Reviewer Notes

<!-- Optional: anything you'd like reviewers to focus on, concerns, etc.
-->

---------

Co-authored-by: Yaxing Cai <[email protected]>
Co-authored-by: Zihao Ye <[email protected]>

Author: cyx-6

.pre-commit-config.yaml

@@ -35,18 +35,6 @@ repos:
       - id: remove-crlf
 
   # Formatters
-  - repo: https://github.com/psf/black-pre-commit-mirror
-    rev: 24.8.0
-    hooks:
-      - id: black
-        exclude: flashinfer/tuning_configs/.*\.py
-
-  - repo: https://github.com/pycqa/isort
-    rev: 5.13.2
-    hooks:
-      - id: isort
-        args: ["--profile=black"] # <-- this one
-
   - repo: https://github.com/pre-commit/mirrors-clang-format
     rev: v19.1.1
     hooks:
@@ -56,6 +44,16 @@ repos:
           (?x)^(3rdparty/.* flashinfer/jit/aot_config.py)$
 
   -   repo: https://github.com/pre-commit/mirrors-mypy
-      rev: ''  # Use the sha / tag you want to point at
+      rev: 'v1.17.1'  # Use the sha / tag you want to point at
       hooks:
       -   id: mypy
+
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    # Ruff version.
+    rev: v0.12.8
+    hooks:
+      # Run the linter.
+      - id: ruff-check
+      # Run the formatter.
+      - id: ruff-format
+        types_or: [ python, pyi ]

benchmarks/bench_append_paged_kv_cache.py

@@ -1,6 +1,6 @@
 import argparse
 import dataclasses
-from typing import Tuple, cast
+from typing import Tuple
 
 import numpy as np
 import torch
@@ -139,10 +139,10 @@ def fn() -> None:
             print(
                 f"model: {model_name:8}",
                 f"seqlens: {seqlens!r:{seqlen_strlen}}",
-                f"convert: {convert_latency_ms*1e3:2.0f}us",
-                f"1layer: {latency_ms*1e3:2.0f}us",
-                f"{model.num_layers}layers: {all_layers_latency_ms*1e3:3.0f}us",
-                f"throughput: {throughput*1e-9:8.3f}GB/s",
+                f"convert: {convert_latency_ms * 1e3:2.0f}us",
+                f"1layer: {latency_ms * 1e3:2.0f}us",
+                f"{model.num_layers}layers: {all_layers_latency_ms * 1e3:3.0f}us",
+                f"throughput: {throughput * 1e-9:8.3f}GB/s",
             )
         print("---")
 

benchmarks/bench_append_paged_mla_kv_cache.py

@@ -1,6 +1,6 @@
 import argparse
 import dataclasses
-from typing import Tuple, cast
+from typing import Tuple
 
 import numpy as np
 import torch
@@ -122,10 +122,10 @@ def fn() -> None:
             print(
                 f"model: {model_name:8}",
                 f"seqlens: {seqlens!r:{seqlen_strlen}}",
-                f"convert: {convert_latency_ms*1e3:2.0f}us",
-                f"1layer: {latency_ms*1e3:2.0f}us",
-                f"{model.num_layers}layers: {all_layers_latency_ms*1e3:3.0f}us",
-                f"throughput: {throughput*1e-9:8.3f}GB/s",
+                f"convert: {convert_latency_ms * 1e3:2.0f}us",
+                f"1layer: {latency_ms * 1e3:2.0f}us",
+                f"{model.num_layers}layers: {all_layers_latency_ms * 1e3:3.0f}us",
+                f"throughput: {throughput * 1e-9:8.3f}GB/s",
             )
         print("---")
 

benchmarks/bench_batch_attention.py

@@ -145,7 +145,6 @@ def main() -> None:
             sweep["num_kv_heads"],
             sweep["num_qo_heads"],
         ):
-
             ms_old, ms_new, mem_MB, bw_old, bw_new = run_bench(
                 kv_lens,
                 qo_lens,

benchmarks/bench_batch_decode.py

@@ -75,7 +75,7 @@ def bench_batch_decode(
         f"batch_size={batch_size}, seq_len={seq_len}, num_qo_heads={num_qo_heads}, num_kv_heads={num_kv_heads}, head_dim={head_dim}, page_block_size={page_block_size}, q_dtype={q_dtype}, kv_dtype={kv_dtype}"
     )
     print(f"execution time: {ms}ms")
-    print(f"memory bandwidth: {io / ms / 1024 / 1024 :.2f} GB/s")
+    print(f"memory bandwidth: {io / ms / 1024 / 1024:.2f} GB/s")
 
 
 if __name__ == "__main__":

benchmarks/bench_blackwell_attention.py

@@ -61,7 +61,7 @@ def bench_fmha_blackwell(
         q_data_type=dtype,
         kv_data_type=dtype,
     )
-    o = wrapper.run(q, k, v)
+    _o = wrapper.run(q, k, v)
     measurements = bench_gpu_time(
         lambda: wrapper.run(q, k, v),
         dry_run_time_ms=100,

benchmarks/bench_cutlass_fused_moe.py

@@ -84,7 +84,6 @@ def bench_cutlass_fused_moe(
     n = intermediate_size
     k = hidden_size
     otype = torch.bfloat16
-    wtype = torch.float8_e4m3fn
     w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=otype) / 10
     w1_cutlass = torch.cat((w1[:, n:, :], w1[:, :n, :]), dim=1).contiguous()
 

benchmarks/bench_fused_add_rmsnorm.py

@@ -1,5 +1,4 @@
 import argparse
-from typing import cast
 
 import numpy as np
 import torch
@@ -54,8 +53,8 @@ def fn() -> None:
                     f"batch_size: {batch_size:3},",
                     f"hidden_size: {hidden_size:5},",
                     f"dtype: {dtype_str:8},",
-                    f"latency: {latency_ms*1e3:2.0f}us,",
-                    f"throughput: {throughput*1e-9:7.3f}GB/s",
+                    f"latency: {latency_ms * 1e3:2.0f}us,",
+                    f"throughput: {throughput * 1e-9:7.3f}GB/s",
                 )
 
         print("---")

benchmarks/bench_mixed_attention.py

@@ -213,8 +213,7 @@ def run_bench(
     for idx, (p_q_lens, p_kv_lens, d_q_len, d_kv_len) in enumerate(
         zip(p_q_configs, p_kv_configs, d_q_len_configs, d_kv_len_configs)
     ):
-
-        print(f"===== Benchmark {idx+1}: (kv_len, qo_len) set =====")
+        print(f"===== Benchmark {idx + 1}: (kv_len, qo_len) set =====")
         run_bench(
             p_q_lens,
             p_kv_lens,

benchmarks/bench_renorm.py

@@ -49,7 +49,7 @@ def main():
                     io = (probs.numel() * probs.element_size()) * 2
                     bandwidth = io * 1e-6 / ms
                     print(
-                        f"vocab_size: {vocab_size}, batch_size: {batch_size}, distrib: {distrib.__name__}, p: {p}, duration: {ms*1e3:.2f} us, effective bandwidth: {bandwidth:.2f} GB/s"
+                        f"vocab_size: {vocab_size}, batch_size: {batch_size}, distrib: {distrib.__name__}, p: {p}, duration: {ms * 1e3:.2f} us, effective bandwidth: {bandwidth:.2f} GB/s"
                     )
 
     print("---")
@@ -75,7 +75,7 @@ def main():
                     io = (probs.numel() * probs.element_size()) * 2
                     bandwidth = io * 1e-6 / ms
                     print(
-                        f"vocab_size: {vocab_size}, batch_size: {batch_size}, distrib: {distrib.__name__}, k: {k}, duration: {ms*1e3:.2f} us, effective bandwidth: {bandwidth:.2f} GB/s"
+                        f"vocab_size: {vocab_size}, batch_size: {batch_size}, distrib: {distrib.__name__}, k: {k}, duration: {ms * 1e3:.2f} us, effective bandwidth: {bandwidth:.2f} GB/s"
                     )
 
     print("---")
@@ -100,7 +100,7 @@ def main():
                     io = (logits.numel() * logits.element_size()) * 2
                     bandwidth = io * 1e-6 / ms
                     print(
-                        f"vocab_size: {vocab_size}, batch_size: {batch_size}, distrib: {distrib.__name__}, k: {k}, duration: {ms*1e3:.2f} us, effective bandwidth: {bandwidth:.2f} GB/s"
+                        f"vocab_size: {vocab_size}, batch_size: {batch_size}, distrib: {distrib.__name__}, k: {k}, duration: {ms * 1e3:.2f} us, effective bandwidth: {bandwidth:.2f} GB/s"
                     )