Use cudnn 9.23 new API to query workspace with override shape (#3291)
<!-- .github/pull_request_template.md --> ## ๐ Description ## Description This MR makes two cuDNN GEMM backend cleanups/improvements: 1. Move effective M-bucket mapper lookup into the cuDNN runners The effective `map_to_tuning_buckets` lookup is now owned by the cuDNN GEMM runners instead of the higher-level GEMM dispatch functions. This keeps the bucket-to-`cache_m` logic local to the backend that uses it, while still respecting active autotune overrides such as custom `tuning_buckets` / `round_up`. 2. Query override-shape workspace size dynamically on cuDNN 9.23+ For cuDNN override-shape GEMM execution, cuDNN 9.23+ can query the workspace requirement for the actual runtime problem shape via `get_workspace_size_plan_at_index(...)` with override shapes and strides. The code now uses this API when available, so workspace allocation matches the executed dynamic problem size. Older cuDNN versions continue to query workspace by execution plan index without override-shape metadata. ## ๐ 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 - [ ] 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 * **Improvements** * Better cuDNN version compatibility checks for more reliable GPU support across driver versions. * More robust and efficient workspace sizing and allocation for matrix-multiply workloads, improving performance and memory use. * Autotuning and runner initialization simplified so tuning behavior is computed automatically, improving reliability across BF16/FP8/FP4/MXFP8 workloads. [](https://app.coderabbit.ai/change-stack/flashinfer-ai/flashinfer/pull/3291) <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: Yanqin Zhai <yanqinz@nvidia.com>
Y
yanqinz2 committed
a1b8a609223fc87f9f8a6570e7602ec7ec4cdf82
Parent: 7d1d46e
Committed by GitHub <noreply@github.com>
on 5/15/2026, 9:30:08 PM