Cap persistent kernel grids at logical PID count

docs/api/config.md

@@ -103,6 +103,14 @@ Configs are typically discovered automatically through autotuning, but can also
    - ``"persistent_blocked"``: Persistent kernels with blocked work distribution
    - ``"persistent_interleaved"``: Persistent kernels with interleaved distribution
 
+.. autoattribute:: Config.num_sm_multiplier
+
+   Controls persistent multi-occupancy. The physical launch grid is capped at
+   ``min(total logical program IDs, num SMs * num_sm_multiplier)`` so the
+   autotuner can select an overprovisioned multiplier without launching workers
+   that cannot own work. Device-dependent bounds retain the SM-derived grid
+   because their logical program count is unavailable to the host launcher.
+
 .. autoattribute:: Config.l2_groupings
 
    Controls reordering of program IDs to improve L2 cache locality.

helion/_compiler/program_id.py

@@ -240,12 +240,26 @@ def codegen_grid(self) -> ast.AST:
         """Generate grid launch expression for kernel execution."""
         raise NotImplementedError
 
+    def codegen_grid_for_total(self, total_pids_expr: str) -> ast.AST:
+        """Generate a grid for an explicit aggregate logical PID count."""
+        return self.codegen_grid()
+
+    def codegen_uncapped_grid(self) -> ast.AST:
+        """Generate a grid when the aggregate logical PID count is unknown."""
+        return self.codegen_grid()
+
     def total_pids_expr(self, *, is_device: bool) -> str:
         """Get total PIDs expression for device or host."""
         return " * ".join(
             f"({pid.num_pids_expr(is_device=is_device)})" for pid in self.pid_info
         )
 
+    def host_total_pids_expr(self) -> str | None:
+        """Return the host-computable PID count, or None for device bounds."""
+        if any(isinstance(pid.numel, str) for pid in self.pid_info):
+            return None
+        return self.total_pids_expr(is_device=False)
+
     def setup_persistent_kernel(
         self, device_function: DeviceFunction, total_pids_expr: str | None = None
     ) -> list[ast.stmt] | None:
@@ -393,6 +407,12 @@ def total_pids_expr(self, *, is_device: bool) -> str:
         cdivs = [pid.total_pids_expr(is_device=is_device) for pid in self.cases]
         return " + ".join(cdivs)
 
+    def host_total_pids_expr(self) -> str | None:
+        totals = [case.host_total_pids_expr() for case in self.cases]
+        if any(total is None for total in totals):
+            return None
+        return " + ".join(total for total in totals if total is not None)
+
     def codegen(self, state: CodegenState) -> None:
         blocks = self._get_cdiv_blocks(state, exclude_last=True)
         if blocks:
@@ -408,8 +428,10 @@ def codegen(self, state: CodegenState) -> None:
     def codegen_grid(self) -> ast.AST:
         # Check if any of the pids is a persistent strategy
         if self.cases[0]._is_persistent():
-            # Use SM count grid for persistent kernels
-            return self.cases[0].codegen_grid()
+            total_pids_expr = self.host_total_pids_expr()
+            if total_pids_expr is not None:
+                return self.cases[0].codegen_grid_for_total(total_pids_expr)
+            return self.cases[0].codegen_uncapped_grid()
 
         # When persistent kernels are not active, use the full grid size
         host_cdivs = [pid.total_pids_expr(is_device=False) for pid in self.cases]
@@ -838,6 +860,14 @@ def codegen_grid(self) -> ast.AST:
         assert self.parent_strategy is not None
         return self.parent_strategy.codegen_grid()
 
+    def codegen_grid_for_total(self, total_pids_expr: str) -> ast.AST:
+        assert self.parent_strategy is not None
+        return self.parent_strategy.codegen_grid_for_total(total_pids_expr)
+
+    def codegen_uncapped_grid(self) -> ast.AST:
+        assert self.parent_strategy is not None
+        return self.parent_strategy.codegen_uncapped_grid()
+
     def setup_persistent_kernel(
         self, device_function: DeviceFunction, total_pids_expr: str | None = None
     ) -> list[ast.stmt] | None:
