fix: resolve CI and Pages workflow errors

- Format test_attention.py and test_gemm.py with ruff to fix CI lint check
- Add missing 'require fileutils' to Ruby script in pages.yml workflow

Fixes:
- CI workflow: Ruff format check failure (exit code 1)
- Pages workflow: FileUtils.mkdir_p NameError in search index generation

Co-authored-by: Qwen-Coder <qwen-coder@alibabacloud.com>

Author: LessUp

.github/workflows/pages.yml

@@ -43,6 +43,7 @@ jobs:
           cat > build_search_index.rb << 'RUBY'
           require 'json'
           require 'yaml'
+          require 'fileutils'
 
           def extract_front_matter(content)
             if content =~ /\A---\s*\n(.*?)\n---\s*\n/m

tests/test_attention.py

@@ -49,15 +49,9 @@ def test_attention_correctness_fp32(self, batch, heads, seq_len, head_dim, devic
             pytest.skip("CUDA kernels not built")
 
         # Generate inputs
-        q = torch.randn(
-            batch, heads, seq_len, head_dim, device=device, dtype=torch.float32
-        )
-        k = torch.randn(
-            batch, heads, seq_len, head_dim, device=device, dtype=torch.float32
-        )
-        v = torch.randn(
-            batch, heads, seq_len, head_dim, device=device, dtype=torch.float32
-        )
+        q = torch.randn(batch, heads, seq_len, head_dim, device=device, dtype=torch.float32)
+        k = torch.randn(batch, heads, seq_len, head_dim, device=device, dtype=torch.float32)
+        v = torch.randn(batch, heads, seq_len, head_dim, device=device, dtype=torch.float32)
 
         # Compute outputs
         output = naive_attention(q, k, v)
@@ -92,15 +86,9 @@ def test_attention_correctness_fp16(self, batch, heads, seq_len, head_dim, devic
         except ImportError:
             pytest.skip("CUDA kernels not built")
 
-        q = torch.randn(
-            batch, heads, seq_len, head_dim, device=device, dtype=torch.float16
-        )
-        k = torch.randn(
-            batch, heads, seq_len, head_dim, device=device, dtype=torch.float16
-        )
-        v = torch.randn(
-            batch, heads, seq_len, head_dim, device=device, dtype=torch.float16
-        )
+        q = torch.randn(batch, heads, seq_len, head_dim, device=device, dtype=torch.float16)
+        k = torch.randn(batch, heads, seq_len, head_dim, device=device, dtype=torch.float16)
+        v = torch.randn(batch, heads, seq_len, head_dim, device=device, dtype=torch.float16)
 
         output = naive_attention(q, k, v)
         reference, _ = compute_attention_reference(q, k, v)
@@ -149,17 +137,16 @@ def test_softmax_invariants(self, batch, heads, seq_len, head_dim, device):
 
         # Property 2: Sum equals 1
         row_sums = softmax_output.sum(dim=-1)
-        assert torch.allclose(
-            row_sums, torch.ones_like(row_sums), rtol=1e-5, atol=1e-5
-        ), "Softmax row sums should equal 1"
+        assert torch.allclose(row_sums, torch.ones_like(row_sums), rtol=1e-5, atol=1e-5), (
+            "Softmax row sums should equal 1"
+        )
 
         # Property 3: Monotonicity (larger input -> larger output)
         for i in range(min(5, seq_len - 1)):  # Check a few pairs
             idx1, idx2 = i, i + 1
             mask = scores[..., idx1] > scores[..., idx2]
             assert (
-                softmax_output[..., idx1][mask]
-                >= softmax_output[..., idx2][mask] - 1e-6
+                softmax_output[..., idx1][mask] >= softmax_output[..., idx2][mask] - 1e-6
             ).all(), "Softmax should preserve relative order"
 
 
@@ -189,15 +176,9 @@ def test_flash_attention_consistency(self, batch, heads, seq_len, head_dim, devi
         except ImportError:
             pytest.skip("CUDA kernels not built")
 
-        q = torch.randn(
-            batch, heads, seq_len, head_dim, device=device, dtype=torch.float32
-        )
-        k = torch.randn(
-            batch, heads, seq_len, head_dim, device=device, dtype=torch.float32
-        )
-        v = torch.randn(
-            batch, heads, seq_len, head_dim, device=device, dtype=torch.float32
-        )
+        q = torch.randn(batch, heads, seq_len, head_dim, device=device, dtype=torch.float32)
+        k = torch.randn(batch, heads, seq_len, head_dim, device=device, dtype=torch.float32)
+        v = torch.randn(batch, heads, seq_len, head_dim, device=device, dtype=torch.float32)
 
         flash_output = flash_attention(q, k, v)
         naive_output = naive_attention(q, k, v)
@@ -233,15 +214,9 @@ def test_causal_mask_correctness(self, batch, heads, seq_len, head_dim, device):
         except ImportError:
             pytest.skip("CUDA kernels not built")
 
-        q = torch.randn(
-            batch, heads, seq_len, head_dim, device=device, dtype=torch.float32
-        )
-        k = torch.randn(
-            batch, heads, seq_len, head_dim, device=device, dtype=torch.float32
-        )
-        v = torch.randn(
-            batch, heads, seq_len, head_dim, device=device, dtype=torch.float32
-        )
+        q = torch.randn(batch, heads, seq_len, head_dim, device=device, dtype=torch.float32)
+        k = torch.randn(batch, heads, seq_len, head_dim, device=device, dtype=torch.float32)
+        v = torch.randn(batch, heads, seq_len, head_dim, device=device, dtype=torch.float32)
 
         # Compute causal attention
         causal_output = flash_attention(q, k, v, is_causal=True)
@@ -260,9 +235,9 @@ def test_causal_mask_correctness(self, batch, heads, seq_len, head_dim, device):
 
         # Verify attention weights are lower triangular
         upper_triangle = torch.triu(attn_weights, diagonal=1)
-        assert torch.allclose(
-            upper_triangle, torch.zeros_like(upper_triangle), atol=1e-6
-        ), "Causal attention weights should be lower triangular"
+        assert torch.allclose(upper_triangle, torch.zeros_like(upper_triangle), atol=1e-6), (
+            "Causal attention weights should be lower triangular"
+        )
 
 
 class TestTiledAttention:
@@ -290,15 +265,9 @@ def test_tiled_attention_consistency(self, batch, heads, seq_len, head_dim, devi
         except ImportError:
             pytest.skip("CUDA kernels not built")
 
-        q = torch.randn(
-            batch, heads, seq_len, head_dim, device=device, dtype=torch.float32
-        )
-        k = torch.randn(
-            batch, heads, seq_len, head_dim, device=device, dtype=torch.float32
-        )
-        v = torch.randn(
-            batch, heads, seq_len, head_dim, device=device, dtype=torch.float32
-        )
+        q = torch.randn(batch, heads, seq_len, head_dim, device=device, dtype=torch.float32)
+        k = torch.randn(batch, heads, seq_len, head_dim, device=device, dtype=torch.float32)
+        v = torch.randn(batch, heads, seq_len, head_dim, device=device, dtype=torch.float32)
 
         tiled_output = tiled_attention(q, k, v)
         naive_output = naive_attention(q, k, v)
@@ -332,15 +301,9 @@ def test_tiled_attention_correctness_fp16(self, batch, heads, seq_len, head_dim,
         except ImportError:
             pytest.skip("CUDA kernels not built")
 
-        q = torch.randn(
-            batch, heads, seq_len, head_dim, device=device, dtype=torch.float16
-        )
-        k = torch.randn(
-            batch, heads, seq_len, head_dim, device=device, dtype=torch.float16
-        )
-        v = torch.randn(
-            batch, heads, seq_len, head_dim, device=device, dtype=torch.float16
-        )
+        q = torch.randn(batch, heads, seq_len, head_dim, device=device, dtype=torch.float16)
+        k = torch.randn(batch, heads, seq_len, head_dim, device=device, dtype=torch.float16)
+        v = torch.randn(batch, heads, seq_len, head_dim, device=device, dtype=torch.float16)
 
         output = tiled_attention(q, k, v)
         reference, _ = compute_attention_reference(q, k, v)
@@ -373,8 +336,9 @@ def test_tiled_attention_scale_parameter(self, device):
         output_custom = tiled_attention(q, k, v, scale=custom_scale)
 
         # Outputs should be different
-        assert not torch.allclose(output_default, output_custom), \
+        assert not torch.allclose(output_default, output_custom), (
             "Custom scale should produce different output"
+        )
 
 
 class TestNaiveAttentionErrorHandling:
@@ -644,15 +608,9 @@ def test_batch_multihead_support(self, batch, heads, seq_len, head_dim, device):
         except ImportError:
             pytest.skip("CUDA kernels not built")
 
-        q = torch.randn(
-            batch, heads, seq_len, head_dim, device=device, dtype=torch.float32
-        )
-        k = torch.randn(
-            batch, heads, seq_len, head_dim, device=device, dtype=torch.float32
-        )
-        v = torch.randn(
-            batch, heads, seq_len, head_dim, device=device, dtype=torch.float32
-        )
+        q = torch.randn(batch, heads, seq_len, head_dim, device=device, dtype=torch.float32)
+        k = torch.randn(batch, heads, seq_len, head_dim, device=device, dtype=torch.float32)
+        v = torch.randn(batch, heads, seq_len, head_dim, device=device, dtype=torch.float32)
 
         output = flash_attention(q, k, v)
 

tests/test_gemm.py

@@ -304,9 +304,7 @@ def test_unaligned_dimensions(self, M, N, K, device):
         output = gemm(a, b)
         reference = torch.matmul(a, b)
 
-        assert output.shape == (M, N), (
-            f"Output shape mismatch: {output.shape} vs {(M, N)}"
-        )
+        assert output.shape == (M, N), f"Output shape mismatch: {output.shape} vs {(M, N)}"
         assert_close(
             output,
             reference,
@@ -390,9 +388,7 @@ def test_int8_gemm_correctness(self, M, N, K, device):
         reference = torch.matmul(a.to(torch.int32), b.to(torch.int32))
 
         assert output.dtype == torch.int32, f"Expected INT32 output, got {output.dtype}"
-        assert output.shape == (M, N), (
-            f"Output shape mismatch: {output.shape} vs {(M, N)}"
-        )
+        assert output.shape == (M, N), f"Output shape mismatch: {output.shape} vs {(M, N)}"
         assert torch.equal(output, reference), (
             f"INT8 GEMM mismatch. Max diff: {(output - reference).abs().max().item()}"
         )
@@ -485,9 +481,7 @@ def test_tall_skinny_matrices(self, device):
         output = gemm(a, b)
         reference = torch.matmul(a, b)
 
-        assert_close(
-            output, reference, rtol=1e-3, atol=1e-3, msg="Tall-skinny GEMM failed"
-        )
+        assert_close(output, reference, rtol=1e-3, atol=1e-3, msg="Tall-skinny GEMM failed")
 
     @pytest.mark.cuda
     def test_alpha_beta_scaling(self, device):
@@ -506,9 +500,7 @@ def test_alpha_beta_scaling(self, device):
         output = gemm(a, b, alpha=alpha)
         reference = alpha * torch.matmul(a, b)
 
-        assert_close(
-            output, reference, rtol=1e-3, atol=1e-3, msg="Alpha scaling failed"
-        )
+        assert_close(output, reference, rtol=1e-3, atol=1e-3, msg="Alpha scaling failed")
 
     @pytest.mark.cuda
     def test_beta_parameter(self, device):
@@ -525,9 +517,7 @@ def test_beta_parameter(self, device):
         # Test with beta=0 (default, should be same as just alpha * A @ B)
         output_beta0 = gemm(a, b, alpha=1.0, beta=0.0)
         reference = torch.matmul(a, b)
-        assert_close(
-            output_beta0, reference, rtol=1e-3, atol=1e-3, msg="Beta=0 failed"
-        )
+        assert_close(output_beta0, reference, rtol=1e-3, atol=1e-3, msg="Beta=0 failed")
 
     @pytest.mark.cuda
     def test_combined_alpha_beta(self, device):
@@ -547,6 +537,4 @@ def test_combined_alpha_beta(self, device):
         output = gemm(a, b, alpha=alpha, beta=beta)
         # Note: beta is currently unused in the implementation, so output should just be alpha * A @ B
         reference = alpha * torch.matmul(a, b)
-        assert_close(
-            output, reference, rtol=1e-3, atol=1e-3, msg="Combined alpha/beta failed"
-        )
+        assert_close(output, reference, rtol=1e-3, atol=1e-3, msg="Combined alpha/beta failed")