ci/cd: bringup flashinfer-jit-cache package

[yzh119]

📌 Description

Move aot modules to a standalone package flashinfer-jit-cache.

This package, together with flashinfer-cubin, are two additional packages to the core flashinfer-python packge.

When only the flashinfer-python, every kernel will be jit compiled or downloaded from artifacory.

# install the core package
pip install flashinfer-python

If these two additional packages are installed, most of the pre-compiled kernels/cubins are be loaded from these wheels, without the need of JIT compile/download.

# optional packages
pip install flashinfer-cubin
pip install flashinfer-jit-cache --index https://flashinfer.ai/whl/cu129

🔍 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

[gemini-code-assist]

Summary of Changes

Hello @yzh119, 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 refactors the Ahead-of-Time (AOT) compilation and packaging of FlashInfer modules by creating a dedicated, standalone Python package. This change aims to improve modularity, streamline the distribution of pre-compiled kernels, and provide a more robust and isolated build process for these performance-critical components. The core compilation logic has been centralized and integrated into the new package's build system.

Highlights

• New Standalone Package: Introduced a new Python package named flashinfer-aot-modules to encapsulate all Ahead-of-Time (AOT) compiled modules, separating them from the main FlashInfer library.
• Centralized AOT Compilation Logic: A new compile_and_package_modules function was added to flashinfer/aot.py to centralize the AOT compilation process, making it reusable and simplifying the build flow.
• Custom Build Commands: Implemented custom setuptools commands (CompileAndBuildPy and CustomSdist) in the new package's setup.py to automatically trigger AOT compilation and include the resulting shared libraries (.so files) during package build and source distribution.
• JIT Environment Context: Added a jit_environment_context context manager in flashinfer/aot.py to temporarily modify JIT environment settings, ensuring isolated and controlled compilation environments.

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 by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description

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 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. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

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.

Footnotes

1. 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 introduces a new package flashinfer-aot-modules to hold the Ahead-Of-Time (AOT) compiled modules, separating them from the main flashinfer package. This is a good step towards better modularity. The changes include a new pyproject.toml and setup.py for the new package with custom build logic to compile the modules. The flashinfer/aot.py script has been refactored to centralize the compilation logic into a new compile_and_package_modules function, which improves code reuse and clarity.

My review focuses on improving code clarity and fixing a potential bug in the command-line usage of the AOT script. I've suggested using more modern pathlib features and ensuring consistency in fallback version strings. Most importantly, I've identified and proposed a fix for an issue where compilation paths might not be set correctly when using a custom build directory, which could break the build.