[Pallas] Downgrade compact_worklist grouping gracefully when kernel doesn't match

helion/_compiler/pallas/backend.py

@@ -1329,8 +1329,16 @@ def _setup_compact_worklist(self, graphs: list[GraphInfo], config: Config) -> No
         grid strategy selects ``WorklistProgramIDs`` and the inner loop remaps its
         begin/end to metadata refs.  Registers the N metadata ref names as
         ``wrapper_only_params`` (kernel-signature-only) and computes the static
-        megablocks ``UPPER``.  ``detect_*`` raises ``exc.InvalidConfig`` on a
-        non-matching kernel (autotuner-skippable).
+        megablocks ``UPPER``.
+
+        ``detect_compact_worklist_plan`` raises ``exc.InvalidConfig`` for
+        kernels whose source doesn't match the compact-worklist pattern
+        (e.g. no owner ``hl.grid``, or an unsupported nest).  Since
+        ``pallas_worklist_grouping`` is a config knob that can be searched
+        independent of the kernel's structure, we catch that here and
+        silently downgrade to no-op grouping rather than propagating the
+        raise past autotune — otherwise a shape whose best-scoring config
+        happens to enable grouping will fail the whole sweep step.
         """
         from ..compile_environment import CompileEnvironment
         from ..device_function import DeviceFunction
@@ -1340,8 +1348,15 @@ def _setup_compact_worklist(self, graphs: list[GraphInfo], config: Config) -> No
 
         env = CompileEnvironment.current()
         host_fn = HostFunction.current()
+        try:
+            detected = detect_compact_worklist_plan(host_fn)
+        except exc.InvalidConfig:
+            # Kernel doesn't support compact-worklist grouping — leave
+            # env.compact_worklist_plan = None so downstream lowering paths
+            # gated on `plan is not None` skip cleanly.
+            return
         plan = dataclasses.replace(
-            detect_compact_worklist_plan(host_fn),
+            detected,
             grouping=cast("int", config.get("pallas_worklist_grouping", 1)),
         )
         env.compact_worklist_plan = plan

test/test_pallas_worklist.py

@@ -470,6 +470,21 @@ def _add_kernel(x, y):
     return out
 
 
+@helion.kernel(backend="pallas", static_shapes=True)
+def _nested_tile_no_grid_kernel(x, y):
+    """Nested ``hl.tile(..., block_size=)`` + ``hl.tile(mb_cta.begin, mb_cta.end)``,
+    no ``hl.grid``. Mirrors ``examples/rms_norm.py::rms_norm_bwd``: the config
+    space offers ``pallas_worklist_grouping`` because the nest is jagged-shaped,
+    but ``detect_compact_worklist_plan`` can't recognise it (needs ``hl.grid``).
+    Used to guard against a raise past the autotuner's skip path."""
+    out = torch.empty_like(x)
+    m_block = hl.register_block_size(x.size(0))
+    for mb_cta in hl.tile(x.size(0), block_size=m_block):
+        for mb in hl.tile(mb_cta.begin, mb_cta.end):
+            out[mb, :] = x[mb, :] + y[mb, :]
+    return out
+
+
 @helion.kernel(backend="pallas", static_shapes=True)
 def _noprep_ordered_kernel(q, k, q_offsets):
     """Jagged ordered reduction whose reused operand (k) has NO transpose prep:
@@ -732,6 +747,30 @@ def test_gating_absent_for_non_jagged(self):
         self.assertNotIn("pallas_loop_type", fields)
         self.assertNotIn("pallas_worklist_grouping", fields)
 
+    def test_grouping_downgrades_on_kernel_without_hl_grid(self):
+        """`pallas_worklist_grouping` is offered for any nested-tile kernel,
+        but ``detect_compact_worklist_plan`` needs an ``hl.grid`` owner. When
+        the kernel uses ``hl.tile(x, block_size=...)`` as the outer loop
+        instead, compiling with ``grouping in (1, 2)`` must silently
+        downgrade to no-op — not raise past the autotuner's skip path."""
+        bk = _nested_tile_no_grid_kernel.bind(
+            (torch.randn(128, 64), torch.randn(128, 64))
+        )
+        fields = bk.env.config_spec._flat_fields()
+        self.assertIn("pallas_worklist_grouping", fields)
+        # Must not raise: compile a config that would have hit
+        # `detect_compact_worklist_plan`'s InvalidConfig raise.
+        for grouping in (1, 2):
+            with self.subTest(grouping=grouping):
+                bk.compile_config(
+                    helion.Config(
+                        block_sizes=[128, 32],
+                        pallas_loop_type="unroll",
+                        pallas_worklist_grouping=grouping,
+                    ),
+                    allow_print=False,
+                )
+
 
 @onlyBackends(["pallas"])
 @unittest.skipUnless(HAS_JAX, "jax not available")
@@ -1070,7 +1109,14 @@ def fn(q, k, v, q_offsets):
         self.assertIn("_wid = pl.program_id(0)", code)
         self.assertIn("work_seq_ref[_wid]", code)
 
-    def test_unsupported_kernel_raises(self):
+    def test_unsupported_kernel_downgrades(self):
+        """A kernel with no owner ``hl.grid`` can't build a compact-worklist
+        plan, but ``pallas_worklist_grouping`` is an independent autotune knob
+        that may still be set on it. Compiling with grouping in (1, 2) must
+        downgrade to a no-op — the generated code contains no compact-worklist
+        builder — rather than raising ``InvalidConfig`` past the autotuner's
+        skip path (which would fail the whole sweep step)."""
+
         def fn(x, y):
             out = torch.empty_like(x)
             for tile in hl.tile(out.size()):
@@ -1081,8 +1127,12 @@ def fn(x, y):
         args = (torch.randn(64, 64), torch.randn(64, 64))
         bound = kernel.bind(args)
         for grouping in (1, 2):
-            with self.subTest(grouping=grouping), self.assertRaises(exc.InvalidConfig):
-                bound.to_triton_code(_worklist_config([16, 16], grouping=grouping))
+            with self.subTest(grouping=grouping):
+                code = bound.to_triton_code(
+                    _worklist_config([16, 16], grouping=grouping)
+                )
+                self.assertNotIn("_compact_build_worklist=", code)
+                self.assertNotIn("def _build_worklist(", code)
 
     def test_invalid_worklist_grouping_raises(self):
         args = (torch.randn(64, 64), torch.randn(64, 64))