@@ -857,6 +887,10 @@ def total_pids_expr(self, *, is_device: bool) -> str:
         assert self.parent_strategy is not None
         return self.parent_strategy.total_pids_expr(is_device=is_device)
 
+    def host_total_pids_expr(self) -> str | None:
+        assert self.parent_strategy is not None
+        return self.parent_strategy.host_total_pids_expr()
+
 
 class PersistentProgramIDs(ProgramIDs):
     """Base class for persistent kernels that use num_sms grid size."""
@@ -911,7 +945,18 @@ def get_device_str(self) -> str:
         return f"torch.{device!r}"
 
     def codegen_grid(self) -> ast.AST:
-        # Use num_sms * multiplier for persistent kernels (multi-occupancy)
+        total_pids_expr = self.host_total_pids_expr()
+        if total_pids_expr is None:
+            return expr_from_string(f"({self.grid_size_expr},)")
+        return self.codegen_grid_for_total(total_pids_expr)
+
+    def codegen_grid_for_total(self, total_pids_expr: str) -> ast.AST:
+        # Avoid launching workers that cannot own a logical PID. Keep the
+        # persistent stride unchanged so overprovisioned configurations still
+        # map virtual PIDs exactly as before when total work exceeds the grid.
+        return expr_from_string(f"(min(({total_pids_expr}), ({self.grid_size_expr})),)")
+
+    def codegen_uncapped_grid(self) -> ast.AST:
         return expr_from_string(f"({self.grid_size_expr},)")
 
     def _persistent_setup_statements(self, total_pids_expr: str) -> list[ast.stmt]:
@@ -1359,6 +1404,14 @@ def codegen_grid(self) -> ast.AST:
         grid_work_clusters = self._tcgen05_grid_work_clusters_expr(total_clusters)
         return expr_from_string(f"({cluster_m}, {cluster_n}, {grid_work_clusters})")
 
+    def codegen_grid_for_total(self, total_pids_expr: str) -> ast.AST:
+        # tcgen05 uses its own cluster-aware scheduler and already caps its
+        # persistent work-cluster grid where required.
+        return self.codegen_grid()
+
+    def codegen_uncapped_grid(self) -> ast.AST:
+        return self.codegen_grid()
+
     def _tcgen05_logical_m_coord_expr(self, coord: str) -> str:
         if self._tcgen05_is_two_cta():
             return f"({coord}) // cutlass.Int32({self._tcgen05_cluster_m()})"

helion/runtime/config.py

@@ -19,7 +19,7 @@
 EvictionPolicyLiteral = Literal["", "first", "last"]
 LoadCacheModifierLiteral = Literal["", ".cg"]
 StoreCacheModifierLiteral = Literal["", ".cs", ".wt"]
-NumSmMultiplierLiteral = Literal[1, 2, 4, 8]
+NumSmMultiplierLiteral = Literal[1, 2, 4, 8, 16, 32, 64, 128]
 MaxnregLiteral = Literal[32, 64, 128, 256] | None
 
 
