[Feature] Disallow FlashMLA on Blackwell #24521
Conversation
Signed-off-by: yewentao256 <[email protected]>
yewentao256
requested review from
WoosukKwon,
alexm-redhat,
comaniac,
njhill,
robertgshaw2-redhat and
ywang96
as code owners
![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 correctly disallows the use of the FlashMLA attention backend on NVIDIA Blackwell GPUs by adding a check for SM 10.0+ device capability. This is a necessary fix due to known support issues and is applied to both v0 and v1 engine implementations. My review suggests changing the exception type from AssertionError to NotImplementedError for better consistency with surrounding code and Python best practices.
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Signed-off-by: Wentao Ye <[email protected]>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Signed-off-by: Wentao Ye <[email protected]>
|
Hi @yewentao256, is it possible to use current_platform.is_cuda() and has_device_capability() instead of CudaPlatform? Will be more friendly for other platforms, thank you |
|
Signed-off-by: yewentao256 <[email protected]> Signed-off-by: Wentao Ye <[email protected]> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Signed-off-by: yewentao256 <[email protected]> Signed-off-by: Wentao Ye <[email protected]> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Signed-off-by: yewentao256 <[email protected]> Signed-off-by: Wentao Ye <[email protected]> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Signed-off-by: xuebwang-amd <[email protected]>
Signed-off-by: yewentao256 <[email protected]> Signed-off-by: Wentao Ye <[email protected]> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Signed-off-by: xuebwang-amd <[email protected]>
3 participants
Purpose
Fixing #24513
FlashMLA is poorly supported on Blackwell, see
deepseek-ai/FlashMLA#83
for more details.