extend FA2 and other cases to XPU, #42536
Conversation
…UDAGraph specific, CUDA compute capability specific and FA3 specific can run XPU. For FA3, we are develioping Signed-off-by: Yao, Matrix <[email protected]>
Signed-off-by: Yao, Matrix <[email protected]>
|
|
||
| MIMI_ATTENTION_CLASSES = { | ||
| "eager": MimiAttention, | ||
| "kernels-community/flash-attn2": MimiFlashAttention2, |
There was a problem hiding this comment.
could you explain this part?
There was a problem hiding this comment.
@ydshieh , sure. in latest design, when users set attn_implementation == "flash_attention_2", there will be 2 branches:
- if
flash_attnpackage is available, it will go directly to use it - else, do not fail as before, but use kernels instead, in this case, the
attn_implementationwill be updated to"kernels-community/flash-attn2", as in code here
For XPU, we go with the kernels path in transformers for FA support, so we need this key.
Thx very much.
There was a problem hiding this comment.
I'd rather not do this, even tho you are correct here. We should rather refactor mimi here with the attention interface and not have these manual registrations. We could infinitely extend these edge cases in the future to FA3 etc which makes this not scalable (without using/refactoring to the interface).
There was a problem hiding this comment.
Let's revert this line 🙏 . We can skip the relevant FA tests if necessary.
| ("xpu", None): { | ||
| "req_1": " 3.5 bolts.\n\nLet's break it down step by step:\n\n- Blue fiber: 2 bolts\n- White fiber: half of 2 bolts = 1 bolt\n\nTotal = ", | ||
| }, | ||
| }).get_expectation() # fmt: skip |
There was a problem hiding this comment.
i need to check why this was {} before, but thank you.
|
@yao-matrix don't forget 🙏 |
Yes, done, thx very much for your always support, :). |
|
[For maintainers] Suggested jobs to run (before merge) run-slow: gemma2, gemma3, glm4v, glm4v_moe, granitemoehybrid, idefics2, kosmos2_5, longcat_flash, mimi, modernbert, musicgen, musicgen_melody, pixtral, qwen2_5_omni, qwen2_5_vl, qwen2_moe |
3 participants
we expect all model cases except CUDAGraph specific, CUDA compute capability specific and FA3 specific can run XPU. For FA3, we are developing.
@ydshieh, pls help review, thx very much.