@@ -83,8 +83,9 @@ def __init__(
             num_warps: Number of warps per block.
             num_stages: Number of stages for software pipelining.
             pid_type: Program ID type strategy ("flat", "xyz", "persistent_blocked", "persistent_interleaved").
-            num_sm_multiplier: Multiplier for the number of SMs in persistent kernels (1, 2, 4, 8).
-                Controls multi-occupancy by launching N * num_sms thread blocks instead of just num_sms.
+            num_sm_multiplier: Multiplier for the number of SMs in persistent kernels.
+                Controls multi-occupancy with a grid capped at the total logical
+                program count: min(total_pids, N * num_sms).
             maxnreg: Maximum number of registers per thread (None, 32, 64, 128, 256).
                 Lower values allow higher occupancy but may hurt performance. Used with persistent kernels
                 to ensure multi-occupancy can be achieved.

test/test_persistent_kernels.py

@@ -683,42 +683,23 @@ def test_kernel(x: torch.Tensor) -> torch.Tensor:
         import re
 
         # Look for _launcher(_kernel_name, (grid_size), ...) pattern
-        # Use a pattern that handles nested parentheses in grid expressions
         grid_pattern = r"_launcher\([^,]+,\s*\((.+?)\)\s*,"
         flat_grid_match = re.search(grid_pattern, code_flat)
-        persistent_blocked_grid_match = re.search(grid_pattern, code_persistent_blocked)
-        persistent_interleaved_grid_match = re.search(
-            grid_pattern, code_persistent_interleaved
-        )
 
         self.assertIsNotNone(flat_grid_match, "Could not find grid size in flat code")
-        self.assertIsNotNone(
-            persistent_blocked_grid_match,
-            "Could not find grid size in persistent blocked code",
-        )
-        self.assertIsNotNone(
-            persistent_interleaved_grid_match,
-            "Could not find grid size in persistent interleaved code",
-        )
 
         flat_grid = flat_grid_match.group(1).rstrip(",")  # Remove trailing comma
-        persistent_blocked_grid = persistent_blocked_grid_match.group(1).rstrip(",")
-        persistent_interleaved_grid = persistent_interleaved_grid_match.group(1).rstrip(
-            ","
-        )
 
         # Flat should use the full grid size calculation (ceiling division)
         self.assertIn("//", flat_grid)
 
-        # Persistent kernels should use NUM_SMS
-        self.assertEqual(
-            persistent_blocked_grid,
-            "_NUM_SM",
-        )
-        self.assertEqual(
-            persistent_interleaved_grid,
-            "_NUM_SM",
+        # Persistent kernels cap the SM-derived grid at the logical PID count.
+        self.assertIn("_launcher(_helion_test_kernel, (min(", code_persistent_blocked)
+        self.assertIn("_NUM_SM", code_persistent_blocked)
+        self.assertIn(
+            "_launcher(_helion_test_kernel, (min(", code_persistent_interleaved
         )
+        self.assertIn("_NUM_SM", code_persistent_interleaved)
 
     def test_persistent_loop_variable_names(self):
         """Test that persistent kernels use correct virtual_pid variable names."""
@@ -1204,6 +1185,39 @@ def data_dependent_kernel(
         self.assertIn("total_pids", code)
         self.assertIn("virtual_pid", code)
 
+    @skipIfRefEager("Code pattern checking not applicable in ref eager mode")
+    def test_mixed_static_and_data_dependent_bounds_use_uncapped_grid(self):
+        """An unknown aggregate PID count must not cap to the static first loop."""
+
+        @helion.kernel(
+            autotune_effort="none",
+            config=helion.Config(pid_type="persistent_interleaved"),
+        )
+        def mixed_bounds_kernel(
+            x: torch.Tensor, num_elements: torch.Tensor
+        ) -> tuple[torch.Tensor, torch.Tensor]:
+            static_result = x.new_zeros(x.size())
+            dynamic_result = x.new_zeros(x.size())
+            for tile in hl.tile(32, block_size=32):
+                static_result[tile] = x[tile] + 1
+            for tile in hl.tile(num_elements, block_size=32):
+                dynamic_result[tile] = x[tile] + 2
+            return static_result, dynamic_result
+
+        x = torch.randn([128], device=DEVICE)
+        num_elements = torch.tensor(64, device=DEVICE)
+        code, (static_result, dynamic_result) = code_and_output(
+            mixed_bounds_kernel, (x, num_elements)
+        )
+
+        expected_static = torch.zeros_like(x)
+        expected_static[:32] = x[:32] + 1
+        expected_dynamic = torch.zeros_like(x)
+        expected_dynamic[:64] = x[:64] + 2
+        torch.testing.assert_close(static_result, expected_static)
+        torch.testing.assert_close(dynamic_result, expected_dynamic)
+        self.assertIn("_launcher(_helion_mixed_bounds_kernel, (_NUM_SM,)", code)
+
 
 @onlyBackends(["triton"])
 class TestNumSmMultiplier(RefEagerTestBase, TestCase):
@@ -1221,21 +1235,24 @@ def test_num_sm_multiplier_blocked_grid_size(self):
         code_m1, result_m1 = code_and_output(
             add_kernel, args, pid_type="persistent_blocked", num_sm_multiplier=1
         )
-        self.assertIn("(_NUM_SM,)", code_m1)
+        self.assertIn("min(", code_m1)
+        self.assertIn("_NUM_SM", code_m1)
         self.assertIn("tl.cdiv(total_pids, _NUM_SM)", code_m1)
 
         # Test with multiplier=2
         code_m2, result_m2 = code_and_output(
             add_kernel, args, pid_type="persistent_blocked", num_sm_multiplier=2
         )
-        self.assertIn("(_NUM_SM * 2,)", code_m2)
+        self.assertIn("min(", code_m2)
+        self.assertIn("_NUM_SM * 2", code_m2)
         self.assertIn("tl.cdiv(total_pids, _NUM_SM * 2)", code_m2)
 
         # Test with multiplier=4
         code_m4, result_m4 = code_and_output(
             add_kernel, args, pid_type="persistent_blocked", num_sm_multiplier=4
         )
-        self.assertIn("(_NUM_SM * 4,)", code_m4)
+        self.assertIn("min(", code_m4)
+        self.assertIn("_NUM_SM * 4", code_m4)
         self.assertIn("tl.cdiv(total_pids, _NUM_SM * 4)", code_m4)
 
         # All should produce the same result
@@ -1256,28 +1273,32 @@ def test_num_sm_multiplier_interleaved_step(self):
         code_m1, result_m1 = code_and_output(
             add_kernel, args, pid_type="persistent_interleaved", num_sm_multiplier=1
         )
-        self.assertIn("(_NUM_SM,)", code_m1)
+        self.assertIn("min(", code_m1)
+        self.assertIn("_NUM_SM", code_m1)
         self.assertIn("tl.range(tl.program_id(0), total_pids, _NUM_SM", code_m1)
 
         # Test with multiplier=2
         code_m2, result_m2 = code_and_output(
             add_kernel, args, pid_type="persistent_interleaved", num_sm_multiplier=2
         )
-        self.assertIn("(_NUM_SM * 2,)", code_m2)
+        self.assertIn("min(", code_m2)
+        self.assertIn("_NUM_SM * 2", code_m2)
         self.assertIn("tl.range(tl.program_id(0), total_pids, _NUM_SM * 2", code_m2)
 
-        # Test with multiplier=8
-        code_m8, result_m8 = code_and_output(
-            add_kernel, args, pid_type="persistent_interleaved", num_sm_multiplier=8
+        # Test with multiplier=16. The host grid is capped at total_pids while
+        # the device loop retains the full persistent stride.
+        code_m16, result_m16 = code_and_output(
+            add_kernel, args, pid_type="persistent_interleaved", num_sm_multiplier=16
         )
-        self.assertIn("(_NUM_SM * 8,)", code_m8)
-        self.assertIn("tl.range(tl.program_id(0), total_pids, _NUM_SM * 8", code_m8)
+        self.assertIn("min(", code_m16)
+        self.assertIn("_NUM_SM * 16", code_m16)
+        self.assertIn("tl.range(tl.program_id(0), total_pids, _NUM_SM * 16", code_m16)
 
         # All should produce the same result
         expected = args[0] + args[1]
         torch.testing.assert_close(result_m1, expected)
         torch.testing.assert_close(result_m2, expected)
-        torch.testing.assert_close(result_m8, expected)
+        torch.testing.assert_close(result_m16, expected)
 
     @skipIfRefEager("Code pattern checking not applicable in ref eager mode")
     def test_num_sm_multiplier_matmul_correctness(self):