Fix kernel cache miss and add RDNA configs

[hyoon1]

• added Navi configurations (Related PR: add RDNA Config triton#640)
• resolved cache miss issue during flash attention calls by fixing max_seqlen_q/k to 0

[maleksan85]

what is the reason to zero seq lens?

[hyoon1]

Below attention fwd kernel is called when we run the model with vllm:

vllm/vllm/attention/ops/triton_flash_attention.py

Line 776 in efb0432

attn_fwd[grid](

However, MAX_SEQLENS_Q/K differs every step, and it occurs different key value and compilation for the triton kernel each step, which leads to the performance degradation.
https://github.com/triton-lang/triton/blob/cf34004b8a67d290a962da166f5aa2fc66751326/python/triton/runtime/jit.py#L620
https://github.com/triton-lang/triton/blob/cf34004b8a67d290a962da166f5aa2fc66751326/python/triton/runtime/jit.py#L660

Currently, VARLEN is always set, and MAX_SEQLENS_Q/K are not used in this case when you look at the kernel in vllm.

vllm/vllm/attention/ops/triton_flash_attention.py

Line 309 in efb0432

def attn_fwd(

Therefore, we just set MAX_SEQLENS_Q/K as a fixed value when we call the kernel for a workaround.

[maleksan85]

seems like not used, right?

[maleksan85]

probably worth to use:

vllm/vllm/utils.py

Line 1620 in 8f3bf8b

def is_navi() -> bool:

[maleksan85]

As per my knowledge AMD has two lines of HW for vllm: MI and Navi. So not navi should work better for future generations of MIs

[maleksan85]

As per @gshtras we need to merge into develop branch instead of main for now. Please correct.

[gshtras]

All this functionality is implemented in a cross-architecture fashion in the platform/rocm.py and its superclasses

[hyoon1]

@maleksan85 @gshtras
Thanks for the comments, I agree with your suggestions. However, I have a few concerns regarding this matter and I'm seeking advice on how to proceed. First, I believe that triton_flash_attention.py in vllm is essentially a copy of the file from ROCm/triton. The modifications in this pull request directly apply the changes from the ROCm/triton repository's pull request #640. While it might be fine to change the relevant functions only in vllm, there is a risk of misalignment later on.

Secondly, our team is using the v0.6.2+rocm release, and I understand that functions like is_navi() are not supported in that version. Implementing them would require significant modifications. Therefore, maintaining backward compatibility is also a concern.

Given these considerations, I would greatly appreciate your advice on how to proceed with the modifications.

[gshtras]

As for your last point, whatever changes will be made here will not have any effect on the previous tags, so v0.6.2+rocm will not get affected.
That's true, our kernel is a snapshot of the one you mentioned, taken in the beginning of 2024. Our attempts to catch up in the past resulted in performance regressions on various models in different configs, which was noted to the team, but I don't believe it was thoroughly investigated. We may take another round of this experiment, but regardless, I think utilizing an existing infrastructure APIs that vllm provides is better here, if nothing else, for uniformity and avoiding code duplication.