Skip to content

Make the Pallas launcher dependency-free (torch + jax only) - #3183

Merged
AmesingFlank merged 1 commit into
mainfrom
AmesingFlank/stack/97
Jul 30, 2026
Merged

Make the Pallas launcher dependency-free (torch + jax only)#3183
AmesingFlank merged 1 commit into
mainfrom
AmesingFlank/stack/97

Conversation

@AmesingFlank

@AmesingFlank AmesingFlank commented Jul 28, 2026

Copy link
Copy Markdown
Contributor

Stacked PRs:


Make the Pallas launcher dependency-free (torch + jax only)

helion/runtime/pallas/launcher.py now imports only torch + jax (+ stdlib) --
no other helion module -- so the precompiler can bulk-export it into a
standalone Pallas kernel with zero Helion runtime dependency. Two helion deps
are severed:

  • _PALLAS_UNSUPPORTED_DTYPES (the {int64,uint64,float64} frozenset TPU can't
    lower) moves from helion/_compiler/backend.py -- which only defined it for
    the launcher's import; nothing else used it -- into the launcher.
  • is_pallas_interpret (from helion.runtime.settings) is replaced by a
    dep-free _pallas_interpret_enabled() env-var read at the launcher's own
    runtime call sites (the two default_pallas_launcher install-path fallbacks),
    where no CompileEnvironment is active so is_pallas_interpret already
    reduces to that env var (Settings.pallas_interpret itself defaults from it).

_get_vmem_limit_bytes is called both at runtime (launcher) and at compile time
(_compiler/pallas/backend.py, where a CompileEnvironment IS active), so
rather than have it read the env var it now takes interpret as a parameter:
the compiler passes env.settings.pallas_interpret (honoring decorator-set
interpret) and the launcher passes _pallas_interpret_enabled(). This keeps the
launcher dependency-free with no behavior change at either site -- mirroring how
the launcher install path already threads _pallas_interpret as a kwarg.

Verified on the TPU pod: non-interpret pallas add/topk/jax_fn and the
compact-worklist window/budget/guard/matmul tests pass; the launcher imports only
torch+jax. ruff/pyrefly clean (83 baseline); Triton add still correct on
H100. (Interpret-mode add tests fail identically on the base commit -- a
pre-existing pod JAX-version issue, not from this change.)

@AmesingFlank
AmesingFlank force-pushed the AmesingFlank/stack/96 branch from 7cc19f5 to aa4871d Compare July 28, 2026 16:02
@AmesingFlank
AmesingFlank force-pushed the AmesingFlank/stack/97 branch from c6b5c2a to 1003016 Compare July 28, 2026 16:02
@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 changed the base branch from AmesingFlank/stack/96 to main July 28, 2026 17:39
@AmesingFlank
AmesingFlank changed the base branch from main to AmesingFlank/stack/96 July 28, 2026 17:39
@AmesingFlank
AmesingFlank marked this pull request as ready for review July 28, 2026 17:50
@AmesingFlank
AmesingFlank marked this pull request as draft July 28, 2026 20:09
@AmesingFlank
AmesingFlank changed the base branch from AmesingFlank/stack/96 to main July 28, 2026 20:09
@AmesingFlank
AmesingFlank force-pushed the AmesingFlank/stack/97 branch from 1003016 to 6b6459d Compare July 28, 2026 20:09
@AmesingFlank
AmesingFlank changed the base branch from main to AmesingFlank/stack/96 July 28, 2026 20:10
@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 changed the base branch from AmesingFlank/stack/96 to main July 28, 2026 21:14
@AmesingFlank
AmesingFlank marked this pull request as draft July 29, 2026 20:37
@AmesingFlank
AmesingFlank changed the base branch from AmesingFlank/stack/96 to main July 29, 2026 20:37
@AmesingFlank
AmesingFlank force-pushed the AmesingFlank/stack/97 branch from 3f6f8c5 to 161fb9d Compare July 29, 2026 20:37
@AmesingFlank
AmesingFlank changed the base branch from main to AmesingFlank/stack/96 July 29, 2026 20:38
@AmesingFlank
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
`helion/runtime/pallas/launcher.py` now imports only torch + jax (+ stdlib) --
no other `helion` module -- so the precompiler can bulk-export it into a
standalone Pallas kernel with zero Helion runtime dependency. Two helion deps
are severed:

