[Bugfix][ROCm] Fix race on AITER MLA FP8 prefill scheduling metadata under async scheduling (#58887)
Signed-off-by: Mehdi Ghanimifard <mehdi.ghanimifard@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
Mehdi Ghanimifard committed
Oct 1, 2026 at 09:02 UTC
aee8fe202a2fbd79a7a33cae70b310edc5960aa5
1 file changed
+49
-11
vllm/v1/attention/backends/mla/rocm_aiter_mla.py
+49
-11
@@ -948,6 +948,31 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
948
device=device,
949
)
950
951
+ # get_ps_metadata_v1 builds the plan on the host and writes its outputs with
952
+ # blocking hipMemcpy that is not ordered on the current stream. Writing the
953
+ # device buffers directly lets step N+1's build (async scheduling) overwrite a
954
+ # plan that step N's prefill kernels are still reading. Build into pinned host
955
+ # staging instead and copy on the current stream. Two staging slots, each
956
+ # reused only after its previous H2D copy has completed.
957
+ self._fp8_ps_device_outputs = (
958
+ self.fp8_ps_work_indptr,
959
+ self.fp8_ps_work_info,
960
+ self.fp8_ps_reduce_indptr,
961
+ self.fp8_ps_reduce_final_map,
962
+ self.fp8_ps_reduce_partial_map,
963
+ )
964
+ self._fp8_ps_staging = [
965
+ tuple(
966
+ torch.empty(t.shape, dtype=t.dtype, device="cpu", pin_memory=True)
967
+ for t in self._fp8_ps_device_outputs
968
+ )
969
+ for _ in range(2)
970
+ ]
971
+ # One event per slot, recorded after that slot's H2D copy. synchronize()
972
+ # on an event that has not been recorded yet returns immediately.
973
+ self._fp8_ps_staging_free = [torch.cuda.Event(), torch.cuda.Event()]
974
+ self._fp8_ps_slot = 0
975
+
976
from vllm.platforms import current_platform
977
from vllm.v1.worker.workspace import current_workspace_manager
978
@@ -1043,6 +1068,20 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
1068
kvlen_granularity = 128
1069
block_size = 1 # non-paged: each "page" is one token
1070
1071
+ slot = self._fp8_ps_slot
1072
+ self._fp8_ps_slot ^= 1
1073
+ staging_free = self._fp8_ps_staging_free[slot]
1074
+ with gpu_sync_allowed():
1075
+ staging_free.synchronize()
1076
+ (
1077
+ work_indptr_host,
1078
+ work_info_host,
1079
+ reduce_indptr_host,
1080
+ reduce_final_map_host,
1081
+ reduce_partial_map_host,
1082
+ ) = self._fp8_ps_staging[slot]
1083
+
1084
+ # work_metadata is not written by the host planner; it stays on device.
1085
get_ps_metadata_v1(
1086
qo_indptr_cpu,
1087
kv_indptr_cpu,
@@ -1050,17 +1089,20 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
1089
gqa_ratio,
1090
num_head_k,
1091
self.fp8_ps_work_metadata,
1053
- self.fp8_ps_work_indptr,
1054
- self.fp8_ps_work_info,
1055
- self.fp8_ps_reduce_indptr,
1056
- self.fp8_ps_reduce_final_map,
1057
- self.fp8_ps_reduce_partial_map,
1092
+ work_indptr_host,
1093
+ work_info_host,
1094
+ reduce_indptr_host,
1095
+ reduce_final_map_host,
1096
+ reduce_partial_map_host,
1097
qhead_granularity=qhead_granularity,
1098
qlen_granularity=qlen_granularity,
1099
kvlen_granularity=kvlen_granularity,
1100
block_size=block_size,
1101
is_causal=True,
1102
)
1103
+ for dst, src in zip(self._fp8_ps_device_outputs, self._fp8_ps_staging[slot]):
1104
+ dst.copy_(src, non_blocking=True)
1105
+ staging_free.record()
1106
1107
total_prefill_tokens = int(qo_indptr_cpu[-1].item())
1108
kv_indices = torch.arange(
@@ -1068,12 +1110,8 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
1110
)
1111
1112
# The actual number of active partial tiles for this batch is the
1071
- # final value of reduce_indptr. Resolving it here (during metadata
1072
- # build) keeps it off the per-layer forward path where a sync would
1073
- # break CUDA Graph capture. Using the device-side reduce_indptr is
1074
- # acceptable since build is allowed to incur an occasional sync.
1075
- with gpu_sync_allowed():
1076
- num_partial_tiles = int(self.fp8_ps_reduce_indptr[-1].item())
1113
+ # final value of reduce_indptr, read from the host plan (no GPU sync).
1114
+ num_partial_tiles = int(reduce_indptr_host[-1])
1115
1116
# Attach PS metadata to the metadata object so forward_mha can read it.
1117
metadata.fp8_prefill_qo_indptr = qo_indptr