Using the hardware specs to derive the cap

Author: umechand-amd

helion/autotuner/config_spec.py

@@ -1051,12 +1051,16 @@ def raise_grid_block_minimums(self) -> None:
     def lower_max_for_imbalanced_grid_dims(self) -> None:
         """Lower max_size for the large grid dimension when the shape is skinny.
 
-        When one grid dimension is much larger than the other (e.g. M=1024, N=8192),
-        the autotuner samples large tile sizes for the big dimension that produce too
-        few grid blocks for good GPU occupancy.  Capping the larger dimension's tile
-        keeps the search in the useful region without hardcoding specific tile values.
+        Mirrors raise_grid_block_minimums: that method raises the minimum tile to
+        prevent too many blocks per dimension; this method lowers the maximum tile
+        to prevent too few blocks when one dimension is much larger than the other.
 
-        Only applied to 2-D grids where max(dim) >= 4 * min(dim).
+        The cap is derived from num_compute_units() so it naturally relaxes on
+        smaller GPUs (fewer SMs/CUs → fewer blocks needed → larger tiles OK).
+
+        Only applied to 2-D grids where max(dim) >= 8 * min(dim).  The 8x threshold
+        is hardware-independent: below it the random sampler has a reasonable chance
+        of finding good balanced-tile configs on its own.
         """
         if len(self.grid_block_ids) != 2:
             return
@@ -1073,16 +1077,24 @@ def lower_max_for_imbalanced_grid_dims(self) -> None:
 
         min_hint = min(hints)
         max_hint = max(hints)
-        if max_hint < min_hint * 4:
-            return  # Square-ish shape — leave the search space alone
+        if max_hint < min_hint * 8:
+            return  # Not severely imbalanced — leave the search space alone
+
+        # Derive cap from actual hardware: require at least 1 block per compute unit
+        # per grid dimension (mirroring raise_grid_block_minimums which uses n_cus*64).
+        n_cus = num_compute_units()
+        n_dims = len(self.grid_block_ids)
+        min_blocks_per_dim = math.ceil(n_cus ** (1.0 / n_dims))
 
-        # Cap the larger dim's tile so that at least 4 blocks cover min_hint,
-        # keeping total blocks comparable across both dims.
-        # e.g. M=1024, N=8192: cap N-tile at max(64, 1024//4) = 256
-        cap = max(64, next_power_of_2(min_hint) // 2)
         for spec, hint in zip(specs, hints, strict=True):
-            if hint == max_hint:
-                spec.update_max(min(spec.max_size, cap))
+            if hint != max_hint:
+                continue
+            max_tile = hint // min_blocks_per_dim
+            if max_tile < 2:
+                continue
+            # Round down to power of two
+            max_tile = 1 << (max_tile.bit_length() - 1)
+            spec.update_max(min(spec.max_size, max_tile))
 
     def create_config_generation(
         self,

test/test_best_available.py

@@ -1100,7 +1100,11 @@ def _make_spec(self, m, n, m_max=None, n_max=None):
         return config_spec
 
     def test_skinny_n_caps_n_tile(self):
-        """Skinny-N (M=1024, N=8192): N-tile max should be capped."""
+        """Skinny-N (M=1024, N=8192): N-tile max should be capped (hardware-derived cap)."""
+        import math
+
+        from helion._compat import num_compute_units
+
         spec = self._make_spec(m=1024, n=8192)
         n_max_before = spec.block_sizes.block_id_lookup(1).max_size
         spec.lower_max_for_imbalanced_grid_dims()
@@ -1110,8 +1114,10 @@ def test_skinny_n_caps_n_tile(self):
         self.assertLess(n_max_after, n_max_before)
         # M tile (the smaller dim) should be unchanged
         self.assertEqual(m_max_after, spec.block_sizes.block_id_lookup(0).max_size)
-        # Cap should allow at least 4 blocks on N dim
-        self.assertGreaterEqual(8192 // n_max_after, 4)
+        # Cap is hardware-derived: at least min_blocks_per_dim blocks on N dim
+        n_cus = num_compute_units()
+        min_blocks_per_dim = math.ceil(n_cus**0.5)
+        self.assertGreaterEqual(8192 // n_max_after, min_blocks_per_dim)
 
     def test_skinny_m_caps_m_tile(self):
         """Skinny-M (M=8192, N=1024): M-tile max should be capped."""
@@ -1131,8 +1137,8 @@ def test_square_shape_unchanged(self):
         self.assertEqual(spec.block_sizes.block_id_lookup(1).max_size, n_max_before)
 
     def test_slightly_imbalanced_unchanged(self):
-        """3:1 ratio (M=1024, N=3072): below the 4x threshold, no change."""
-        spec = self._make_spec(m=1024, n=3072)
+        """6.4x ratio (M=1280, N=8192): below the 8x threshold, no change (covers int4_gemm regression)."""
+        spec = self._make_spec(m=1280, n=8192)
         n_max_before = spec.block_sizes.block_id_lookup(1).max_size
         spec.lower_max_for_imbalanced_grid_dims()
         self.assertEqual(spec.block_sizes.block_id_lookup(1).max_size, n_max_before)