A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 10, 15:15 UTC
[Bugfix] Fix race condition in async-scheduling for vlm model (#31841)
Signed-off-by: Tianshu Yu <tianshuyu.formal@gmail.com>
tianshu-Michael-yu committed
Jan 6, 2026 at 22:45 UTC
efeaac92f22f8a0a26c6bb9b9182f316210bb19c
1 file changed
+14
-3
vllm/v1/worker/gpu_model_runner.py
+14
-3
@@ -554,7 +554,13 @@ class GPUModelRunner(
554
555
# Only relevant for multimodal models
556
if self.supports_mm_inputs:
557
- self.is_mm_embed = self._make_buffer(self.max_num_tokens, dtype=torch.bool)
557
+ # Double buffer to avoid race condition: previous iteration's async
558
+ # copy may still be reading from CPU while current iteration writes.
559
+ self.is_mm_embed_buffers = [
560
+ self._make_buffer(self.max_num_tokens, dtype=torch.bool),
561
+ self._make_buffer(self.max_num_tokens, dtype=torch.bool),
562
+ ]
563
+ self.is_mm_embed_idx = 0
564
565
# Only relevant for models using M-RoPE (e.g, Qwen2-VL)
566
if self.uses_mrope:
@@ -2337,8 +2343,13 @@ class GPUModelRunner(
2343
) -> tuple[list[torch.Tensor], torch.Tensor]:
2344
total_num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
2345
2346
+ # Swap to the other buffer to avoid race condition with previous
2347
+ # iteration's async copy that may still be reading from CPU.
2348
+ self.is_mm_embed_idx = 1 - self.is_mm_embed_idx
2349
+ is_mm_embed_buf = self.is_mm_embed_buffers[self.is_mm_embed_idx]
2350
+
2351
mm_embeds = list[torch.Tensor]()
2341
- is_mm_embed = self.is_mm_embed.cpu
2352
+ is_mm_embed = is_mm_embed_buf.cpu
2353
is_mm_embed[:total_num_scheduled_tokens] = False
2354
2355
req_start_idx = 0
@@ -2416,7 +2427,7 @@ class GPUModelRunner(
2427
mm_embeds.extend(mm_embeds_req)
2428
req_start_idx += num_scheduled_tokens
2429
2419
- is_mm_embed = self.is_mm_embed.copy_to_gpu(total_num_scheduled_tokens)
2430
+ is_mm_embed = is_mm_embed_buf.copy_to_gpu(total_num_scheduled_tokens)
2431
2432
if should_sync_mrope_positions:
2433
self._calc_mrope_positions(scheduler_output)