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, 21:30 UTC
[BugFix] Re-fix async multimodal cpu tensor race condition (#31373)
Signed-off-by: Nick Hill <nickhill123@gmail.com> Signed-off-by: njhill <nickhill123@gmail.com>
Nick Hill committed
Dec 28, 2025 at 03:05 UTC
094fcce250110b7987f36468a3b4e28cc4858378
1 file changed
+112
-114
vllm/v1/worker/gpu_model_runner.py
+112
-114
@@ -3058,131 +3058,129 @@ class GPUModelRunner(
3058
scheduler_output = deepcopy(scheduler_output)
3059
3060
num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
3061
- with record_function_or_nullcontext("gpu_model_runner: preprocess"):
3062
- with self.synchronize_input_prep():
3063
- # Update persistent batch states.
3064
- self._update_states(scheduler_output)
3065
-
3066
- if has_ec_transfer() and get_ec_transfer().is_producer:
3067
- with self.maybe_get_ec_connector_output(
3068
- scheduler_output,
3069
- encoder_cache=self.encoder_cache,
3070
- ) as ec_connector_output:
3071
- self._execute_mm_encoder(scheduler_output)
3072
- return make_empty_encoder_model_runner_output(scheduler_output)
3073
-
3074
- if not num_scheduled_tokens:
3075
- if (
3076
- self.parallel_config.distributed_executor_backend
3077
- == "external_launcher"
3078
- and self.parallel_config.data_parallel_size > 1
3079
- ):
3080
- # this is a corner case when both external launcher
3081
- # and DP are enabled, num_scheduled_tokens could be
3082
- # 0, and has_unfinished_requests in the outer loop
3083
- # returns True. before returning early here we call
3084
- # dummy run to ensure coordinate_batch_across_dp
3085
- # is called into to avoid out of sync issues.
3086
- self._dummy_run(1)
3087
- if not has_kv_transfer_group():
3088
- # Return empty ModelRunnerOutput if no work to do.
3089
- return EMPTY_MODEL_RUNNER_OUTPUT
3090
- return self.kv_connector_no_forward(
3091
- scheduler_output, self.vllm_config
3092
- )
3093
- if self.cache_config.kv_sharing_fast_prefill:
3094
- assert not self.num_prompt_logprobs, (
3095
- "--kv-sharing-fast-prefill produces incorrect "
3096
- "logprobs for prompt tokens, tokens, please disable "
3097
- "it when the requests need prompt logprobs"
3098
- )
3099
-
3100
- num_reqs = self.input_batch.num_reqs
3101
- req_ids = self.input_batch.req_ids
3102
- tokens = [scheduler_output.num_scheduled_tokens[i] for i in req_ids]
3103
- num_scheduled_tokens_np = np.array(tokens, dtype=np.int32)
3104
- max_num_scheduled_tokens = int(num_scheduled_tokens_np.max())
3105
- num_tokens_unpadded = scheduler_output.total_num_scheduled_tokens
3061
+ with (
3062
+ record_function_or_nullcontext("gpu_model_runner: preprocess"),
3063
+ self.synchronize_input_prep(),
3064
+ ):
3065
+ # Update persistent batch states.
3066
+ self._update_states(scheduler_output)
3067
3107
- (
3108
- logits_indices,
3109
- spec_decode_metadata,
3110
- ) = self._prepare_inputs(
3068
+ if has_ec_transfer() and get_ec_transfer().is_producer:
3069
+ with self.maybe_get_ec_connector_output(
3070
scheduler_output,
3112
- num_scheduled_tokens_np,
3071
+ encoder_cache=self.encoder_cache,
3072
+ ) as ec_connector_output:
3073
+ self._execute_mm_encoder(scheduler_output)
3074
+ return make_empty_encoder_model_runner_output(scheduler_output)
3075
+
3076
+ if not num_scheduled_tokens:
3077
+ if (
3078
+ self.parallel_config.distributed_executor_backend
3079
+ == "external_launcher"
3080
+ and self.parallel_config.data_parallel_size > 1
3081
+ ):
3082
+ # this is a corner case when both external launcher
3083
+ # and DP are enabled, num_scheduled_tokens could be
3084
+ # 0, and has_unfinished_requests in the outer loop
3085
+ # returns True. before returning early here we call
3086
+ # dummy run to ensure coordinate_batch_across_dp
3087
+ # is called into to avoid out of sync issues.
3088
+ self._dummy_run(1)
3089
+ if not has_kv_transfer_group():
3090
+ # Return empty ModelRunnerOutput if no work to do.
3091
+ return EMPTY_MODEL_RUNNER_OUTPUT
3092
+ return self.kv_connector_no_forward(scheduler_output, self.vllm_config)
3093
+
3094
+ if self.cache_config.kv_sharing_fast_prefill:
3095
+ assert not self.num_prompt_logprobs, (
3096
+ "--kv-sharing-fast-prefill produces incorrect "
3097
+ "logprobs for prompt tokens, tokens, please disable "
3098
+ "it when the requests need prompt logprobs"
3099
)
3100
3115
- cascade_attn_prefix_lens = None
3116
- # Disable cascade attention when using microbatching (DBO)
3117
- if self.cascade_attn_enabled and not self.parallel_config.use_ubatching:
3118
- # Pre-compute cascade attention prefix lengths
3119
- cascade_attn_prefix_lens = self._compute_cascade_attn_prefix_lens(
3120
- num_scheduled_tokens_np,
3121
- self.input_batch.num_computed_tokens_cpu[:num_reqs],
3122
- scheduler_output.num_common_prefix_blocks,
3123
- )
3124
-
3125
- (
3126
- cudagraph_mode,
3127
- batch_desc,
3128
- should_ubatch,
3129
- num_tokens_across_dp,
3130
- cudagraph_stats,
3131
- ) = self._determine_batch_execution_and_padding(
3132
- num_tokens=num_tokens_unpadded,
3133
- num_reqs=num_reqs,
3134
- num_scheduled_tokens_np=num_scheduled_tokens_np,
3135
- max_num_scheduled_tokens=max_num_scheduled_tokens,
3136
- use_cascade_attn=cascade_attn_prefix_lens is not None,
3137
- num_encoder_reqs=len(scheduler_output.scheduled_encoder_inputs),
3138
- )
3101
+ num_reqs = self.input_batch.num_reqs
3102
+ req_ids = self.input_batch.req_ids
3103
+ tokens = [scheduler_output.num_scheduled_tokens[i] for i in req_ids]
3104
+ num_scheduled_tokens_np = np.array(tokens, dtype=np.int32)
3105
+ max_num_scheduled_tokens = int(num_scheduled_tokens_np.max())
3106
+ num_tokens_unpadded = scheduler_output.total_num_scheduled_tokens
3107
3140
- logger.debug(
3141
- "Running batch with cudagraph_mode: %s, batch_descriptor: %s, "
3142
- "should_ubatch: %s, num_tokens_across_dp: %s",
3143
- cudagraph_mode,
3144
- batch_desc,
3145
- should_ubatch,
3146
- num_tokens_across_dp,
3147
- )
3108
+ logits_indices, spec_decode_metadata = self._prepare_inputs(
3109
+ scheduler_output,
3110
+ num_scheduled_tokens_np,
3111
+ )
3112
3149
- num_tokens_padded = batch_desc.num_tokens
3150
- num_reqs_padded = (
3151
- batch_desc.num_reqs if batch_desc.num_reqs is not None else num_reqs
3152
- )
3153
- ubatch_slices, ubatch_slices_padded = maybe_create_ubatch_slices(
3154
- should_ubatch,
3113
+ cascade_attn_prefix_lens = None
3114
+ # Disable cascade attention when using microbatching (DBO)
3115
+ if self.cascade_attn_enabled and not self.parallel_config.use_ubatching:
3116
+ # Pre-compute cascade attention prefix lengths
3117
+ cascade_attn_prefix_lens = self._compute_cascade_attn_prefix_lens(
3118
num_scheduled_tokens_np,
3156
- num_tokens_padded,
3157
- num_reqs_padded,
3158
- self.parallel_config.num_ubatches,
3119
+ self.input_batch.num_computed_tokens_cpu[:num_reqs],
3120
+ scheduler_output.num_common_prefix_blocks,
3121
)
3122
3161
- logger.debug(
3162
- "ubatch_slices: %s, ubatch_slices_padded: %s",
3163
- ubatch_slices,
3164
- ubatch_slices_padded,
3165
- )
3123
+ (
3124
+ cudagraph_mode,
3125
+ batch_desc,
3126
+ should_ubatch,
3127
+ num_tokens_across_dp,
3128
+ cudagraph_stats,
3129
+ ) = self._determine_batch_execution_and_padding(
3130
+ num_tokens=num_tokens_unpadded,
3131
+ num_reqs=num_reqs,
3132
+ num_scheduled_tokens_np=num_scheduled_tokens_np,
3133
+ max_num_scheduled_tokens=max_num_scheduled_tokens,
3134
+ use_cascade_attn=cascade_attn_prefix_lens is not None,
3135
+ num_encoder_reqs=len(scheduler_output.scheduled_encoder_inputs),
3136
+ )
3137
3167
- pad_attn = cudagraph_mode == CUDAGraphMode.FULL
3168
-
3169
- use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0
3170
- ubatch_slices_attn = ubatch_slices_padded if pad_attn else ubatch_slices
3171
-
3172
- (attn_metadata, spec_decode_common_attn_metadata) = (
3173
- self._build_attention_metadata(
3174
- num_tokens=num_tokens_unpadded,
3175
- num_tokens_padded=num_tokens_padded if pad_attn else None,
3176
- num_reqs=num_reqs,
3177
- num_reqs_padded=num_reqs_padded if pad_attn else None,
3178
- max_query_len=max_num_scheduled_tokens,
3179
- ubatch_slices=ubatch_slices_attn,
3180
- logits_indices=logits_indices,
3181
- use_spec_decode=use_spec_decode,
3182
- num_scheduled_tokens=scheduler_output.num_scheduled_tokens,
3183
- cascade_attn_prefix_lens=cascade_attn_prefix_lens,
3184
- )
3138
+ logger.debug(
3139
+ "Running batch with cudagraph_mode: %s, batch_descriptor: %s, "
3140
+ "should_ubatch: %s, num_tokens_across_dp: %s",
3141
+ cudagraph_mode,
3142
+ batch_desc,
3143
+ should_ubatch,
3144
+ num_tokens_across_dp,
3145
+ )
3146
+
3147
+ num_tokens_padded = batch_desc.num_tokens
3148
+ num_reqs_padded = (
3149
+ batch_desc.num_reqs if batch_desc.num_reqs is not None else num_reqs
3150
+ )
3151
+ ubatch_slices, ubatch_slices_padded = maybe_create_ubatch_slices(
3152
+ should_ubatch,
3153
+ num_scheduled_tokens_np,
3154
+ num_tokens_padded,
3155
+ num_reqs_padded,
3156
+ self.parallel_config.num_ubatches,
3157
+ )
3158
+
3159
+ logger.debug(
3160
+ "ubatch_slices: %s, ubatch_slices_padded: %s",
3161
+ ubatch_slices,
3162
+ ubatch_slices_padded,
3163
+ )
3164
+
3165
+ pad_attn = cudagraph_mode == CUDAGraphMode.FULL
3166
+
3167
+ use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0
3168
+ ubatch_slices_attn = ubatch_slices_padded if pad_attn else ubatch_slices
3169
+
3170
+ attn_metadata, spec_decode_common_attn_metadata = (
3171
+ self._build_attention_metadata(
3172
+ num_tokens=num_tokens_unpadded,
3173
+ num_tokens_padded=num_tokens_padded if pad_attn else None,
3174
+ num_reqs=num_reqs,
3175
+ num_reqs_padded=num_reqs_padded if pad_attn else None,
3176
+ max_query_len=max_num_scheduled_tokens,
3177
+ ubatch_slices=ubatch_slices_attn,
3178
+ logits_indices=logits_indices,
3179
+ use_spec_decode=use_spec_decode,
3180
+ num_scheduled_tokens=scheduler_output.num_scheduled_tokens,
3181
+ cascade_attn_prefix_lens=cascade_attn_prefix_lens,
3182
)
3183
+ )
3184
3185
(
3186
input_ids,