use CK FA for glm-4v on navi3 #281
Conversation
| @@ -79,6 +80,31 @@ def __init__( | |||
| self.output_dropout = torch.nn.Dropout(config.dropout_prob) | |||
|
|
|||
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |||
| _ON_NAVI3 = "gfx11" in torch.cuda.get_device_properties("cuda").gcnArchName | |||
There was a problem hiding this comment.
Calling torch.cuda directly in vllm will not work on platforms other than ROCm. Please refer to platform/rocm.py and its superclass. Also is_navi() in utils.py
There was a problem hiding this comment.
The CK FA implementation is specify to navi3 only. but is_navi() include all navi GPUs. how about :
_ON_NAVI3 = current_platform.is_rocm() and "gfx11" in torch.cuda.get_device_properties("cuda").gcnArchName
There was a problem hiding this comment.
First checking for is_rocm is the right way to go, just please move this alongside the general is_navi, in utils.py for visibility and others to use as well
There was a problem hiding this comment.
Thanks for the suggestion. I have added is_navi3 function in utils.py
| _ON_NAVI3 = "gfx11" in torch.cuda.get_device_properties("cuda").gcnArchName | ||
| if _ON_NAVI3: | ||
| try: | ||
| # git clone -b howiejay/navi_support https://github.com/ROCm/flash-attention.git |
There was a problem hiding this comment.
Will this fail without flash_attn built from this particular branch?
There was a problem hiding this comment.
yes, only this branch support navi3.
There was a problem hiding this comment.
Will any other FA branch trigger a ModuleNotFoundError?
Is there an ETA on when this branch is going to be accepted into mainline FA?
There was a problem hiding this comment.
The other FA branch doesn't support navi3. It will be fail to install on gfx1100. I think there isn't plan to upstream this branch.
SDPA is very slow on navi3. Replace it with CK FA for glm-4v on navi3