diff --git a/helion/_compiler/pallas/backend.py b/helion/_compiler/pallas/backend.py index 53ce08625..89e60f133 100644 --- a/helion/_compiler/pallas/backend.py +++ b/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 diff --git a/test/test_pallas_worklist.py b/test/test_pallas_worklist.py index 6026ed8b9..d042c6de7 100644 --- a/test/test_pallas_worklist.py +++ b/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,10 @@ 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_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))