support sinks

vllm/v1/attention/backends/flash_attn.py

@@ -453,8 +453,9 @@ class FlashAttentionImpl(AttentionImpl):
 
         self.sinks = sinks
         if self.sinks is not None:
-            assert self.vllm_flash_attn_version == 3, (
-                "Sinks are only supported in FlashAttention 3")
+            if not current_platform.is_rocm():
+                assert self.vllm_flash_attn_version == 3, (
+                    "Sinks are only supported in FlashAttention 3")
             assert self.sinks.shape[0] == num_heads, (
                 "Sinks must have the same number of heads as the number of "
                 "heads in the layer")