chore: allow custom paths for external dependencies like CUTLASS #1827
Conversation
Summary of ChangesHello @yzh119, I'm Gemini Code Assist1! 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 enhances the flexibility of FlashInfer's build system by allowing users to define custom installation paths for critical external dependencies such as CUTLASS, SPDLOG, and NVSHMEM. This change simplifies integration into diverse development environments and system configurations, reducing potential conflicts and improving overall usability for advanced users who manage their own dependency installations. The update also ensures that configuration details are easily inspectable via the Highlights
Using Gemini Code AssistThe 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
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 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
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a valuable feature by allowing users to specify custom paths for external dependencies like CUTLASS, spdlog, and nvshmem via environment variables. This greatly improves the flexibility of the build system, especially for users who have these dependencies installed system-wide. The implementation is clean and maintains backward compatibility by falling back to older environment variables or default paths. I have one suggestion to improve performance by caching the results of the new path-getter functions to avoid redundant work.
| SPDLOG_INCLUDE_DIR = _package_root / "data" / "spdlog" / "include" | ||
|
|
||
|
|
||
| def get_cutlass_include_dirs(): |
There was a problem hiding this comment.
For performance and consistency, it's a good practice to cache the results of this getter function. Since environment variables are not expected to change during the program's execution, caching will prevent redundant lookups and path processing. You can use @functools.cache for this, which is already used elsewhere in the codebase (e.g., in cpp_ext.py).
Please apply this caching to get_spdlog_include_dir, get_nvshmem_include_dirs, and get_nvshmem_lib_dirs as well. You'll need to add import functools at the top of the file.
| def get_cutlass_include_dirs(): | |
| @functools.cache | |
| def get_cutlass_include_dirs(): |
| def gen_nvshmem_module() -> JitSpec: | ||
| lib_dirs = jit_env.get_nvshmem_lib_dirs() | ||
| # Check new environment variable first, then fall back to old one for backward compatibility | ||
| ldflags_env = os.environ.get("FLASHINFER_NVSHMEM_LDFLAGS") or os.environ.get( |
There was a problem hiding this comment.
shall we have documentation collecting these FLASHINFER* vars in some .md file, and briefly mention what they may affect and the order in which related vars are consulted. this could help both users and developers understand the expected behavior (which could otherwise run away after a few ppl modifying the same)
There was a problem hiding this comment.
currently user can run
python -m flashinfer show-config
do display these environment variables, but yes let's clearly document it.
📌 Description
Implement features mentioned in #1806 , now user can specify custom dependencies path (e.g. if they are installed as system library) via environment variables
FLASHINFER_CUTLASS_INCLUDE_PATHetc.User can run
python -m flashinfer show-configto display these environment variables.cc @BwL1289
🔍 Related Issues
#1806
🚀 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
pre-commitby runningpip install pre-commit(or used your preferred method).pre-commit install.pre-commit run --all-filesand fixed any reported issues.🧪 Tests
unittest, etc.).Reviewer Notes