Fix the FA2 logic in the longcat_flash model #42549
Conversation
|
@vasqu , please help review. Thanks! |
| if config.qk_head_dim != config.v_head_dim: | ||
| self.skipTest( | ||
| reason="Flash Attention 2 requires qk_head_dim == v_head_dim, but got " | ||
| f"qk_head_dim={config.qk_head_dim}, v_head_dim={config.v_head_dim}" | ||
| ) |
There was a problem hiding this comment.
Will this not skip all tests here? I doubt that the classes will have different head dims.
I'd rather we properly adjust the sizes than to skip - should always be the last resort.
There was a problem hiding this comment.
Yes, this test only involves this single model_class. After digging deeper into LongcatFlashForCausalLM, I found that it already implements the padding pre-processing for FA2 internally. However, for the fallback FA2 path(kernels), it failed to correctly match the parameter naming. This PR updates that part with the fix. Please review it again!
YangKai0616
force-pushed
the
fix-FA2-UT
branch
from
13038ae to
c492337
Compare
YangKai0616
force-pushed
the
fix-FA2-UT
branch
2 times, most recently
from
13038ae to
11e4ac9
Compare
YangKai0616
changed the title
vasqu
left a comment
vasqu
left a comment
There was a problem hiding this comment.
Thx for iterating! I think we can generalize this some more to include all attentions? They will probably face similar issues
| uses_flash_attention_2 = ( | ||
| "flash" in self.config._attn_implementation and self.config._attn_implementation.endswith("2") | ||
| ) | ||
| if uses_flash_attention_2 and self.qk_head_dim != self.v_head_dim: |
There was a problem hiding this comment.
| uses_flash_attention_2 = ( | |
| "flash" in self.config._attn_implementation and self.config._attn_implementation.endswith("2") | |
| ) | |
| if uses_flash_attention_2 and self.qk_head_dim != self.v_head_dim: | |
| if "flash" in self.config._attn_implementation and self.qk_head_dim != self.v_head_dim: |
I think we should generalize this here to check for all flavors. FA3 etc would face the same issue
| ) | ||
|
|
||
| if self.config._attn_implementation == "flash_attention_2" and self.qk_head_dim != self.v_head_dim: | ||
| if uses_flash_attention_2 and self.qk_head_dim != self.v_head_dim: |
There was a problem hiding this comment.
| if uses_flash_attention_2 and self.qk_head_dim != self.v_head_dim: | |
| if "flash" in self.config._attn_implementation and self.qk_head_dim != self.v_head_dim: |
Same here then
|
[For maintainers] Suggested jobs to run (before merge) run-slow: longcat_flash |
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
* Matching FA2 naming under kernels * make style * convert model * Follow the comments
3 participants
What does this PR do?
FA2 does not support MLA (i.e., cases where the dimensions of Q, K, and V heads are inconsistent), so skip this test.