diff --git a/.buildkite/test_areas/plugins.yaml b/.buildkite/test_areas/plugins.yaml index ecaf287956..4ba3392071 100644 --- a/.buildkite/test_areas/plugins.yaml +++ b/.buildkite/test_areas/plugins.yaml @@ -75,6 +75,9 @@ steps: - vllm/model_executor/layers/quantization - tests/plugins_tests/bitsandbytes commands: + # bitsandbytes' int8 quant syncs internally, and it is out-of-tree, so + # there is nothing to wrap on the vLLM side. + - unset VLLM_GPU_SYNC_CHECK - pip install "vllm-bnb-plugin >= 0.0.1" - pytest -v -s plugins_tests/bitsandbytes -m 'not distributed' @@ -89,5 +92,8 @@ steps: - vllm/model_executor/layers/quantization - tests/plugins_tests/bitsandbytes commands: + # bitsandbytes' int8 quant syncs internally, and it is out-of-tree, so + # there is nothing to wrap on the vLLM side. + - unset VLLM_GPU_SYNC_CHECK - pip install "vllm-bnb-plugin >= 0.0.1" - pytest -v -s plugins_tests/bitsandbytes -m 'distributed(num_gpus=2)' diff --git a/docker/Dockerfile b/docker/Dockerfile index 0b68f89dde..804a3a4bb3 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -965,6 +965,9 @@ ENV HF_XET_HIGH_PERFORMANCE 1 # increase timeout for hf downloads (for testing) ENV HF_HUB_DOWNLOAD_TIMEOUT 60 +# Catch GPU<->CPU syncs in execute_model/sample_tokens +ENV VLLM_GPU_SYNC_CHECK=error + # Copy in the v1 package for testing (it isn't distributed yet) COPY vllm/v1 /usr/local/lib/python${PYTHON_VERSION}/dist-packages/vllm/v1 diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index 837b1b652e..d3504d1a3d 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -820,6 +820,9 @@ COPY --from=export_vllm /vllm_v1 /usr/local/lib/python${PYTHON_VERSION}/dist-pac # Hide source under src/ so it won't shadow the installed package in tests. RUN mkdir src && mv vllm src/vllm +# Catch GPU<->CPU syncs in execute_model/sample_tokens during tests. +ENV VLLM_GPU_SYNC_CHECK=error + # ----------------------- # Final vLLM image FROM mori_base AS final diff --git a/tests/models/multimodal/generation/test_mm_prefix_lm.py b/tests/models/multimodal/generation/test_mm_prefix_lm.py index 8d3f5b77b7..10ff63f6c0 100644 --- a/tests/models/multimodal/generation/test_mm_prefix_lm.py +++ b/tests/models/multimodal/generation/test_mm_prefix_lm.py @@ -8,6 +8,7 @@ import torch from transformers import AutoModelForImageTextToText from vllm.platforms import current_platform +from vllm.utils.gpu_sync_debug import gpu_sync_allowed from ....conftest import HfRunner, ImageTestAssets, VllmRunner from .vlm_utils import model_utils @@ -30,7 +31,9 @@ def _install_prefill_hidden_capture(model): def forward(*args, **kwargs): hidden_states = original_forward(*args, **kwargs) if model._prefill_hidden is None and torch.is_tensor(hidden_states): - model._prefill_hidden = hidden_states.detach().float().cpu() + # Capturing hidden states for comparison is a deliberate D2H. + with gpu_sync_allowed(): + model._prefill_hidden = hidden_states.detach().float().cpu() return hidden_states language_model.forward = forward diff --git a/tests/v1/e2e/general/test_mamba_prefix_cache.py b/tests/v1/e2e/general/test_mamba_prefix_cache.py index 559d687928..d8b42b9abf 100644 --- a/tests/v1/e2e/general/test_mamba_prefix_cache.py +++ b/tests/v1/e2e/general/test_mamba_prefix_cache.py @@ -19,6 +19,7 @@ from vllm.distributed import cleanup_dist_env_and_memory from vllm.model_executor.layers.mamba.mamba_utils import MambaStateCopyFunc from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors +from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.torch_utils import async_tensor_h2d from vllm.v1.attention.backends.utils import CommonAttentionMetadata from vllm.v1.core.kv_cache_manager import KVCacheBlocks, KVCacheManager @@ -981,7 +982,10 @@ def _run_mamba_prefix_cache_mrv2( yield forward_context[layer_name].kv_cache[-1], block_table def temporal_block(temporal_state, block_table, col): - return temporal_state[int(block_table[0, col].item())] + # Resolving the block id for assertions is a deliberate D2H. + with gpu_sync_allowed(): + block_id = int(block_table[0, col].item()) + return temporal_state[block_id] def wrapped_preprocess_state( self: MambaHybridModelState, @@ -1005,19 +1009,24 @@ def _run_mamba_prefix_cache_mrv2( self, input_batch, block_tables, kv_cache_config, num_computed_tokens ) if cur_step_action is not None: - req_idx = int(input_batch.idx_mapping[0].item()) - src_col = int(self._mamba_src_col_gpu[req_idx].item()) - off = int(self._mamba_src_off_gpu[req_idx].item()) - dst = int(self._mamba_state_idx_gpu[req_idx].item()) + # Reading the GPU-side copy state back to assert on it is a + # deliberate D2H. + with gpu_sync_allowed(): + req_idx = int(input_batch.idx_mapping[0].item()) + src_col = int(self._mamba_src_col_gpu[req_idx].item()) + off = int(self._mamba_src_off_gpu[req_idx].item()) + dst = int(self._mamba_state_idx_gpu[req_idx].item()) actual = (-1, -1) if src_col < 0 or src_col == dst else (src_col + off, dst) assert actual == expected, ( f"V2 align preprocess copy: expected={expected}, " f"actual={actual}, {cur_step_action=}" ) - for temporal, bt, src_state in snapshots: - torch.testing.assert_close( - temporal_block(temporal, bt, expected[1]), src_state - ) + # Comparing device tensors for the assertion is a deliberate D2H. + with gpu_sync_allowed(): + for temporal, bt, src_state in snapshots: + torch.testing.assert_close( + temporal_block(temporal, bt, expected[1]), src_state + ) return ret def wrapped_postprocess_state( @@ -1048,10 +1057,12 @@ def _run_mamba_prefix_cache_mrv2( ret = original_postprocess_state( self, idx_mapping, num_sampled, num_computed_tokens ) - for temporal, bt, src_state in snapshots: - torch.testing.assert_close( - temporal_block(temporal, bt, expected[1]), src_state - ) + # Comparing device tensors for the assertion is a deliberate D2H. + with gpu_sync_allowed(): + for temporal, bt, src_state in snapshots: + torch.testing.assert_close( + temporal_block(temporal, bt, expected[1]), src_state + ) return ret def wrapped_execute_model( @@ -1072,10 +1083,10 @@ def _run_mamba_prefix_cache_mrv2( ret = original_execute_model(self, scheduler_output, *args, **kwargs) if cur_step_action is not None and self.execute_model_state is not None: input_batch = self.execute_model_state.input_batch - assert ( - cur_step_action.num_computed_tokens_start - == input_batch.positions[input_batch.query_start_loc[0]].item() - ) + # Reading positions back to assert on them is a deliberate D2H. + with gpu_sync_allowed(): + start_pos = input_batch.positions[input_batch.query_start_loc[0]].item() + assert cur_step_action.num_computed_tokens_start == start_pos return ret def fake_sample( diff --git a/vllm/distributed/eplb/eplb_communicator.py b/vllm/distributed/eplb/eplb_communicator.py index f9a9a8a90a..3f04b7bc91 100644 --- a/vllm/distributed/eplb/eplb_communicator.py +++ b/vllm/distributed/eplb/eplb_communicator.py @@ -33,6 +33,7 @@ from vllm.distributed.stateless_coordinator import StatelessGroupCoordinator from vllm.distributed.utils import is_weak_contiguous from vllm.logger import init_logger from vllm.platforms import current_platform +from vllm.utils.gpu_sync_debug import gpu_sync_allowed logger = init_logger(__name__) @@ -222,10 +223,11 @@ class TorchDistGlooStagedEplbCommunicator(EplbCommunicator): # Wait for all D2H copies to finish # before issuing gloo batch_isend_irecv operations. - if self._cuda_stream is not None: - self._cuda_stream.synchronize() - else: - torch.cuda.current_stream().synchronize() + with gpu_sync_allowed(): + if self._cuda_stream is not None: + self._cuda_stream.synchronize() + else: + torch.cuda.current_stream().synchronize() reqs = batch_isend_irecv(p2p_ops) for req in reqs: diff --git a/vllm/distributed/eplb/eplb_state.py b/vllm/distributed/eplb/eplb_state.py index 267beda9f1..5b2bb05ed0 100644 --- a/vllm/distributed/eplb/eplb_state.py +++ b/vllm/distributed/eplb/eplb_state.py @@ -47,6 +47,7 @@ from vllm.distributed.utils import StatelessProcessGroup from vllm.logger import init_logger from vllm.model_executor.models.interfaces import MixtureOfExperts from vllm.platforms import current_platform +from vllm.utils.gpu_sync_debug import gpu_sync_allowed from .async_worker import start_async_worker from .eplb_communicator import EplbCommunicator, create_eplb_communicator @@ -596,10 +597,11 @@ class EplbState: num_tokens_per_rank ) - # Just to make type checker happy - tokens_tensors: list[float] = torch.stack( - [avg_tokens_tensor, max_tokens_tensor] - ).tolist() + # This gpu/cpu sync only happens with expert stat logging enabled. + with gpu_sync_allowed(): + tokens_tensors: list[float] = torch.stack( + [avg_tokens_tensor, max_tokens_tensor] + ).tolist() avg_tokens, max_tokens = tokens_tensors balancedness = avg_tokens / max_tokens if max_tokens > 0 else 0.0 @@ -825,15 +827,17 @@ class EplbState: self.model_states.values(), global_expert_load_windows ): if not self.is_async or is_profile: - # Get new expert mappings for the model - new_physical_to_logical_map = self.policy.rebalance_experts( - global_expert_load_window.cpu(), - num_replicas, - num_groups, - num_nodes, - num_gpus, - eplb_model_state.physical_to_logical_map.cpu(), - ) + # Get new expert mappings for the model. The policy runs on the + # host, so the load window and current map have to come back. + with gpu_sync_allowed(): + new_physical_to_logical_map = self.policy.rebalance_experts( + global_expert_load_window.cpu(), + num_replicas, + num_groups, + num_nodes, + num_gpus, + eplb_model_state.physical_to_logical_map.cpu(), + ) skip_rearrange = False if ( @@ -1248,7 +1252,9 @@ def _pad_out_tensor(src: torch.Tensor, dst: torch.Tensor) -> None: src_padding = dst.shape[-1] - src.shape[-1] assert src_padding >= 0 new_src = torch.nn.functional.pad(src, (0, src_padding), value=-1) - dst.copy_(new_src) + # The map is committed from the host once per layer per rearrangement. + with gpu_sync_allowed(): + dst.copy_(new_src) def _commit_eplb_maps_for_layer( @@ -1276,10 +1282,7 @@ def _commit_eplb_maps_for_layer( num_logical_experts = model_state.logical_to_physical_map.shape[1] new_logical, new_replica_count = compute_logical_maps(src, num_logical_experts) # Commit logical_to_physical_map - _pad_out_tensor( - src=new_logical, - dst=model_state.logical_to_physical_map[layer], - ) + _pad_out_tensor(src=new_logical, dst=model_state.logical_to_physical_map[layer]) # Commit logical_replica_count src = new_replica_count diff --git a/vllm/distributed/eplb/rebalance_execute.py b/vllm/distributed/eplb/rebalance_execute.py index 53b0356dcd..cbd5132945 100644 --- a/vllm/distributed/eplb/rebalance_execute.py +++ b/vllm/distributed/eplb/rebalance_execute.py @@ -16,6 +16,7 @@ from torch.distributed import ProcessGroup, all_gather from vllm.distributed.eplb.eplb_communicator import EplbCommunicator from vllm.distributed.eplb.eplb_utils import CpuGpuEvent from vllm.logger import init_logger +from vllm.utils.gpu_sync_debug import gpu_sync_allowed logger = init_logger(__name__) @@ -590,8 +591,11 @@ def rearrange_expert_weights_inplace( weights_buffer = list(expert_buffer) - old_global_expert_indices_cpu = old_global_expert_indices.cpu().numpy() - new_global_expert_indices_cpu = new_global_expert_indices.cpu().numpy() + # The per-layer transfer plan is built in Python, so both maps have to + # come back to the host. Once per rearrangement. + with gpu_sync_allowed(): + old_global_expert_indices_cpu = old_global_expert_indices.cpu().numpy() + new_global_expert_indices_cpu = new_global_expert_indices.cpu().numpy() for layer_idx in range(num_moe_layers): transfer_metadata = move_to_buffer( diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/example_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/example_connector.py index 93b7f7c3c1..782972b03a 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/example_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/example_connector.py @@ -15,6 +15,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.base import ( ) from vllm.logger import init_logger from vllm.model_executor.layers.attention.mla_attention import MLACommonMetadata +from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.hashing import safe_hash from vllm.v1.attention.backend import AttentionMetadata from vllm.v1.core.sched.output import SchedulerOutput @@ -243,7 +244,8 @@ class ExampleConnector(KVConnectorBase_V1): layer_name, request.token_ids, request.mm_hashes ) kv_cache = extract_kv_from_layer(kv_layer, request.slot_mapping) - tensors = {"kv_cache": kv_cache.detach().cpu()} + with gpu_sync_allowed(): + tensors = {"kv_cache": kv_cache.detach().cpu()} safetensors.torch.save_file(tensors, filename) def wait_for_save(self): diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py index 7d1f9466a8..d2c073676d 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py @@ -67,6 +67,7 @@ from vllm.distributed.parallel_state import ( ) from vllm.logger import init_logger from vllm.platforms import current_platform +from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.network_utils import make_zmq_path from vllm.utils.torch_utils import async_tensor_h2d from vllm.v1.attention.backends.utils import get_kv_cache_layout @@ -1882,14 +1883,16 @@ class NixlBaseConnectorWorker: local_block_ids = meta.local_physical_block_ids # TODO (NickLucche) D2H<>H2D ops could benefit from coalescing io across groups - for group_block_ids in local_block_ids: - self.copy_blocks( - self.host_xfer_buffers, - self.device_kv_caches, - group_block_ids, - group_block_ids, - "h2d", - ) + # The h2d block copies below are intentionally synchronous. + with gpu_sync_allowed(): + for group_block_ids in local_block_ids: + self.copy_blocks( + self.host_xfer_buffers, + self.device_kv_caches, + group_block_ids, + group_block_ids, + "h2d", + ) if logger.isEnabledFor(logging.DEBUG): logger.debug( "synced recved kv of request[%s] to device kv buffer," @@ -1903,26 +1906,28 @@ class NixlBaseConnectorWorker: assert self.use_host_buffer assert self.copy_blocks is not None - for req_id, meta in metadata.reqs_to_save.items(): - meta.local_physical_block_ids = self._logical_to_kernel_block_ids( - meta.local_block_ids, self._physical_blocks_per_logical_kv_block - ) - if logger.isEnabledFor(logging.DEBUG): - logger.debug( - "save_load_kv for request[%s] to host xfer buffer." - "local_block_ids: %s. ", - req_id, - ",".join(map(str, meta.local_physical_block_ids)), - ) - # blocking - for group_block_ids in meta.local_physical_block_ids: - self.copy_blocks( - self.device_kv_caches, - self.host_xfer_buffers, - group_block_ids, - group_block_ids, - "d2h", + # The d2h block copies below are intentionally synchronous. + with gpu_sync_allowed(): + for req_id, meta in metadata.reqs_to_save.items(): + meta.local_physical_block_ids = self._logical_to_kernel_block_ids( + meta.local_block_ids, self._physical_blocks_per_logical_kv_block ) + if logger.isEnabledFor(logging.DEBUG): + logger.debug( + "save_load_kv for request[%s] to host xfer buffer." + "local_block_ids: %s. ", + req_id, + ",".join(map(str, meta.local_physical_block_ids)), + ) + # blocking + for group_block_ids in meta.local_physical_block_ids: + self.copy_blocks( + self.device_kv_caches, + self.host_xfer_buffers, + group_block_ids, + group_block_ids, + "d2h", + ) @cached_property def _attention_kv_caches(self) -> list[torch.Tensor]: diff --git a/vllm/lora/punica_wrapper/punica_gpu.py b/vllm/lora/punica_wrapper/punica_gpu.py index 18272354b4..2f746f967e 100644 --- a/vllm/lora/punica_wrapper/punica_gpu.py +++ b/vllm/lora/punica_wrapper/punica_gpu.py @@ -14,6 +14,7 @@ import torch from vllm.lora.layers import LoRAMapping from vllm.lora.utils import get_captured_lora_counts from vllm.triton_utils import HAS_TRITON, triton +from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.math_utils import round_up if HAS_TRITON: @@ -83,9 +84,11 @@ class PunicaWrapperGPU(PunicaWrapperBase): self.is_prefill = mapping.is_prefill self._update_base_metadata(mapping, lora_index_to_id, max_loras, vocab_size) - # Prepare cuda kernel metadata tensors - self.token_mapping_meta.prepare_tensors(self.token_lora_indices) - self.prompt_mapping_meta.prepare_tensors(self.sampler_indices) + # TODO avoid gpu<->cpu sync here + with gpu_sync_allowed(): + # Prepare cuda kernel metadata tensors + self.token_mapping_meta.prepare_tensors(self.token_lora_indices) + self.prompt_mapping_meta.prepare_tensors(self.sampler_indices) def add_shrink( self, diff --git a/vllm/lora/worker_manager.py b/vllm/lora/worker_manager.py index 7082b7287d..105a998072 100644 --- a/vllm/lora/worker_manager.py +++ b/vllm/lora/worker_manager.py @@ -19,6 +19,7 @@ from vllm.lora.model_manager import ( from vllm.lora.peft_helper import PEFTHelper from vllm.lora.request import LoRARequest from vllm.lora.utils import get_adapter_absolute_path +from vllm.utils.gpu_sync_debug import gpu_sync_allowed logger = init_logger(__name__) @@ -223,9 +224,11 @@ class WorkerLoRAManager: def add_adapter(self, adapter_request: Any) -> bool: if adapter_request.adapter_id in self.list_adapters(): return False - loaded_adapter = self._load_adapter(adapter_request) - loaded = self._adapter_manager.add_adapter(loaded_adapter) - self._adapter_manager.activate_adapter(loaded_adapter.id) + # One-time per adapter + with gpu_sync_allowed(): + loaded_adapter = self._load_adapter(adapter_request) + loaded = self._adapter_manager.add_adapter(loaded_adapter) + self._adapter_manager.activate_adapter(loaded_adapter.id) return loaded def remove_adapter(self, adapter_id: int) -> bool: @@ -288,32 +291,34 @@ class LRUCacheWorkerLoRAManager(WorkerLoRAManager): # This is ok because it's currently only called from # the single-threaded core engine loop. - if ( - lora_request.lora_int_id not in self.list_adapters() - or lora_request.load_inplace - ): - # Load the new adapter first to ensure it is actually valid, before - # evicting any existing adapters. - # This may cause the # of loaded lora adapters to very temporarily - # exceed `--max-cpu-loras`. - lora = self._load_adapter(lora_request) + with gpu_sync_allowed(): + if ( + lora_request.lora_int_id not in self.list_adapters() + or lora_request.load_inplace + ): + # Load the new adapter first to ensure it is actually valid, before + # evicting any existing adapters. + # This may cause the # of loaded lora adapters to very temporarily + # exceed `--max-cpu-loras`. + lora = self._load_adapter(lora_request) - # Remove the existing adapter if it exists - # Use case for LoRA inplace - self._adapter_manager.remove_adapter(lora.id) + # Remove the existing adapter if it exists + # Use case for LoRA inplace + self._adapter_manager.remove_adapter(lora.id) - # Loading succeeded, now check if we will exceed cache capacity and - # evict if the oldest adapter if so - if len(self._adapter_manager) + 1 > self._adapter_manager.capacity: - assert isinstance(self._adapter_manager, LRUCacheLoRAModelManager) - self._adapter_manager.remove_oldest_adapter() - # Then add the new adapter to the cache - loaded = self._adapter_manager.add_adapter(lora) - else: - # If the lora is already loaded, just touch it to - # update its position in the caches - loaded = ( - self._adapter_manager.get_adapter(lora_request.lora_int_id) is not None - ) - self._adapter_manager.activate_adapter(lora_request.lora_int_id) + # Loading succeeded, now check if we will exceed cache capacity and + # evict if the oldest adapter if so + if len(self._adapter_manager) + 1 > self._adapter_manager.capacity: + assert isinstance(self._adapter_manager, LRUCacheLoRAModelManager) + self._adapter_manager.remove_oldest_adapter() + # Then add the new adapter to the cache + loaded = self._adapter_manager.add_adapter(lora) + else: + # If the lora is already loaded, just touch it to + # update its position in the caches + loaded = ( + self._adapter_manager.get_adapter(lora_request.lora_int_id) + is not None + ) + self._adapter_manager.activate_adapter(lora_request.lora_int_id) return loaded diff --git a/vllm/model_executor/layers/mamba/mamba_mixer2.py b/vllm/model_executor/layers/mamba/mamba_mixer2.py index 7538a2b6b4..14b1121389 100644 --- a/vllm/model_executor/layers/mamba/mamba_mixer2.py +++ b/vllm/model_executor/layers/mamba/mamba_mixer2.py @@ -48,6 +48,7 @@ from vllm.model_executor.model_loader.weight_utils import ( from vllm.model_executor.parameter import BasevLLMParameter from vllm.model_executor.utils import set_weight_attrs from vllm.platforms import current_platform +from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.torch_utils import ( LayerNameType, _encode_layer_name, @@ -904,17 +905,21 @@ class MambaMixer2(MambaBase, PluggableLayer): # then chunk_stride = 2 chunk_stride = mamba_block_size // chunk_size + # The per-sequence loop below uses these as Python scalars. + # TODO avoid sync here? + with gpu_sync_allowed(): + block_idx_first_cpu = block_idx_first_scheduled_token_p.tolist() + block_idx_last_cpu = block_idx_last_scheduled_token_p.tolist() + num_computed_tokens_p_cpu = num_computed_tokens_p.tolist() + last_chunk_indices_p_cpu = last_chunk_indices_p.tolist() + # Save state for sequences with more than just final state for seq_idx in range(num_prefills): # Block index for the first scheduled token - block_idx_first_scheduled_token = block_idx_first_scheduled_token_p[ - seq_idx - ] + block_idx_first_scheduled_token = block_idx_first_cpu[seq_idx] # Block index for the last scheduled token - block_idx_last_scheduled_token = block_idx_last_scheduled_token_p[ - seq_idx - ] + block_idx_last_scheduled_token = block_idx_last_cpu[seq_idx] # Number of blocks that need to be written n_blocks_to_fill = ( @@ -935,7 +940,7 @@ class MambaMixer2(MambaBase, PluggableLayer): if seq_idx == 0: first_chunk = 0 else: - first_chunk = 1 + last_chunk_indices_p[seq_idx - 1] + first_chunk = 1 + last_chunk_indices_p_cpu[seq_idx - 1] # First chunk that is aligned on the mamba block boundary first_aligned_chunk = first_chunk + chunk_stride - 1 @@ -943,7 +948,7 @@ class MambaMixer2(MambaBase, PluggableLayer): # Calculate the number of computed tokens that were not # already cached num_unaligned_computed_tokens = ( - num_computed_tokens_p[seq_idx] % mamba_block_size + num_computed_tokens_p_cpu[seq_idx] % mamba_block_size ) if num_unaligned_computed_tokens > 0: diff --git a/vllm/model_executor/models/audioflamingo3.py b/vllm/model_executor/models/audioflamingo3.py index 4fa03261c4..d8099c311a 100644 --- a/vllm/model_executor/models/audioflamingo3.py +++ b/vllm/model_executor/models/audioflamingo3.py @@ -54,6 +54,7 @@ from vllm.multimodal.processing import ( PromptUpdateDetails, ) from vllm.sequence import IntermediateTensors +from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.tensor_schema import TensorSchema, TensorShape from .interfaces import ( @@ -306,7 +307,10 @@ def _flatten_valid_audio_embeddings( < output_lengths[:, None] ) - return audio_embeddings[valid_mask], output_lengths + # Boolean-mask indexing has a data-dependent output shape, so the count + # has to come back to the host. + with gpu_sync_allowed(): + return audio_embeddings[valid_mask], output_lengths def _count_audio_tokens_from_mask( @@ -630,10 +634,10 @@ class AudioFlamingo3ForConditionalGeneration( audio_features, feature_attention_mask, ) - chunk_embeddings = torch.split( - masked_audio_features, - audio_output_lengths.tolist(), - ) + # `split` needs Python int sizes. + with gpu_sync_allowed(): + split_sizes = audio_output_lengths.tolist() + chunk_embeddings = torch.split(masked_audio_features, split_sizes) grouped_embeddings = [] current_idx = 0 diff --git a/vllm/model_executor/models/cohere_asr.py b/vllm/model_executor/models/cohere_asr.py index b124f30d52..0f9bd86dd0 100644 --- a/vllm/model_executor/models/cohere_asr.py +++ b/vllm/model_executor/models/cohere_asr.py @@ -56,6 +56,7 @@ from vllm.transformers_utils.processors.cohere_asr import ( CohereASRProcessor, ) from vllm.utils.collection_utils import is_list_of +from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.v1.attention.backend import ( AttentionType, ) @@ -1770,10 +1771,9 @@ class CohereASRModel(nn.Module): out = self.encoder_decoder_proj(out) # Convert padded tensor to packed - outs = [] - for i, feat in enumerate(out): - feat_len = encoder_output_length[i] - outs.append(feat[:feat_len, :]) + with gpu_sync_allowed(): + lengths = encoder_output_length.tolist() + outs = [feat[:length, :] for feat, length in zip(out, lengths)] return outs else: diff --git a/vllm/model_executor/models/ernie45_vl.py b/vllm/model_executor/models/ernie45_vl.py index f3ef38329b..0c1781b04f 100644 --- a/vllm/model_executor/models/ernie45_vl.py +++ b/vllm/model_executor/models/ernie45_vl.py @@ -71,6 +71,7 @@ from vllm.multimodal.processing import ( PromptUpdate, ) from vllm.sequence import IntermediateTensors +from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.tensor_schema import TensorSchema, TensorShape from vllm.utils.torch_utils import async_tensor_h2d from vllm.v1.attention.backends.registry import AttentionBackendEnum @@ -788,7 +789,10 @@ class VariableResolutionResamplerModel(nn.Module): return x def fwd_placeholder(x, grid_thw, to_tensor=False): - grid_thw_cpu = grid_thw.cpu().numpy() + # The per-image grids drive the Python-level offset arithmetic + # below; the same read is wrapped on the cudagraph path. + with gpu_sync_allowed(): + grid_thw_cpu = grid_thw.cpu().numpy() grid_t, grid_hw = grid_thw_cpu[:, 0], grid_thw_cpu[:, 1:] grid_hw_after_conv = grid_hw.prod(-1) // (self.spatial_conv_size**2) @@ -1697,8 +1701,11 @@ class Ernie4_5_VLMoeForConditionalGeneration( EncoderCudaGraphReplayBuffers, ) + # The per-image grids are needed as Python ints to size the buffers. + with gpu_sync_allowed(): + grid_thw_list = mm_kwargs["image_grid_thw"].tolist() metadata = self.vision_model.prepare_encoder_metadata( - mm_kwargs["image_grid_thw"].tolist(), max_batch_size=max_batch_size + grid_thw_list, max_batch_size=max_batch_size ) values = metadata | { "pixel_values": mm_kwargs["pixel_values"], @@ -1738,7 +1745,10 @@ class Ernie4_5_VLMoeForConditionalGeneration( # Ernie only uses the single "default" encoder path. output = outputs["default"] grid_thw = batch_mm_kwargs["image_grid_thw"].to(output.device) - num_valid = int((grid_thw[:, 0] * grid_thw[:, 1] * grid_thw[:, 2]).sum()) + # The valid token count slices the graph output for the eager + # resampler call, so it has to come back to the host. + with gpu_sync_allowed(): + num_valid = int((grid_thw[:, 0] * grid_thw[:, 1] * grid_thw[:, 2]).sum()) image_embeds = self.resampler_model(output[:num_valid], grid_thw) scatter_output_slices(image_embeds, indices, per_item_out_tokens, dest, clone) diff --git a/vllm/model_executor/models/gemma4_mm.py b/vllm/model_executor/models/gemma4_mm.py index 1bfc0c0b7f..7263d77896 100644 --- a/vllm/model_executor/models/gemma4_mm.py +++ b/vllm/model_executor/models/gemma4_mm.py @@ -69,6 +69,7 @@ from vllm.multimodal.processing.processor import ( PromptUpdateDetails, ) from vllm.sequence import IntermediateTensors +from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.tensor_schema import TensorSchema, TensorShape from vllm.utils.torch_utils import async_tensor_h2d @@ -1365,11 +1366,14 @@ class Gemma4ForConditionalGeneration( pp_tensor, pad_tensor, ).to(self.model_dtype) - encoder_outputs = vt.encoder( - inputs_embeds=inputs_embeds, - attention_mask=~pad_tensor, - pixel_position_ids=pp_tensor, - ) + # HuggingFace's mask builder probes `padding_mask.all()` to + # decide whether the mask can be skipped, which syncs. + with gpu_sync_allowed(): + encoder_outputs = vt.encoder( + inputs_embeds=inputs_embeds, + attention_mask=~pad_tensor, + pixel_position_ids=pp_tensor, + ) hidden_states = encoder_outputs.last_hidden_state for i, (orig_idx, _, _) in enumerate(chunk_items): @@ -1387,13 +1391,17 @@ class Gemma4ForConditionalGeneration( single_pos_ids = pool_position_ids[orig_idx].unsqueeze(0) padding_positions = (single_pos_ids == -1).all(dim=-1) - pooled_states, valid_mask = vt.pooler( - hidden_states=single_hidden, - pixel_position_ids=single_pos_ids, - padding_positions=padding_positions, - output_length=output_length, - ) - valid_states = pooled_states[valid_mask] + # The pooler goes through HuggingFace's mask builder, which probes + # `padding_mask.all()`, and the mask indexing below needs the + # selected count on the host. + with gpu_sync_allowed(): + pooled_states, valid_mask = vt.pooler( + hidden_states=single_hidden, + pixel_position_ids=single_pos_ids, + padding_positions=padding_positions, + output_length=output_length, + ) + valid_states = pooled_states[valid_mask] if getattr(vt.config, "standardize", False): valid_states = (valid_states - vt.std_bias) * vt.std_scale @@ -1444,7 +1452,9 @@ class Gemma4ForConditionalGeneration( pooling_k2 = vision_cfg.pooling_kernel_size**2 if isinstance(frame_counts, torch.Tensor): - fc_list = frame_counts.tolist() + # Per-video frame counts drive the Python-level batching below. + with gpu_sync_allowed(): + fc_list = frame_counts.tolist() else: fc_list = list(frame_counts) @@ -1474,11 +1484,13 @@ class Gemma4ForConditionalGeneration( pp_chunk, pad_chunk, ).to(self.model_dtype) - encoder_outputs = vt.encoder( - inputs_embeds=inputs_embeds, - attention_mask=~pad_chunk, - pixel_position_ids=pp_chunk, - ) + # HuggingFace's mask builder probes `padding_mask.all()`. + with gpu_sync_allowed(): + encoder_outputs = vt.encoder( + inputs_embeds=inputs_embeds, + attention_mask=~pad_chunk, + pixel_position_ids=pp_chunk, + ) last_hidden_states_list.append(encoder_outputs.last_hidden_state) last_hidden_states = torch.cat(last_hidden_states_list, dim=0) @@ -1493,13 +1505,15 @@ class Gemma4ForConditionalGeneration( single_pos_ids = pixel_position_ids[i].unsqueeze(0) single_pad_pos = padding_positions[i].unsqueeze(0) - pooled_states, valid_mask = vt.pooler( - hidden_states=single_hidden, - pixel_position_ids=single_pos_ids, - padding_positions=single_pad_pos, - output_length=output_length, - ) - valid_states = pooled_states[valid_mask] + # As above, plus mask indexing that needs the count on the host. + with gpu_sync_allowed(): + pooled_states, valid_mask = vt.pooler( + hidden_states=single_hidden, + pixel_position_ids=single_pos_ids, + padding_positions=single_pad_pos, + output_length=output_length, + ) + valid_states = pooled_states[valid_mask] if getattr(vt.config, "standardize", False): valid_states = (valid_states - vt.std_bias) * vt.std_scale @@ -1553,9 +1567,11 @@ class Gemma4ForConditionalGeneration( # Strip padding per-batch element: only keep valid (non-padding) # tokens. + # Boolean-mask indexing needs the selected count on the host. per_audio = [] - for enc, mask in zip(audio_features, audio_mask, strict=True): - per_audio.append(enc[mask]) # [num_real, hidden_size] + with gpu_sync_allowed(): + for enc, mask in zip(audio_features, audio_mask, strict=True): + per_audio.append(enc[mask]) # [num_real, hidden_size] return per_audio diff --git a/vllm/model_executor/models/granite_speech.py b/vllm/model_executor/models/granite_speech.py index 262eda4c25..0e10121274 100644 --- a/vllm/model_executor/models/granite_speech.py +++ b/vllm/model_executor/models/granite_speech.py @@ -60,6 +60,7 @@ from vllm.multimodal.processing import ( from vllm.sequence import IntermediateTensors from vllm.tokenizers import cached_tokenizer_from_config from vllm.transformers_utils.processor import cached_processor_from_config +from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.tensor_schema import TensorSchema, TensorShape from vllm.utils.torch_utils import PIN_MEMORY @@ -788,8 +789,9 @@ class GraniteSpeechForConditionalGeneration( encoder_embeds = self.encoder(audio_input["input_features"]) # [bsz, , 4096] projected_embeds = self.projector(encoder_embeds) - # Apply mask on variable length audio features - masked_embeds = projected_embeds[audio_input["input_features_mask"]] + # Apply mask on variable length audio features. + with gpu_sync_allowed(): + masked_embeds = projected_embeds[audio_input["input_features_mask"]] # Split variable length features into a tuple return torch.split(masked_embeds, audio_input["audio_embed_sizes"]) diff --git a/vllm/model_executor/models/idefics2_vision_model.py b/vllm/model_executor/models/idefics2_vision_model.py index b6f5531074..72a26993a5 100644 --- a/vllm/model_executor/models/idefics2_vision_model.py +++ b/vllm/model_executor/models/idefics2_vision_model.py @@ -39,6 +39,7 @@ from vllm.model_executor.layers.linear import ( RowParallelLinear, ) from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.utils.gpu_sync_debug import gpu_sync_allowed from .utils import AutoWeightsLoader, WeightsMapper from .vision import is_vit_use_data_parallel, run_dp_sharded_vision_model @@ -96,10 +97,14 @@ class Idefics2VisionEmbeddings(nn.Module): size=(batch_size, max_nb_patches_h * max_nb_patches_w), fill_value=0 ) - for batch_idx, p_attn_mask in enumerate(patch_attention_mask): - if tgt_sizes is not None: - nb_patches_h = tgt_sizes[batch_idx][0] - nb_patches_w = tgt_sizes[batch_idx][1] + with gpu_sync_allowed(): + patch_attention_mask_cpu = patch_attention_mask.cpu() + tgt_sizes_cpu = tgt_sizes.cpu() if tgt_sizes is not None else None + + for batch_idx, p_attn_mask in enumerate(patch_attention_mask_cpu): + if tgt_sizes_cpu is not None: + nb_patches_h = tgt_sizes_cpu[batch_idx][0] + nb_patches_w = tgt_sizes_cpu[batch_idx][1] else: nb_patches_h = p_attn_mask[:, 0].sum() nb_patches_w = p_attn_mask[0].sum() @@ -114,8 +119,10 @@ class Idefics2VisionEmbeddings(nn.Module): pos_ids = ( bucket_coords_h[:, None] * self.num_patches_per_side + bucket_coords_w ).flatten() - position_ids[batch_idx][p_attn_mask.view(-1).cpu()] = pos_ids - position_ids = position_ids.to(self.position_embedding.weight.device) + position_ids[batch_idx][p_attn_mask.view(-1)] = pos_ids + position_ids = position_ids.to( + self.position_embedding.weight.device, non_blocking=True + ) embeddings += self.position_embedding(position_ids) return embeddings diff --git a/vllm/model_executor/models/idefics3.py b/vllm/model_executor/models/idefics3.py index 7f26e01c0c..7b3e552f0c 100644 --- a/vllm/model_executor/models/idefics3.py +++ b/vllm/model_executor/models/idefics3.py @@ -51,6 +51,7 @@ from vllm.multimodal.processing import ( PromptUpdateDetails, ) from vllm.sequence import IntermediateTensors +from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.tensor_schema import TensorSchema, TensorShape from .idefics2_vision_model import ( @@ -443,11 +444,12 @@ class Idefics3Model(nn.Module): real_images_inds = (pixel_values == 0.0).sum( dim=(-1, -2, -3) ) != nb_values_per_image - pixel_values = pixel_values[real_images_inds].contiguous() + with gpu_sync_allowed(): + pixel_values = pixel_values[real_images_inds].contiguous() - # Handle the vision attention mask - # Remove padding images from the mask - pixel_attention_mask = pixel_attention_mask[real_images_inds].contiguous() + # Handle the vision attention mask + # Remove padding images from the mask + pixel_attention_mask = pixel_attention_mask[real_images_inds].contiguous() patch_size = self.config.vision_config.patch_size patches_subgrid = pixel_attention_mask.unfold( diff --git a/vllm/model_executor/models/jina.py b/vllm/model_executor/models/jina.py index 5fae303154..0595df8e3a 100644 --- a/vllm/model_executor/models/jina.py +++ b/vllm/model_executor/models/jina.py @@ -16,6 +16,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead from vllm.sequence import IntermediateTensors from vllm.tasks import PoolingTask from vllm.transformers_utils.repo_utils import get_hf_file_bytes +from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.v1.pool.metadata import PoolingMetadata from ..layers.pooler import DispatchPooler @@ -110,19 +111,24 @@ class JinaForRankingPool(StepPool): prompt_token_ids = pooling_metadata.get_prompt_token_ids() embeds_list = list[torch.Tensor | None]() - for data, token_ids in zip(pooled_data_lst, prompt_token_ids): - # for unfinished chunked prefill - if data is None: - embeds_list.append(None) - else: - docs_indexes = torch.where(torch.eq(token_ids, self.doc_token_id))[0] - query_indexes = torch.where(torch.eq(token_ids, self.query_token_id))[0] + # `torch.where` resolves the match count on the host. + with gpu_sync_allowed(): + for data, token_ids in zip(pooled_data_lst, prompt_token_ids): + # for unfinished chunked prefill + if data is None: + embeds_list.append(None) + else: + doc_match = torch.eq(token_ids, self.doc_token_id) + docs_indexes = torch.where(doc_match)[0] + query_indexes = torch.where( + torch.eq(token_ids, self.query_token_id) + )[0] - # The JinaForRanking model concatenates docs first, then query. - # Let's stay consistent with this novel design. - indexes = torch.cat([docs_indexes, query_indexes]) - embeds = self.projector(data[indexes]) - embeds_list.append(embeds) + # The JinaForRanking model concatenates docs first, then query. + # Let's stay consistent with this novel design. + indexes = torch.cat([docs_indexes, query_indexes]) + embeds = self.projector(data[indexes]) + embeds_list.append(embeds) return embeds_list diff --git a/vllm/model_executor/models/minicpmv.py b/vllm/model_executor/models/minicpmv.py index ff34a4fb36..d9358fed42 100644 --- a/vllm/model_executor/models/minicpmv.py +++ b/vllm/model_executor/models/minicpmv.py @@ -98,6 +98,7 @@ from vllm.transformers_utils.processor import ( ) from vllm.transformers_utils.utils import convert_model_repo_to_path from vllm.utils.collection_utils import flatten_2d_lists +from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.tensor_schema import TensorSchema, TensorShape from vllm.utils.torch_utils import set_default_torch_dtype @@ -1548,7 +1549,8 @@ class MiniCPMV2_5(MiniCPMVBaseModel, SupportsLoRA): all_pixel_values[i, ..., :L_item] = pixel_values_item num_patches = tgt_sizes.prod(-1) - max_patches = num_patches.max().item() + with gpu_sync_allowed(): + max_patches = num_patches.max().item() assert isinstance(max_patches, int) patch_attn_mask = torch.zeros((B, max_patches), dtype=torch.bool, device=device) @@ -1640,7 +1642,9 @@ class MiniCPMV2_6(MiniCPMVBaseModel, SupportsLoRA): all_pixel_values[i, ..., :L_item] = pixel_values_item num_patches = tgt_sizes.prod(-1) - max_patches = num_patches.max().item() + # Needed as a Python int to size the mask below. + with gpu_sync_allowed(): + max_patches = num_patches.max().item() assert isinstance(max_patches, int) patch_attn_mask = torch.zeros((B, max_patches), dtype=torch.bool, device=device) @@ -1737,7 +1741,8 @@ class MiniCPMV4_0(MiniCPMVBaseModel, SupportsLoRA): all_pixel_values[i, ..., :L_item] = pixel_values_item num_patches = tgt_sizes.prod(-1) - max_patches = num_patches.max().item() + with gpu_sync_allowed(): + max_patches = num_patches.max().item() assert isinstance(max_patches, int) patch_attn_mask = torch.zeros((B, max_patches), dtype=torch.bool, device=device) @@ -1839,7 +1844,8 @@ class MiniCPMV4_5(MiniCPMVBaseModel, SupportsLoRA): all_pixel_values[i, ..., :L_item] = pixel_values_item num_patches = tgt_sizes.prod(-1) - max_patches = num_patches.max().item() + with gpu_sync_allowed(): + max_patches = num_patches.max().item() assert isinstance(max_patches, int) patch_attn_mask = torch.zeros((B, max_patches), dtype=torch.bool, device=device) diff --git a/vllm/model_executor/models/moonvit.py b/vllm/model_executor/models/moonvit.py index 56204dd3c6..29400f597b 100644 --- a/vllm/model_executor/models/moonvit.py +++ b/vllm/model_executor/models/moonvit.py @@ -66,6 +66,7 @@ from vllm.model_executor.models.utils import maybe_prefix from vllm.model_executor.models.vision import is_vit_use_data_parallel from vllm.platforms import current_platform from vllm.transformers_utils.configs.moonvit import MoonViTConfig +from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.torch_utils import async_tensor_h2d @@ -811,9 +812,14 @@ class MoonVitPretrainedModel(PreTrainedModel): merge_kernel_size=self.merge_kernel_size, ) - hidden_states = self.patch_embed(pixel_values, grid_hw) - hidden_states = self.encoder(hidden_states, grid_hw) - hidden_states = patch_merger( - hidden_states, grid_hw, merge_kernel_size=self.merge_kernel_size - ) + # Legacy path: patch_embed, encoder and patch_merger each iterate the + # per-image grids in Python, so all three read `grid_hw` back to the + # host. The `encoder_metadata` path above precomputes that outside the + # graph and does not sync. + with gpu_sync_allowed(): + hidden_states = self.patch_embed(pixel_values, grid_hw) + hidden_states = self.encoder(hidden_states, grid_hw) + hidden_states = patch_merger( + hidden_states, grid_hw, merge_kernel_size=self.merge_kernel_size + ) return hidden_states diff --git a/vllm/model_executor/models/moss_audio.py b/vllm/model_executor/models/moss_audio.py index 878eabc0e3..bfcebe7f92 100644 --- a/vllm/model_executor/models/moss_audio.py +++ b/vllm/model_executor/models/moss_audio.py @@ -55,6 +55,7 @@ from vllm.multimodal.processing import ( ) from vllm.sequence import IntermediateTensors from vllm.transformers_utils.repo_utils import get_hf_file_to_dict +from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.tensor_schema import TensorSchema, TensorShape from .interfaces import ( @@ -1667,15 +1668,19 @@ class MossAudioModel(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA): """ audio_data = audio_input["audio_data"] audio_data_seqlens = audio_input["audio_data_seqlens"] - last_hidden_state, deepstack = self.audio_encoder( - audio_data.to(self.audio_encoder.dtype), - feature_lens=audio_data_seqlens, - output_deepstack_hidden_states=len(self.deepstack_audio_merger_list) > 0, - ) - audio_embeds = self.audio_adapter(last_hidden_state) - audio_lengths = MossAudioEncoder._compute_downsampled_length( - audio_data_seqlens.to(device=audio_embeds.device, dtype=torch.long) - ).tolist() + # The encoder chunks the input by per-audio feature lengths, which + # needs Python ints for `split`/`pad_sequence`. + want_deepstack = len(self.deepstack_audio_merger_list) > 0 + with gpu_sync_allowed(): + last_hidden_state, deepstack = self.audio_encoder( + audio_data.to(self.audio_encoder.dtype), + feature_lens=audio_data_seqlens, + output_deepstack_hidden_states=want_deepstack, + ) + audio_embeds = self.audio_adapter(last_hidden_state) + audio_lengths = MossAudioEncoder._compute_downsampled_length( + audio_data_seqlens.to(device=audio_embeds.device, dtype=torch.long) + ).tolist() main_embeddings = tuple(audio_embeds.squeeze(0).split(audio_lengths, dim=0)) deepstack_embeddings: list[tuple[torch.Tensor, ...]] = [] diff --git a/vllm/model_executor/models/phi4mm.py b/vllm/model_executor/models/phi4mm.py index 0fc10398a4..9a8da5ed32 100644 --- a/vllm/model_executor/models/phi4mm.py +++ b/vllm/model_executor/models/phi4mm.py @@ -49,6 +49,7 @@ from vllm.multimodal.processing.processor import ( ResolvedPromptUpdate, ) from vllm.sequence import IntermediateTensors +from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.tensor_schema import TensorSchema, TensorShape from .idefics2_vision_model import Idefics2VisionTransformer @@ -396,12 +397,23 @@ class Phi4MMImageEncoder(nn.Module): w * base_feat_width // base_feat_height_reduction, ) ) - useful_height = int(reshaped_image_attention_mask[0, :, 0].sum().item()) - useful_width = int(reshaped_image_attention_mask[0, 0, :].sum().item()) + # The mask stays on device for the encoder above, so these + # per-image counts have to come back to the host to drive the + # slicing and the Python-level length arithmetic. + with gpu_sync_allowed(): + useful_height = int( + reshaped_image_attention_mask[0, :, 0].sum().item() + ) + useful_width = int( + reshaped_image_attention_mask[0, 0, :].sum().item() + ) + mask_token_count = int( + image_attention_mask[_bs, : B_ + 1, 0::2, 0::2].sum().item() + ) sub_img = sub_img[:, :useful_height, :useful_width] temp_sub_GN = self.sub_GN.repeat(1, useful_height, 1, 1) temp_len = ( - int(image_attention_mask[_bs, : B_ + 1, 0::2, 0::2].sum().item()) + mask_token_count + (useful_height + 1) + base_feat_height // base_feat_height_reduction ) diff --git a/vllm/model_executor/models/qwen2_5_omni_thinker.py b/vllm/model_executor/models/qwen2_5_omni_thinker.py index 0715235ebb..96a58df1dd 100644 --- a/vllm/model_executor/models/qwen2_5_omni_thinker.py +++ b/vllm/model_executor/models/qwen2_5_omni_thinker.py @@ -89,6 +89,7 @@ from vllm.multimodal.processing.processor import ( PromptUpdateDetails, ) from vllm.sequence import IntermediateTensors +from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.tensor_schema import TensorSchema, TensorShape from .interfaces import ( @@ -1036,12 +1037,14 @@ class Qwen2_5OmniConditionalGenerationMixin: self.audio_tower._get_feat_extract_output_lengths(audio_feature_lengths) ) - audio_outputs = self.audio_tower( - input_features.to(self.audio_tower.dtype), - feature_lens=audio_feature_lengths, - aftercnn_lens=audio_feat_lengths, - ) - return audio_outputs.last_hidden_state.split(audio_output_lengths.tolist()) + with gpu_sync_allowed(): + audio_outputs = self.audio_tower( + input_features.to(self.audio_tower.dtype), + feature_lens=audio_feature_lengths, + aftercnn_lens=audio_feat_lengths, + ) + split_sizes = audio_output_lengths.tolist() + return audio_outputs.last_hidden_state.split(split_sizes) def _process_image_input( self, image_input: Qwen2_5_VLImageInputs @@ -1526,12 +1529,13 @@ class Qwen2_5OmniThinkerForConditionalGeneration( video_token_id = self.config.video_token_index audio_token_id = self.config.audio_token_index - input_ids_cpu = input_ids.cpu() - is_video = is_multimodal & (input_ids_cpu == video_token_id) - is_audio = is_multimodal & (input_ids_cpu == audio_token_id) + with gpu_sync_allowed(): + input_ids_cpu = input_ids.cpu() + is_video = is_multimodal & (input_ids_cpu == video_token_id) + is_audio = is_multimodal & (input_ids_cpu == audio_token_id) - num_video = is_video.sum().item() - num_audio = is_audio.sum().item() + num_video = is_video.sum().item() + num_audio = is_audio.sum().item() if check_interleaved_audio_video(is_video, is_audio, num_video, num_audio): inputs_embeds = self._embed_text_input_ids( diff --git a/vllm/model_executor/models/qwen2_5_vl.py b/vllm/model_executor/models/qwen2_5_vl.py index aa105ed8e8..d7d0d1fa5a 100644 --- a/vllm/model_executor/models/qwen2_5_vl.py +++ b/vllm/model_executor/models/qwen2_5_vl.py @@ -87,6 +87,7 @@ from vllm.utils.torch_utils import PIN_MEMORY, async_tensor_h2d from vllm.v1.attention.backends.registry import AttentionBackendEnum from vllm.v1.worker.encoder_cudagraph_defs import EncoderCudaGraphReplayBuffers +from ...utils.gpu_sync_debug import gpu_sync_allowed from .interfaces import ( MultiModalEmbeddings, SupportsEagle, @@ -1569,8 +1570,9 @@ class Qwen2_5_VLForConditionalGeneration( video_second_per_grid=video_second_per_grid_t.item(), ).to(emb.device, non_blocking=True) - emb = emb[retention_mask] - positions = positions[retention_mask] + with gpu_sync_allowed(): + emb = emb[retention_mask] + positions = positions[retention_mask] emb = torch.cat([emb, positions], dim=1) video_embeds_out.append(emb) return tuple(video_embeds_out) @@ -1624,15 +1626,16 @@ class Qwen2_5_VLForConditionalGeneration( mm[:, -4:].permute(1, 0).long() for mm in multimodal_embeddings ] - positions, mrope_positions_delta = recompute_mrope_positions( - input_ids_t, - mm_embeddings_pos, - mrope_positions, - num_computed_tokens, - vision_start_token_id, - image_token_id, - video_token_id, - ) + with gpu_sync_allowed(): + positions, mrope_positions_delta = recompute_mrope_positions( + input_ids_t, + mm_embeddings_pos, + mrope_positions, + num_computed_tokens, + vision_start_token_id, + image_token_id, + video_token_id, + ) return mm_embeddings_out, positions, mrope_positions_delta diff --git a/vllm/model_executor/models/qwen2_audio.py b/vllm/model_executor/models/qwen2_audio.py index 9692d9615b..115a7f7f79 100644 --- a/vllm/model_executor/models/qwen2_audio.py +++ b/vllm/model_executor/models/qwen2_audio.py @@ -60,6 +60,7 @@ from vllm.multimodal.processing import ( PromptUpdate, ) from vllm.sequence import IntermediateTensors +from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.tensor_schema import TensorSchema, TensorShape from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP @@ -444,12 +445,15 @@ class Qwen2AudioForConditionalGeneration(nn.Module, SupportsMultiModal, Supports ) < audio_output_lengths ) - masked_audio_features = audio_features[audio_features_mask].view(-1, embed_dim) + with gpu_sync_allowed(): + masked_audio_features = audio_features[audio_features_mask].view( + -1, embed_dim + ) - # Split to tuple of embeddings for individual audio input. - return torch.split( - masked_audio_features, audio_output_lengths.flatten().tolist() - ) + # Split to tuple of embeddings for individual audio input. + return torch.split( + masked_audio_features, audio_output_lengths.flatten().tolist() + ) def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings: audio_input = self._parse_and_validate_audio_input(**kwargs) diff --git a/vllm/model_executor/models/qwen3_asr.py b/vllm/model_executor/models/qwen3_asr.py index b9224adc6c..f5ec1a4526 100644 --- a/vllm/model_executor/models/qwen3_asr.py +++ b/vllm/model_executor/models/qwen3_asr.py @@ -92,6 +92,7 @@ from vllm.transformers_utils.processor import cached_processor_from_config from vllm.transformers_utils.processors.qwen3_asr import ( Qwen3ASRProcessor, ) +from vllm.utils.gpu_sync_debug import gpu_sync_allowed logger = init_logger(__name__) _ASR_TEXT_TAG = "" @@ -467,7 +468,9 @@ class Qwen3ASRForConditionalGeneration( feature_lens=audio_feature_lengths, aftercnn_lens=audio_output_lengths, ) - return audio_features.split(audio_output_lengths.tolist()) + with gpu_sync_allowed(): + split_sizes = audio_output_lengths.tolist() + return audio_features.split(split_sizes) def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings | None: mm_input_by_modality = self._parse_and_validate_multimodal_inputs(**kwargs) diff --git a/vllm/model_executor/models/qwen3_omni_moe_thinker.py b/vllm/model_executor/models/qwen3_omni_moe_thinker.py index 2f84f12834..48a16e97d1 100755 --- a/vllm/model_executor/models/qwen3_omni_moe_thinker.py +++ b/vllm/model_executor/models/qwen3_omni_moe_thinker.py @@ -77,6 +77,7 @@ from vllm.multimodal.processing.processor import ( ) from vllm.sequence import IntermediateTensors from vllm.transformers_utils.processor import cached_processor_from_config +from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.torch_utils import async_tensor_h2d from vllm.v1.attention.backends.registry import AttentionBackendEnum @@ -434,11 +435,14 @@ class Qwen3OmniMoeAudioEncoder(nn.Module): feature_lens: torch.Tensor, aftercnn_lens: torch.Tensor, ): - # Compute chunk information + # Compute chunk information. chunk_num = torch.ceil(feature_lens / (self.n_window * 2)).long() + with gpu_sync_allowed(): + total_chunks = int(chunk_num.sum()) + chunk_lengths = async_tensor_h2d( - [self.n_window * 2] * chunk_num.sum(), + [self.n_window * 2] * total_chunks, dtype=torch.long, device=feature_lens.device, ) @@ -447,15 +451,18 @@ class Qwen3OmniMoeAudioEncoder(nn.Module): chunk_lengths[chunk_lengths == 0] = self.n_window * 2 # Split input features into chunks and pad - chunk_list = input_features.T.split(chunk_lengths.tolist(), dim=0) + with gpu_sync_allowed(): + chunk_lengths_list = chunk_lengths.tolist() + chunk_list = input_features.T.split(chunk_lengths_list, dim=0) padded_feature = nn.utils.rnn.pad_sequence( chunk_list, batch_first=True ).transpose(1, 2) # Compute feature lengths after CNN feature_lens_after_cnn = self._get_cnn_output_lengths(chunk_lengths) - # Vectorized mask creation: avoid creating many small tensors - max_len_after_cnn = feature_lens_after_cnn.max().item() + # Vectorized mask creation: avoid creating many small tensors. + with gpu_sync_allowed(): + max_len_after_cnn = feature_lens_after_cnn.max().item() indices = torch.arange(max_len_after_cnn, device=padded_feature.device) padded_mask_after_cnn = indices.unsqueeze(0) < feature_lens_after_cnn.unsqueeze( 1 @@ -494,16 +501,17 @@ class Qwen3OmniMoeAudioEncoder(nn.Module): ) padded_embed = padded_embed + positional_embedding - # Extract valid hidden states and compute cu_seqlens - hidden_states = padded_embed[padded_mask_after_cnn] + with gpu_sync_allowed(): + hidden_states = padded_embed[padded_mask_after_cnn] + # Use tolist() for efficient batch conversion from tensor to Python. + aftercnn_lens_list = aftercnn_lens.tolist() # Compute cumulative sequence lengths for chunked attention cu_chunk_lens = [0] window_aftercnn = padded_mask_after_cnn.shape[-1] * ( self.n_window_infer // (self.n_window * 2) ) - # Use tolist() for efficient batch conversion from tensor to Python - for cnn_len in aftercnn_lens.tolist(): + for cnn_len in aftercnn_lens_list: num_full_chunks = cnn_len // window_aftercnn remainder = cnn_len % window_aftercnn cu_chunk_lens.extend([window_aftercnn] * num_full_chunks) diff --git a/vllm/model_executor/models/qwen3_vl.py b/vllm/model_executor/models/qwen3_vl.py index 89495ff949..de341d302a 100644 --- a/vllm/model_executor/models/qwen3_vl.py +++ b/vllm/model_executor/models/qwen3_vl.py @@ -115,6 +115,7 @@ from vllm.utils.math_utils import round_up from vllm.utils.torch_utils import PIN_MEMORY from vllm.v1.worker.encoder_cudagraph_defs import EncoderCudaGraphReplayBuffers +from ...utils.gpu_sync_debug import gpu_sync_allowed from ...utils.torch_utils import async_tensor_h2d from .interfaces import ( MultiModalEmbeddings, @@ -2368,19 +2369,17 @@ class Qwen3VLForConditionalGeneration( spatial_merge_size=self.visual.spatial_merge_size, q=self.video_pruning_rate, ) - # Apply retention mask. - emb = emb[retention_mask] - # Calculate the actual number of retained tokens per frame. - num_frames, rows, cols = ( - t, - h // merge_size, - w // merge_size, - ) - retention_mask_thw = retention_mask.reshape(num_frames, rows, cols) - num_tokens_per_frame = ( - retention_mask_thw.sum(dim=(1, 2)).long().tolist() - ) + with gpu_sync_allowed(): + # Apply retention mask. + emb = emb[retention_mask] + + # Calculate the actual number of retained tokens per frame. + num_frames, rows, cols = t, h // merge_size, w // merge_size + retention_mask_thw = retention_mask.reshape(num_frames, rows, cols) + num_tokens_per_frame = ( + retention_mask_thw.sum(dim=(1, 2)).long().tolist() + ) else: feature_size = emb.shape[0] // num_frames num_tokens_per_frame = [feature_size] * num_frames @@ -2548,10 +2547,14 @@ class Qwen3VLForConditionalGeneration( .permute(1, 0) ) full_is_video_embed = unpruned_token_ids_tensor == embed_token_id - expanded_positions[is_video_embed, :3] = original_mrope[full_is_video_embed][ - retention_mask - ] - expanded_positions[~is_video_embed, :3] = original_mrope[~full_is_video_embed] + + with gpu_sync_allowed(): + expanded_positions[is_video_embed, :3] = original_mrope[ + full_is_video_embed + ][retention_mask] + expanded_positions[~is_video_embed, :3] = original_mrope[ + ~full_is_video_embed + ] expanded_positions[..., 3] = is_vision_start expanded_positions[..., 4] = is_video_embed @@ -2814,15 +2817,16 @@ class Qwen3VLForConditionalGeneration( torch.empty(5, 0, device=device, dtype=torch.long) ) - positions, mrope_positions_delta = recompute_mrope_positions( - input_ids_t, - mm_embeddings_pos, - mrope_positions, - num_computed_tokens, - vision_start_token_id, - image_token_id, - video_token_id, - ) + with gpu_sync_allowed(): + positions, mrope_positions_delta = recompute_mrope_positions( + input_ids_t, + mm_embeddings_pos, + mrope_positions, + num_computed_tokens, + vision_start_token_id, + image_token_id, + video_token_id, + ) return mm_embeddings_out, positions, mrope_positions_delta diff --git a/vllm/model_executor/models/siglip.py b/vllm/model_executor/models/siglip.py index 5808c9539b..6d924a2865 100644 --- a/vllm/model_executor/models/siglip.py +++ b/vllm/model_executor/models/siglip.py @@ -54,6 +54,7 @@ from vllm.multimodal.processing import ( TimingContext, ) from vllm.sequence import IntermediateTensors +from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.tensor_schema import TensorSchema, TensorShape from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsQuant @@ -998,27 +999,31 @@ class SiglipEmbeddingModel(nn.Module, SupportsMultiModal, SupportsQuant): position_diffs = position_ids[1:] - position_ids[:-1] boundary_mask = position_diffs <= 0 - boundary_indices = torch.cat( - [ - torch.tensor([0], device=features.device), - torch.where(boundary_mask)[0] + 1, - torch.tensor([len(features)], device=features.device), - ] - ) + with gpu_sync_allowed(): + boundary_mid_cpu = torch.where(boundary_mask.cpu())[0] + 1 + + zero = torch.zeros(1, dtype=boundary_mid_cpu.dtype) + end = torch.full((1,), len(features), dtype=boundary_mid_cpu.dtype) + boundary_indices_cpu = torch.cat([zero, boundary_mid_cpu, end]) # For each sequence [start, end), position i flips to: start + end - 1 - i - lengths = boundary_indices[1:] - boundary_indices[:-1] - starts = boundary_indices[:-1] - ends = boundary_indices[1:] + lengths_cpu = boundary_indices_cpu[1:] - boundary_indices_cpu[:-1] + starts_cpu = boundary_indices_cpu[:-1] + ends_cpu = boundary_indices_cpu[1:] # Assign sequence ID to each element - sequence_ids = torch.arange( - len(lengths), device=features.device - ).repeat_interleave(lengths) + sequence_ids_cpu = torch.arange( + len(lengths_cpu), dtype=boundary_mid_cpu.dtype + ).repeat_interleave(lengths_cpu) # Calculate flipped indices for all positions at once - current_positions = torch.arange(len(features), device=features.device) - flip_indices = starts[sequence_ids] + ends[sequence_ids] - 1 - current_positions + current_positions_cpu = torch.arange( + len(features), dtype=boundary_mid_cpu.dtype + ) + flip_indices_cpu = ( + starts_cpu[sequence_ids_cpu] + ends_cpu[sequence_ids_cpu] + ) - (1 + current_positions_cpu) + flip_indices = flip_indices_cpu.to(features.device, non_blocking=True) return features[flip_indices] diff --git a/vllm/model_executor/models/terratorch.py b/vllm/model_executor/models/terratorch.py index b49f7827da..a160d5bd17 100644 --- a/vllm/model_executor/models/terratorch.py +++ b/vllm/model_executor/models/terratorch.py @@ -61,6 +61,7 @@ from vllm.multimodal.processing import ( TimingContext, ) from vllm.sequence import IntermediateTensors +from vllm.utils.gpu_sync_debug import gpu_sync_allowed from .interfaces import IsAttentionFree, MultiModalEmbeddings, SupportsMultiModal from .interfaces_base import attn_type @@ -280,7 +281,9 @@ class Terratorch(nn.Module, IsAttentionFree, SupportsMultiModal): inputs_embeds: torch.Tensor | None = None, **kwargs: object, ): - model_output = self.inference_runner.forward(**kwargs) + # terratorch's forward has internal GPU syncs. + with gpu_sync_allowed(): + model_output = self.inference_runner.forward(**kwargs) return model_output.output def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: diff --git a/vllm/model_executor/models/transformers/multimodal.py b/vllm/model_executor/models/transformers/multimodal.py index 84d6813c8e..f328a18a3c 100644 --- a/vllm/model_executor/models/transformers/multimodal.py +++ b/vllm/model_executor/models/transformers/multimodal.py @@ -54,6 +54,7 @@ from vllm.multimodal.processing import ( TimingContext, ) from vllm.sequence import IntermediateTensors +from vllm.utils.gpu_sync_debug import gpu_sync_allowed if TYPE_CHECKING: from transformers import BatchFeature, PreTrainedModel @@ -726,12 +727,17 @@ class MultiModalMixin(SupportsMultiModal, SupportsMRoPE): kwargs.pop("token_type_ids", None) kwargs.pop("mm_token_type_ids", None) - audio_output = self.model.get_audio_features( - input_features, return_dict=True, **kwargs - ) + # HuggingFace's `get_audio_features` implementations branch on + # per-sample feature lengths internally. + with gpu_sync_allowed(): + audio_output = self.model.get_audio_features( + input_features, return_dict=True, **kwargs + ) audio_embeddings = audio_output.pooler_output - split_sizes = num_audio_tokens.flatten().tolist() + # Per-audio token counts are needed as Python ints to split. + with gpu_sync_allowed(): + split_sizes = num_audio_tokens.flatten().tolist() return self._split_embeddings(audio_embeddings, split_sizes) def _process_image_input(self, **kwargs) -> list[torch.Tensor] | None: @@ -749,7 +755,12 @@ class MultiModalMixin(SupportsMultiModal, SupportsMRoPE): num_image_patches = kwargs.pop("num_image_patches") - vision_embeddings = self.model.get_image_features(pixel_values, **kwargs) + # The underlying HuggingFace `get_image_features` implementations + # contain model-internal syncs (e.g. Idefics3 filters all-zero + # padding images via boolean-mask indexing, LlavaOnevision + # branches on per-sample batch counts). + with gpu_sync_allowed(): + vision_embeddings = self.model.get_image_features(pixel_values, **kwargs) # Transformers `v5`, `self.get_image_features` returns a tuple # containing the features and optionally attentions/hidden_states diff --git a/vllm/model_executor/models/ultravox.py b/vllm/model_executor/models/ultravox.py index a2c051cf74..551c2c5694 100644 --- a/vllm/model_executor/models/ultravox.py +++ b/vllm/model_executor/models/ultravox.py @@ -43,6 +43,7 @@ from vllm.multimodal.processing import ( from vllm.renderers import TokenizeParams from vllm.sequence import IntermediateTensors from vllm.transformers_utils.configs.ultravox import UltravoxConfig +from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.tensor_schema import TensorSchema, TensorShape from .interfaces import ( @@ -841,15 +842,19 @@ class UltravoxModel(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA): embeddings.shape[0], -1 ) mask = indices < audio_token_len[:, None] - # Apply mask and flatten - flattened_embeddings = embeddings[mask] - # Return one tensor per input audio - embed_lens = [ - chunk_lens.sum().item() - for chunk_lens in audio_token_len.split(audio_input["num_chunks"].tolist()) - ] - return flattened_embeddings.split(embed_lens) + with gpu_sync_allowed(): + # Apply mask and flatten + flattened_embeddings = embeddings[mask] + + # Return one tensor per input audio + embed_lens = [ + chunk_lens.sum().item() + for chunk_lens in audio_token_len.split( + audio_input["num_chunks"].tolist() + ) + ] + return flattened_embeddings.split(embed_lens) def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings: audio_input = self._parse_and_validate_audio_input(**kwargs) diff --git a/vllm/third_party/flash_linear_attention/ops/index.py b/vllm/third_party/flash_linear_attention/ops/index.py index 8dc5f75f0f..cc6f54d8af 100644 --- a/vllm/third_party/flash_linear_attention/ops/index.py +++ b/vllm/third_party/flash_linear_attention/ops/index.py @@ -9,6 +9,7 @@ # ruff: noqa: E501 import torch +from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.triton_utils import triton from .utils import tensor_cache @@ -21,12 +22,10 @@ def prepare_lens(cu_seqlens: torch.Tensor) -> torch.Tensor: @tensor_cache def prepare_chunk_indices(cu_seqlens: torch.Tensor, chunk_size: int) -> torch.Tensor: - indices = torch.cat( - [ - torch.arange(n) - for n in triton.cdiv(prepare_lens(cu_seqlens), chunk_size).tolist() - ] - ) + # This will be fixed by https://github.com/vllm-project/vllm/pull/51540. + with gpu_sync_allowed(): + chunk_counts = triton.cdiv(prepare_lens(cu_seqlens), chunk_size).tolist() + indices = torch.cat([torch.arange(n) for n in chunk_counts]) chunk_indices = torch.stack([indices.eq(0).cumsum(0) - 1, indices], 1) return chunk_indices.to( device=cu_seqlens.device, dtype=cu_seqlens.dtype, non_blocking=True diff --git a/vllm/v1/attention/backends/flashinfer.py b/vllm/v1/attention/backends/flashinfer.py index 2898e51c24..de51e2198b 100755 --- a/vllm/v1/attention/backends/flashinfer.py +++ b/vllm/v1/attention/backends/flashinfer.py @@ -46,6 +46,7 @@ from vllm.utils.flashinfer import ( supports_trtllm_attention, use_trtllm_attention, ) +from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.math_utils import cdiv from vllm.utils.torch_utils import ( PIN_MEMORY, @@ -1408,13 +1409,15 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]): # seq_lens_cpu is not needed since TRTLLM paths use GPU tensors # (block_tables, seq_lens) directly. needs_seq_lens_cpu = self.use_dcp or use_cascade or not all_uses_trtllm - seq_lens_cpu = common_attn_metadata.seq_lens_cpu if needs_seq_lens_cpu else None - seq_lens_np = seq_lens_cpu.numpy() if seq_lens_cpu is not None else None - num_blocks_np = ( - (seq_lens_np + (page_size - 1)) // page_size - if seq_lens_np is not None - else None - ) + if needs_seq_lens_cpu: + with gpu_sync_allowed(): + seq_lens_cpu = common_attn_metadata.seq_lens_cpu + seq_lens_np = seq_lens_cpu.numpy() + num_blocks_np = (seq_lens_np + (page_size - 1)) // page_size + else: + seq_lens_cpu = None + seq_lens_np = None + num_blocks_np = None # Adjust seq_lens_cpu for DCP if self.use_dcp: diff --git a/vllm/v1/attention/ops/vit_attn_wrappers.py b/vllm/v1/attention/ops/vit_attn_wrappers.py index 86a0b6ee9d..b7aeb8c6a7 100644 --- a/vllm/v1/attention/ops/vit_attn_wrappers.py +++ b/vllm/v1/attention/ops/vit_attn_wrappers.py @@ -20,6 +20,7 @@ import torch.nn.functional as F from vllm._aiter_ops import rocm_aiter_ops from vllm.platforms import current_platform +from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.torch_utils import direct_register_custom_op @@ -50,7 +51,12 @@ def flash_attn_maxseqlen_wrapper( cu_seqlens = torch.arange( 0, (batch_size + 1) * q_len, step=q_len, dtype=torch.int32, device=q.device ) - max_seqlen = q_len if max_seqlen is None else max_seqlen.item() + if max_seqlen is None: + max_seqlen = q_len + else: + # `flash_attn_varlen_func` needs a Python int for kernel launch bounds. + with gpu_sync_allowed(): + max_seqlen = max_seqlen.item() q, k, v = (einops.rearrange(x, "b s ... -> (b s) ...") for x in [q, k, v]) output = flash_attn_varlen_func( @@ -159,7 +165,12 @@ def triton_attn_wrapper( cu_seqlens = torch.arange( 0, (batch_size + 1) * q_len, step=q_len, dtype=torch.int32, device=q.device ) - max_seqlen = q_len if max_seqlen is None else max_seqlen.item() + if max_seqlen is None: + max_seqlen = q_len + else: + # `context_attention_fwd` needs a Python int. + with gpu_sync_allowed(): + max_seqlen = max_seqlen.item() q, k, v = (einops.rearrange(x, "b s ... -> (b s) ...") for x in [q, k, v]) output = torch.empty_like(q) @@ -261,7 +272,9 @@ def torch_sdpa_wrapper( outputs = [] - lens = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist() + # `torch.split` needs Python int sizes. + with gpu_sync_allowed(): + lens = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist() q_chunks = torch.split(q, lens, dim=1) k_chunks = torch.split(k, lens, dim=1) v_chunks = torch.split(v, lens, dim=1) @@ -337,7 +350,10 @@ def flashinfer_wrapper( batch_offsets_qko = cu_seqlens[:cu_seqlength].view(-1, 1, 1, 1) batch_offsets_v = cu_seqlens[cu_seqlength:].view(-1, 1, 1, 1) sequence_lengths = sequence_lengths.view(-1, 1, 1, 1) - max_seqlen = max_seqlen.item() + # `cudnn_batch_prefill_with_kv_cache` needs Python ints for the + # max-token-per-seq bounds. + with gpu_sync_allowed(): + max_seqlen = max_seqlen.item() output, _ = cudnn_batch_prefill_with_kv_cache( q, diff --git a/vllm/v1/sample/ops/topk_topp_triton.py b/vllm/v1/sample/ops/topk_topp_triton.py index 7d593e69cd..d8f8c8956d 100755 --- a/vllm/v1/sample/ops/topk_topp_triton.py +++ b/vllm/v1/sample/ops/topk_topp_triton.py @@ -12,6 +12,7 @@ using Pivot-based Truncation and Selection" By Park et al. import torch from vllm.triton_utils import tl, triton +from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.math_utils import next_power_of_2 from vllm.utils.platform_utils import num_compute_units @@ -919,12 +920,13 @@ def apply_top_k_top_p_triton( # Cache lookup table entries on each device. tables = _TRITON_TABLE_CACHE.get(logits.device) if tables is None: - normal_cdf_to_sigma_table = logits.new_tensor(_NORMAL_CDF_TO_SIGMA_TABLE) - percentile_to_std_table = logits.new_tensor(_PERCENTILE_TO_STD_TABLE) - _TRITON_TABLE_CACHE[logits.device] = ( - normal_cdf_to_sigma_table, - percentile_to_std_table, - ) + with gpu_sync_allowed(): + normal_cdf_to_sigma_table = logits.new_tensor(_NORMAL_CDF_TO_SIGMA_TABLE) + percentile_to_std_table = logits.new_tensor(_PERCENTILE_TO_STD_TABLE) + _TRITON_TABLE_CACHE[logits.device] = ( + normal_cdf_to_sigma_table, + percentile_to_std_table, + ) else: normal_cdf_to_sigma_table, percentile_to_std_table = tables diff --git a/vllm/v1/sample/sampler.py b/vllm/v1/sample/sampler.py index ac8512f22c..3efaf33d98 100644 --- a/vllm/v1/sample/sampler.py +++ b/vllm/v1/sample/sampler.py @@ -6,6 +6,7 @@ import torch import torch.nn as nn from vllm.config.model import LogprobsMode +from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.torch_utils import PIN_MEMORY from vllm.v1.outputs import LogprobsTensors, SamplerOutput from vllm.v1.sample.metadata import SamplingMetadata @@ -342,9 +343,10 @@ class Sampler(nn.Module): # of the compiled batched_count_greater_than. mark_unbacked makes # the size fully symbolic so dynamo doesn't specialize when # batch_size transitions from 1 to >=2. - torch._dynamo.decorators.mark_unbacked(logprobs, 0) - torch._dynamo.decorators.mark_unbacked(token_logprobs, 0) - token_ranks = batched_count_greater_than(logprobs, token_logprobs) + with gpu_sync_allowed(first_only=True): + torch._dynamo.decorators.mark_unbacked(logprobs, 0) + torch._dynamo.decorators.mark_unbacked(token_logprobs, 0) + token_ranks = batched_count_greater_than(logprobs, token_logprobs) # Concatenate together with the topk. indices = torch.cat((token_ids, topk_indices), dim=1) diff --git a/vllm/v1/worker/dp_utils.py b/vllm/v1/worker/dp_utils.py index e7c6d81a99..a7476bffa4 100644 --- a/vllm/v1/worker/dp_utils.py +++ b/vllm/v1/worker/dp_utils.py @@ -1,12 +1,15 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from contextlib import nullcontext + import torch import torch.distributed as dist from vllm.config import ParallelConfig from vllm.distributed.parallel_state import get_dp_group from vllm.logger import init_logger +from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.v1.worker.ubatch_utils import ( check_ubatch_thresholds, is_last_ubatch_empty, @@ -136,27 +139,33 @@ def _synchronize_dp_ranks( parallel_config=parallel_config, ) - # Synchronize cudagraph_mode across ranks first (take min). - # This is needed before DP padding decision since we use the synced - # cudagraph mode to determine whether DP padding is needed. - synced_cudagraph_mode = _post_process_cudagraph_mode(tensor) + # Only the NCCL path leaves `tensor` on device. With Gloo -- the default + # under async scheduling -- the all-reduce runs on CPU, so the reads below + # are host-side and the check should stay armed. + with ( + nullcontext() + if parallel_config.disable_nccl_for_dp_synchronization + else gpu_sync_allowed() + ): + # Synchronize cudagraph_mode across ranks first (take min). + # This is needed before DP padding decision since we use the synced + # cudagraph mode to determine whether DP padding is needed. + synced_cudagraph_mode = _post_process_cudagraph_mode(tensor) - # Check conditions for microbatching - should_ubatch = _post_process_ubatch(tensor, parallel_config.num_ubatches) + # Check conditions for microbatching + should_ubatch = _post_process_ubatch(tensor, parallel_config.num_ubatches) - # DP padding is needed when cudagraph is enabled (synced across ranks) - # or when ubatching/DBO is active (ubatching requires uniform batch - # sizes across DP ranks currently). - # Use the synced runtime cudagraph mode rather than the compilation config - # so we can avoid padding when cudagraph is not enabled for this step. - should_dp_pad = synced_cudagraph_mode != 0 or should_ubatch + # DP padding is needed when cudagraph is enabled (synced across ranks) + # or when ubatching/DBO is active (ubatching requires uniform batch + # sizes across DP ranks currently). + # Use the synced runtime cudagraph mode rather than the compilation + # config so we can avoid padding when cudagraph is not enabled for + # this step. + should_dp_pad = synced_cudagraph_mode != 0 or should_ubatch - # Pad all DP ranks up to the maximum token count across ranks if - # should_dp_pad is True - num_tokens_after_padding = _post_process_dp_padding( - tensor, - should_dp_pad, - ) + # Pad all DP ranks up to the maximum token count across ranks if + # should_dp_pad is True + num_tokens_after_padding = _post_process_dp_padding(tensor, should_dp_pad) return should_ubatch, num_tokens_after_padding, synced_cudagraph_mode diff --git a/vllm/v1/worker/encoder_cudagraph.py b/vllm/v1/worker/encoder_cudagraph.py index 9f61b704cb..04a54a39a0 100644 --- a/vllm/v1/worker/encoder_cudagraph.py +++ b/vllm/v1/worker/encoder_cudagraph.py @@ -19,6 +19,7 @@ from vllm.model_executor.models.interfaces import ( ) from vllm.model_executor.models.utils import scatter_output_slices from vllm.model_executor.models.vision import get_load_balance_assignment +from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.v1.worker.encoder_cudagraph_defs import ( EncoderCudaGraphConfig, EncoderItemSpec, @@ -298,7 +299,19 @@ class EncoderCudaGraphManager: def _get_item_specs(self, mm_kwargs: dict[str, Any]) -> list[EncoderItemSpec]: """Get item specs from the model.""" - return self.model.get_encoder_cudagraph_item_specs(mm_kwargs) + # Implementations read per-item grid/patch counts off device tensors + # to size the cudagraph buffers, so the D2H is inherent here. + with gpu_sync_allowed(): + return self.model.get_encoder_cudagraph_item_specs(mm_kwargs) + + def _select_items( + self, mm_kwargs: dict[str, Any], indices: list[int] + ) -> dict[str, Any]: + """Select the mm kwargs for `indices` from the model.""" + # Same as `_get_item_specs`: implementations re-read the per-item + # grid/patch counts to slice the batch, so the D2H is inherent. + with gpu_sync_allowed(): + return self.model.select_encoder_cudagraph_items(mm_kwargs, indices) def _get_per_item_out_tokens(self, mm_kwargs: dict[str, Any]) -> list[int]: """Get per-item output token counts as plain ints.""" @@ -415,9 +428,7 @@ class EncoderCudaGraphManager: outputs_by_orig_idx: dict[int, torch.Tensor] = {} for batch_indices, path_budgets in batches: - batch_mm_kwargs = self.model.select_encoder_cudagraph_items( - mm_kwargs, batch_indices - ) + batch_mm_kwargs = self._select_items(mm_kwargs, batch_indices) graph_outputs: dict[str, torch.Tensor] = {} all_eager = True @@ -491,11 +502,9 @@ class EncoderCudaGraphManager: ] if len(local_indices) > 0: - local_mm_kwargs = self.model.select_encoder_cudagraph_items( - mm_kwargs, local_indices - ) + local_mm_kwargs = self._select_items(mm_kwargs, local_indices) else: - local_mm_kwargs = self.model.select_encoder_cudagraph_items(mm_kwargs, []) + local_mm_kwargs = self._select_items(mm_kwargs, []) max_output_tokens_per_rank = ( max( diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index cd1679bdd9..9087f87854 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -120,6 +120,7 @@ from vllm.sequence import IntermediateTensors from vllm.tasks import GenerationTask, PoolingTask, SupportedTask from vllm.tracing import instrument from vllm.utils import length_from_prompt_token_ids_or_embeds +from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.math_utils import cdiv, round_up from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib from vllm.utils.nvtx_pytorch_hooks import PytHooks @@ -3848,26 +3849,28 @@ class GPUModelRunner( self.routed_experts_slot_mapping_device[:total], non_blocking=True, ) + with gpu_sync_allowed(): + # Get the valid generated tokens. + max_gen_len = sampled_token_ids.shape[-1] + if max_gen_len == 1: + # No spec decode tokens. + valid_sampled_token_ids = self._to_list(sampled_token_ids) + # Mask out the sampled tokens that should not be sampled. + for i in discard_sampled_tokens_req_indices: + valid_sampled_token_ids[int(i)].clear() - # Get the valid generated tokens. - max_gen_len = sampled_token_ids.shape[-1] - if max_gen_len == 1: - # No spec decode tokens. - valid_sampled_token_ids = self._to_list(sampled_token_ids) - # Mask out the sampled tokens that should not be sampled. - for i in discard_sampled_tokens_req_indices: - valid_sampled_token_ids[int(i)].clear() - - if logprobs_tensors is not None: - logprobs_lists = logprobs_tensors.tolists() - else: - # Includes spec decode tokens. - valid_sampled_token_ids, logprobs_lists = RejectionSampler.parse_output( - sampled_token_ids, - self.input_batch.vocab_size, - discard_sampled_tokens_req_indices, - logprobs_tensors=logprobs_tensors, - ) + if logprobs_tensors is not None: + logprobs_lists = logprobs_tensors.tolists() + else: + # Includes spec decode tokens. + valid_sampled_token_ids, logprobs_lists = ( + RejectionSampler.parse_output( + sampled_token_ids, + self.input_batch.vocab_size, + discard_sampled_tokens_req_indices, + logprobs_tensors=logprobs_tensors, + ) + ) else: valid_sampled_token_ids = [] invalid_req_indices = discard_sampled_tokens_req_indices.tolist() @@ -5843,7 +5846,10 @@ class GPUModelRunner( on device instead; see`AsyncGPUModelRunnerOutput`. """ try: - counts = [] if logits is None else count_nans_per_row(logits).tolist() + # Reporting per-request NaN counts requires them on the host; this + # path is opt-in diagnostics, so the D2H is intended. + with gpu_sync_allowed(): + counts = [] if logits is None else count_nans_per_row(logits).tolist() num_nans_in_logits = nans_to_dict(counts, self.input_batch.req_id_to_index) if envs.VLLM_RAISE_ON_LOGIT_NANS: raise_if_nan_logits(num_nans_in_logits) diff --git a/vllm/v1/worker/mamba_utils.py b/vllm/v1/worker/mamba_utils.py index f87a34a364..9b89af0812 100644 --- a/vllm/v1/worker/mamba_utils.py +++ b/vllm/v1/worker/mamba_utils.py @@ -16,6 +16,7 @@ from vllm.model_executor.layers.mamba.mamba_utils import ( is_conv_state_dim_first, ) from vllm.triton_utils import tl, triton +from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.math_utils import cdiv from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec @@ -754,7 +755,22 @@ class MambaSpecDecodeGPUContext: """ if self.is_initialized: return + # This only runs once per worker. + with gpu_sync_allowed(): + self._populate_metadata( + kv_cache_config, + forward_context, + mamba_state_copy_funcs, + block_tables, + ) + def _populate_metadata( + self, + kv_cache_config: KVCacheConfig, + forward_context: dict[str, Any], + mamba_state_copy_funcs: tuple[MambaStateCopyFunc, ...], + block_tables: list[torch.Tensor], + ) -> None: idx = 0 for group_local_idx, mamba_group_id in enumerate(self.mamba_group_ids): layer_names = kv_cache_config.kv_cache_groups[mamba_group_id].layer_names