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 14, 05:30 UTC
[BugFix] Fix async scheduling CPU tensor race take 2 (#25279)
Signed-off-by: Nick Hill <nhill@redhat.com>
Nick Hill committed
Sep 19, 2025 at 16:34 UTC
14c1432789c9c1b66308481b2c37439d3ee6661a
1 file changed
+32
-22
vllm/v1/worker/gpu_model_runner.py
+32
-22
@@ -1903,7 +1903,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1903
**self._init_model_kwargs(num_scheduled_tokens),
1904
**self._extract_mm_kwargs(scheduler_output),
1905
}
1906
- elif (self.enable_prompt_embeds and get_pp_group().is_first_rank):
1906
+ elif self.enable_prompt_embeds and get_pp_group().is_first_rank:
1907
# Get the input embeddings for the tokens that are not input embeds,
1908
# then put them into the appropriate positions.
1909
# TODO(qthequartermasterman): Since even when prompt embeds are
@@ -2125,6 +2125,21 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2125
invalid_req_indices,
2126
)
2127
2128
+ @contextmanager
2129
+ def synchronize_input_prep(self):
2130
+ if self.prepare_inputs_event is None:
2131
+ yield
2132
+ return
2133
+
2134
+ # Ensure prior step has finished with reused CPU tensors.
2135
+ # This is required in the async scheduling case because
2136
+ # the CPU->GPU transfer happens async.
2137
+ self.prepare_inputs_event.synchronize()
2138
+ try:
2139
+ yield
2140
+ finally:
2141
+ self.prepare_inputs_event.record()
2142
+
2143
@torch.inference_mode()
2144
def execute_model(
2145
self,
@@ -2132,33 +2147,28 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2147
intermediate_tensors: Optional[IntermediateTensors] = None,
2148
) -> Union[ModelRunnerOutput, AsyncModelRunnerOutput, IntermediateTensors]:
2149
with record_function_or_nullcontext("Preprocess"):
2135
- self._update_states(scheduler_output)
2136
- if not scheduler_output.total_num_scheduled_tokens:
2137
- if not has_kv_transfer_group():
2138
- # Return empty ModelRunnerOutput if there's no work to do.
2139
- return EMPTY_MODEL_RUNNER_OUTPUT
2140
- return self.kv_connector_no_forward(scheduler_output,
2141
- self.vllm_config)
2142
- if self.cache_config.kv_sharing_fast_prefill:
2143
- assert not self.input_batch.num_prompt_logprobs, (
2144
- "--kv-sharing-fast-prefill produces incorrect logprobs for "
2145
- "prompt tokens, tokens, please disable it when the requests"
2146
- " need prompt logprobs")
2147
-
2148
- if self.prepare_inputs_event is not None:
2149
- # Ensure prior step has finished with reused CPU tensors.
2150
- self.prepare_inputs_event.synchronize()
2151
- try:
2150
+ with self.synchronize_input_prep():
2151
+ # Update persistent batch states.
2152
+ self._update_states(scheduler_output)
2153
+
2154
+ if not scheduler_output.total_num_scheduled_tokens:
2155
+ if not has_kv_transfer_group():
2156
+ # Return empty ModelRunnerOutput if no work to do.
2157
+ return EMPTY_MODEL_RUNNER_OUTPUT
2158
+ return self.kv_connector_no_forward(
2159
+ scheduler_output, self.vllm_config)
2160
+ if self.cache_config.kv_sharing_fast_prefill:
2161
+ assert not self.input_batch.num_prompt_logprobs, (
2162
+ "--kv-sharing-fast-prefill produces incorrect "
2163
+ "logprobs for prompt tokens, tokens, please disable "
2164
+ "it when the requests need prompt logprobs")
2165
+
2166
# Prepare the decoder inputs.
2167
(attn_metadata, logits_indices, spec_decode_metadata,
2168
num_scheduled_tokens_np, spec_decode_common_attn_metadata,
2169
max_query_len, ubatch_slices, num_tokens_after_padding
2170
) = self._prepare_inputs(scheduler_output)
2171
2158
- finally:
2159
- if self.prepare_inputs_event is not None:
2160
- self.prepare_inputs_event.record()
2161
-
2172
(
2173
num_scheduled_tokens,
2174
num_input_tokens,