fix fa and triton tests

tests/kernels/test_attention.py

@@ -359,6 +359,8 @@ def test_multi_query_kv_attention(
         attn_bias=attn_bias,
         p=0.0,
         scale=scale,
+        op=xops.fmha.MemoryEfficientAttentionFlashAttentionOp[0] if
+        (is_hip()) else None,
     )
     output = output.squeeze(0)
 

vllm/model_executor/layers/ops/rand.py

@@ -3,6 +3,7 @@ from typing import Optional, Union
 import torch
 import triton
 import triton.language as tl
+from vllm.utils import is_hip
 
 
 def seeded_uniform(
@@ -69,8 +70,14 @@ def seeded_uniform(
     # Manual tuning. This seems to give best performance on A100 for
     # simple kernels like this.
     if philox_block_size >= 8192:
+        if is_hip():
+            num_warps = 16
+        else:
             num_warps = 32
     elif philox_block_size >= 4096:
+        if is_hip():
+            num_warps = 8
+        else:
             num_warps = 16
     elif philox_block_size >= 2048:
         num_warps = 8

vllm/model_executor/layers/ops/sample.py

@@ -6,6 +6,7 @@ import triton
 import triton.language as tl
 
 from vllm.model_executor.layers.ops.rand import seeded_uniform
+from vllm.utils import is_hip
 
 _EPS = 1e-6
 
@@ -278,8 +279,14 @@ def _sample(probs: torch.Tensor,
     # Manual tuning. This seems to give best performance on A100 for
     # simple kernels like this.
     if block_size >= 8192:
+        if is_hip():
+            num_warps = 16
+        else:
             num_warps = 32
     elif block_size >= 4096:
+        if is_hip():
+            num_warps = 8
+        else:
             num_warps = 16
     elif block_size >= 2048:
         num_warps = 8