Support OutputCodeOptions(allow_helion_deps=False) for the Pallas backend (TorchTPU) - #3184
Merged
Conversation
AmesingFlank
force-pushed
the
AmesingFlank/stack/98
branch
from
July 28, 2026 16:02
e2f682b to
ab11374
Compare
AmesingFlank
force-pushed
the
AmesingFlank/stack/97
branch
from
July 28, 2026 16:02
c6b5c2a to
1003016
Compare
This was referenced Jul 28, 2026
meta-cla
Bot
added
the
CLA Signed
AmesingFlank
force-pushed
the
AmesingFlank/stack/98
branch
from
July 28, 2026 20:09
ab11374 to
2b10a7a
Compare
AmesingFlank
force-pushed
the
AmesingFlank/stack/98
branch
from
July 28, 2026 21:14
2b10a7a to
55b6703
Compare
AmesingFlank
added a commit
that referenced
this pull request
Jul 28, 2026
Add the Pallas precompiler beside its launcher, in `helion/runtime/pallas/precompile.py` (`PallasPrecompiler`), and register it in the orchestrator's per-backend dispatch. It uses the backend-neutral default flow -- inline the dependency-free Pallas launcher (`helion/runtime/pallas/launcher.py`) -- so a precompiled Pallas kernel runs on TPU with torch + jax and no helion at runtime. (`jax_fn=True` is added next.) Verified on the TPU pod: precompiling a Pallas `add` produces a standalone with no `import helion`, and running it on TPU returns correct results. `ruff`/`pyrefly` clean (83 baseline). stack-info: PR: #3184, branch: AmesingFlank/stack/98
AmesingFlank
force-pushed
the
AmesingFlank/stack/98
branch
from
July 28, 2026 22:27
55b6703 to
83fa975
Compare
AmesingFlank
force-pushed
the
AmesingFlank/stack/98
branch
from
July 30, 2026 03:47
ad120fb to
45ebee7
Compare
AmesingFlank
added a commit
that referenced
this pull request
Jul 30, 2026
Add the Pallas precompiler beside its launcher, in `helion/runtime/pallas/precompile.py` (`PallasPrecompiler`), and register it in the orchestrator's per-backend dispatch. It uses the backend-neutral default flow -- inline the dependency-free Pallas launcher (`helion/runtime/pallas/launcher.py`) -- so a precompiled Pallas kernel runs on TPU with torch + jax and no helion at runtime. (`jax_fn=True` is added next.) Verified on the TPU pod: precompiling a Pallas `add` produces a standalone with no `import helion`, and running it on TPU returns correct results. `ruff`/`pyrefly` clean (83 baseline). stack-info: PR: #3184, branch: AmesingFlank/stack/98
AmesingFlank
force-pushed
the
AmesingFlank/stack/98
branch
from
July 30, 2026 08:15
45ebee7 to
cdf0bc3
Compare
AmesingFlank
changed the title
to_code(allow_helion_deps=False) for the Pallas backend (TorchTPU)
AmesingFlank
force-pushed
the
AmesingFlank/stack/98
branch
from
July 30, 2026 15:31
cdf0bc3 to
012f471
Compare
AmesingFlank
changed the title
to_code(allow_helion_deps=False) for the Pallas backend (TorchTPU)OutputCodeOptions(allow_helion_deps=False) for the Pallas backend (TorchTPU)
AmesingFlank
added a commit
that referenced
this pull request
Jul 30, 2026
…ackend (TorchTPU) Implement `PallasBackend.dependency_free_launcher_info` so `to_code(options=OutputCodeOptions(allow_helion_deps=False))` works for Pallas (TorchTPU) kernels: the emitted module inlines the dependency-free Pallas launcher as a local `helion.runtime` shim, and its only runtime deps are `torch` + `jax`. Pallas kernels call two helion-defined pure-`jax` helpers in-kernel -- `divide_filter_topk` (aten.topk lowering) and `flatten_worklist` (compact-worklist builder). Regular `to_code` imports them from helion (added as conditional `library_imports` entries); the dependency-free path drops those imports and embeds the helper source instead (`embedded_helper_source`, called only by the standalone builders). Their modules carry a note to stay free of any `helion` import so the embed stays self-contained. This folds in the previously separate "Embed in-kernel Pallas runtime helpers" change: embedding is now a dependency-free-only step rather than applied to every `to_code`, so regular output is import-based and only standalones embed. stack-info: PR: #3184, branch: AmesingFlank/stack/98
AmesingFlank
force-pushed
the
AmesingFlank/stack/98
branch
from
July 30, 2026 17:37
012f471 to
9e7c129
Compare
AmesingFlank
force-pushed
the
AmesingFlank/stack/98
branch
from
July 30, 2026 18:07
9e7c129 to
9da8bc7
Compare
AmesingFlank
marked this pull request as ready for review
July 30, 2026 18:08
AmesingFlank
marked this pull request as draft
July 30, 2026 19:10
AmesingFlank
force-pushed
the
AmesingFlank/stack/98
branch
from
July 30, 2026 19:10
9da8bc7 to
d7a3a2c
Compare
jansel
approved these changes
Jul 31, 2026
…ackend (TorchTPU)
Implement `PallasBackend.dependency_free_launcher_info` so
`to_code(options=OutputCodeOptions(allow_helion_deps=False))` works for Pallas
(TorchTPU) kernels: the emitted module inlines the dependency-free Pallas
launcher as a local `helion.runtime` shim, and its only runtime deps are
`torch` + `jax`.
Example (kernel bound with backend="pallas"):
code = bound.to_code(cfg, options=OutputCodeOptions(allow_helion_deps=False))
emits a module whose deps are `torch` + `jax` (no `helion` import):
import torch, types
helion = types.SimpleNamespace(runtime=_make_helion_runtime()) # inlined Pallas
# launcher; topk/worklist helpers embedded (not imported)
def _helion_add(x, y, out): ...
def add(x, y, *, _launcher=_default_pallas_launcher): # torch.Tensor -> torch.Tensor
Pallas kernels call two helion-defined pure-`jax` helpers in-kernel --
`divide_filter_topk` (aten.topk lowering) and `flatten_worklist`
(compact-worklist builder). Regular `to_code` imports them from helion (added as
conditional `library_imports` entries); the dependency-free path drops those
imports and embeds the helper source instead (`embedded_helper_source`, called
only by the standalone builders). Their modules carry a note to stay free of any
`helion` import so the embed stays self-contained.
This folds in the previously separate "Embed in-kernel Pallas runtime helpers"
change: embedding is now a dependency-free-only step rather than applied to every
`to_code`, so regular output is import-based and only standalones embed.
stack-info: PR: #3184, branch: AmesingFlank/stack/98
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) #3184Support
OutputCodeOptions(allow_helion_deps=False)for the Pallas backend (TorchTPU)Implement
PallasBackend.dependency_free_launcher_infosoto_code(options=OutputCodeOptions(allow_helion_deps=False))works for Pallas(TorchTPU) kernels: the emitted module inlines the dependency-free Pallas
launcher as a local
helion.runtimeshim, and its only runtime deps aretorch+jax.Example (kernel bound with backend="pallas"):
emits a module whose deps are
torch+jax(nohelionimport):Pallas kernels call two helion-defined pure-
jaxhelpers in-kernel --divide_filter_topk(aten.topk lowering) andflatten_worklist(compact-worklist builder). Regular
to_codeimports them from helion (added asconditional
library_importsentries); the dependency-free path drops thoseimports and embeds the helper source instead (
embedded_helper_source, calledonly by the standalone builders). Their modules carry a note to stay free of any
helionimport so the embed stays self-contained.This folds in the previously separate "Embed in-kernel Pallas runtime helpers"
change: embedding is now a dependency-free-only step rather than applied to every
to_code, so regular output is import-based and only standalones embed.