Embed in-kernel Pallas runtime helpers into generated code

helion/_compiler/backend.py

@@ -703,6 +703,17 @@ def library_imports(self) -> dict[str, str]:
         """
         ...
 
+    def embedded_helper_source(self, body: str) -> str:
+        """Source of in-kernel runtime helpers to inline into the generated module.
+
+        Backends that call a helion-defined helper from inside the generated
+        kernel can return its source here (instead of importing it) so the output
+        is self-contained -- which lets the precompiler produce a helion-free
+        standalone. Only helpers actually referenced in ``body`` should be
+        emitted. Injected between the imports and the kernel body. Default: none.
+        """
+        return ""
+
     def launcher_keyword_args(self, config: Config, *, has_barrier: bool) -> list[str]:
         return []
 

helion/_compiler/pallas/backend.py

@@ -6,6 +6,7 @@
 import ast
 import dataclasses
 import enum
+import inspect
 import math
 from typing import TYPE_CHECKING
 from typing import Any
@@ -35,6 +36,35 @@
     InductorOpOverrides = OpsHandler[Any]
 
 
+def _embed_source(source: str) -> str:
+    """Return a module's source ready to inline: its module docstring and
+    ``from __future__`` lines stripped (leading comments -- e.g. an SPDX header --
+    and everything else preserved), so the docstring prose can't leak into the
+    generated code and the mid-module ``from __future__`` (a SyntaxError) is gone.
+
+    The docstring span is located via ``ast`` (not a quote scan) so a docstring
+    whose prose contains a triple-quote can't corrupt the output, and a module
+    that opens with code rather than a docstring is handled correctly.
+    """
+    lines = source.split("\n")
+    doc_lines: set[int] = set()
+    tree = ast.parse(source)
+    if (
+        tree.body
+        and isinstance(first := tree.body[0], ast.Expr)
+        and isinstance(first.value, ast.Constant)
+        and isinstance(first.value.value, str)
+    ):
+        # ast line numbers are 1-based; end_lineno is the closing-quote line.
+        doc_lines = set(range(first.lineno - 1, (first.end_lineno or first.lineno)))
+    kept = [
+        line
+        for idx, line in enumerate(lines)
+        if idx not in doc_lines and not line.strip().startswith("from __future__")
+    ]
+    return "\n".join(kept).strip("\n")
+
+
 # Mapping from torch dtype to JAX dtype string (e.g., "jnp.float32")
 _TORCH_TO_JAX_DTYPE: dict[str, str] = {
     "torch.float16": "jnp.float16",
@@ -169,9 +199,33 @@ def library_imports(self) -> dict[str, str]:
             "lax": "import jax.lax as lax",
             "pltpu": "from jax.experimental.pallas import tpu as pltpu",
             "_default_pallas_launcher": "from helion.runtime import default_pallas_launcher as _default_pallas_launcher",
-            "_helion_divide_filter_topk": "from helion._compiler.pallas.topk_impl import divide_filter_topk as _helion_divide_filter_topk",
         }
 
+    def embedded_helper_source(self, body: str) -> str:
+        """Inline the in-kernel Pallas helpers referenced by ``body``.
+
+        ``divide_filter_topk`` (aten.topk lowering) and ``flatten_worklist``
+        (compact-worklist builder) are pure-``jax`` helpers the generated kernel
+        calls. Embedding their source -- instead of importing them from helion --
+        makes the generated module self-contained (so it precompiles to a
+        helion-free standalone) and applies to regular kernels too.
+        """
+        blocks: list[str] = []
+        if "_helion_divide_filter_topk" in body:
+            from . import topk_impl
+
+            blocks.extend(
+                [
+                    _embed_source(inspect.getsource(topk_impl)),
+                    "_helion_divide_filter_topk = divide_filter_topk",
+                ]
+            )
+        if "flatten_worklist" in body:
+            from ...runtime import compact_worklist
+
+            blocks.append(_embed_source(inspect.getsource(compact_worklist)))
+        return "\n\n\n".join(blocks)
+
     # Config keys that Pallas actually uses.  Everything else
     # (pid_type, num_warps, num_stages, maxnreg, indexing, etc.)
     # is GPU-specific and should not be tuned.

helion/_compiler/pallas/compact_worklist.py

@@ -703,7 +703,6 @@ def render_build_worklist(
     lines = [
         f"def {builder_name}({', '.join(offset_params)}):",
         "    import jax.numpy as jnp",
-        "    from helion.runtime.compact_worklist import flatten_worklist",
         f"    {owner_array} = jnp.arange({num_owners_expr}, dtype=jnp.int32)",
         f"    base = {base_src}",
         f"    length = {length_src}",

helion/_compiler/pallas/topk_impl.py

@@ -37,6 +37,11 @@
 if TYPE_CHECKING:
     import jax
 
+# NOTE: this module is embedded verbatim into helion.precompile standalones (via
+# PallasBackend.embedded_helper_source), so it must not import anything from the
+# `helion` package (nor mention such an import in a comment) -- the precompiler's
+# helion-free guard is a substring check and would reject any topk kernel.
+
 _NUM_LANES = 128
 NUM_LANES = 128
 NUM_SUBLANES = 8

helion/runtime/compact_worklist.py

@@ -27,6 +27,12 @@
 if TYPE_CHECKING:
     from jax import Array
 
+# NOTE: this module is embedded verbatim into helion.precompile standalones (via
+# PallasBackend.embedded_helper_source), so its non-docstring code must not import
+# anything from the `helion` package (nor mention such an import in a comment) --
+# the precompiler's helion-free guard is a substring check and would reject any
+# compact-worklist kernel.
+
 
 class CompactWorkMetadata(NamedTuple):
     """Scalar-prefetch metadata describing one compact worklist.

