update fa interface

vllm/attention/backends/dual_chunk_flash_attn.py

@@ -1221,6 +1221,7 @@ class DualChunkFlashAttentionImpl(FlashAttentionImpl):
                 s_lse = s_lse.view(q_len, q_heads, 1).transpose(0, 2).float()
             return res, s_lse
 
+        if not current_platform.is_rocm():
             output, softmax_lse = flash_attn_varlen_func(
                 q=query_states,
                 k=key_states,
@@ -1238,6 +1239,24 @@ class DualChunkFlashAttentionImpl(FlashAttentionImpl):
                 block_table=block_table.unsqueeze(0),
                 return_softmax_lse=True,
             )
+        else:
+            output, softmax_lse = flash_attn_varlen_func(
+                q=query_states,
+                k=key_states,
+                v=value_states,
+                softmax_scale=softmax_scale,
+                cu_seqlens_q=torch.tensor([0, query_states.shape[0]],
+                                        dtype=torch.int32,
+                                        device=query_states.device),
+                max_seqlen_q=query_states.shape[0],
+                cu_seqlens_k=torch.tensor([0, max_seqlen_k],
+                                        dtype=torch.int32,
+                                        device=query_states.device),
+                max_seqlen_k=max_seqlen_k,
+                causal=causal,
+                block_table=block_table.unsqueeze(0),
+                return_attn_probs=True,
+            )
         softmax_lse = softmax_lse.view(q_len, q_heads, 1).transpose(0,
                                                                     2).float()
         return output, softmax_lse