feat: Restore convenience `FLASHINFER_ENABLE_AOT` option #1235

mgorny · 2025-07-08T09:03:23Z

📌 Description

Restore the minimal support for FLASHINFER_ENABLE_AOT environment variable that was removed as part of the refactoring in #1075. This option was quite useful for downstreams like Red Hat, since it provided for convenient integration of AOT mode in regular PEP 517 workflows.

This change aims to be absolutely minimal -- it does not change anything for the default workflow, merely readds FLASHINFER_ENABLE_AOT that automatically invokes flashinfer.aot.main(). This makes it possible to perform an AOT build in a single step such as:

TORCH_CUDA_ARCH_LIST="7.5 8.0 8.9 9.0a" FLASHINFER_ENABLE_AOT=1 python -m build -w

or:

TORCH_CUDA_ARCH_LIST="7.5 8.0 8.9 9.0a" FLASHINFER_ENABLE_AOT=1 pip install .

🔍 Related Issues

🚀 Pull Request Checklist

Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete.

✅ Pre-commit Checks

I have installed pre-commit by running pip install pre-commit (or used your preferred method).
I have installed the hooks with pre-commit install.
I have run the hooks manually with pre-commit run --all-files and fixed any reported issues.

If you are unsure about how to set up pre-commit, see the pre-commit documentation.

🧪 Tests

Tests have been added or updated as needed.
All tests are passing (unittest, etc.).

Reviewer Notes

Restore the minimal support for `FLASHINFER_ENABLE_AOT` environment variable that was removed as part of the refactoring in flashinfer-ai#1075. This option was quite useful for downstreams like Red Hat, since it provided for convenient integration of AOT mode in regular PEP 517 workflows. This change aims to be absolutely minimal -- it does not change anything for the default workflow, merely readds `FLASHINFER_ENABLE_AOT` that automatically invokes `flashinfer.aot.main()`. This makes it possible to perform an AOT build in a single step such as: ``` TORCH_CUDA_ARCH_LIST="7.5 8.0 8.9 9.0a" FLASHINFER_ENABLE_AOT=1 python -m build -w ``` or: ``` TORCH_CUDA_ARCH_LIST="7.5 8.0 8.9 9.0a" FLASHINFER_ENABLE_AOT=1 pip install . ```

gemini-code-assist

Summary of Changes

Hello @mgorny, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request reintroduces the FLASHINFER_ENABLE_AOT environment variable, providing a convenient mechanism to trigger Ahead-of-Time (AOT) compilation for flashinfer within standard Python packaging workflows (PEP 517). This restores functionality previously removed, making AOT builds more accessible for downstream users without altering the default build process.

Highlights

AOT Build Convenience: Reinstates the FLASHINFER_ENABLE_AOT environment variable, allowing users to easily enable Ahead-of-Time (AOT) compilation during standard Python package builds (e.g., pip install or python -m build). This restores a feature previously removed in #1075, making AOT integration more convenient for downstream users.
Integration with Build Process: Modifies custom_backend.py to detect the FLASHINFER_ENABLE_AOT variable within the build_wheel function. If the variable is set to '1', it automatically invokes the flashinfer.aot.main() function, triggering the AOT compilation process during wheel creation.
AOT Main Function Adaptability: Updates the flashinfer.aot.main function to accept command-line arguments programmatically. This change enables its direct invocation from the build backend (custom_backend.py) without relying on sys.argv directly, making the AOT entry point more flexible.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

Feature	Command	Description
Code Review	`/gemini review`	Performs a code review for the current pull request in its current state.
Pull Request Summary	`/gemini summary`	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in issue comments and review comments.
Help	`/gemini help`	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist is currently in preview and may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments to provide feedback.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

gemini-code-assist

Code Review

This pull request restores the FLASHINFER_ENABLE_AOT environment variable, providing a convenient way to trigger Ahead-of-Time (AOT) compilation. I suggested accepting more truthy values for the environment variable and ensuring the main function in flashinfer/aot.py maintains its original CLI behavior.

gemini-code-assist · 2025-07-08T09:05:09Z

flashinfer/aot.py

@@ -426,7 +427,7 @@ def main():
        type=parse_bool,
        help="Add kernels for Gemma Model (head_dim=256, use_sliding_window, use_logits_soft_cap)",
    )


Pass args to parse_args only when args is not None to maintain the original CLI behavior when the function is called directly from the command line.

args = parser.parse_args() if args is None else parser.parse_args(args)

gemini-code-assist · 2025-07-08T09:05:10Z

custom_backend.py

@@ -109,5 +110,10 @@ def build_sdist(sdist_directory, config_settings=None):


 def build_wheel(wheel_directory, config_settings=None, metadata_directory=None):
+    if os.environ.get("FLASHINFER_ENABLE_AOT", "0") == "1":


Consider accepting more truthy values for the FLASHINFER_ENABLE_AOT environment variable to improve user experience. The current implementation only accepts "1", which might be less intuitive for users accustomed to boolean flags with values like true or yes.

Suggested change

if os.environ.get("FLASHINFER_ENABLE_AOT", "0") == "1":

if os.environ.get("FLASHINFER_ENABLE_AOT", "0").lower() in ("1", "true", "yes", "on", "t", "y"):

yzh119 · 2025-07-08T18:17:50Z

Hi @mgorny sure this is something we can do.
Our recent goal is to totally forsaken the AOT wheel, and have a unified package for JIT/AOT mode:

pip install flashinfer-python[all]

or

pip install flashinfer-python
flashinfer --download all # or flashinfer --compile all

to get all modules, WDYT?

mgorny · 2025-07-08T18:33:20Z

Well, we're using source distributions to build our own AOT wheel, and that's basically what we need to work. Our goal is specifically not to require our customers to wait for kernels to compile and to avoid persistent cache in Kubernetes / OpenShift. So anything that lets us get AOT mode from sdist readily works.

pip install flashinfer-python[all]

I presume this means having a second package that installs the precompiled kernels then? I think that's going to work for us, provided that package also has a sdist. I'll ask to make sure.

gemini-code-assist bot reviewed Jul 8, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

feat: Restore convenience `FLASHINFER_ENABLE_AOT` option #1235

feat: Restore convenience `FLASHINFER_ENABLE_AOT` option #1235

Uh oh!

mgorny commented Jul 8, 2025

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

gemini-code-assist bot Jul 8, 2025

Uh oh!

gemini-code-assist bot Jul 8, 2025

Uh oh!

yzh119 commented Jul 8, 2025

Uh oh!

mgorny commented Jul 8, 2025

Uh oh!

Uh oh!

		@@ -109,5 +110,10 @@ def build_sdist(sdist_directory, config_settings=None):


		def build_wheel(wheel_directory, config_settings=None, metadata_directory=None):
		if os.environ.get("FLASHINFER_ENABLE_AOT", "0") == "1":

feat: Restore convenience FLASHINFER_ENABLE_AOT option #1235

Are you sure you want to change the base?

feat: Restore convenience FLASHINFER_ENABLE_AOT option #1235

Uh oh!

Conversation

mgorny commented Jul 8, 2025

📌 Description

🔍 Related Issues

🚀 Pull Request Checklist

✅ Pre-commit Checks

🧪 Tests

Reviewer Notes

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist bot Jul 8, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Jul 8, 2025

Choose a reason for hiding this comment

Uh oh!

yzh119 commented Jul 8, 2025

Uh oh!

mgorny commented Jul 8, 2025

Uh oh!

Uh oh!

feat: Restore convenience `FLASHINFER_ENABLE_AOT` option #1235

feat: Restore convenience `FLASHINFER_ENABLE_AOT` option #1235