Commit 292f9be
authored
fix: include fp8_blockscale_gemm_90 in AOT jit-cache (#2533)
## Summary
- Add fp8_blockscale_gemm_90 (gen_fp8_blockscale_gemm_sm90_module) to
the AOT build list when SM90 is enabled.
- Avoid runtime JIT compilation for fp8_blockscale_gemm_sm90 in
environments without CUDA dev headers, which can fail with cublasLt.h
not found.
## Changes
- flashinfer/aot.py: append gen_fp8_blockscale_gemm_sm90_module() under
add_moe + has_sm90 gating.
## Related Issues
- Fixes #2527
- #2527
## Tests
<!-- .github/pull_request_template.md -->
## 📌 Description
<!-- What does this PR do? Briefly describe the changes and why they’re
needed. -->
## 🔍 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
- [x] I have installed `pre-commit` by running `pip install pre-commit`
(or used your preferred method).
- [x] I have installed the hooks with `pre-commit install`.
- [x] 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
- [x] 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.
-->
<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit
* **New Features**
* Added support for FP8 blockscale matrix multiplication operations on
SM90 GPU architecture.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->1 parent c5b8a2e commit 292f9be
1 file changed
+3
-0
lines changed| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
52 | 52 | | |
53 | 53 | | |
54 | 54 | | |
| 55 | + | |
55 | 56 | | |
56 | 57 | | |
57 | 58 | | |
| |||
477 | 478 | | |
478 | 479 | | |
479 | 480 | | |
| 481 | + | |
| 482 | + | |
480 | 483 | | |
481 | 484 | | |
482 | 485 | | |
| |||
0 commit comments