Move Pallas launchers into helion/runtime/pallas/launcher.py - #3182
Merged
Conversation
AmesingFlank
force-pushed
the
AmesingFlank/stack/95
branch
from
July 28, 2026 16:02
301565b to
e64ad7c
Compare
AmesingFlank
force-pushed
the
AmesingFlank/stack/96
branch
from
July 28, 2026 16:02
7cc19f5 to
aa4871d
Compare
This was referenced Jul 28, 2026
meta-cla
Bot
added
the
CLA Signed
This was referenced Jul 28, 2026
AmesingFlank
marked this pull request as ready for review
July 28, 2026 17:51
AmesingFlank
marked this pull request as draft
July 28, 2026 20:09
AmesingFlank
force-pushed
the
AmesingFlank/stack/96
branch
from
July 28, 2026 20:09
aa4871d to
3742da2
Compare
AmesingFlank
marked this pull request as ready for review
July 28, 2026 20:10
AmesingFlank
marked this pull request as draft
July 28, 2026 21:14
AmesingFlank
force-pushed
the
AmesingFlank/stack/96
branch
from
July 28, 2026 21:14
3742da2 to
8b3ec21
Compare
AmesingFlank
marked this pull request as ready for review
July 29, 2026 09:04
norx1991
approved these changes
Jul 29, 2026
AmesingFlank
added a commit
that referenced
this pull request
Jul 29, 2026
Relocate the entire Pallas launch + compile core (`default_pallas_launcher`, the ~40 `_pallas_*` block-spec / compile / cache helpers, `_DirectCallKernel`, `_LauncherFastPath`, `_PallasCompileResult`, `_PallasInterpretCallable`, the compact-worklist launcher path, and the torch<->jax converters) out of the 2500-line catch-all `helion/runtime/__init__.py` into a dedicated `helion/runtime/pallas/launcher.py` (2388 lines moved). Pure move: the only edits are re-basing the two relative imports one level deeper (`from .settings` -> `from ..settings`, `from .._compiler.backend` -> `from ..._compiler.backend`); every function body is unchanged. `helion.runtime` re-exports `default_pallas_launcher` (generated Pallas code imports it), and every direct importer of a moved symbol is repointed to the new home: `pallas_jax_export` (7 `_pallas_*`/`_BlockSpecInfo` helpers), `_compiler/pallas/backend.py` (`_get_vmem_limit_bytes`), `_compiler/pallas/compact_worklist.py` (`compact_ordered_physical_window`), `test_pallas_compact_worklist` (`_get_vmem_limit_bytes`, `compact_ordered_*`, `_compact_raise_if_range_exceeds_window`), and `test_pallas` (`_build_matmul_dot_general_jit_fn`, imported as the launcher module so the `patch.object` spy still targets the launcher's namespace). Two stale `runtime._compact_raise_if_range_exceeds_window` doc comments are updated. The now-unused `dataclass`/`Literal`/`TYPE_CHECKING`/`jax` imports are dropped from `runtime/__init__.py`. Isolating this into one file (still importing a couple of `helion` helpers for now) sets up making it torch+jax-only so the precompiler can bulk-export it. Verified on the TPU pod: pallas add/topk, all 5 `test_pallas.py` `jax_fn` tests, compact-worklist vmem/budget, resident-window-guard, and matmul-dot-general tests pass; `ruff`/`pyrefly` clean (83 baseline); Triton `add` still correct on H100. (One pre-existing TPU failure, `test_jax_fn_under_jit_matches_eager`, and the interpret-mode `add` tests fail identically on the base commit -- pod JAX-version skew, unrelated to this move.) stack-info: PR: #3182, branch: AmesingFlank/stack/96
AmesingFlank
marked this pull request as draft
July 29, 2026 20:37
AmesingFlank
force-pushed
the
AmesingFlank/stack/96
branch
from
July 29, 2026 20:37
065a798 to
c7a6ab1
Compare
AmesingFlank
marked this pull request as ready for review
July 29, 2026 20:38
AmesingFlank
added a commit
that referenced
this pull request
Jul 30, 2026
Relocate the entire Pallas launch + compile core (`default_pallas_launcher`, the ~40 `_pallas_*` block-spec / compile / cache helpers, `_DirectCallKernel`, `_LauncherFastPath`, `_PallasCompileResult`, `_PallasInterpretCallable`, the compact-worklist launcher path, and the torch<->jax converters) out of the 2500-line catch-all `helion/runtime/__init__.py` into a dedicated `helion/runtime/pallas/launcher.py` (2388 lines moved). Pure move: the only edits are re-basing the two relative imports one level deeper (`from .settings` -> `from ..settings`, `from .._compiler.backend` -> `from ..._compiler.backend`); every function body is unchanged. `helion.runtime` re-exports `default_pallas_launcher` (generated Pallas code imports it), and every direct importer of a moved symbol is repointed to the new home: `pallas_jax_export` (7 `_pallas_*`/`_BlockSpecInfo` helpers), `_compiler/pallas/backend.py` (`_get_vmem_limit_bytes`), `_compiler/pallas/compact_worklist.py` (`compact_ordered_physical_window`), `test_pallas_compact_worklist` (`_get_vmem_limit_bytes`, `compact_ordered_*`, `_compact_raise_if_range_exceeds_window`), and `test_pallas` (`_build_matmul_dot_general_jit_fn`, imported as the launcher module so the `patch.object` spy still targets the launcher's namespace). Two stale `runtime._compact_raise_if_range_exceeds_window` doc comments are updated. The now-unused `dataclass`/`Literal`/`TYPE_CHECKING`/`jax` imports are dropped from `runtime/__init__.py`. Isolating this into one file (still importing a couple of `helion` helpers for now) sets up making it torch+jax-only so the precompiler can bulk-export it. Verified on the TPU pod: pallas add/topk, all 5 `test_pallas.py` `jax_fn` tests, compact-worklist vmem/budget, resident-window-guard, and matmul-dot-general tests pass; `ruff`/`pyrefly` clean (83 baseline); Triton `add` still correct on H100. (One pre-existing TPU failure, `test_jax_fn_under_jit_matches_eager`, and the interpret-mode `add` tests fail identically on the base commit -- pod JAX-version skew, unrelated to this move.) stack-info: PR: #3182, branch: AmesingFlank/stack/96
AmesingFlank
marked this pull request as draft
July 30, 2026 01:55
AmesingFlank
force-pushed
the
AmesingFlank/stack/96
branch
from
July 30, 2026 01:56
c7a6ab1 to
7a96ae6
Compare
AmesingFlank
marked this pull request as ready for review
July 30, 2026 01:56
AmesingFlank
added a commit
that referenced
this pull request
Jul 30, 2026
Relocate the entire Pallas launch + compile core (`default_pallas_launcher`, the ~40 `_pallas_*` block-spec / compile / cache helpers, `_DirectCallKernel`, `_LauncherFastPath`, `_PallasCompileResult`, `_PallasInterpretCallable`, the compact-worklist launcher path, and the torch<->jax converters) out of the 2500-line catch-all `helion/runtime/__init__.py` into a dedicated `helion/runtime/pallas/launcher.py` (2388 lines moved). Pure move: the only edits are re-basing the two relative imports one level deeper (`from .settings` -> `from ..settings`, `from .._compiler.backend` -> `from ..._compiler.backend`); every function body is unchanged. `helion.runtime` re-exports `default_pallas_launcher` (generated Pallas code imports it), and every direct importer of a moved symbol is repointed to the new home: `pallas_jax_export` (7 `_pallas_*`/`_BlockSpecInfo` helpers), `_compiler/pallas/backend.py` (`_get_vmem_limit_bytes`), `_compiler/pallas/compact_worklist.py` (`compact_ordered_physical_window`), `test_pallas_compact_worklist` (`_get_vmem_limit_bytes`, `compact_ordered_*`, `_compact_raise_if_range_exceeds_window`), and `test_pallas` (`_build_matmul_dot_general_jit_fn`, imported as the launcher module so the `patch.object` spy still targets the launcher's namespace). Two stale `runtime._compact_raise_if_range_exceeds_window` doc comments are updated. The now-unused `dataclass`/`Literal`/`TYPE_CHECKING`/`jax` imports are dropped from `runtime/__init__.py`. Isolating this into one file (still importing a couple of `helion` helpers for now) sets up making it torch+jax-only so the precompiler can bulk-export it. Verified on the TPU pod: pallas add/topk, all 5 `test_pallas.py` `jax_fn` tests, compact-worklist vmem/budget, resident-window-guard, and matmul-dot-general tests pass; `ruff`/`pyrefly` clean (83 baseline); Triton `add` still correct on H100. (One pre-existing TPU failure, `test_jax_fn_under_jit_matches_eager`, and the interpret-mode `add` tests fail identically on the base commit -- pod JAX-version skew, unrelated to this move.) stack-info: PR: #3182, branch: AmesingFlank/stack/96
AmesingFlank
marked this pull request as draft
July 30, 2026 02:43
AmesingFlank
force-pushed
the
AmesingFlank/stack/96
branch
from
July 30, 2026 02:43
7a96ae6 to
889a66f
Compare
AmesingFlank
marked this pull request as ready for review
July 30, 2026 02:43
AmesingFlank
marked this pull request as draft
July 30, 2026 03:17
AmesingFlank
force-pushed
the
AmesingFlank/stack/96
branch
from
July 30, 2026 03:17
889a66f to
307dfec
Compare
AmesingFlank
marked this pull request as ready for review
July 30, 2026 03:17
Relocate the entire Pallas launch + compile core (`default_pallas_launcher`, the ~40 `_pallas_*` block-spec / compile / cache helpers, `_DirectCallKernel`, `_LauncherFastPath`, `_PallasCompileResult`, `_PallasInterpretCallable`, the compact-worklist launcher path, and the torch<->jax converters) out of the 2500-line catch-all `helion/runtime/__init__.py` into a dedicated `helion/runtime/pallas/launcher.py` (2388 lines moved). Pure move: the only edits are re-basing the two relative imports one level deeper (`from .settings` -> `from ..settings`, `from .._compiler.backend` -> `from ..._compiler.backend`); every function body is unchanged. `helion.runtime` re-exports `default_pallas_launcher` (generated Pallas code imports it), and every direct importer of a moved symbol is repointed to the new home: `pallas_jax_export` (7 `_pallas_*`/`_BlockSpecInfo` helpers), `_compiler/pallas/backend.py` (`_get_vmem_limit_bytes`), `_compiler/pallas/compact_worklist.py` (`compact_ordered_physical_window`), `test_pallas_compact_worklist` (`_get_vmem_limit_bytes`, `compact_ordered_*`, `_compact_raise_if_range_exceeds_window`), and `test_pallas` (`_build_matmul_dot_general_jit_fn`, imported as the launcher module so the `patch.object` spy still targets the launcher's namespace). Two stale `runtime._compact_raise_if_range_exceeds_window` doc comments are updated. The now-unused `dataclass`/`Literal`/`TYPE_CHECKING`/`jax` imports are dropped from `runtime/__init__.py`. Isolating this into one file (still importing a couple of `helion` helpers for now) sets up making it torch+jax-only so the precompiler can bulk-export it. Verified on the TPU pod: pallas add/topk, all 5 `test_pallas.py` `jax_fn` tests, compact-worklist vmem/budget, resident-window-guard, and matmul-dot-general tests pass; `ruff`/`pyrefly` clean (83 baseline); Triton `add` still correct on H100. (One pre-existing TPU failure, `test_jax_fn_under_jit_matches_eager`, and the interpret-mode `add` tests fail identically on the base commit -- pod JAX-version skew, unrelated to this move.) stack-info: PR: #3182, branch: AmesingFlank/stack/96
AmesingFlank
marked this pull request as draft
July 30, 2026 03:28
AmesingFlank
force-pushed
the
AmesingFlank/stack/96
branch
from
July 30, 2026 03:29
307dfec to
45cc490
Compare
AmesingFlank
marked this pull request as ready for review
July 30, 2026 03:29
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Stacked PRs:
OutputCodeOptions(jax_fn=True)for Pallas (JAX-array entrypoints) #3186OutputCodeOptions(allow_helion_deps=False)for the Pallas backend (TorchTPU) #3184OutputCodeOptions(allow_helion_deps=False)toBoundKernel.to_codefor standalone Triton kernels #3181helion/runtime/pallas/launcher.py#3182Move Pallas launchers into
helion/runtime/pallas/launcher.pyRelocate the entire Pallas launch + compile core (
default_pallas_launcher, the~40
_pallas_*block-spec / compile / cache helpers,_DirectCallKernel,_LauncherFastPath,_PallasCompileResult,_PallasInterpretCallable, thecompact-worklist launcher path, and the torch<->jax converters) out of the
2500-line catch-all
helion/runtime/__init__.pyinto a dedicatedhelion/runtime/pallas/launcher.py(2388 lines moved).Pure move: the only edits are re-basing the two relative imports one level deeper
(
from .settings->from ..settings,from .._compiler.backend->from ..._compiler.backend); every function body is unchanged.helion.runtimere-exportsdefault_pallas_launcher(generated Pallas codeimports it), and every direct importer of a moved symbol is repointed to the new
home:
pallas_jax_export(7_pallas_*/_BlockSpecInfohelpers),_compiler/pallas/backend.py(_get_vmem_limit_bytes),_compiler/pallas/compact_worklist.py(compact_ordered_physical_window),test_pallas_compact_worklist(_get_vmem_limit_bytes,compact_ordered_*,_compact_raise_if_range_exceeds_window), andtest_pallas(
_build_matmul_dot_general_jit_fn, imported as the launcher module so thepatch.objectspy still targets the launcher's namespace). Two staleruntime._compact_raise_if_range_exceeds_windowdoc comments are updated. Thenow-unused
dataclass/Literal/TYPE_CHECKING/jaximports are dropped fromruntime/__init__.py.Isolating this into one file (still importing a couple of
helionhelpers fornow) sets up making it torch+jax-only so the precompiler can bulk-export it.
Verified on the TPU pod: pallas add/topk, all 5
test_pallas.pyjax_fntests,compact-worklist vmem/budget, resident-window-guard, and matmul-dot-general tests
pass;
ruff/pyreflyclean (83 baseline); Tritonaddstill correct on H100.(One pre-existing TPU failure,
test_jax_fn_under_jit_matches_eager, and theinterpret-mode
addtests fail identically on the base commit -- pod JAX-versionskew, unrelated to this move.)