Autotune between tma gather and cp.async in SM100 (Dao-AILab#100)

* autotune tma gather

* fix comment

* ruff format

Author: GarlGuo

quack/gemm_act.py

@@ -286,6 +286,7 @@ def _compile_gemm_act(
     gemm_cls_name,
     rounding_mode=RoundingMode.RN,
     sr_seed_mode=0,
+    use_tma_gather=False,
 ):
     sm_to_cls = {
         "act": {9: GemmActSm90, 10: GemmActSm100, 11: GemmActSm100, 12: GemmActSm120},
@@ -360,6 +361,7 @@ def fake_scalar(mode, dtype=Int32):
         epi_args,
         scheduler_args,
         varlen_args,
+        use_tma_gather=use_tma_gather,
     )
 
 
@@ -385,6 +387,7 @@ def gemm_act(
     A_idx: Optional[Tensor] = None,  # (total_m,) if gather_A with varlen_m
     rounding_mode: int = RoundingMode.RN,
     sr_seed: int | Tensor = 0,
+    use_tma_gather: bool = False,
 ) -> None:
     if activation in gate_fn_map:
         gemm_cls_name = "gated"
@@ -462,6 +465,7 @@ def gemm_act(
         gemm_cls_name,
         rounding_mode=rounding_mode,
         sr_seed_mode=sr_seed_mode,
+        use_tma_gather=use_tma_gather,
     )
 
     from quack.cache_utils import COMPILE_ONLY

quack/gemm_config.py

@@ -18,6 +18,8 @@ class GemmConfig:
     # raster_order: int = 1
     max_swizzle_size: int = 8
     device_capacity: int = 9
+    # whether to use TMA gather (vs normal cp.async) for gather_A on SM100
+    use_tma_gather: bool = False
 
 
 def _get_sm90_configs(
@@ -58,6 +60,7 @@ def _get_sm90_configs(
             swap_ab=swap_ab,
             device_capacity=9,
             is_dynamic_persistent=False,  # default to not use dynamic persistent on SM90
+            use_tma_gather=False,  # TMA gather not supported on SM90
         )
         for (tile_m, tile_n, pingpong), (cluster_m, cluster_n), swap_ab in itertools.product(
             tile_mn_vals,
@@ -87,6 +90,7 @@ def _get_sm100_configs(
         GemmConfig, pingpong=False, device_capacity=10
     )  # There's no pingpong on Sm100
     use_clc_vals = [True, False]
+    use_tma_gather_vals = [True, False]
     return [
         GemmConfigCls(
             tile_m=m,
@@ -96,9 +100,10 @@ def _get_sm100_configs(
             swap_ab=sab,
             max_swizzle_size=8,
             is_dynamic_persistent=use_clc,
+            use_tma_gather=use_tma_gather,
         )
-        for (m, n, (cm, cn)), sab, use_clc in itertools.product(
-            tile_mn_cluster_vals, swap_ab_vals, use_clc_vals
+        for (m, n, (cm, cn)), sab, use_clc, use_tma_gather in itertools.product(
+            tile_mn_cluster_vals, swap_ab_vals, use_clc_vals, use_tma_gather_vals
         )
     ]
 
@@ -126,6 +131,7 @@ def _get_sm120_configs(
             swap_ab=swap_ab,
             device_capacity=12,
             is_dynamic_persistent=True,
+            use_tma_gather=False,  # TMA gather not supported on SM120
         )
         for (tile_m, tile_n, pingpong), swap_ab in itertools.product(tile_mn_vals, swap_ab_vals)
     ]

quack/gemm_dact.py

@@ -259,6 +259,7 @@ def _compile_gemm_dact(
     gather_A,
     device_capacity,
     gemm_cls_name,
+    use_tma_gather=False,
 ):
     is_dgated = gemm_cls_name == "dgated"
     sm_to_cls = {
@@ -352,6 +353,7 @@ def _set_implicit_dtype(gemm_obj):
         scheduler_args,
         varlen_args,
         post_init=post_init,
+        use_tma_gather=use_tma_gather,
     )
 
 
@@ -376,7 +378,7 @@ def gemm_dact(
     colvec_reduce: Optional[Tensor] = None,
     cu_seqlens_m: Optional[Tensor] = None,  # (l+1,) cumulative sum of m values for variable length
     A_idx: Optional[Tensor] = None,  # (total_m,) if gather_A with varlen_m
-    use_clc_persistence: bool = False,
+    use_tma_gather: bool = False,
 ) -> None:
     is_dgated = activation in dgate_fn_map
     if not is_dgated:
@@ -463,6 +465,7 @@ def gemm_dact(
         gather_A,
         device_capacity,
         gemm_cls_name,
+        use_tma_gather=use_tma_gather,
     )
 
     from quack.cache_utils import COMPILE_ONLY

quack/gemm_interface.py

@@ -114,6 +114,9 @@ def prune_invalid_gemm_configs(configs, named_args: dict, **kwargs):
         if device_capacity == 9:
             configs = [conf for conf in configs if conf.kwargs["config"].tile_n != 208]
             configs = [conf for conf in configs if not conf.kwargs["config"].is_dynamic_persistent]
+    # use_tma_gather only valid when gather_A is active on SM100/SM110
+    if not gather_A or device_capacity not in [10, 11]:
+        configs = [conf for conf in configs if not conf.kwargs["config"].use_tma_gather]
     return configs
 
 
@@ -215,6 +218,7 @@ def gemm_tuned(
         add_to_output=add_to_output,
         rounding_mode=rounding_mode,
         sr_seed=sr_seed,
+        use_tma_gather=config.use_tma_gather,
     )
 
 
@@ -286,6 +290,7 @@ def gemm_act_tuned(
         colvec_bias=bias if config.swap_ab else None,
         cu_seqlens_m=cu_seqlens_m,
         A_idx=A_idx,
+        use_tma_gather=config.use_tma_gather,
     )
 
 
@@ -351,6 +356,7 @@ def gemm_dact_tuned(
         max_swizzle_size=config.max_swizzle_size,
         cu_seqlens_m=cu_seqlens_m,
         A_idx=A_idx,
+        use_tma_gather=config.use_tma_gather,
     )
 
 
@@ -1198,6 +1204,7 @@ def gemm_gated_tuned(
         colvec_bias=bias if config.swap_ab else None,
         cu_seqlens_m=cu_seqlens_m,
         A_idx=A_idx,
+        use_tma_gather=config.use_tma_gather,
     )
 
 
@@ -1292,6 +1299,7 @@ def gemm_dgated_tuned(
         colvec_reduce=colvec_reduce_partial,
         cu_seqlens_m=cu_seqlens_m,
         A_idx=A_idx,
+        use_tma_gather=config.use_tma_gather,
     )
     if colvec_reduce:
         colvec_reduce_final = colvec_reduce_partial.sum(dim=-1)