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 8, 01:15 UTC
[BUGFIX] Fix `_dummy_run` missing `prepare_inputs_event` synchronization (#34866)
Signed-off-by: Vadim Gimpelson <vadim.gimpelson@gmail.com>
Vadim Gimpelson committed
Feb 20, 2026 at 17:54 UTC
59965affbd6e652a3c8ed229b66ef34a681e5693
1 file changed
+31
-26
vllm/v1/worker/gpu_model_runner.py
+31
-26
@@ -4771,34 +4771,39 @@ class GPUModelRunner(
4771
ubatch_slices=ubatch_slices_padded,
4772
)
4773
4774
- # If force_attention is True, we always capture attention. Otherwise,
4775
- # it only happens for cudagraph_runtime_mode=FULL.
4776
- if force_attention or cudagraph_runtime_mode == CUDAGraphMode.FULL:
4777
- if create_mixed_batch:
4778
- # In the mixed batch mode (used for FI warmup), we use
4779
- # shorter sequence lengths to run faster.
4780
- # TODO(luka) better system for describing dummy batches
4781
- seq_lens = [1] * num_decode_tokens + [num_prefill_tokens + 1]
4782
- else:
4783
- seq_lens = max_query_len # type: ignore[assignment]
4784
- self.seq_lens.np[:num_reqs] = seq_lens
4785
- self.seq_lens.np[num_reqs:] = 0
4786
- self.seq_lens.copy_to_gpu()
4774
+ # _dummy_run shares pinned CPU buffers (seq_lens, query_start_loc,
4775
+ # etc.) with execute_model. It must participate in the same event
4776
+ # protocol so that back-to-back dummy/real steps don't overwrite
4777
+ # pinned memory while a prior non_blocking H2D DMA is still reading.
4778
+ with self.synchronize_input_prep():
4779
+ # If force_attention is True, we always capture attention.
4780
+ # Otherwise, it only happens for cudagraph_runtime_mode=FULL.
4781
+ if force_attention or cudagraph_runtime_mode == CUDAGraphMode.FULL:
4782
+ if create_mixed_batch:
4783
+ # In the mixed batch mode (used for FI warmup), we use
4784
+ # shorter sequence lengths to run faster.
4785
+ # TODO(luka) better system for describing dummy batches
4786
+ seq_lens = [1] * num_decode_tokens + [num_prefill_tokens + 1]
4787
+ else:
4788
+ seq_lens = max_query_len # type: ignore[assignment]
4789
+ self.seq_lens.np[:num_reqs] = seq_lens
4790
+ self.seq_lens.np[num_reqs:] = 0
4791
+ self.seq_lens.copy_to_gpu()
4792
4788
- cum_num_tokens, _ = self._get_cumsum_and_arange(num_scheduled_tokens)
4789
- self.query_start_loc.np[1 : num_reqs + 1] = cum_num_tokens
4790
- self.query_start_loc.copy_to_gpu()
4793
+ cum_num_tokens, _ = self._get_cumsum_and_arange(num_scheduled_tokens)
4794
+ self.query_start_loc.np[1 : num_reqs + 1] = cum_num_tokens
4795
+ self.query_start_loc.copy_to_gpu()
4796
4792
- pad_attn = cudagraph_runtime_mode == CUDAGraphMode.FULL
4793
- attn_metadata, _ = self._build_attention_metadata(
4794
- num_tokens=num_tokens_unpadded,
4795
- num_tokens_padded=num_tokens_padded if pad_attn else None,
4796
- num_reqs=num_reqs_padded,
4797
- max_query_len=max_query_len,
4798
- ubatch_slices=ubatch_slices_padded if pad_attn else ubatch_slices,
4799
- for_cudagraph_capture=is_graph_capturing,
4800
- slot_mappings=slot_mappings_by_group,
4801
- )
4797
+ pad_attn = cudagraph_runtime_mode == CUDAGraphMode.FULL
4798
+ attn_metadata, _ = self._build_attention_metadata(
4799
+ num_tokens=num_tokens_unpadded,
4800
+ num_tokens_padded=num_tokens_padded if pad_attn else None,
4801
+ num_reqs=num_reqs_padded,
4802
+ max_query_len=max_query_len,
4803
+ ubatch_slices=(ubatch_slices_padded if pad_attn else ubatch_slices),
4804
+ for_cudagraph_capture=is_graph_capturing,
4805
+ slot_mappings=slot_mappings_by_group,
4806
+ )
4807
4808
with self.maybe_dummy_run_with_lora(
4809
self.lora_config,