From f00103da4af3b2404002630473cf0e4ee12d76d9 Mon Sep 17 00:00:00 2001 From: Yifei Xu Date: Wed, 29 Jul 2026 15:49:26 -0700 Subject: [PATCH 1/3] [Pallas] Downgrade compact_worklist grouping gracefully when kernel doesn't match MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `pallas_worklist_grouping` became a search-space knob independent of `pallas_loop_type` in #3131. Kernels whose source doesn't match the compact-worklist pattern (e.g. `examples/rms_norm.py::rms_norm_bwd`, which uses a nested `hl.tile` and no `hl.grid`) now hit `detect_compact_worklist_plan`'s `exc.InvalidConfig` raise whenever autotune picks `pallas_worklist_grouping in (1, 2)`, and that raise escapes past autotune's skip path, FAILing the whole sweep step. Catch `InvalidConfig` in `_setup_compact_worklist` and leave `env.compact_worklist_plan = None`. Downstream lowering paths already gate on `plan is not None`, so grouping silently downgrades to a no-op for kernels that can't use it. Regression seen in the 2026-07-28 TPU sweep: `rms_norm-bwd[8192,8192]` FAILed with "compact_worklist could not locate the owner hl.grid loop in source." after #3131 landed. Verified fixed on TPU v7 — both rms_norm-bwd shapes PASS full autotune, [8192,8192] gets a 2.34x speedup vs eager (matching pre-regression 2.62x within noise). --- helion/_compiler/pallas/backend.py | 21 +++++++++++++--- test/test_pallas_worklist.py | 39 ++++++++++++++++++++++++++++++ 2 files changed, 57 insertions(+), 3 deletions(-) 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..094ec483d 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") From 2438a9636a59cf97b152e4ce6190ffc690608e4b Mon Sep 17 00:00:00 2001 From: Yifei Xu Date: Thu, 30 Jul 2026 15:06:37 -0700 Subject: [PATCH 2/3] Update test_unsupported_kernel_raises for graceful downgrade The raise this test asserted is exactly the behavior this PR removes: _setup_compact_worklist now catches InvalidConfig and downgrades grouping to a no-op instead of propagating past the autotuner. Convert it to test_unsupported_kernel_downgrades, asserting to_triton_code succeeds and emits no compact-worklist builder. --- test/test_pallas_worklist.py | 17 ++++++++++++++--- 1 file changed, 14 insertions(+), 3 deletions(-) diff --git a/test/test_pallas_worklist.py b/test/test_pallas_worklist.py index 094ec483d..93e6f8c31 100644 --- a/test/test_pallas_worklist.py +++ b/test/test_pallas_worklist.py @@ -1109,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()): @@ -1120,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)) From e5e1b117decd704a1101f7b0b8c339dfeb4cda07 Mon Sep 17 00:00:00 2001 From: Yifei Xu Date: Fri, 31 Jul 2026 15:07:26 -0700 Subject: [PATCH 3/3] test_unsupported_kernel_downgrades: use to_code, not the to_triton_code alias The Pallas backend generates JAX/Pallas code, not Triton. The Triton-flavored method name is a backward-compat alias for to_code (see helion/runtime/kernel.py) that pre-dates multi-backend support; using it in a Pallas test is misleading. --- test/test_pallas_worklist.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/test/test_pallas_worklist.py b/test/test_pallas_worklist.py index 93e6f8c31..d042c6de7 100644 --- a/test/test_pallas_worklist.py +++ b/test/test_pallas_worklist.py @@ -1128,9 +1128,7 @@ def fn(x, y): bound = kernel.bind(args) for grouping in (1, 2): with self.subTest(grouping=grouping): - code = bound.to_triton_code( - _worklist_config([16, 16], grouping=grouping) - ) + code = bound.to_code(_worklist_config([16, 16], grouping=grouping)) self.assertNotIn("_compact_build_worklist=", code) self.assertNotIn("def _build_worklist(", code)