Fix examples and pytest, run ruff (#3230)

Co-authored-by: dePaul Miller <23461061+depaulmillz@users.noreply.github.com>

test/examples/CuTeDSL/sm_100a/conftest.py

@@ -26,5 +26,6 @@
 # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
+
 def pytest_configure(config):
-    config.default_SMs[__file__] = "100f"
+    config.default_SMs[__file__] = "100f"

test/examples/CuTeDSL/sm_100a/test_dense_blockscaled_gemm_persistent_prefetch.py

@@ -41,29 +41,33 @@ from typing import Tuple, Type, Optional
 
 import pytest
 
-from blackwell.dense_blockscaled_gemm_persistent_prefetch import (
+from blackwell.kernel.blockscaled_gemm.dense_blockscaled_gemm_persistent_prefetch import (
     Sm100BlockScaledPersistentDenseGemmKernel,
     run,
 )
 
 import cutlass
+
 pytestmark = [pytest.mark.arch(["100a"])]
 
+
 @pytest.mark.invalid_case(
-    lambda: not Sm100BlockScaledPersistentDenseGemmKernel.can_implement(
-        ab_dtype,
-        sf_dtype,
-        sf_vec_size,
-        c_dtype,
-        mma_tiler_mn,
-        cluster_shape_mn,
-        mnkl[0],
-        mnkl[1],
-        mnkl[2],
-        mnkl[3],
-        a_major,
-        b_major,
-        c_major,
+    lambda: (
+        not Sm100BlockScaledPersistentDenseGemmKernel.can_implement(
+            ab_dtype,
+            sf_dtype,
+            sf_vec_size,
+            c_dtype,
+            mma_tiler_mn,
+            cluster_shape_mn,
+            mnkl[0],
+            mnkl[1],
+            mnkl[2],
+            mnkl[3],
+            a_major,
+            b_major,
+            c_major,
+        )
     )
 )
 @pytest.mark.parametrize(
@@ -110,7 +114,7 @@ pytestmark = [pytest.mark.arch(["100a"])]
     "prefetch_dist",
     [
         None,  # Default: auto (uses num_ab_stage)
-        0,     # Disabled
+        0,  # Disabled
     ],
 )
 @pytest.mark.parametrize("tolerance", [1e-01])
@@ -145,20 +149,22 @@ def test_dense_blockscaled_gemm_prefetch(
 
 
 @pytest.mark.invalid_case(
-    lambda: not Sm100BlockScaledPersistentDenseGemmKernel.can_implement(
-        ab_dtype,
-        sf_dtype,
-        sf_vec_size,
-        c_dtype,
-        mma_tiler_mn,
-        cluster_shape_mn,
-        mnkl[0],
-        mnkl[1],
-        mnkl[2],
-        mnkl[3],
-        a_major,
-        b_major,
-        c_major,
+    lambda: (
+        not Sm100BlockScaledPersistentDenseGemmKernel.can_implement(
+            ab_dtype,
+            sf_dtype,
+            sf_vec_size,
+            c_dtype,
+            mma_tiler_mn,
+            cluster_shape_mn,
+            mnkl[0],
+            mnkl[1],
+            mnkl[2],
+            mnkl[3],
+            a_major,
+            b_major,
+            c_major,
+        )
     )
 )
 @pytest.mark.parametrize(
@@ -190,7 +196,7 @@ def test_dense_blockscaled_gemm_prefetch(
     "prefetch_dist",
     [
         None,  # Default: auto (uses num_ab_stage)
-        4,     # Explicit distance
+        4,  # Explicit distance
     ],
 )
 @pytest.mark.parametrize("tolerance", [1e-01])
@@ -228,15 +234,15 @@ def test_dense_blockscaled_gemm_prefetch_L0(
     "prefetch_dist",
     [
         None,  # Auto: uses num_ab_stage
-        0,     # Disabled
-        2,     # Small distance
-        4,     # Medium distance
+        0,  # Disabled
+        2,  # Small distance
+        4,  # Medium distance
     ],
 )
 def test_prefetch_dist_configurations(prefetch_dist: Optional[int]):
     """
     Test different prefetch_dist configurations specifically for blockscaled GEMM.
-    
+
     - None: Auto mode, uses num_ab_stage as prefetch distance
     - 0: Prefetch disabled
     - >0: Explicit prefetch distance
@@ -451,4 +457,3 @@ def test_invalid_tensor_alignment(
             cluster_shape_mn,
             tolerance,
         )
-

test/examples/CuTeDSL/sm_100a/test_dense_gemm_persistent_prefetch.py

@@ -40,8 +40,7 @@ from typing import Tuple, Type, Optional
 
 import pytest
 
-from blackwell.dense_gemm_persistent_prefetch import (
-    PersistentDenseGemmKernel,
+from blackwell.kernel.dense_gemm.dense_gemm_persistent_prefetch import (
     run,
 )
 
@@ -92,8 +91,8 @@ import cutlass.cute.testing as testing
     "prefetch_dist",
     [
         None,  # Default: auto (uses num_ab_stage)
-        0,     # Disabled
-        2,     # Explicit distance
+        0,  # Disabled
+        2,  # Explicit distance
     ],
 )
 @pytest.mark.parametrize("tolerance", [1e-01])
@@ -168,7 +167,7 @@ def test_dense_gemm_prefetch(
     "prefetch_dist",
     [
         None,  # Default: auto (uses num_ab_stage)
-        4,     # Explicit distance
+        4,  # Explicit distance
     ],
 )
 def test_dense_gemm_prefetch_L0(
@@ -215,15 +214,15 @@ def test_dense_gemm_prefetch_L0(
     "prefetch_dist",
     [
         None,  # Auto: uses num_ab_stage
-        0,     # Disabled
-        2,     # Small distance
-        4,     # Medium distance
+        0,  # Disabled
+        2,  # Small distance
+        4,  # Medium distance
     ],
 )
 def test_prefetch_dist_configurations(prefetch_dist: Optional[int]):
     """
     Test different prefetch_dist configurations specifically.
-    
+
     - None: Auto mode, uses num_ab_stage as prefetch distance
     - 0: Prefetch disabled
     - >0: Explicit prefetch distance
@@ -259,4 +258,3 @@ def test_prefetch_dist_configurations(prefetch_dist: Optional[int]):
         )
     except testing.CantImplementError:
         pytest.skip(f"Skip unsupported testcase with prefetch_dist={prefetch_dist}")
-

test/examples/CuTeDSL/sm_100a/test_rmsnorm.py

@@ -38,11 +38,10 @@ Tests various configurations of:
 """
 
 import pytest
-import torch
 
 import cutlass
 
-from blackwell.rmsnorm import (
+from blackwell.kernel.rmsnorm.rmsnorm import (
     run,
     get_sm_version,
     supports_cluster,
@@ -104,7 +103,9 @@ class TestRMSNormCorrectness:
 class TestRMSNormClusterPath:
     """Test the cluster path for large N (SM90+/SM100 only)."""
 
-    @pytest.mark.skipif(not supports_cluster(), reason="Cluster not supported on this GPU")
+    @pytest.mark.skipif(
+        not supports_cluster(), reason="Cluster not supported on this GPU"
+    )
     @pytest.mark.parametrize("N", [32768, 65536])
     def test_cluster_path_correctness(self, N):
         """Test cluster path produces correct results."""
@@ -119,6 +120,7 @@ class TestRMSNormClusterPath:
             benchmark=False,
         )
 
+
 class TestRMSNormLargeN:
     """Test RMSNorm with large N values."""
 
@@ -151,7 +153,6 @@ class TestRMSNormLargeN:
         )
 
 
-
 class TestRMSNormEdgeCases:
     """Test edge cases for RMSNorm."""
 
@@ -197,4 +198,4 @@ class TestRMSNormFloat32:
             tolerance=1e-4,  # Tighter tolerance for FP32
             skip_ref_check=False,
             benchmark=False,
-        )
+        )

test/examples/CuTeDSL/sm_100a/test_tutorial_gemm.py

@@ -26,14 +26,14 @@
 # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
-from blackwell.tutorial_gemm import fp16_gemm_0
-from blackwell.tutorial_gemm import fp16_gemm_1
-from blackwell.tutorial_gemm import fp16_gemm_2
-from blackwell.tutorial_gemm import fp16_gemm_3
-from blackwell.tutorial_gemm import fp16_gemm_3_1
-from blackwell.tutorial_gemm import fp16_gemm_4
-from blackwell.tutorial_gemm import fp16_gemm_5
-from blackwell.tutorial_gemm import fp16_gemm_6
+from blackwell.tutorial.tutorial_gemm import fp16_gemm_0
+from blackwell.tutorial.tutorial_gemm import fp16_gemm_1
+from blackwell.tutorial.tutorial_gemm import fp16_gemm_2
+from blackwell.tutorial.tutorial_gemm import fp16_gemm_3
+from blackwell.tutorial.tutorial_gemm import fp16_gemm_3_1
+from blackwell.tutorial.tutorial_gemm import fp16_gemm_4
+from blackwell.tutorial.tutorial_gemm import fp16_gemm_5
+from blackwell.tutorial.tutorial_gemm import fp16_gemm_6
 
 import pytest
 from typing import Tuple
@@ -63,7 +63,6 @@ def test_fp16_gemm_1(
     fp16_gemm_1.run_dense_gemm(mnk, tolerance)
 
 
-
 @pytest.mark.parametrize(
     "mnk",
     [(512, 512, 256)],