Add native cudnn_decode for improved cudnn decode performance

[Anerudhan]

📌 Description

This PR tries to integrate cudnn decode by calling the cudnn kernels directly instead of through the cubin path.

Also, enabled nvidia-cudnn-frontend on all platforms

🔍 Related Issues

🚀 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.

🧪 Tests

• Tests have been added or updated as needed.
• All tests are passing (unittest, etc.).

Reviewer Notes

[gemini-code-assist]

Summary of Changes

Hello @Anerudhan, I'm Gemini Code Assist¹! 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 integrates a native cuDNN implementation for the decode operation, aiming to leverage direct cuDNN kernel calls for potentially improved performance and more granular control over attention mechanisms. The changes include the core logic for building and executing cuDNN graphs, a conditional dispatch system, and updated tests to cover the new functionality and larger scale scenarios.

Highlights

• Native cuDNN Integration: I've introduced a new, native cuDNN-based implementation for the decode operation. This involves defining cuDNN graph building logic (_build_decode_graph) and a helper function (_batch_decode_with_kv_cache) to map PyTorch tensors to cuDNN graph inputs and execute the graph directly, bypassing the previous cubin path.
• Conditional Dispatch: The main cudnn_batch_decode_with_kv_cache function now conditionally dispatches to the new native cuDNN path if the cudnn library is available. If cudnn is not found, it gracefully falls back to the existing JIT-compiled decode path, ensuring backward compatibility.
• Enhanced Attention Features: The new cuDNN graph building explicitly supports advanced attention features such as batch_offsets_q and batch_offsets_o for ragged tensors, and block_tables for paged attention. This provides greater flexibility for handling varying sequence lengths and efficient memory management.
• Test Suite Updates: The test_cudnn_decode suite has been updated to reflect the new capabilities. This includes adjusting the parameter ranges for batch_size, s_kv, page_size, num_kv_heads, and num_qo_heads to test larger scenarios, and adding batch_offsets_q and batch_offsets_o to the test calls to validate the new ragged tensor support.

Using Gemini Code Assist

The 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 in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in issue comments and review comments.
Help	/gemini help	Displays a list of available commands.

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 .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist is currently in preview and 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 to provide feedback.

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

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

The code changes introduce the ability to call the cudnn kernels directly instead of through the cubin path. I found a few critical issues that need to be addressed, along with some suggestions for improving maintainability and correctness. The main concerns are around potential race conditions with the global cuDNN handle in multi-GPU scenarios, an UnboundLocalError if block_tables is not provided, and what appear to be incorrect tensor strides which could lead to memory access errors.

[yzh119]

For package level dependency, it's important to track the version of dependencies. Can you specify a minimal version number of nvidia-cudnn-cu12 and nvidia-cudnn-frontend:

flashinfer/setup.py

Lines 63 to 64 in 513c613

	"nvidia-cudnn-cu12",
	"nvidia-cudnn-frontend",

[Anerudhan]

Can you help flashinfer/flashinfer-ci:latest to install the pytorch such that the cudnn dependency is not hard coded ? (probably with custom constraint.txt and --no-deps option) @yzh119

[yzh119]

Sure I'll double check the container

[Anerudhan]

Rebased to resolve merge conflicts.

[yzh119]

LGTM, thank you @Anerudhan !