docs: Fix inaccurate API docstrings for attention prefill by bkryu · Pull Request #2196 · flashinfer-ai/flashinfer

bkryu · 2025-12-10T22:55:33Z

📌 Description

This PR addresses the documentation issues and adds proper validation for unsupported configurations in the attention prefill wrappers raised in #1709.
Changes

BatchPrefillWithPagedKVCacheWrapper:
- Updated docstring to include trtllm-gen as a valid backend option, which was
- Added validation for non-causal attention with trtllm-gen backend: The kernel only supports causal attention. Previously, using causal=False would silently run causal attention anyway (causing incorrect outputs). Now it raises a clear NotImplementedError.
  supported but undocumented.
BatchPrefillWithRaggedKVCacheWrapper: Fixed docstring to list cutlass instead of trtllm-gen as a backend option.
- The cutlass backend is supported but was missing from docs, while trtllm-gen raises NotImplementedError for ragged KV cache.

🔍 Related Issues

#1709

🚀 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

Summary by CodeRabbit

New Features
- Added support for trtllm-gen and cutlass backend options
Documentation
- Clarified available backends and default configurations
- Improved error guidance for unsupported backend combinations

_{✏️ Tip: You can customize this high-level summary in your review settings.}

coderabbitai · 2025-12-10T22:55:45Z

Walkthrough

The changes expand backend option documentation for prefill wrapper classes, introducing trtllm-gen and cutlass as valid backends, while enforcing a runtime constraint preventing non-causal attention with trtllm-gen backend and paged KV cache.

Changes

Cohort / File(s)	Summary
Backend documentation and constraint enforcement `flashinfer/prefill.py`	Updated docstrings for `BatchPrefillWithPagedKVCacheWrapper` and `BatchPrefillWithRaggedKVCacheWrapper` to include trtllm-gen and cutlass as valid backend options; added runtime guard in `plan()` method raising `NotImplementedError` when non-causal attention is used with trtllm-gen backend and paged KV cache.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~8 minutes

Documentation updates comprise the majority of changes (docstring modifications for backend descriptions and defaults)
Single runtime constraint check is straightforward logic with clear error messaging
All changes confined to a single file with no complex interactions

Possibly related PRs

fix: Fix trtllm-gen prefill IMA when batch_size==1 #1912: Fixes trtllm-gen prefill logic for batch_size==1 in ragged attention context, directly related to trtllm-gen backend behavior addressed in this PR.

Suggested reviewers

aleozlx
yzh119
cyx-6
wenscarl
nvmbreughe

Poem

🐰 Docs shine bright with backends new,
trtllm-gen joins the crew,
Guards protect the paged cache flow,
Causal constraints steal the show! ✨

Pre-merge checks and finishing touches

✅ Passed checks (3 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title 'docs: Fix inaccurate API docstrings for attention prefill' accurately and concisely describes the main change in the PR: updating documentation for the attention prefill API.
Description check	✅ Passed	The PR description includes the required sections with sufficient detail: a clear description of changes addressing issue #1709, related issues link, and a checklist (though items are unchecked). The description clearly explains what was changed and why.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.

✨ Finishing touches

📝 Generate docstrings

🧪 Generate unit tests (beta)

Create PR with unit tests
Post copyable unit tests in a comment

📜 Recent review details

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 6bb01d1 and b697506.

📒 Files selected for processing (1)

flashinfer/prefill.py (3 hunks)

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)

GitHub Check: Deploy Docs

🔇 Additional comments (3)

flashinfer/prefill.py (3)

1474-1475: LGTM!

The docstring correctly documents trtllm-gen as a valid backend option for paged KV cache, aligning documentation with actual supported functionality.

1923-1928: LGTM!

The validation correctly prevents silent incorrect outputs by raising NotImplementedError early in plan() when non-causal attention is requested with the trtllm-gen backend. The error message provides clear guidance on alternatives.

2466-2469: LGTM!

The docstring correctly documents cutlass as a valid backend option for ragged KV cache. This aligns with the implementation where trtllm-gen is explicitly unsupported for ragged operations (raises NotImplementedError at line 265-267) and cutlass is properly handled via get_fmha_module.

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

gemini-code-assist · 2025-12-10T22:55:48Z

Summary of Changes

Hello @bkryu, 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 improves the accuracy and robustness of the attention prefill wrappers by correcting API docstrings and implementing crucial validation checks. It ensures that users are properly informed about supported backends and configurations, preventing silent errors and clarifying expected behavior for paged and ragged KV cache implementations.

Highlights