helion/runtime/kernel.py

@@ -843,6 +843,11 @@ def to_code(
                 ast.fix_missing_locations(body_root)
                 imports = "\n".join(import_lines)
                 body = unparse(body_root, output_origin_lines=output_origin_lines)
+                # Inline any in-kernel runtime helpers (e.g. Pallas topk /
+                # compact-worklist) so the generated module is self-contained.
+                embedded = self.env.backend.embedded_helper_source(body)
+                if embedded:
+                    body = f"{embedded}\n\n\n{body}"
                 if imports:
                     return f"from __future__ import annotations\n\n{imports}\n\n{body}"
                 return f"from __future__ import annotations\n\n{body}"

test/test_pallas.py

@@ -900,6 +900,13 @@ def test_topk_divide_and_filter_lowering(self) -> None:
         # (1) the lowering emits the divide-and-filter helper, not jax.lax.top_k
         self.assertIn("_helion_divide_filter_topk", code)
         self.assertNotIn("lax.top_k", code)
+        # (1b) the helper source is embedded (no helion import) so the kernel is
+        # self-contained / precompilable, and the embed doesn't corrupt the module.
+        self.assertIn("def divide_filter_topk(", code)
+        self.assertIn("_helion_divide_filter_topk = divide_filter_topk", code)
+        self.assertNotIn("from helion._compiler.pallas.topk_impl import", code)
+        self.assertNotIn("import helion", code)
+        ast.parse(code)
         # (2) correctness vs the exact top-k
         ref_v, ref_i = torch.topk(x, _TOPK_TEST_K, dim=-1, largest=True)
         idx_c = idx.cpu()

test/test_pallas_worklist.py

@@ -749,7 +749,10 @@ def _render_and_exec(self, plan, offset_arrays, upper):
         src, offset_params = render_build_worklist(
             plan, block_expr=str(self.BLOCK), upper_expr=str(upper)
         )
-        namespace: dict = {}
+        # The rendered builder now calls a module-level ``flatten_worklist``
+        # (provided by the backend's embedded-helper inlining in real generated
+        # modules) rather than importing it inline, so supply it here.
+        namespace: dict = {"flatten_worklist": flatten_worklist}
         exec(compile(src, "<build_worklist>", "exec"), namespace)
         builder = namespace["_build_worklist"]
         return src, offset_params, builder(*offset_arrays)
@@ -881,7 +884,10 @@ def test_lo_hi_builder_includes_both_args(self):
         self.assertEqual(offset_params, ["lo", "hi"])
         self.assertIn("jnp.arange(lo.shape[0]", src)
 
-        namespace: dict = {}
+        # The rendered builder now calls a module-level ``flatten_worklist``
+        # (supplied by the backend's embedded-helper inlining in real modules)
+        # rather than importing it inline, so provide it here.
+        namespace: dict = {"flatten_worklist": flatten_worklist}
         exec(compile(src, "<bw>", "exec"), namespace)
         meta = namespace["_build_worklist"](
             jnp.asarray(lo.numpy()), jnp.asarray(hi.numpy())
@@ -1064,6 +1070,13 @@ def fn(q, k, v, q_offsets):
         # builder kwargs; there is no separate launcher name.
         self.assertIn("_compact_build_worklist=_build_worklist", code)
         self.assertIn("def _build_worklist(", code)
+        # ``flatten_worklist`` is embedded at module scope (no Helion import),
+        # so the generated module is self-contained for precompilation rather
+        # than importing the runtime helper inside ``_build_worklist``.
+        self.assertIn("def flatten_worklist(", code)
+        self.assertNotIn("from helion.runtime.compact_worklist import", code)
+        # The embedded helper source must not corrupt the generated module.
+        ast.parse(code)
         # Offsets arg index is non-empty (q_offsets feeds the builder).
         self.assertRegex(code, r"_compact_offset_arg_indices=\[\d")
         self.assertIn("_compact_num_scalar_prefetch=3", code)