Update early prune to support any N (#4734)

renganxu · facebook-github-bot · commit 6a7e5d80c3f3 · 2025-08-19T15:56:50.000-07:00
Summary: Pull Request resolved: #4734 X-link: facebookresearch/FBGEMM#1754 When N is small and not multiple of BLOCK_N, e.g. 221, and M=121, the early prune will prune all configs and then an empty config is returned. This will lead to error "ValueError: min() arg is an empty sequence". This diff fixes this issue. Reviewed By: sunfish2010 Differential Revision: D80542648 fbshipit-source-id: 8e29d90ac87978b1bb4c69002b3583d10ba5f8a3
diff --git a/fbgemm_gpu/experimental/gemm/triton_gemm/grouped_gemm.py b/fbgemm_gpu/experimental/gemm/triton_gemm/grouped_gemm.py
@@ -184,7 +184,7 @@ def early_config_prune(configs, named_args, dtsize=None, dtype=None, **kwargs):
         num_sm = driver.active.utils.get_device_properties(device)[
             "multiprocessor_count"
         ]
-        N_TILES = N // BLOCK_N
+        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 (

Original file line number	Diff line number	Diff line change
`@@ -184,7 +184,7 @@ def early_config_prune(configs, named_args, dtsize=None, dtype=None, **kwargs):`
`184`	`184`	`num_sm = driver.active.utils.get_device_properties(device)[`
`185`	`185`	`"multiprocessor_count"`
`186`	`186`	`]`
`187`		`- N_TILES = N // BLOCK_N`
	`187`	`+ N_TILES = (N + BLOCK_N - 1) // BLOCK_N`
`188`	`188`	`MIN_N_TILES = 32 if torch.version.hip else 64`
`189`	`189`	`# 4. make sure we don't load N tiles that are too big`
`190`	`190`	`if (`