SIGN IN SIGN UP

FlashInfer: Kernel Library for LLM Serving

0 0 188 Python

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.

[![Review Change
Stack](https://storage.googleapis.com/coderabbit_public_assets/review-stack-in-coderabbit-ui.svg)](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