Move Triton launcher into helion/runtime/triton/launcher.py - #3178
Merged
Conversation
AmesingFlank
force-pushed
the
AmesingFlank/stack/92
branch
from
July 28, 2026 16:02
b1ef8d0 to
df47047
Compare
This was referenced Jul 28, 2026
AmesingFlank
marked this pull request as ready for review
July 28, 2026 16:29
AmesingFlank
marked this pull request as draft
July 28, 2026 17:38
AmesingFlank
marked this pull request as ready for review
July 28, 2026 17:40
Relocate the runtime symbols that Helion's *generated* Triton code depends on at execution time out of the catch-all `helion/runtime/__init__.py` into a dedicated `helion/runtime/triton/launcher.py`: - `default_launcher` — runs a compiled `triton.jit` kernel - `get_num_sm` — persistent-kernel grid size (host statement) - `set_triton_allocator` (+ `_alloc_fn`) — TMA / tensor-descriptor scratch allocator `helion.runtime` re-exports all three, so `helion.runtime.default_launcher`, `helion.runtime.get_num_sm`, `helion.runtime.set_triton_allocator`, and the generated `from helion.runtime import default_launcher` are unchanged. This is a pure move with no behavior change — the module still reaches back into a few `helion` helpers (`exc`, `settings`, `_utils`). Isolating these into one small module is the first step toward making them dependency-free, so the ahead-of-time precompiler can bulk-export exactly this file into a standalone kernel. Because `get_num_sm` now resolves its globals in the new module, `test/test_runtime.py` is updated to patch `helion.runtime.triton.launcher._module_is_pallas_interpret` (the moved function's real lookup site) instead of `helion.runtime._module_is_pallas_interpret`. Verified: `ruff`/`pyrefly` clean, `test_runtime` and `test_persistent_kernels` pass, and the generated Triton for `examples/add.py` is byte-identical. stack-info: PR: #3178, branch: AmesingFlank/stack/92
AmesingFlank
marked this pull request as draft
July 28, 2026 20:09
AmesingFlank
force-pushed
the
AmesingFlank/stack/92
branch
from
July 28, 2026 20:09
df47047 to
a82000d
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
marked this pull request as ready for review
July 28, 2026 21:15
AmesingFlank
marked this pull request as draft
July 28, 2026 22:27
AmesingFlank
marked this pull request as ready for review
July 28, 2026 22:28
AmesingFlank
marked this pull request as draft
July 29, 2026 00:03
AmesingFlank
marked this pull request as ready for review
July 29, 2026 00:04
AmesingFlank
marked this pull request as draft
July 29, 2026 05:05
AmesingFlank
marked this pull request as ready for review
July 29, 2026 05:06
AmesingFlank
marked this pull request as draft
July 29, 2026 09:03
AmesingFlank
marked this pull request as ready for review
July 29, 2026 09:04
yushangdi
approved these changes
Jul 29, 2026
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
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:
pretuned_kernel,jax_fn, helper embedding #3189aot_kerneltopretuned_kernel#3188OutputCodeOptions(jax_fn=True)for Pallas (JAX-array entrypoints) #3186helion.precompilefor standalone kernels #3185OutputCodeOptions(allow_helion_deps=False)for the Pallas backend (TorchTPU) #3184helion/runtime/pallas/launcher.py#3182OutputCodeOptions(allow_helion_deps=False)toBoundKernel.to_codefor standalone Triton kernels #3181get_num_xcdinto the dependency-free Triton launcher #3180helion/runtime/triton/launcher.py#3178Move Triton launcher into
helion/runtime/triton/launcher.pyRelocate the runtime symbols that Helion's generated Triton code depends on
at execution time out of the catch-all
helion/runtime/__init__.pyinto adedicated
helion/runtime/triton/launcher.py:default_launcher— runs a compiledtriton.jitkernelget_num_sm— persistent-kernel grid size (host statement)set_triton_allocator(+_alloc_fn) — TMA / tensor-descriptor scratch allocatorhelion.runtimere-exports all three, sohelion.runtime.default_launcher,helion.runtime.get_num_sm,helion.runtime.set_triton_allocator, and thegenerated
from helion.runtime import default_launcherare unchanged.This is a pure move with no behavior change — the module still reaches back into
a few
helionhelpers (exc,settings,_utils). Isolating these into onesmall module is the first step toward making them dependency-free, so the
ahead-of-time precompiler can bulk-export exactly this file into a standalone
kernel.
Because
get_num_smnow resolves its globals in the new module,test/test_runtime.pyis updated to patchhelion.runtime.triton.launcher._module_is_pallas_interpret(the movedfunction's real lookup site) instead of
helion.runtime._module_is_pallas_interpret.Verified:
ruff/pyreflyclean,test_runtimeandtest_persistent_kernelspass, and the generated Triton for
examples/add.pyis byte-identical.