[Bugfix] Fix ChunkedLocalAttention CUDA Graph setting #28739
Conversation
Signed-off-by: Benjamin Chislett <[email protected]>
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request provides a crucial bugfix for ChunkedLocalAttention. The root cause, an incorrect inheritance of get_cudagraph_support from the underlying attention backend's builder, led to CUDA graphs being improperly enabled and causing incorrect model outputs. The proposed solution correctly overrides the get_cudagraph_support method to always return AttentionCGSupport.NEVER, effectively disabling CUDA graphs for this attention mechanism as intended. This change is robust, well-targeted, and prevents a critical correctness issue. The addition of an issubclass assertion is a good defensive measure. The implementation is clean and I have no further suggestions for improvement.
mgoin
left a comment
mgoin
left a comment
There was a problem hiding this comment.
LGTM as an immediate fix, thank you
LucasWilkinson
left a comment
LucasWilkinson
left a comment
There was a problem hiding this comment.
LGTM; thanks for the quick fix!
…8739) Signed-off-by: Benjamin Chislett <[email protected]> Signed-off-by: George D. Torres <[email protected]>
…8739) Signed-off-by: Benjamin Chislett <[email protected]> Signed-off-by: Bram Wasti <[email protected]>
|
@benchislett does this fix mean we can re-enable llama4 E2E fusion tests using the FI attention backend on Blackwell? |
|
I'm not sure about the overall status of Llama4 support, but this fix definitely should clear up a major blocker for blackwell. Seems like it's worth trying it out to see where it's at. CC @pavanimajety @xinli-sw who might have some context? |
…8739) Signed-off-by: Benjamin Chislett <[email protected]>
…8739) Signed-off-by: Benjamin Chislett <[email protected]>
…8739) Signed-off-by: Benjamin Chislett <[email protected]>
5 participants
Purpose
Bugfix for incorrect outputs on llama4 caused by the refactor in #28479: ChunkedLocalAttentionBuilder incorrectly inherits FlashInfer's
get_cudagraph_supportmethod, causing the_cudagraph_supportto be ignored and CUDA graphs to always get used.Test Plan
Breaking command:
Test Result
Now works. Ongoing discussion to make ChunkedLocalAttention compatible with CUDA graphs, but for now it should stay as NEVER to avoid errors.
FIX #28604