update test_prefix_prefill.py

tests/kernels/test_prefix_prefill.py

@@ -20,7 +20,7 @@ CUDA_DEVICES = [
     f"cuda:{i}" for i in range(1 if torch.cuda.device_count() == 1 else 2)
 ]
 SLIDING_WINDOW = [0, 16, 64, 128, 256, 512, 2048]
-KV_CACHE_DTYPES = ["auto", "fp8", "fp8_e5m2"]
+KV_CACHE_DTYPES = ["auto", "fp8", "fp8_e5m2"] if not is_hip() else ["auto"]
 
 
 @pytest.mark.parametrize("num_heads", NUM_HEADS)
@@ -432,54 +432,6 @@ def test_contexted_kv_attention_alibi(
         key = key.unsqueeze(0)
         value = value.unsqueeze(0)
 
-        attn_bias = _make_alibi_bias(alibi_slopes, num_kv_heads, dtype, seq_lens)
-        output_ref = torch.empty_like(output)
-        seq_start = 0
-        query_start = 0
-        start_time = time.time()
-        # Attention with alibi slopes.
-        # FIXME(DefTruth): Because xformers does not support dynamic sequence
-        # lengths with custom attention bias, we process each prompt one by
-        # one. This is inefficient, especially when we have many short prompts.
-        # modified from: vllm/attention/backends/xformers.py#L343
-        for i, (query_len, seq_len) in enumerate(zip(query_lens, seq_lens)):
-            seq_end = seq_start + seq_len
-            query_end = query_start + query_len
-            out = xops.memory_efficient_attention_forward(query[:,
-                                                                seq_start:seq_end],
-                                                        key[:,
-                                                            seq_start:seq_end],
-                                                        value[:,
-                                                                seq_start:seq_end],
-                                                        attn_bias=attn_bias[i],
-                                                        p=0.0,
-                                                        scale=scale)
-            out = out.view_as(query[:, seq_start:seq_end]).view(
-                seq_len, num_heads, head_size)
-            output_ref[query_start:query_end, ...].copy_(out[seq_len - query_len:,
-                                                            ...])
-            seq_start += seq_len
-            query_start += query_len
-        query = query_pad
-
-    if num_kv_heads != num_heads:
-        # As of Nov 2023, xformers only supports MHA. For MQA/GQA,
-        # project the key and value tensors to the desired number of
-        # heads.
-        #
-        # see also: vllm/model_executor/layers/attention.py
-        query = query.view(query.shape[0], num_kv_heads, num_queries_per_kv,
-                           query.shape[-1])
-        key = key[:, :, None, :].expand(key.shape[0], num_kv_heads,
-                                        num_queries_per_kv, key.shape[-1])
-        value = value[:, :,
-                      None, :].expand(value.shape[0], num_kv_heads,
-                                      num_queries_per_kv, value.shape[-1])
-
-    query = query.unsqueeze(0)
-    key = key.unsqueeze(0)
-    value = value.unsqueeze(0)
-
         attn_bias = _make_alibi_bias(alibi_slopes, num_kv_heads, dtype, seq_lens)
         output_ref = torch.empty_like(output)
         seq_start = 0