- `_PALLAS_UNSUPPORTED_DTYPES` (the {int64,uint64,float64} frozenset TPU can't
  lower) moves from `helion/_compiler/backend.py` -- which only defined it for
  the launcher's import; nothing else used it -- into the launcher.
- `is_pallas_interpret` (from `helion.runtime.settings`) is replaced by a
  dep-free `_pallas_interpret_enabled()` env-var read at the launcher's own
  runtime call sites (the two `default_pallas_launcher` install-path fallbacks),
  where no `CompileEnvironment` is active so `is_pallas_interpret` already
  reduces to that env var (`Settings.pallas_interpret` itself defaults from it).

`_get_vmem_limit_bytes` is called both at runtime (launcher) and at compile time
(`_compiler/pallas/backend.py`, where a `CompileEnvironment` IS active), so
rather than have it read the env var it now takes `interpret` as a parameter:
the compiler passes `env.settings.pallas_interpret` (honoring decorator-set
interpret) and the launcher passes `_pallas_interpret_enabled()`. This keeps the
launcher dependency-free with no behavior change at either site -- mirroring how
the launcher install path already threads `_pallas_interpret` as a kwarg.

Verified on the TPU pod: non-interpret pallas add/topk/jax_fn and the
compact-worklist window/budget/guard/matmul tests pass; the launcher imports only
torch+jax. `ruff`/`pyrefly` clean (83 baseline); Triton `add` still correct on
H100. (Interpret-mode `add` tests fail identically on the base commit -- a
pre-existing pod JAX-version issue, not from this change.)

stack-info: PR: #3183, branch: AmesingFlank/stack/97
@AmesingFlank
AmesingFlank marked this pull request as draft July 30, 2026 01:55
@AmesingFlank
AmesingFlank changed the base branch from AmesingFlank/stack/96 to main July 30, 2026 01:55
@AmesingFlank
AmesingFlank force-pushed the AmesingFlank/stack/97 branch from 161fb9d to c87972c Compare July 30, 2026 01:56
@AmesingFlank
AmesingFlank changed the base branch from main to AmesingFlank/stack/96 July 30, 2026 01:56
@AmesingFlank
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
`helion/runtime/pallas/launcher.py` now imports only torch + jax (+ stdlib) --
no other `helion` module -- so the precompiler can bulk-export it into a
standalone Pallas kernel with zero Helion runtime dependency. Two helion deps
are severed:

- `_PALLAS_UNSUPPORTED_DTYPES` (the {int64,uint64,float64} frozenset TPU can't
  lower) moves from `helion/_compiler/backend.py` -- which only defined it for
  the launcher's import; nothing else used it -- into the launcher.
- `is_pallas_interpret` (from `helion.runtime.settings`) is replaced by a
  dep-free `_pallas_interpret_enabled()` env-var read at the launcher's own
  runtime call sites (the two `default_pallas_launcher` install-path fallbacks),
  where no `CompileEnvironment` is active so `is_pallas_interpret` already
  reduces to that env var (`Settings.pallas_interpret` itself defaults from it).

`_get_vmem_limit_bytes` is called both at runtime (launcher) and at compile time
(`_compiler/pallas/backend.py`, where a `CompileEnvironment` IS active), so
rather than have it read the env var it now takes `interpret` as a parameter:
the compiler passes `env.settings.pallas_interpret` (honoring decorator-set
interpret) and the launcher passes `_pallas_interpret_enabled()`. This keeps the
launcher dependency-free with no behavior change at either site -- mirroring how
the launcher install path already threads `_pallas_interpret` as a kwarg.

Verified on the TPU pod: non-interpret pallas add/topk/jax_fn and the
compact-worklist window/budget/guard/matmul tests pass; the launcher imports only
torch+jax. `ruff`/`pyrefly` clean (83 baseline); Triton `add` still correct on
H100. (Interpret-mode `add` tests fail identically on the base commit -- a
pre-existing pod JAX-version issue, not from this change.)

