wip

Author: Amir-19

flashinfer/cute_dsl/blockscaled_gemm.py

@@ -50,26 +50,20 @@
     Uint64,
     T,
     Integer,
+    dsl_user_op,
+    extract_mlir_values,
+    new_from_mlir_values,
+)
+
+from cutlass.cute.typing import (
+    Int32,
     Float16,
     BFloat16,
     Float32,
     Float8E4M3FN,
     Float8E5M2,
     Tensor,
-    dsl_user_op,
-    extract_mlir_values,
-    new_from_mlir_values,
 )
-
-# from cutlass.cute.typing import (
-#     Int32,
-#     Float16,
-#     BFloat16,
-#     Float32,
-#     Float8E4M3FN,
-#     Float8E5M2,
-#     Tensor,
-# )
 from cutlass._mlir.dialects import llvm
 from flashinfer.utils import get_compute_capability
 from cutlass.utils.static_persistent_tile_scheduler import WorkTileInfo

tests/test_cute_dsl_blockscaled_gemm_allreduce_two_shot.py

@@ -86,7 +86,7 @@ def create_barrier_flags(m, n, l, mma_tiler_mn, cluster_shape_mn, sm_count):
     barrier_size = Sm100BlockScaledPersistentDenseGemmKernel.compute_barrier_flag_size(
         m, n, l, mma_tiler_mn, cluster_shape_mn, sm_count
     )
-    print("LOOK HERE", (barrier_size,))
+    #print("LOOK HERE", (barrier_size,))
     # NOTE: use_2cta_instrs from blockedscaled_gemm logic
 
     # use_2cta_instrs = mma_tiler_mn[0] == 256
@@ -481,23 +481,23 @@ def multi_process_parallel(
 @pytest.mark.parametrize(
     "ab_dtype,sf_dtype,c_dtype,sf_vec_size",
     [
-        ("float8_e5m2", "float8_e8m0fnu", "bfloat16", 32)
-        ("float4_e2m1fn", "float8_e8m0fnu", "float16", 16),
-        ("float4_e2m1fn", "float8_e8m0fnu", "bfloat16", 16),
-        ("float4_e2m1fn", "float8_e8m0fnu", "float32", 16),
-        ("float4_e2m1fn", "float8_e4m3fn", "float16", 16),
-        ("float4_e2m1fn", "float8_e4m3fn", "bfloat16", 16),
-        ("float4_e2m1fn", "float8_e4m3fn", "float32", 16),
-        ("float8_e4m3fn", "float8_e8m0fnu", "bfloat16", 32),
-        ("float8_e4m3fn", "float8_e8m0fnu", "float16", 32),
-        ("float8_e4m3fn", "float8_e8m0fnu", "float32", 32),
+        # ("float8_e5m2", "float8_e8m0fnu", "bfloat16", 32)
+        # ("float4_e2m1fn", "float8_e8m0fnu", "float16", 16),
+        # ("float4_e2m1fn", "float8_e8m0fnu", "bfloat16", 16),
+        # ("float4_e2m1fn", "float8_e8m0fnu", "float32", 16),
+        # ("float4_e2m1fn", "float8_e4m3fn", "float16", 16),
+        # ("float4_e2m1fn", "float8_e4m3fn", "bfloat16", 16),
+        # ("float4_e2m1fn", "float8_e4m3fn", "float32", 16),
+        # ("float8_e4m3fn", "float8_e8m0fnu", "bfloat16", 32),
+        # ("float8_e4m3fn", "float8_e8m0fnu", "float16", 32),
+        # ("float8_e4m3fn", "float8_e8m0fnu", "float32", 32),
         ("float8_e4m3fn", "float8_e8m0fnu", "float8_e4m3fn", 32),
-        ("float8_e4m3fn", "float8_e8m0fnu", "float8_e5m2", 32),
-        ("float8_e5m2", "float8_e8m0fnu", "bfloat16", 32),
-        ("float8_e5m2", "float8_e8m0fnu", "float16", 32),
-        ("float8_e5m2", "float8_e8m0fnu", "float32", 32),
-        ("float8_e5m2", "float8_e8m0fnu", "float8_e4m3fn", 32),
-        ("float8_e5m2", "float8_e8m0fnu", "float8_e5m2", 32),
+        # ("float8_e4m3fn", "float8_e8m0fnu", "float8_e5m2", 32),
+        # ("float8_e5m2", "float8_e8m0fnu", "bfloat16", 32),
+        # ("float8_e5m2", "float8_e8m0fnu", "float16", 32),
+        # ("float8_e5m2", "float8_e8m0fnu", "float32", 32),
+        # ("float8_e5m2", "float8_e8m0fnu", "float8_e4m3fn", 32),
+        # ("float8_e5m2", "float8_e8m0fnu", "float8_e5m2", 32),
     ],
 )
 @pytest.mark.parametrize("a_major", ["k"])