Fix env for chunked prefill

Signed-off-by: fsx950223 <[email protected]>

Author: fsx950223

vllm/attention/backends/rocm_flash_attn.py

@@ -943,7 +943,8 @@ def forward(
             use_custom = use_rocm_custom_paged_attention(
                 decode_query.dtype, head_size, block_size, gqa_ratio,
                 decode_meta.max_decode_seq_len, self.sliding_window,
-                self.kv_cache_dtype, self.alibi_slopes)
+                self.kv_cache_dtype,
+                self.alibi_slopes) and not is_rocm_aiter_paged_attn_enabled()
 
             if use_custom:
                 max_seq_len = (decode_meta.max_decode_seq_len if self.attn_type

vllm/platforms/rocm.py

@@ -150,9 +150,8 @@ def use_rocm_custom_paged_attention(
                 and (block_size == 16 or block_size == 32)
                 and (gqa_ratio >= 1 and gqa_ratio <= 16)
                 and max_seq_len <= 128 * 1024
-                and (envs.VLLM_ROCM_CUSTOM_PAGED_ATTN)
-                and not (envs.VLLM_ROCM_USE_AITER_PAGED_ATTN
-                         and envs.VLLM_ROCM_USE_AITER) and sinks is None)
+                and (envs.VLLM_ROCM_CUSTOM_PAGED_ATTN) and sinks is None)
+
     else:
         return (ON_GFX11_GFX12 and (not envs.VLLM_USE_V1 or sliding_window == 0
                                     or sliding_window == (-1, -1))