[Bugfix] Fix race condition in d2h transfer (#3372)
### What this PR does / why we need it? Using non-blocking operations for device-to-host transfers can lead to data corruption in later steps. The CPU tensor is accessed right after the transfer is triggered, but the transfer might not be complete yet. As a result, the data could be wrong. This problem was seen in the A3 environment during `profile_run`. ### How was this patch tested? CI pass. - vLLM version: v0.11.0rc3 - vLLM main: https://github.com/vllm-project/vllm/commit/v0.11.0 Signed-off-by: Jade Zheng <zheng.shoujian@outlook.com>
This commit is contained in:
@@ -459,7 +459,7 @@ def torchair_fused_experts_with_all2all(
|
||||
token_counts_combined = token_counts_combined.view(
|
||||
2, ep_group.world_size, -1).sum(dim=2)
|
||||
token_counts_combined_cpu = token_counts_combined.to(
|
||||
torch.device("cpu"), non_blocking=True).numpy()
|
||||
torch.device("cpu"), non_blocking=False).numpy()
|
||||
all_tokens = gather_sizes.sum()
|
||||
|
||||
gathered_tokens = quantized_tokens.new_empty(all_tokens.item(),
|
||||
|
||||
Reference in New Issue
Block a user