fix(ci): repair python CI and lint for llm-speed

Made-with: Cursor

Author: CI

.github/workflows/ci.yml

@@ -68,7 +68,10 @@ jobs:
           cache: pip
 
       - name: Install test dependencies
-        run: pip install pytest hypothesis torch --index-url https://download.pytorch.org/whl/cpu
+        run: |
+          pip install -U pip
+          pip install pytest hypothesis
+          pip install torch --index-url https://download.pytorch.org/whl/cpu
 
       - name: Run tests (skip CUDA)
         run: pytest tests/ -v -m "not cuda" --tb=short || true

benchmarks/benchmark_attention.py

@@ -7,7 +7,7 @@
 import argparse
 import json
 import torch
-from typing import List, Dict, Tuple
+from typing import List, Dict
 import sys
 from pathlib import Path
 sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
@@ -205,7 +205,7 @@ def main():
 
     dtype = torch.float16 if args.dtype == 'fp16' else torch.float32
 
-    print(f"Configuration:")
+    print("Configuration:")
     print(f"  Batch size: {args.batch_size}")
     print(f"  Num heads: {args.num_heads}")
     print(f"  Head dim: {args.head_dim}")

benchmarks/benchmark_gemm.py

@@ -48,8 +48,12 @@ def compute_gemm_flops(M: int, N: int, K: int) -> int:
 
 
 def benchmark_gemm(
-    sizes: List[Tuple[int, int, int]] = [(1024, 1024, 1024), (2048, 2048, 2048), 
-                                          (4096, 4096, 4096), (8192, 8192, 8192)],
+    sizes: List[Tuple[int, int, int]] = [
+        (1024, 1024, 1024),
+        (2048, 2048, 2048),
+        (4096, 4096, 4096),
+        (8192, 8192, 8192),
+    ],
     dtype: torch.dtype = torch.float16,
     warmup: int = 10,
     iterations: int = 100
@@ -213,7 +217,7 @@ def main():
 
     dtype = torch.float16 if args.dtype == 'fp16' else torch.float32
 
-    print(f"Configuration:")
+    print("Configuration:")
     print(f"  Sizes: {sizes}")
     print(f"  Dtype: {args.dtype}")
     print(f"  Warmup: {args.warmup}")

python/profiler.py

@@ -5,7 +5,6 @@
 from __future__ import annotations
 
 import torch
-import time
 from dataclasses import dataclass
 from typing import Callable
 from enum import Enum

tests/conftest.py

@@ -5,7 +5,6 @@
 import pytest
 import torch
 import numpy as np
-from typing import Tuple
 
 
 def pytest_configure(config):

tests/test_attention.py

@@ -5,8 +5,7 @@
 
 import pytest
 import torch
-import numpy as np
-from hypothesis import given, settings, strategies as st, assume
+from hypothesis import given, settings, strategies as st
 
 from conftest import assert_close, compute_attention_reference
 

tests/test_gemm.py

@@ -5,8 +5,7 @@
 
 import pytest
 import torch
-import numpy as np
-from hypothesis import given, settings, strategies as st, assume
+from hypothesis import given, settings, strategies as st
 
 from conftest import assert_close