update fa interface

vllm/attention/backends/dual_chunk_flash_attn.py

@@ -1221,23 +1221,42 @@ class DualChunkFlashAttentionImpl(FlashAttentionImpl):
                 s_lse = s_lse.view(q_len, q_heads, 1).transpose(0, 2).float()
             return res, s_lse
 
-        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_softmax_lse=True,
-        )
+        if not current_platform.is_rocm():
+            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_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