stack-info: PR: #3183, branch: AmesingFlank/stack/97
@AmesingFlank
AmesingFlank marked this pull request as draft July 30, 2026 02:43
@AmesingFlank
AmesingFlank changed the base branch from AmesingFlank/stack/96 to main July 30, 2026 02:43
@AmesingFlank
AmesingFlank force-pushed the AmesingFlank/stack/97 branch from c87972c to a611585 Compare July 30, 2026 02:43
@AmesingFlank
AmesingFlank changed the base branch from main to AmesingFlank/stack/96 July 30, 2026 02:43
@AmesingFlank
AmesingFlank marked this pull request as ready for review July 30, 2026 02:43
@AmesingFlank
AmesingFlank marked this pull request as draft July 30, 2026 03:17
@AmesingFlank
AmesingFlank changed the base branch from AmesingFlank/stack/96 to main July 30, 2026 03:17
@AmesingFlank
AmesingFlank force-pushed the AmesingFlank/stack/97 branch from a611585 to 44c2433 Compare July 30, 2026 03:17
@AmesingFlank
AmesingFlank changed the base branch from main to AmesingFlank/stack/96 July 30, 2026 03:17
@AmesingFlank
AmesingFlank marked this pull request as ready for review July 30, 2026 03:17
@AmesingFlank
AmesingFlank marked this pull request as draft July 30, 2026 03:28
@AmesingFlank
AmesingFlank changed the base branch from AmesingFlank/stack/96 to main July 30, 2026 03:28
@AmesingFlank
AmesingFlank force-pushed the AmesingFlank/stack/97 branch from 44c2433 to 0cb409a Compare July 30, 2026 03:29
@AmesingFlank
AmesingFlank changed the base branch from main to AmesingFlank/stack/96 July 30, 2026 03:29
@AmesingFlank
AmesingFlank marked this pull request as ready for review July 30, 2026 03:29
`helion/runtime/pallas/launcher.py` now imports only torch + jax (+ stdlib) --
no other `helion` module -- so the precompiler can bulk-export it into a
standalone Pallas kernel with zero Helion runtime dependency. Two helion deps
are severed:

- `_PALLAS_UNSUPPORTED_DTYPES` (the {int64,uint64,float64} frozenset TPU can't
  lower) moves from `helion/_compiler/backend.py` -- which only defined it for
  the launcher's import; nothing else used it -- into the launcher.
- `is_pallas_interpret` (from `helion.runtime.settings`) is replaced by a
  dep-free `_pallas_interpret_enabled()` env-var read at the launcher's own
  runtime call sites (the two `default_pallas_launcher` install-path fallbacks),
  where no `CompileEnvironment` is active so `is_pallas_interpret` already
  reduces to that env var (`Settings.pallas_interpret` itself defaults from it).

`_get_vmem_limit_bytes` is called both at runtime (launcher) and at compile time
(`_compiler/pallas/backend.py`, where a `CompileEnvironment` IS active), so
rather than have it read the env var it now takes `interpret` as a parameter:
the compiler passes `env.settings.pallas_interpret` (honoring decorator-set
interpret) and the launcher passes `_pallas_interpret_enabled()`. This keeps the
launcher dependency-free with no behavior change at either site -- mirroring how
the launcher install path already threads `_pallas_interpret` as a kwarg.

Verified on the TPU pod: non-interpret pallas add/topk/jax_fn and the
compact-worklist window/budget/guard/matmul tests pass; the launcher imports only
torch+jax. `ruff`/`pyrefly` clean (83 baseline); Triton `add` still correct on
H100. (Interpret-mode `add` tests fail identically on the base commit -- a
pre-existing pod JAX-version issue, not from this change.)

stack-info: PR: #3183, branch: AmesingFlank/stack/97
@AmesingFlank
AmesingFlank marked this pull request as draft July 30, 2026 03:47
@AmesingFlank
AmesingFlank changed the base branch from AmesingFlank/stack/96 to main July 30, 2026 03:47
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