sync flashinfer and update sgl-kernel tests (#3081)

.github/workflows/pr-test-sgl-kernel.yml

@@ -47,7 +47,7 @@ jobs:
           pip3 list | grep sgl-kernel
 
       - name: Run test
-        timeout-minutes: 10
+        timeout-minutes: 30
         run: |
           cd sgl-kernel
           find tests -name "test_*.py" | xargs -n 1 python3

sgl-kernel/Makefile

@@ -19,7 +19,7 @@ clean:
 	@rm -rf build dist *.egg-info
 
 test:
-	@find tests -name "test_*.py" | xargs -n 1 python3 && pytest tests/test_norm.py && pytest tests/test_activation.py
+	@find tests -name "test_*.py" | xargs -n 1 python3
 
 format:
 	@find src tests -name '*.cc' -o -name '*.cu' -o -name '*.cuh' -o -name '*.h' -o -name '*.hpp' | xargs clang-format -i && find src tests -name '*.py' | xargs isort && find src tests -name '*.py' | xargs black

sgl-kernel/tests/test_activation.py

@@ -35,4 +35,5 @@ def test_fused_gelu_mul(dim, batch_size, seq_len):
     torch.testing.assert_close(y_ref, y, rtol=1e-3, atol=1e-3)
 
 
-test_fused_silu_mul(128, 1, 1)
+if __name__ == "__main__":
+    pytest.main([__file__])

sgl-kernel/tests/test_lightning_attention_decode.py

@@ -82,3 +82,7 @@ def test_lightning_attention_decode(dtype, batch_size, num_heads, dim, embed_dim
         msg=f"New KV mismatch for batch_size={batch_size}, num_heads={num_heads}, "
         f"dim={dim}, embed_dim={embed_dim}, dtype={dtype}",
     )
+
+
+if __name__ == "__main__":
+    pytest.main([__file__])

sgl-kernel/tests/test_norm.py

@@ -127,3 +127,7 @@ def test_gemma_fused_add_rmsnorm(batch_size, hidden_size, dtype):
 
     torch.testing.assert_close(x_fused, x_native, rtol=1e-3, atol=1e-3)
     torch.testing.assert_close(residual_fused, residual_native, rtol=1e-3, atol=1e-3)
+
+
+if __name__ == "__main__":
+    pytest.main([__file__])