[Core] Check for GPU<->CPU syncs during CI (#43107)

Signed-off-by: Nick Hill <nickhill123@gmail.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

.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)'

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
 

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

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

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(

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:

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

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(

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):

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]:

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,

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

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:

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

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:

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)
 

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
 

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, <max feature size>, 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"])
 

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
 

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(

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
 

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)

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

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, ...]] = []

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
                 )

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(

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
 

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)

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 = "<asr_text>"
@@ -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)

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)

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
 

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]
 

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]:

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

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)

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

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:

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,

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
 

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)

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
 

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(

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)

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