BatchPrefillWithPagedKVCacheWrapper Docstring Update: The documentation for BatchPrefillWithPagedKVCacheWrapper now correctly lists trtllm-gen as a supported backend option, which was previously missing.
BatchPrefillWithPagedKVCacheWrapper Validation: Added a NotImplementedError check to BatchPrefillWithPagedKVCacheWrapper when using the trtllm-gen backend with non-causal attention, as this configuration is not supported and previously led to incorrect silent behavior.
BatchPrefillWithRaggedKVCacheWrapper Docstring Correction: The docstring for BatchPrefillWithRaggedKVCacheWrapper has been updated to accurately reflect cutlass as a supported backend, replacing the incorrect mention of trtllm-gen.

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 by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

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 pull request 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. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

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 effectively addresses documentation inconsistencies and adds crucial validation for attention prefill wrappers. The docstring updates for BatchPrefillWithPagedKVCacheWrapper and BatchPrefillWithRaggedKVCacheWrapper accurately reflect the supported backends, improving clarity for users. The addition of a NotImplementedError for non-causal attention with the trtllm-gen backend is a significant improvement, preventing silent failures and incorrect outputs. The changes are well-executed and enhance the overall quality and usability of the library. The code is clean and correct. I have no further suggestions.

bkryu · 2025-12-11T00:06:01Z

/bot run

flashinfer-bot · 2025-12-11T00:06:45Z

GitLab MR !189 has been created, and the CI pipeline #39990613 is currently running. I'll report back once the pipeline job completes.

…-ai#2196)  ## 📌 Description This PR addresses the documentation issues and adds proper validation for unsupported configurations in the attention prefill wrappers raised in flashinfer-ai#1709. Changes * `BatchPrefillWithPagedKVCacheWrapper`: * Updated docstring to include `trtllm-gen` as a valid backend option, which was * Added validation for non-causal attention with `trtllm-gen` backend: The kernel only supports causal attention. Previously, using `causal=False` would silently run causal attention anyway (causing incorrect outputs). Now it raises a clear NotImplementedError. supported but undocumented. * `BatchPrefillWithRaggedKVCacheWrapper`: Fixed docstring to list `cutlass` instead of `trtllm-gen` as a backend option. * The `cutlass` backend is supported but was missing from docs, while `trtllm-gen` raises NotImplementedError for ragged KV cache.  ## 🔍 Related Issues flashinfer-ai#1709  ## 🚀 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   ## Summary by CodeRabbit * **New Features** * Added support for trtllm-gen and cutlass backend options * **Documentation** * Clarified available backends and default configurations * Improved error guidance for unsupported backend combinations <sub>✏️ Tip: You can customize this high-level summary in your review settings.</sub>

Fix prefill docstrings and add causal check for trtllm-gen

b697506

bkryu requested review from aleozlx, cyx-6, nvmbreughe, wenscarl and yzh119 as code owners December 10, 2025 22:55

gemini-code-assist Bot reviewed Dec 10, 2025

View reviewed changes

bkryu self-assigned this Dec 10, 2025

bkryu added this to the 2025 Dec milestone Dec 11, 2025

yzh119 approved these changes Dec 11, 2025

View reviewed changes

yzh119 enabled auto-merge (squash) December 11, 2025 07:33

yzh119 disabled auto-merge December 11, 2025 07:33

yzh119 merged commit 3a29167 into flashinfer-ai:main Dec 11, 2025
3 checks passed

bkryu mentioned this pull request Dec 11, 2025

Inaccurate API Docstrings for Attention Prefill #1709

Closed

bkryu deleted the prefill_docstring branch December 11, 2025 17:55

coderabbitai Bot mentioned this pull request Dec 22, 2025

Incorrect trtllm_batch_decode_with_kv_cache_mla API doc string #2258

Closed

coderabbitai Bot mentioned this pull request Feb 24, 2026

allow cudnn to be chosen for prefill #2622

Open

5 tasks

This was referenced Apr 9, 2026

trtllm non causal support #3020

Merged

feat: Integrate CuTe DSL FMHA prefill kernels by loading cubin #3039

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

docs: Fix inaccurate API docstrings for attention prefill#2196

docs: Fix inaccurate API docstrings for attention prefill#2196
yzh119 merged 1 commit intoflashinfer-ai:mainfrom
bkryu:prefill_docstring

bkryu commented Dec 10, 2025 •

edited by coderabbitai Bot

Loading

Uh oh!

coderabbitai Bot commented Dec 10, 2025 •

edited

Loading

Uh oh!

gemini-code-assist Bot commented Dec 10, 2025

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

bkryu commented Dec 11, 2025

Uh oh!

flashinfer-bot commented Dec 11, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

bkryu commented Dec 10, 2025 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

📌 Description

🔍 Related Issues

🚀 Pull Request Checklist

✅ Pre-commit Checks

🧪 Tests

Reviewer Notes

Summary by CodeRabbit

Uh oh!

coderabbitai Bot commented Dec 10, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Estimated code review effort

Possibly related PRs

Suggested reviewers

Poem

Pre-merge checks and finishing touches

Uh oh!

gemini-code-assist Bot commented Dec 10, 2025

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!

bkryu commented Dec 11, 2025

Uh oh!

flashinfer-bot commented Dec 11, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

bkryu commented Dec 10, 2025 •

edited by coderabbitai Bot

Loading

coderabbitai Bot commented Dec 10, 2025 •

edited

Loading