feat: add an env to allow specify cubins directory (#1462)

<!-- .github/pull_request_template.md -->

## 📌 Description

Add a env `FLASHINFER_CUBIN_DIR` for cubins directory which gives more
flexibility for aot deployment. If it's not specified, will fall back to
`FLASHINFER_CACHE_DIR / "cubins"` which is what we currently have.

## 🔍 Related Issues

<!-- Link any related issues here -->

## 🚀 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](https://pre-commit.com/).

## 🧪 Tests

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

## Reviewer Notes

<!-- Optional: anything you'd like reviewers to focus on, concerns, etc.
-->

Author: 842974287

flashinfer/jit/cubin_loader.py

@@ -23,7 +23,7 @@
 import filelock
 
 from .core import logger
-from .env import FLASHINFER_CACHE_DIR
+from .env import FLASHINFER_CUBIN_DIR
 
 # This is the storage path for the cubins, it can be replaced
 # with a local path for testing.
@@ -145,7 +145,7 @@ def get_cubin(name, sha256, file_extension=".cubin"):
     None on failure.
     """
     cubin_fname = name + file_extension
-    cubin_path = FLASHINFER_CACHE_DIR / "cubins" / cubin_fname
+    cubin_path = FLASHINFER_CUBIN_DIR / cubin_fname
     cubin = load_cubin(cubin_path, sha256)
     if cubin:
         return cubin

flashinfer/jit/env.py

@@ -30,6 +30,9 @@
 )
 
 FLASHINFER_CACHE_DIR = FLASHINFER_BASE_DIR / ".cache" / "flashinfer"
+FLASHINFER_CUBIN_DIR = pathlib.Path(
+    os.getenv("FLASHINFER_CUBIN_DIR", (FLASHINFER_CACHE_DIR / "cubins").as_posix())
+)
 
 
 def _get_workspace_dir_name() -> pathlib.Path: