[AMD][BugFix] Add gpu_sync_allowed to ROCm AITER FA backend (#53641)
Signed-off-by: Randall Smith <Randall.Smith@amd.com>
rasmith committed
Aug 26, 2026 at 20:23 UTC
2dc53bae174397e763995120c4095a23d9ad5ca3
1 file changed
+9
-7
vllm/v1/attention/backends/rocm_aiter_fa.py
+9
-7
@@ -14,6 +14,7 @@ from vllm.logger import init_logger
14
from vllm.model_executor.layers.attention import Attention
15
from vllm.platforms import current_platform
16
from vllm.platforms.interface import DeviceCapability
17
+from vllm.utils.gpu_sync_debug import gpu_sync_allowed
18
from vllm.utils.math_utils import cdiv
19
from vllm.utils.platform_utils import num_compute_units
20
from vllm.utils.torch_utils import is_quantized_kv_cache
@@ -496,13 +497,14 @@ class AiterFlashAttentionMetadataBuilder(
497
498
query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu
499
499
- # Only copy seq_lens to CPU when prefill or extend is present to avoid a
500
- # blocking device→host transfer.
501
- seq_lens = (
502
- common_attn_metadata.seq_lens.cpu()
503
- if num_prefills > 0 or num_extends > 0
504
- else None
505
- )
500
+ with gpu_sync_allowed():
501
+ # Only copy seq_lens to CPU when prefill or extend is present to avoid a
502
+ # blocking device→host transfer.
503
+ seq_lens = (
504
+ common_attn_metadata.seq_lens.cpu()
505
+ if num_prefills > 0 or num_extends > 0
506
+ else None
507
+ )
508
509
query_lens_cpu = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]
510