Skip to content

Move Triton launcher into helion/runtime/triton/launcher.py - #3178

Merged
AmesingFlank merged 1 commit into
mainfrom
AmesingFlank/stack/92
Jul 29, 2026
Merged

Move Triton launcher into helion/runtime/triton/launcher.py#3178
AmesingFlank merged 1 commit into
mainfrom
AmesingFlank/stack/92

Conversation

@AmesingFlank

@AmesingFlank AmesingFlank commented Jul 28, 2026

Copy link
Copy Markdown
Contributor

Stacked PRs:


Move Triton launcher into helion/runtime/triton/launcher.py

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.

@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Meta Open Source bot. label Jul 28, 2026
@AmesingFlank
AmesingFlank force-pushed the AmesingFlank/stack/92 branch from b1ef8d0 to df47047 Compare July 28, 2026 16:02
@AmesingFlank
AmesingFlank marked this pull request as ready for review July 28, 2026 16:29
@AmesingFlank
AmesingFlank marked this pull request as draft July 28, 2026 17:38
@AmesingFlank
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
AmesingFlank marked this pull request as draft July 28, 2026 20:09
@AmesingFlank
AmesingFlank force-pushed the AmesingFlank/stack/92 branch from df47047 to a82000d Compare July 28, 2026 20:09
@AmesingFlank
AmesingFlank marked this pull request as ready for review July 28, 2026 20:10
@AmesingFlank
AmesingFlank marked this pull request as draft July 28, 2026 21:14
@AmesingFlank
AmesingFlank marked this pull request as ready for review July 28, 2026 21:15
@AmesingFlank
AmesingFlank requested review from choijon5 and ethche July 28, 2026 21:52
@AmesingFlank
AmesingFlank marked this pull request as draft July 28, 2026 22:27
@AmesingFlank
AmesingFlank marked this pull request as ready for review July 28, 2026 22:28
@AmesingFlank
AmesingFlank marked this pull request as draft July 29, 2026 00:03
@AmesingFlank
AmesingFlank marked this pull request as ready for review July 29, 2026 00:04
@AmesingFlank
AmesingFlank marked this pull request as draft July 29, 2026 05:05
@AmesingFlank
AmesingFlank marked this pull request as ready for review July 29, 2026 05:06
@AmesingFlank
AmesingFlank marked this pull request as draft July 29, 2026 09:03
@AmesingFlank
AmesingFlank marked this pull request as ready for review July 29, 2026 09:04
@AmesingFlank
AmesingFlank merged commit 0762fda into main Jul 29, 2026
39 of 40 checks passed
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