Allow BatchPrefillPagedWrapper to call cudnn API #1384

Anerudhan · 2025-08-05T05:04:34Z

📌 Description

🔍 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).
[x ] 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!

I've implemented fixes to the benchmark attention wrapper, specifically targeting the CUDNN backend. My changes ensure that the batch prefill operations are correctly handled by accurately calculating and passing batch offset pointers and by using device-specific sequence length tensors. This improves the robustness and correctness of the CUDNN attention benchmark.

Highlights

Batch Offset Calculation: I introduced a new variable, "qo_indptr_cudnn", which is crucial for correctly calculating the batch offsets required by the CUDNN attention wrapper. This ensures proper indexing for batched prefill operations.
Device-Specific Tensor Usage: I updated the call to the CUDNN backend within "run_backend_wrapper" to use the device-specific sequence length tensors, "actual_seq_lens_q_device" and "actual_seq_lens_kv_device". This prevents potential issues with host-device mismatches.
CUDNN Parameter Update: I added the "batch_offsets_q" and "batch_offsets_o" parameters to the CUDNN attention wrapper call, passing the newly computed "qo_indptr_cudnn". This enables the CUDNN backend to correctly handle the batching of queries and outputs.

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 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 or fill out our survey 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.

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

This pull request introduces fixes for the cudnn backend in the testBatchPrefillWithPagedKVCacheWrapper benchmark. It correctly computes and passes the ragged tensor offsets for query and output. My review identifies a potential integer overflow issue in the offset calculation that could lead to incorrect benchmark results for large inputs. A fix is suggested to ensure 64-bit integer precision for the offsets.

benchmarks/routines/attention.py

flashinfer/prefill.py

yzh119

Thank you @Anerudhan , the PR looks good to me!

weireweire · 2025-08-12T02:42:36Z

flashinfer/prefill.py

-                self._max_q_len,
-                self._max_kv_len,
-                self._batch_size,
+            assert self._plan_info is not None, "plan info is not initialized"


This assertion break trtllm-gen backend, pytest tests/test_trtllm_gen_context.py failed

fix in #1375 by check backend

See https://github.com/flashinfer-ai/flashinfer/pull/1375/files#r2268464907

@weireweire thanks for the fix!

Edenzzzz · 2025-08-12T21:47:33Z

flashinfer/prefill.py

+            paged_kv_indptr_host = paged_kv_indptr.to("cpu")
+            paged_kv_last_page_len_host = paged_kv_last_page_len.to("cpu")
+            if seq_lens is None:
+                kv_lens_arr_host = get_seq_lens(
+                    paged_kv_indptr_host, paged_kv_last_page_len_host, page_size
+                )
+            else:
+                kv_lens_arr_host = seq_lens.cpu().flatten()


Apply non_blocking for these too?

Edenzzzz · 2025-08-12T21:52:07Z

flashinfer/prefill.py

+        seq_lens_q: Optional[torch.Tensor]
+            A uint32 1D tensor indicating the q sequence length of each prompt. shape: ``[batch_size]``.
+            If not provided, will be set to the same value as ``seq_lens``.


Comment that it's only used for cudnn attention?

## 📌 Description  ## 🔍 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](https://pre-commit.com/). ## 🧪 Tests - [ ] Tests have been added or updated as needed. - [ ] All tests are passing (`unittest`, etc.). ## Reviewer Notes

gemini-code-assist bot reviewed Aug 5, 2025

View reviewed changes

benchmarks/routines/attention.py Outdated Show resolved Hide resolved

Anerudhan force-pushed the update/BatchPrefillPagedWrapper branch 3 times, most recently from 538af25 to 2b5977e Compare August 7, 2025 06:53

Anerudhan marked this pull request as ready for review August 7, 2025 06:54

Anerudhan changed the title ~~Fixes to the benchmark attention wrapper~~ Allow BatchPrefillPagedWrapper to call cudnn API Aug 7, 2025

Anerudhan force-pushed the update/BatchPrefillPagedWrapper branch from 2b5977e to 220041c Compare August 9, 2025 05:45

yzh119 reviewed Aug 10, 2025

View reviewed changes

flashinfer/prefill.py Show resolved Hide resolved

flashinfer/prefill.py Show resolved Hide resolved

flashinfer/prefill.py Show resolved Hide resolved

Fixes to the benchmark attention wrapper

90d00f0

Anerudhan force-pushed the update/BatchPrefillPagedWrapper branch from c27d1fa to bd2ec06 Compare August 11, 2025 16:51

Add cudnn to the BatchPrefillWithPagedKVCacheWrapper

5ad0187

Anerudhan force-pushed the update/BatchPrefillPagedWrapper branch from bd2ec06 to 5ad0187 Compare August 11, 2025 20:13

yzh119 approved these changes Aug 11, 2025

View reviewed changes

yzh119 merged commit eebb92d into flashinfer-ai:main Aug 11, 2025
2 checks passed

weireweire reviewed Aug 12, 2025

View reviewed changes

Edenzzzz reviewed Aug 12, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Allow BatchPrefillPagedWrapper to call cudnn API #1384

Allow BatchPrefillPagedWrapper to call cudnn API #1384

Anerudhan commented Aug 5, 2025

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

yzh119 left a comment

Uh oh!

Uh oh!

weireweire Aug 12, 2025

Uh oh!

weireweire Aug 12, 2025

Uh oh!

Anerudhan Aug 12, 2025

Uh oh!

yzh119 Aug 13, 2025

Uh oh!

Edenzzzz Aug 12, 2025

Uh oh!

Edenzzzz Aug 12, 2025

Uh oh!

Uh oh!

Allow BatchPrefillPagedWrapper to call cudnn API #1384

Allow BatchPrefillPagedWrapper to call cudnn API #1384

Conversation

Anerudhan commented Aug 5, 2025

📌 Description

🔍 Related Issues

🚀 Pull Request Checklist

✅ Pre-commit Checks

🧪 Tests

Reviewer Notes

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

yzh119 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

weireweire Aug 12, 2025

Choose a reason for hiding this comment

Uh oh!

weireweire Aug 12, 2025

Choose a reason for hiding this comment

Uh oh!

Anerudhan Aug 12, 2025

Choose a reason for hiding this comment

Uh oh!

yzh119 Aug 13, 2025

Choose a reason for hiding this comment

Uh oh!

Edenzzzz Aug 12, 2025

Choose a reason for hiding this comment

Uh oh!

Edenzzzz Aug 12, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!