Triton early prune config fix (#4917)

Summary:
Pull Request resolved: #4917

X-link: facebookresearch/FBGEMM#1941

Fix early prune config for cases not using warp-specializations.

Reviewed By: renganxu, jasonjk-park

Differential Revision: D83019498

fbshipit-source-id: f35883eefe8e437020722a1a738c4f2ef12e8f11

Author: sunfish2010

fbgemm_gpu/experimental/gemm/triton_gemm/grouped_gemm.py

@@ -129,6 +129,75 @@ def early_config_prune(configs, named_args, dtsize=None, dtype=None, **kwargs):
     if dtype is None:
         dtype = named_args["c_ptr"].dtype
 
+    pruned_configs = []
+    for config in configs:
+        kw = config.kwargs
+        (
+            BLOCK_M,
+            BLOCK_N,
+            BLOCK_K,
+            num_stages,
+            use_tma_load_on_scales,
+        ) = (
+            kw["BLOCK_SIZE_M"],
+            kw["BLOCK_SIZE_N"],
+            kw["BLOCK_SIZE_K"],
+            config.num_stages,
+            kw.get("USE_TMA_LOAD_ON_SCALES", False),
+        )
+        G, M, N = (
+            named_args["G"],
+            named_args["M_BUCKET"],
+            named_args["N"],
+        )
+
+        # 1. make sure we have enough smem
+        max_shared_memory = driver.active.utils.get_device_properties(device)[
+            "max_shared_mem"
+        ]
+        if torch.version.hip:
+            required_shared_memory = BLOCK_N * BLOCK_K * num_stages * dtsize
+        else:
+            required_shared_memory = (BLOCK_M + BLOCK_N) * BLOCK_K * num_stages * dtsize
+        if required_shared_memory > max_shared_memory:
+            continue
+
+        M_PER_GROUP = M // G
+        MIN_M_TILES = 32 if torch.version.hip else 64
+        # 2. make sure we don't load M tiles that are too big
+        if BLOCK_M > MIN_M_TILES and BLOCK_M > (M_PER_GROUP * 2):
+            continue
+        # 3. make sure we don't load N tiles that are too small
+        if BLOCK_M < 128 and BLOCK_M < (M_PER_GROUP // 2):
+            continue
+
+        num_sm = driver.active.utils.get_device_properties(device)[
+            "multiprocessor_count"
+        ]
+        N_TILES = (N + BLOCK_N - 1) // BLOCK_N
+        MIN_N_TILES = 32 if torch.version.hip else 64
+        # 4. make sure we don't load N tiles that are too big
+        if BLOCK_N > MIN_N_TILES and M * N_TILES < num_sm:
+            continue
+        # 5. make sure we don't load N tiles that are too small
+        if BLOCK_N < 128 and M * N_TILES > 2 * num_sm:
+            continue
+        if dtsize >= 2:
+            if use_tma_load_on_scales:
+                continue
+        pruned_configs.append(config)
+
+    return pruned_configs
+
+
+def early_config_prune_ws(configs, named_args, dtsize=None, dtype=None, **kwargs):
+    device = torch.cuda.current_device()
+    # BLOCK_M, BLOCK_N, BLOCK_K, SPLIT_K, num_warps, num_stages
+    if dtsize is None:
+        dtsize = named_args["c_ptr"].element_size()
+    if dtype is None:
+        dtype = named_args["c_ptr"].dtype
+
     pruned_configs = []
     for config in configs:
         kw = config.kwargs
@@ -384,7 +453,7 @@ def _fbgemm_grouped_gemm(
 @triton.autotune(
     configs=_NV_WS_CONFIGS,
     key=["G", "M_BUCKET", "N", "K"],
-    prune_configs_by={"early_config_prune": early_config_prune},
+    prune_configs_by={"early_config_prune": early_config_prune_ws},
     restore_value=["c_ptr"],  # restore for scatter_add fusion
 )
 @triton.jit
@@ -712,7 +781,7 @@ def _fbgemm_grouped_gemm_fp8_rowwise(
     key=["G", "M_BUCKET", "N", "K"],
     prune_configs_by={
         "early_config_prune": functools.partial(
-            early_config_prune, dtype=TT_FP8_DTYPE, dtsize=1
+            early_config_prune_ws, dtype=TT_FP8_DTYPE, dtsize=1
         )
     },
     restore_value=["c_ptr"],  # restore for scatter_add fusion