Skip to content

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

Open
norx1991 wants to merge 3 commits into
mainfrom
yifeixu/pallas-worklist-graceful-downgrade
Open

[Pallas] Downgrade compact_worklist grouping gracefully when kernel doesn't match#3206
norx1991 wants to merge 3 commits into
mainfrom
yifeixu/pallas-worklist-graceful-downgrade

Conversation

@norx1991

@norx1991 norx1991 commented Jul 29, 2026

Copy link
Copy Markdown
Contributor

Summary

pallas_worklist_grouping became an autotune knob independent of pallas_loop_type in #3131. Its config-space gating (offer whenever the kernel has a nested hl.tile shape) is looser than detect_compact_worklist_plan's requirement (must have an hl.grid owner), so autotune can pick pallas_worklist_grouping in (1, 2) for kernels the compact-worklist detector can't handle. In that case detect_compact_worklist_plan raises exc.InvalidConfig("compact_worklist could not locate the owner hl.grid loop in source.") — and that raise escapes past the autotuner's skip path (it's called via the pre-codegen setup, not via a config-compile-check), FAILing the whole sweep step.

This regressed rms_norm-bwd[8192,8192] in the 07-28 TPU sweep (was 2.62× on 07-07, now FAIL). Reproduces on any kernel with the nested hl.tile(x, block_size=…) + hl.tile(mb_cta.begin, mb_cta.end) shape but no hl.grid — e.g. examples/rms_norm.py::rms_norm_bwd.

Catch exc.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. Matches the "autotuner-skippable" semantic the raise's docstring already documented.

Verification

  • Regression test added (test_grouping_downgrades_on_kernel_without_hl_grid): confirmed FAILs on main with the exact InvalidConfig message from the sweep and PASSes with this fix.
  • Pre-existing test updated: test_unsupported_kernel_raisestest_unsupported_kernel_downgrades. It asserted the raise this PR intentionally removes; it now asserts to_triton_code succeeds and emits no compact-worklist builder. test_invalid_worklist_grouping_raises still passes, confirming the try/except doesn't swallow the separate invalid-value validation. Full test_pallas_worklist.py: 86 passed on TPU v7.
  • End-to-end on TPU v7 (full autotune + --subprocess-per-shape):
Shape 07-07 07-28 (main) With fix
rms_norm-bwd[2048,4096] 1.37× 1.60× 1.58×
rms_norm-bwd[8192,8192] 2.62× FAIL 2.85×

@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Meta Open Source bot. label Jul 29, 2026
…oesn't match

`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).
@norx1991
norx1991 force-pushed the yifeixu/pallas-worklist-graceful-downgrade branch from b934298 to f00103d Compare July 30, 2026 00:06
@norx1991
norx1991 marked this pull request as ready for review July 30, 2026 21:03
@norx1991
norx1991 requested review from AmesingFlank and ethche and removed request for AmesingFlank July 30, 2026 21:18
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.
norx1991 added a commit that referenced this pull request Jul 30, 2026
welford fails accuracy on the tritonbench TPU bridge for every implementation
(Helion and torch.compile alike) while passing under run_tpu.py — a torch_tpu
harness issue, not the kernel (a direct check produces correct output). Until
that is understood, run welford via the run_tpu.py holdout (kernels_tpu) where
it measures correctly; the bridge keeps everything else.

rms_norm-bwd stays on the bridge: its bridge failures are fixed at the source
(the backward harness fed each implementation a different random grad_output on
TPU, meta-pytorch/tritonbench#1207; the pallas autotune crash is #3206), so it
no longer needs the holdout.
…de 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.
@norx1991
norx1991 force-pushed the yifeixu/pallas-worklist-graceful-downgrade branch from 7f49323 to e5e1b11 Compare July 31, 2026 22:12
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Meta Open Source bot.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants