615d4cfade · tduncan/Nvidia-Optimization-Stack-VLLM

.buildkite/test_areas/plugins.yaml

+6

@@ -75,6 +75,9 @@ steps:

75 - vllm/model_executor/layers/quantization

76 - tests/plugins_tests/bitsandbytes

77 commands:

78 + # bitsandbytes' int8 quant syncs internally, and it is out-of-tree, so

79 + # there is nothing to wrap on the vLLM side.

80 + - unset VLLM_GPU_SYNC_CHECK

81 - pip install "vllm-bnb-plugin >= 0.0.1"

82 - pytest -v -s plugins_tests/bitsandbytes -m 'not distributed'

83

@@ -89,5 +92,8 @@ steps:

92 - vllm/model_executor/layers/quantization

93 - tests/plugins_tests/bitsandbytes

94 commands:

95 + # bitsandbytes' int8 quant syncs internally, and it is out-of-tree, so

96 + # there is nothing to wrap on the vLLM side.

97 + - unset VLLM_GPU_SYNC_CHECK

98 - pip install "vllm-bnb-plugin >= 0.0.1"

99 - pytest -v -s plugins_tests/bitsandbytes -m 'distributed(num_gpus=2)'

docker/Dockerfile

+3

@@ -965,6 +965,9 @@ ENV HF_XET_HIGH_PERFORMANCE 1

965 # increase timeout for hf downloads (for testing)

966 ENV HF_HUB_DOWNLOAD_TIMEOUT 60

967

968 +# Catch GPU<->CPU syncs in execute_model/sample_tokens

969 +ENV VLLM_GPU_SYNC_CHECK=error

970 +

971 # Copy in the v1 package for testing (it isn't distributed yet)

972 COPY vllm/v1 /usr/local/lib/python${PYTHON_VERSION}/dist-packages/vllm/v1

973

docker/Dockerfile.rocm

+3

@@ -820,6 +820,9 @@ COPY --from=export_vllm /vllm_v1 /usr/local/lib/python${PYTHON_VERSION}/dist-pac

820 # Hide source under src/ so it won't shadow the installed package in tests.

821 RUN mkdir src && mv vllm src/vllm

822

823 +# Catch GPU<->CPU syncs in execute_model/sample_tokens during tests.

824 +ENV VLLM_GPU_SYNC_CHECK=error

825 +

826 # -----------------------

827 # Final vLLM image

828 FROM mori_base AS final

tests/models/multimodal/generation/test_mm_prefix_lm.py

+4 -1

@@ -8,6 +8,7 @@ import torch

8 from transformers import AutoModelForImageTextToText

9

10 from vllm.platforms import current_platform

11 +from vllm.utils.gpu_sync_debug import gpu_sync_allowed

12

13 from ....conftest import HfRunner, ImageTestAssets, VllmRunner

14 from .vlm_utils import model_utils

@@ -30,7 +31,9 @@ def _install_prefill_hidden_capture(model):

31 def forward(*args, **kwargs):

32 hidden_states = original_forward(*args, **kwargs)

33 if model._prefill_hidden is None and torch.is_tensor(hidden_states):

33 - model._prefill_hidden = hidden_states.detach().float().cpu()

34 + # Capturing hidden states for comparison is a deliberate D2H.

35 + with gpu_sync_allowed():

36 + model._prefill_hidden = hidden_states.detach().float().cpu()

37 return hidden_states

38

39 language_model.forward = forward

tests/v1/e2e/general/test_mamba_prefix_cache.py

+28 -17

@@ -19,6 +19,7 @@ from vllm.distributed import cleanup_dist_env_and_memory

19 from vllm.model_executor.layers.mamba.mamba_utils import MambaStateCopyFunc

20 from vllm.platforms import current_platform

21 from vllm.sequence import IntermediateTensors

22 +from vllm.utils.gpu_sync_debug import gpu_sync_allowed

23 from vllm.utils.torch_utils import async_tensor_h2d

24 from vllm.v1.attention.backends.utils import CommonAttentionMetadata

25 from vllm.v1.core.kv_cache_manager import KVCacheBlocks, KVCacheManager

@@ -981,7 +982,10 @@ def _run_mamba_prefix_cache_mrv2(

982 yield forward_context[layer_name].kv_cache[-1], block_table

983

984 def temporal_block(temporal_state, block_table, col):

984 - return temporal_state[int(block_table[0, col].item())]

985 + # Resolving the block id for assertions is a deliberate D2H.

986 + with gpu_sync_allowed():

987 + block_id = int(block_table[0, col].item())

988 + return temporal_state[block_id]

989

990 def wrapped_preprocess_state(

991 self: MambaHybridModelState,

@@ -1005,19 +1009,24 @@ def _run_mamba_prefix_cache_mrv2(

1009 self, input_batch, block_tables, kv_cache_config, num_computed_tokens

1010 )

1011 if cur_step_action is not None:

1008 - req_idx = int(input_batch.idx_mapping[0].item())

1009 - src_col = int(self._mamba_src_col_gpu[req_idx].item())

1010 - off = int(self._mamba_src_off_gpu[req_idx].item())

1011 - dst = int(self._mamba_state_idx_gpu[req_idx].item())

1012 + # Reading the GPU-side copy state back to assert on it is a

1013 + # deliberate D2H.

1014 + with gpu_sync_allowed():

1015 + req_idx = int(input_batch.idx_mapping[0].item())

1016 + src_col = int(self._mamba_src_col_gpu[req_idx].item())

1017 + off = int(self._mamba_src_off_gpu[req_idx].item())

1018 + dst = int(self._mamba_state_idx_gpu[req_idx].item())

1019 actual = (-1, -1) if src_col < 0 or src_col == dst else (src_col + off, dst)

1020 assert actual == expected, (

1021 f"V2 align preprocess copy: expected={expected}, "

1022 f"actual={actual}, {cur_step_action=}"

1023 )

1017 - for temporal, bt, src_state in snapshots:

1018 - torch.testing.assert_close(

1019 - temporal_block(temporal, bt, expected[1]), src_state

1020 - )

1024 + # Comparing device tensors for the assertion is a deliberate D2H.

1025 + with gpu_sync_allowed():

1026 + for temporal, bt, src_state in snapshots:

1027 + torch.testing.assert_close(

1028 + temporal_block(temporal, bt, expected[1]), src_state

1029 + )

1030 return ret

1031

1032 def wrapped_postprocess_state(

@@ -1048,10 +1057,12 @@ def _run_mamba_prefix_cache_mrv2(

1057 ret = original_postprocess_state(

1058 self, idx_mapping, num_sampled, num_computed_tokens

1059 )

1051 - for temporal, bt, src_state in snapshots:

1052 - torch.testing.assert_close(

1053 - temporal_block(temporal, bt, expected[1]), src_state

1054 - )

1060 + # Comparing device tensors for the assertion is a deliberate D2H.

1061 + with gpu_sync_allowed():

1062 + for temporal, bt, src_state in snapshots:

1063 + torch.testing.assert_close(

1064 + temporal_block(temporal, bt, expected[1]), src_state

1065 + )

1066 return ret

1067

1068 def wrapped_execute_model(

@@ -1072,10 +1083,10 @@ def _run_mamba_prefix_cache_mrv2(

1083 ret = original_execute_model(self, scheduler_output, *args, **kwargs)

1084 if cur_step_action is not None and self.execute_model_state is not None:

1085 input_batch = self.execute_model_state.input_batch

1075 - assert (

1076 - cur_step_action.num_computed_tokens_start

1077 - == input_batch.positions[input_batch.query_start_loc[0]].item()

1078 - )

1086 + # Reading positions back to assert on them is a deliberate D2H.

1087 + with gpu_sync_allowed():

1088 + start_pos = input_batch.positions[input_batch.query_start_loc[0]].item()

1089 + assert cur_step_action.num_computed_tokens_start == start_pos

1090 return ret

1091

1092 def fake_sample(

vllm/distributed/eplb/eplb_communicator.py

+6 -4

@@ -33,6 +33,7 @@ from vllm.distributed.stateless_coordinator import StatelessGroupCoordinator

33 from vllm.distributed.utils import is_weak_contiguous

34 from vllm.logger import init_logger

35 from vllm.platforms import current_platform

36 +from vllm.utils.gpu_sync_debug import gpu_sync_allowed

37

38 logger = init_logger(__name__)

39

@@ -222,10 +223,11 @@ class TorchDistGlooStagedEplbCommunicator(EplbCommunicator):

223

224 # Wait for all D2H copies to finish

225 # before issuing gloo batch_isend_irecv operations.

225 - if self._cuda_stream is not None:

226 - self._cuda_stream.synchronize()

227 - else:

228 - torch.cuda.current_stream().synchronize()

226 + with gpu_sync_allowed():

227 + if self._cuda_stream is not None:

228 + self._cuda_stream.synchronize()

229 + else:

230 + torch.cuda.current_stream().synchronize()

231

232 reqs = batch_isend_irecv(p2p_ops)

233 for req in reqs:

vllm/distributed/eplb/eplb_state.py

+21 -18

@@ -47,6 +47,7 @@ from vllm.distributed.utils import StatelessProcessGroup

47 from vllm.logger import init_logger

48 from vllm.model_executor.models.interfaces import MixtureOfExperts

49 from vllm.platforms import current_platform

50 +from vllm.utils.gpu_sync_debug import gpu_sync_allowed

51

52 from .async_worker import start_async_worker

53 from .eplb_communicator import EplbCommunicator, create_eplb_communicator

@@ -596,10 +597,11 @@ class EplbState:

597 num_tokens_per_rank

598 )

599

599 - # Just to make type checker happy

600 - tokens_tensors: list[float] = torch.stack(

601 - [avg_tokens_tensor, max_tokens_tensor]

602 - ).tolist()

600 + # This gpu/cpu sync only happens with expert stat logging enabled.

601 + with gpu_sync_allowed():

602 + tokens_tensors: list[float] = torch.stack(

603 + [avg_tokens_tensor, max_tokens_tensor]

604 + ).tolist()

605 avg_tokens, max_tokens = tokens_tensors

606 balancedness = avg_tokens / max_tokens if max_tokens > 0 else 0.0

607

@@ -825,15 +827,17 @@ class EplbState:

827 self.model_states.values(), global_expert_load_windows

828 ):

829 if not self.is_async or is_profile:

828 - # Get new expert mappings for the model

829 - new_physical_to_logical_map = self.policy.rebalance_experts(

830 - global_expert_load_window.cpu(),

831 - num_replicas,

832 - num_groups,

833 - num_nodes,

834 - num_gpus,

835 - eplb_model_state.physical_to_logical_map.cpu(),

836 - )

830 + # Get new expert mappings for the model. The policy runs on the

831 + # host, so the load window and current map have to come back.

832 + with gpu_sync_allowed():

833 + new_physical_to_logical_map = self.policy.rebalance_experts(

834 + global_expert_load_window.cpu(),

835 + num_replicas,

836 + num_groups,

837 + num_nodes,

838 + num_gpus,

839 + eplb_model_state.physical_to_logical_map.cpu(),

840 + )

841

842 skip_rearrange = False

843 if (

@@ -1248,7 +1252,9 @@ def _pad_out_tensor(src: torch.Tensor, dst: torch.Tensor) -> None:

1252 src_padding = dst.shape[-1] - src.shape[-1]

1253 assert src_padding >= 0

1254 new_src = torch.nn.functional.pad(src, (0, src_padding), value=-1)

1251 - dst.copy_(new_src)

1255 + # The map is committed from the host once per layer per rearrangement.

1256 + with gpu_sync_allowed():

1257 + dst.copy_(new_src)

1258

1259

1260 def _commit_eplb_maps_for_layer(

@@ -1276,10 +1282,7 @@ def _commit_eplb_maps_for_layer(

1282 num_logical_experts = model_state.logical_to_physical_map.shape[1]

1283 new_logical, new_replica_count = compute_logical_maps(src, num_logical_experts)

1284 # Commit logical_to_physical_map

1279 - _pad_out_tensor(

1280 - src=new_logical,

1281 - dst=model_state.logical_to_physical_map[layer],

1282 - )

1285 + _pad_out_tensor(src=new_logical, dst=model_state.logical_to_physical_map[layer])

1286

1287 # Commit logical_replica_count

1288 src = new_replica_count

vllm/distributed/eplb/rebalance_execute.py

+6 -2

@@ -16,6 +16,7 @@ from torch.distributed import ProcessGroup, all_gather

16 from vllm.distributed.eplb.eplb_communicator import EplbCommunicator

17 from vllm.distributed.eplb.eplb_utils import CpuGpuEvent

18 from vllm.logger import init_logger

19 +from vllm.utils.gpu_sync_debug import gpu_sync_allowed

20

21 logger = init_logger(__name__)

22

@@ -590,8 +591,11 @@ def rearrange_expert_weights_inplace(

591

592 weights_buffer = list(expert_buffer)

593

593 - old_global_expert_indices_cpu = old_global_expert_indices.cpu().numpy()

594 - new_global_expert_indices_cpu = new_global_expert_indices.cpu().numpy()

594 + # The per-layer transfer plan is built in Python, so both maps have to

595 + # come back to the host. Once per rearrangement.

596 + with gpu_sync_allowed():

597 + old_global_expert_indices_cpu = old_global_expert_indices.cpu().numpy()

598 + new_global_expert_indices_cpu = new_global_expert_indices.cpu().numpy()

599

600 for layer_idx in range(num_moe_layers):

601 transfer_metadata = move_to_buffer(

vllm/distributed/kv_transfer/kv_connector/v1/example_connector.py

+3 -1

@@ -15,6 +15,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.base import (

15 )

16 from vllm.logger import init_logger

17 from vllm.model_executor.layers.attention.mla_attention import MLACommonMetadata

18 +from vllm.utils.gpu_sync_debug import gpu_sync_allowed

19 from vllm.utils.hashing import safe_hash

20 from vllm.v1.attention.backend import AttentionMetadata

21 from vllm.v1.core.sched.output import SchedulerOutput

@@ -243,7 +244,8 @@ class ExampleConnector(KVConnectorBase_V1):

244 layer_name, request.token_ids, request.mm_hashes

245 )

246 kv_cache = extract_kv_from_layer(kv_layer, request.slot_mapping)

246 - tensors = {"kv_cache": kv_cache.detach().cpu()}

247 + with gpu_sync_allowed():

248 + tensors = {"kv_cache": kv_cache.detach().cpu()}

249 safetensors.torch.save_file(tensors, filename)

250

251 def wait_for_save(self):

vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py

+32 -27

@@ -67,6 +67,7 @@ from vllm.distributed.parallel_state import (

67 )

68 from vllm.logger import init_logger

69 from vllm.platforms import current_platform

70 +from vllm.utils.gpu_sync_debug import gpu_sync_allowed

71 from vllm.utils.network_utils import make_zmq_path

72 from vllm.utils.torch_utils import async_tensor_h2d

73 from vllm.v1.attention.backends.utils import get_kv_cache_layout

@@ -1882,14 +1883,16 @@ class NixlBaseConnectorWorker:

1883

1884 local_block_ids = meta.local_physical_block_ids

1885 # TODO (NickLucche) D2H<>H2D ops could benefit from coalescing io across groups

1885 - for group_block_ids in local_block_ids:

1886 - self.copy_blocks(

1887 - self.host_xfer_buffers,

1888 - self.device_kv_caches,

1889 - group_block_ids,

1890 - group_block_ids,

1891 - "h2d",

1892 - )

1886 + # The h2d block copies below are intentionally synchronous.

1887 + with gpu_sync_allowed():

1888 + for group_block_ids in local_block_ids:

1889 + self.copy_blocks(

1890 + self.host_xfer_buffers,

1891 + self.device_kv_caches,

1892 + group_block_ids,

1893 + group_block_ids,

1894 + "h2d",

1895 + )

1896 if logger.isEnabledFor(logging.DEBUG):

1897 logger.debug(

1898 "synced recved kv of request[%s] to device kv buffer,"

@@ -1903,26 +1906,28 @@ class NixlBaseConnectorWorker:

1906 assert self.use_host_buffer

1907 assert self.copy_blocks is not None

1908

1906 - for req_id, meta in metadata.reqs_to_save.items():

1907 - meta.local_physical_block_ids = self._logical_to_kernel_block_ids(

1908 - meta.local_block_ids, self._physical_blocks_per_logical_kv_block

1909 - )

1910 - if logger.isEnabledFor(logging.DEBUG):

1911 - logger.debug(

1912 - "save_load_kv for request[%s] to host xfer buffer."

1913 - "local_block_ids: %s. ",

1914 - req_id,

1915 - ",".join(map(str, meta.local_physical_block_ids)),

1916 - )

1917 - # blocking

1918 - for group_block_ids in meta.local_physical_block_ids:

1919 - self.copy_blocks(

1920 - self.device_kv_caches,

1921 - self.host_xfer_buffers,

1922 - group_block_ids,

1923 - group_block_ids,

1924 - "d2h",

1909 + # The d2h block copies below are intentionally synchronous.

1910 + with gpu_sync_allowed():

1911 + for req_id, meta in metadata.reqs_to_save.items():

1912 + meta.local_physical_block_ids = self._logical_to_kernel_block_ids(

1913 + meta.local_block_ids, self._physical_blocks_per_logical_kv_block

1914 )

1915 + if logger.isEnabledFor(logging.DEBUG):

1916 + logger.debug(

1917 + "save_load_kv for request[%s] to host xfer buffer."

1918 + "local_block_ids: %s. ",

1919 + req_id,

1920 + ",".join(map(str, meta.local_physical_block_ids)),

1921 + )

1922 + # blocking

1923 + for group_block_ids in meta.local_physical_block_ids:

1924 + self.copy_blocks(

1925 + self.device_kv_caches,

1926 + self.host_xfer_buffers,

1927 + group_block_ids,

1928 + group_block_ids,

1929 + "d2h",

1930 + )

1931

1932 @cached_property

1933 def _attention_kv_caches(self) -> list[torch.Tensor]:

vllm/lora/punica_wrapper/punica_gpu.py

+6 -3

@@ -14,6 +14,7 @@ import torch

14 from vllm.lora.layers import LoRAMapping

15 from vllm.lora.utils import get_captured_lora_counts

16 from vllm.triton_utils import HAS_TRITON, triton

17 +from vllm.utils.gpu_sync_debug import gpu_sync_allowed

18 from vllm.utils.math_utils import round_up

19

20 if HAS_TRITON:

@@ -83,9 +84,11 @@ class PunicaWrapperGPU(PunicaWrapperBase):

84 self.is_prefill = mapping.is_prefill

85 self._update_base_metadata(mapping, lora_index_to_id, max_loras, vocab_size)

86

86 - # Prepare cuda kernel metadata tensors

87 - self.token_mapping_meta.prepare_tensors(self.token_lora_indices)

88 - self.prompt_mapping_meta.prepare_tensors(self.sampler_indices)

87 + # TODO avoid gpu<->cpu sync here

88 + with gpu_sync_allowed():

89 + # Prepare cuda kernel metadata tensors

90 + self.token_mapping_meta.prepare_tensors(self.token_lora_indices)

91 + self.prompt_mapping_meta.prepare_tensors(self.sampler_indices)

92

93 def add_shrink(

94 self,

vllm/lora/worker_manager.py

+36 -31

@@ -19,6 +19,7 @@ from vllm.lora.model_manager import (

19 from vllm.lora.peft_helper import PEFTHelper

20 from vllm.lora.request import LoRARequest

21 from vllm.lora.utils import get_adapter_absolute_path

22 +from vllm.utils.gpu_sync_debug import gpu_sync_allowed

23

24 logger = init_logger(__name__)

25

@@ -223,9 +224,11 @@ class WorkerLoRAManager:

224 def add_adapter(self, adapter_request: Any) -> bool:

225 if adapter_request.adapter_id in self.list_adapters():

226 return False

226 - loaded_adapter = self._load_adapter(adapter_request)

227 - loaded = self._adapter_manager.add_adapter(loaded_adapter)

228 - self._adapter_manager.activate_adapter(loaded_adapter.id)

227 + # One-time per adapter

228 + with gpu_sync_allowed():

229 + loaded_adapter = self._load_adapter(adapter_request)

230 + loaded = self._adapter_manager.add_adapter(loaded_adapter)

231 + self._adapter_manager.activate_adapter(loaded_adapter.id)

232 return loaded

233

234 def remove_adapter(self, adapter_id: int) -> bool:

@@ -288,32 +291,34 @@ class LRUCacheWorkerLoRAManager(WorkerLoRAManager):

291 # This is ok because it's currently only called from

292 # the single-threaded core engine loop.

293

291 - if (

292 - lora_request.lora_int_id not in self.list_adapters()

293 - or lora_request.load_inplace

294 - ):

295 - # Load the new adapter first to ensure it is actually valid, before

296 - # evicting any existing adapters.

297 - # This may cause the # of loaded lora adapters to very temporarily

298 - # exceed `--max-cpu-loras`.

299 - lora = self._load_adapter(lora_request)

300 -

301 - # Remove the existing adapter if it exists

302 - # Use case for LoRA inplace

303 - self._adapter_manager.remove_adapter(lora.id)

304 -

305 - # Loading succeeded, now check if we will exceed cache capacity and

306 - # evict if the oldest adapter if so

307 - if len(self._adapter_manager) + 1 > self._adapter_manager.capacity:

308 - assert isinstance(self._adapter_manager, LRUCacheLoRAModelManager)

309 - self._adapter_manager.remove_oldest_adapter()

310 - # Then add the new adapter to the cache

311 - loaded = self._adapter_manager.add_adapter(lora)

312 - else:

313 - # If the lora is already loaded, just touch it to

314 - # update its position in the caches

315 - loaded = (

316 - self._adapter_manager.get_adapter(lora_request.lora_int_id) is not None

317 - )

318 - self._adapter_manager.activate_adapter(lora_request.lora_int_id)

294 + with gpu_sync_allowed():

295 + if (

296 + lora_request.lora_int_id not in self.list_adapters()

297 + or lora_request.load_inplace

298 + ):

299 + # Load the new adapter first to ensure it is actually valid, before

300 + # evicting any existing adapters.

301 + # This may cause the # of loaded lora adapters to very temporarily

302 + # exceed `--max-cpu-loras`.

303 + lora = self._load_adapter(lora_request)

304 +

305 + # Remove the existing adapter if it exists

306 + # Use case for LoRA inplace

307 + self._adapter_manager.remove_adapter(lora.id)

308 +

309 + # Loading succeeded, now check if we will exceed cache capacity and

310 + # evict if the oldest adapter if so

311 + if len(self._adapter_manager) + 1 > self._adapter_manager.capacity:

312 + assert isinstance(self._adapter_manager, LRUCacheLoRAModelManager)

313 + self._adapter_manager.remove_oldest_adapter()

314 + # Then add the new adapter to the cache

315 + loaded = self._adapter_manager.add_adapter(lora)

316 + else:

317 + # If the lora is already loaded, just touch it to

318 + # update its position in the caches

319 + loaded = (

320 + self._adapter_manager.get_adapter(lora_request.lora_int_id)

321 + is not None

322 + )

323 + self._adapter_manager.activate_adapter(lora_request.lora_int_id)

324 return loaded

vllm/model_executor/layers/mamba/mamba_mixer2.py

+13 -8

@@ -48,6 +48,7 @@ from vllm.model_executor.model_loader.weight_utils import (

48 from vllm.model_executor.parameter import BasevLLMParameter

49 from vllm.model_executor.utils import set_weight_attrs

50 from vllm.platforms import current_platform

51 +from vllm.utils.gpu_sync_debug import gpu_sync_allowed

52 from vllm.utils.torch_utils import (

53 LayerNameType,

54 _encode_layer_name,

@@ -904,17 +905,21 @@ class MambaMixer2(MambaBase, PluggableLayer):

905 # then chunk_stride = 2

906 chunk_stride = mamba_block_size // chunk_size

907

908 + # The per-sequence loop below uses these as Python scalars.

909 + # TODO avoid sync here?

910 + with gpu_sync_allowed():

911 + block_idx_first_cpu = block_idx_first_scheduled_token_p.tolist()

912 + block_idx_last_cpu = block_idx_last_scheduled_token_p.tolist()

913 + num_computed_tokens_p_cpu = num_computed_tokens_p.tolist()

914 + last_chunk_indices_p_cpu = last_chunk_indices_p.tolist()

915 +

916 # Save state for sequences with more than just final state

917 for seq_idx in range(num_prefills):

918 # Block index for the first scheduled token

910 - block_idx_first_scheduled_token = block_idx_first_scheduled_token_p[

911 - seq_idx

912 - ]

919 + block_idx_first_scheduled_token = block_idx_first_cpu[seq_idx]

920

921 # Block index for the last scheduled token

915 - block_idx_last_scheduled_token = block_idx_last_scheduled_token_p[

916 - seq_idx

917 - ]

922 + block_idx_last_scheduled_token = block_idx_last_cpu[seq_idx]

923

924 # Number of blocks that need to be written

925 n_blocks_to_fill = (

@@ -935,7 +940,7 @@ class MambaMixer2(MambaBase, PluggableLayer):

940 if seq_idx == 0:

941 first_chunk = 0

942 else:

938 - first_chunk = 1 + last_chunk_indices_p[seq_idx - 1]

943 + first_chunk = 1 + last_chunk_indices_p_cpu[seq_idx - 1]

944

945 # First chunk that is aligned on the mamba block boundary

946 first_aligned_chunk = first_chunk + chunk_stride - 1

@@ -943,7 +948,7 @@ class MambaMixer2(MambaBase, PluggableLayer):

948 # Calculate the number of computed tokens that were not

949 # already cached

950 num_unaligned_computed_tokens = (

946 - num_computed_tokens_p[seq_idx] % mamba_block_size

951 + num_computed_tokens_p_cpu[seq_idx] % mamba_block_size

952 )

953

954 if num_unaligned_computed_tokens > 0:

vllm/model_executor/models/audioflamingo3.py

+9 -5

@@ -54,6 +54,7 @@ from vllm.multimodal.processing import (

54 PromptUpdateDetails,

55 )

56 from vllm.sequence import IntermediateTensors

57 +from vllm.utils.gpu_sync_debug import gpu_sync_allowed

58 from vllm.utils.tensor_schema import TensorSchema, TensorShape

59

60 from .interfaces import (

@@ -306,7 +307,10 @@ def _flatten_valid_audio_embeddings(

307 < output_lengths[:, None]

308 )

309

309 - return audio_embeddings[valid_mask], output_lengths

310 + # Boolean-mask indexing has a data-dependent output shape, so the count

311 + # has to come back to the host.

312 + with gpu_sync_allowed():

313 + return audio_embeddings[valid_mask], output_lengths

314

315

316 def _count_audio_tokens_from_mask(

@@ -630,10 +634,10 @@ class AudioFlamingo3ForConditionalGeneration(

634 audio_features,

635 feature_attention_mask,

636 )

633 - chunk_embeddings = torch.split(

634 - masked_audio_features,

635 - audio_output_lengths.tolist(),

636 - )

637 + # `split` needs Python int sizes.

638 + with gpu_sync_allowed():

639 + split_sizes = audio_output_lengths.tolist()

640 + chunk_embeddings = torch.split(masked_audio_features, split_sizes)

641

642 grouped_embeddings = []

643 current_idx = 0

vllm/model_executor/models/cohere_asr.py

+4 -4

@@ -56,6 +56,7 @@ from vllm.transformers_utils.processors.cohere_asr import (

56 CohereASRProcessor,

57 )

58 from vllm.utils.collection_utils import is_list_of

59 +from vllm.utils.gpu_sync_debug import gpu_sync_allowed

60 from vllm.v1.attention.backend import (

61 AttentionType,

62 )

@@ -1770,10 +1771,9 @@ class CohereASRModel(nn.Module):

1771 out = self.encoder_decoder_proj(out)

1772

1773 # Convert padded tensor to packed

1773 - outs = []

1774 - for i, feat in enumerate(out):

1775 - feat_len = encoder_output_length[i]

1776 - outs.append(feat[:feat_len, :])

1774 + with gpu_sync_allowed():

1775 + lengths = encoder_output_length.tolist()

1776 + outs = [feat[:length, :] for feat, length in zip(out, lengths)]

1777

1778 return outs

1779 else:

vllm/model_executor/models/ernie45_vl.py

+13 -3

@@ -71,6 +71,7 @@ from vllm.multimodal.processing import (

71 PromptUpdate,

72 )

73 from vllm.sequence import IntermediateTensors

74 +from vllm.utils.gpu_sync_debug import gpu_sync_allowed

75 from vllm.utils.tensor_schema import TensorSchema, TensorShape

76 from vllm.utils.torch_utils import async_tensor_h2d

77 from vllm.v1.attention.backends.registry import AttentionBackendEnum

@@ -788,7 +789,10 @@ class VariableResolutionResamplerModel(nn.Module):

789 return x

790

791 def fwd_placeholder(x, grid_thw, to_tensor=False):

791 - grid_thw_cpu = grid_thw.cpu().numpy()

792 + # The per-image grids drive the Python-level offset arithmetic

793 + # below; the same read is wrapped on the cudagraph path.

794 + with gpu_sync_allowed():

795 + grid_thw_cpu = grid_thw.cpu().numpy()

796 grid_t, grid_hw = grid_thw_cpu[:, 0], grid_thw_cpu[:, 1:]

797 grid_hw_after_conv = grid_hw.prod(-1) // (self.spatial_conv_size**2)

798

@@ -1697,8 +1701,11 @@ class Ernie4_5_VLMoeForConditionalGeneration(

1701 EncoderCudaGraphReplayBuffers,

1702 )

1703

1704 + # The per-image grids are needed as Python ints to size the buffers.

1705 + with gpu_sync_allowed():

1706 + grid_thw_list = mm_kwargs["image_grid_thw"].tolist()

1707 metadata = self.vision_model.prepare_encoder_metadata(

1701 - mm_kwargs["image_grid_thw"].tolist(), max_batch_size=max_batch_size

1708 + grid_thw_list, max_batch_size=max_batch_size

1709 )

1710 values = metadata | {

1711 "pixel_values": mm_kwargs["pixel_values"],

@@ -1738,7 +1745,10 @@ class Ernie4_5_VLMoeForConditionalGeneration(

1745 # Ernie only uses the single "default" encoder path.

1746 output = outputs["default"]

1747 grid_thw = batch_mm_kwargs["image_grid_thw"].to(output.device)

1741 - num_valid = int((grid_thw[:, 0] * grid_thw[:, 1] * grid_thw[:, 2]).sum())

1748 + # The valid token count slices the graph output for the eager

1749 + # resampler call, so it has to come back to the host.

1750 + with gpu_sync_allowed():

1751 + num_valid = int((grid_thw[:, 0] * grid_thw[:, 1] * grid_thw[:, 2]).sum())

1752 image_embeds = self.resampler_model(output[:num_valid], grid_thw)

1753 scatter_output_slices(image_embeds, indices, per_item_out_tokens, dest, clone)

1754

vllm/model_executor/models/gemma4_mm.py

+43 -27

@@ -69,6 +69,7 @@ from vllm.multimodal.processing.processor import (

69 PromptUpdateDetails,

70 )

71 from vllm.sequence import IntermediateTensors

72 +from vllm.utils.gpu_sync_debug import gpu_sync_allowed

73 from vllm.utils.tensor_schema import TensorSchema, TensorShape

74 from vllm.utils.torch_utils import async_tensor_h2d

75

@@ -1365,11 +1366,14 @@ class Gemma4ForConditionalGeneration(

1366 pp_tensor,

1367 pad_tensor,

1368 ).to(self.model_dtype)

1368 - encoder_outputs = vt.encoder(

1369 - inputs_embeds=inputs_embeds,

1370 - attention_mask=~pad_tensor,

1371 - pixel_position_ids=pp_tensor,

1372 - )

1369 + # HuggingFace's mask builder probes `padding_mask.all()` to

1370 + # decide whether the mask can be skipped, which syncs.

1371 + with gpu_sync_allowed():

1372 + encoder_outputs = vt.encoder(

1373 + inputs_embeds=inputs_embeds,

1374 + attention_mask=~pad_tensor,

1375 + pixel_position_ids=pp_tensor,

1376 + )

1377 hidden_states = encoder_outputs.last_hidden_state

1378

1379 for i, (orig_idx, _, _) in enumerate(chunk_items):

@@ -1387,13 +1391,17 @@ class Gemma4ForConditionalGeneration(

1391 single_pos_ids = pool_position_ids[orig_idx].unsqueeze(0)

1392 padding_positions = (single_pos_ids == -1).all(dim=-1)

1393

1390 - pooled_states, valid_mask = vt.pooler(

1391 - hidden_states=single_hidden,

1392 - pixel_position_ids=single_pos_ids,

1393 - padding_positions=padding_positions,

1394 - output_length=output_length,

1395 - )

1396 - valid_states = pooled_states[valid_mask]

1394 + # The pooler goes through HuggingFace's mask builder, which probes

1395 + # `padding_mask.all()`, and the mask indexing below needs the

1396 + # selected count on the host.

1397 + with gpu_sync_allowed():

1398 + pooled_states, valid_mask = vt.pooler(

1399 + hidden_states=single_hidden,

1400 + pixel_position_ids=single_pos_ids,

1401 + padding_positions=padding_positions,

1402 + output_length=output_length,

1403 + )

1404 + valid_states = pooled_states[valid_mask]

1405

1406 if getattr(vt.config, "standardize", False):

1407 valid_states = (valid_states - vt.std_bias) * vt.std_scale

@@ -1444,7 +1452,9 @@ class Gemma4ForConditionalGeneration(

1452 pooling_k2 = vision_cfg.pooling_kernel_size**2

1453

1454 if isinstance(frame_counts, torch.Tensor):

1447 - fc_list = frame_counts.tolist()

1455 + # Per-video frame counts drive the Python-level batching below.

1456 + with gpu_sync_allowed():

1457 + fc_list = frame_counts.tolist()

1458 else:

1459 fc_list = list(frame_counts)

1460

@@ -1474,11 +1484,13 @@ class Gemma4ForConditionalGeneration(

1484 pp_chunk,

1485 pad_chunk,

1486 ).to(self.model_dtype)

1477 - encoder_outputs = vt.encoder(

1478 - inputs_embeds=inputs_embeds,

1479 - attention_mask=~pad_chunk,

1480 - pixel_position_ids=pp_chunk,

1481 - )

1487 + # HuggingFace's mask builder probes `padding_mask.all()`.

1488 + with gpu_sync_allowed():

1489 + encoder_outputs = vt.encoder(

1490 + inputs_embeds=inputs_embeds,

1491 + attention_mask=~pad_chunk,

1492 + pixel_position_ids=pp_chunk,

1493 + )

1494 last_hidden_states_list.append(encoder_outputs.last_hidden_state)

1495

1496 last_hidden_states = torch.cat(last_hidden_states_list, dim=0)

@@ -1493,13 +1505,15 @@ class Gemma4ForConditionalGeneration(

1505 single_pos_ids = pixel_position_ids[i].unsqueeze(0)

1506 single_pad_pos = padding_positions[i].unsqueeze(0)

1507

1496 - pooled_states, valid_mask = vt.pooler(

1497 - hidden_states=single_hidden,

1498 - pixel_position_ids=single_pos_ids,

1499 - padding_positions=single_pad_pos,

1500 - output_length=output_length,

1501 - )

1502 - valid_states = pooled_states[valid_mask]

1508 + # As above, plus mask indexing that needs the count on the host.

1509 + with gpu_sync_allowed():

1510 + pooled_states, valid_mask = vt.pooler(

1511 + hidden_states=single_hidden,

1512 + pixel_position_ids=single_pos_ids,

1513 + padding_positions=single_pad_pos,

1514 + output_length=output_length,

1515 + )

1516 + valid_states = pooled_states[valid_mask]

1517

1518 if getattr(vt.config, "standardize", False):

1519 valid_states = (valid_states - vt.std_bias) * vt.std_scale

@@ -1553,9 +1567,11 @@ class Gemma4ForConditionalGeneration(

1567

1568 # Strip padding per-batch element: only keep valid (non-padding)

1569 # tokens.

1570 + # Boolean-mask indexing needs the selected count on the host.

1571 per_audio = []

1557 - for enc, mask in zip(audio_features, audio_mask, strict=True):

1558 - per_audio.append(enc[mask]) # [num_real, hidden_size]

1572 + with gpu_sync_allowed():

1573 + for enc, mask in zip(audio_features, audio_mask, strict=True):

1574 + per_audio.append(enc[mask]) # [num_real, hidden_size]

1575

1576 return per_audio

1577

vllm/model_executor/models/granite_speech.py

+4 -2

@@ -60,6 +60,7 @@ from vllm.multimodal.processing import (

60 from vllm.sequence import IntermediateTensors

61 from vllm.tokenizers import cached_tokenizer_from_config

62 from vllm.transformers_utils.processor import cached_processor_from_config

63 +from vllm.utils.gpu_sync_debug import gpu_sync_allowed

64 from vllm.utils.tensor_schema import TensorSchema, TensorShape

65 from vllm.utils.torch_utils import PIN_MEMORY

66

@@ -788,8 +789,9 @@ class GraniteSpeechForConditionalGeneration(

789 encoder_embeds = self.encoder(audio_input["input_features"])

790 # [bsz, <max feature size>, 4096]

791 projected_embeds = self.projector(encoder_embeds)

791 - # Apply mask on variable length audio features

792 - masked_embeds = projected_embeds[audio_input["input_features_mask"]]

792 + # Apply mask on variable length audio features.

793 + with gpu_sync_allowed():

794 + masked_embeds = projected_embeds[audio_input["input_features_mask"]]

795 # Split variable length features into a tuple

796 return torch.split(masked_embeds, audio_input["audio_embed_sizes"])

797

vllm/model_executor/models/idefics2_vision_model.py

+13 -6

@@ -39,6 +39,7 @@ from vllm.model_executor.layers.linear import (

39 RowParallelLinear,

40 )

41 from vllm.model_executor.layers.quantization import QuantizationConfig

42 +from vllm.utils.gpu_sync_debug import gpu_sync_allowed

43

44 from .utils import AutoWeightsLoader, WeightsMapper

45 from .vision import is_vit_use_data_parallel, run_dp_sharded_vision_model

@@ -96,10 +97,14 @@ class Idefics2VisionEmbeddings(nn.Module):

97 size=(batch_size, max_nb_patches_h * max_nb_patches_w), fill_value=0

98 )

99

99 - for batch_idx, p_attn_mask in enumerate(patch_attention_mask):

100 - if tgt_sizes is not None:

101 - nb_patches_h = tgt_sizes[batch_idx][0]

102 - nb_patches_w = tgt_sizes[batch_idx][1]

100 + with gpu_sync_allowed():

101 + patch_attention_mask_cpu = patch_attention_mask.cpu()

102 + tgt_sizes_cpu = tgt_sizes.cpu() if tgt_sizes is not None else None

103 +

104 + for batch_idx, p_attn_mask in enumerate(patch_attention_mask_cpu):

105 + if tgt_sizes_cpu is not None:

106 + nb_patches_h = tgt_sizes_cpu[batch_idx][0]

107 + nb_patches_w = tgt_sizes_cpu[batch_idx][1]

108 else:

109 nb_patches_h = p_attn_mask[:, 0].sum()

110 nb_patches_w = p_attn_mask[0].sum()

@@ -114,8 +119,10 @@ class Idefics2VisionEmbeddings(nn.Module):

119 pos_ids = (

120 bucket_coords_h[:, None] * self.num_patches_per_side + bucket_coords_w

121 ).flatten()

117 - position_ids[batch_idx][p_attn_mask.view(-1).cpu()] = pos_ids

118 - position_ids = position_ids.to(self.position_embedding.weight.device)

122 + position_ids[batch_idx][p_attn_mask.view(-1)] = pos_ids

123 + position_ids = position_ids.to(

124 + self.position_embedding.weight.device, non_blocking=True

125 + )

126 embeddings += self.position_embedding(position_ids)

127 return embeddings

128

vllm/model_executor/models/idefics3.py

+6 -4

@@ -51,6 +51,7 @@ from vllm.multimodal.processing import (

51 PromptUpdateDetails,

52 )

53 from vllm.sequence import IntermediateTensors

54 +from vllm.utils.gpu_sync_debug import gpu_sync_allowed

55 from vllm.utils.tensor_schema import TensorSchema, TensorShape

56

57 from .idefics2_vision_model import (

@@ -443,11 +444,12 @@ class Idefics3Model(nn.Module):

444 real_images_inds = (pixel_values == 0.0).sum(

445 dim=(-1, -2, -3)

446 ) != nb_values_per_image

446 - pixel_values = pixel_values[real_images_inds].contiguous()

447 + with gpu_sync_allowed():

448 + pixel_values = pixel_values[real_images_inds].contiguous()

449

448 - # Handle the vision attention mask

449 - # Remove padding images from the mask

450 - pixel_attention_mask = pixel_attention_mask[real_images_inds].contiguous()

450 + # Handle the vision attention mask

451 + # Remove padding images from the mask

452 + pixel_attention_mask = pixel_attention_mask[real_images_inds].contiguous()

453

454 patch_size = self.config.vision_config.patch_size

455 patches_subgrid = pixel_attention_mask.unfold(

vllm/model_executor/models/jina.py

+19 -13

@@ -16,6 +16,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead

16 from vllm.sequence import IntermediateTensors

17 from vllm.tasks import PoolingTask

18 from vllm.transformers_utils.repo_utils import get_hf_file_bytes

19 +from vllm.utils.gpu_sync_debug import gpu_sync_allowed

20 from vllm.v1.pool.metadata import PoolingMetadata

21

22 from ..layers.pooler import DispatchPooler

@@ -110,19 +111,24 @@ class JinaForRankingPool(StepPool):

111 prompt_token_ids = pooling_metadata.get_prompt_token_ids()

112

113 embeds_list = list[torch.Tensor | None]()

113 - for data, token_ids in zip(pooled_data_lst, prompt_token_ids):

114 - # for unfinished chunked prefill

115 - if data is None:

116 - embeds_list.append(None)

117 - else:

118 - docs_indexes = torch.where(torch.eq(token_ids, self.doc_token_id))[0]

119 - query_indexes = torch.where(torch.eq(token_ids, self.query_token_id))[0]

120 -

121 - # The JinaForRanking model concatenates docs first, then query.

122 - # Let's stay consistent with this novel design.

123 - indexes = torch.cat([docs_indexes, query_indexes])

124 - embeds = self.projector(data[indexes])

125 - embeds_list.append(embeds)

114 + # `torch.where` resolves the match count on the host.

115 + with gpu_sync_allowed():

116 + for data, token_ids in zip(pooled_data_lst, prompt_token_ids):

117 + # for unfinished chunked prefill

118 + if data is None:

119 + embeds_list.append(None)

120 + else:

121 + doc_match = torch.eq(token_ids, self.doc_token_id)

122 + docs_indexes = torch.where(doc_match)[0]

123 + query_indexes = torch.where(

124 + torch.eq(token_ids, self.query_token_id)

125 + )[0]

126 +

127 + # The JinaForRanking model concatenates docs first, then query.

128 + # Let's stay consistent with this novel design.

129 + indexes = torch.cat([docs_indexes, query_indexes])

130 + embeds = self.projector(data[indexes])

131 + embeds_list.append(embeds)

132

133 return embeds_list

134

vllm/model_executor/models/minicpmv.py

+10 -4

@@ -98,6 +98,7 @@ from vllm.transformers_utils.processor import (

98 )

99 from vllm.transformers_utils.utils import convert_model_repo_to_path

100 from vllm.utils.collection_utils import flatten_2d_lists

101 +from vllm.utils.gpu_sync_debug import gpu_sync_allowed

102 from vllm.utils.tensor_schema import TensorSchema, TensorShape

103 from vllm.utils.torch_utils import set_default_torch_dtype

104

@@ -1548,7 +1549,8 @@ class MiniCPMV2_5(MiniCPMVBaseModel, SupportsLoRA):

1549 all_pixel_values[i, ..., :L_item] = pixel_values_item

1550

1551 num_patches = tgt_sizes.prod(-1)

1551 - max_patches = num_patches.max().item()

1552 + with gpu_sync_allowed():

1553 + max_patches = num_patches.max().item()

1554 assert isinstance(max_patches, int)

1555

1556 patch_attn_mask = torch.zeros((B, max_patches), dtype=torch.bool, device=device)

@@ -1640,7 +1642,9 @@ class MiniCPMV2_6(MiniCPMVBaseModel, SupportsLoRA):

1642 all_pixel_values[i, ..., :L_item] = pixel_values_item

1643

1644 num_patches = tgt_sizes.prod(-1)

1643 - max_patches = num_patches.max().item()

1645 + # Needed as a Python int to size the mask below.

1646 + with gpu_sync_allowed():

1647 + max_patches = num_patches.max().item()

1648 assert isinstance(max_patches, int)

1649

1650 patch_attn_mask = torch.zeros((B, max_patches), dtype=torch.bool, device=device)

@@ -1737,7 +1741,8 @@ class MiniCPMV4_0(MiniCPMVBaseModel, SupportsLoRA):

1741 all_pixel_values[i, ..., :L_item] = pixel_values_item

1742

1743 num_patches = tgt_sizes.prod(-1)

1740 - max_patches = num_patches.max().item()

1744 + with gpu_sync_allowed():

1745 + max_patches = num_patches.max().item()

1746 assert isinstance(max_patches, int)

1747

1748 patch_attn_mask = torch.zeros((B, max_patches), dtype=torch.bool, device=device)

@@ -1839,7 +1844,8 @@ class MiniCPMV4_5(MiniCPMVBaseModel, SupportsLoRA):

1844 all_pixel_values[i, ..., :L_item] = pixel_values_item

1845

1846 num_patches = tgt_sizes.prod(-1)

1842 - max_patches = num_patches.max().item()

1847 + with gpu_sync_allowed():

1848 + max_patches = num_patches.max().item()

1849 assert isinstance(max_patches, int)

1850

1851 patch_attn_mask = torch.zeros((B, max_patches), dtype=torch.bool, device=device)

vllm/model_executor/models/moonvit.py

+11 -5

@@ -66,6 +66,7 @@ from vllm.model_executor.models.utils import maybe_prefix

66 from vllm.model_executor.models.vision import is_vit_use_data_parallel

67 from vllm.platforms import current_platform

68 from vllm.transformers_utils.configs.moonvit import MoonViTConfig

69 +from vllm.utils.gpu_sync_debug import gpu_sync_allowed

70 from vllm.utils.torch_utils import async_tensor_h2d

71

72

@@ -811,9 +812,14 @@ class MoonVitPretrainedModel(PreTrainedModel):

812 merge_kernel_size=self.merge_kernel_size,

813 )

814

814 - hidden_states = self.patch_embed(pixel_values, grid_hw)

815 - hidden_states = self.encoder(hidden_states, grid_hw)

816 - hidden_states = patch_merger(

817 - hidden_states, grid_hw, merge_kernel_size=self.merge_kernel_size

818 - )

815 + # Legacy path: patch_embed, encoder and patch_merger each iterate the

816 + # per-image grids in Python, so all three read `grid_hw` back to the

817 + # host. The `encoder_metadata` path above precomputes that outside the

818 + # graph and does not sync.

819 + with gpu_sync_allowed():

820 + hidden_states = self.patch_embed(pixel_values, grid_hw)

821 + hidden_states = self.encoder(hidden_states, grid_hw)

822 + hidden_states = patch_merger(

823 + hidden_states, grid_hw, merge_kernel_size=self.merge_kernel_size

824 + )

825 return hidden_states

vllm/model_executor/models/moss_audio.py

+14 -9

@@ -55,6 +55,7 @@ from vllm.multimodal.processing import (

55 )

56 from vllm.sequence import IntermediateTensors

57 from vllm.transformers_utils.repo_utils import get_hf_file_to_dict

58 +from vllm.utils.gpu_sync_debug import gpu_sync_allowed

59 from vllm.utils.tensor_schema import TensorSchema, TensorShape

60

61 from .interfaces import (

@@ -1667,15 +1668,19 @@ class MossAudioModel(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA):

1668 """

1669 audio_data = audio_input["audio_data"]

1670 audio_data_seqlens = audio_input["audio_data_seqlens"]

1670 - last_hidden_state, deepstack = self.audio_encoder(

1671 - audio_data.to(self.audio_encoder.dtype),

1672 - feature_lens=audio_data_seqlens,

1673 - output_deepstack_hidden_states=len(self.deepstack_audio_merger_list) > 0,

1674 - )

1675 - audio_embeds = self.audio_adapter(last_hidden_state)

1676 - audio_lengths = MossAudioEncoder._compute_downsampled_length(

1677 - audio_data_seqlens.to(device=audio_embeds.device, dtype=torch.long)

1678 - ).tolist()

1671 + # The encoder chunks the input by per-audio feature lengths, which

1672 + # needs Python ints for `split`/`pad_sequence`.

1673 + want_deepstack = len(self.deepstack_audio_merger_list) > 0

1674 + with gpu_sync_allowed():

1675 + last_hidden_state, deepstack = self.audio_encoder(

1676 + audio_data.to(self.audio_encoder.dtype),

1677 + feature_lens=audio_data_seqlens,

1678 + output_deepstack_hidden_states=want_deepstack,

1679 + )

1680 + audio_embeds = self.audio_adapter(last_hidden_state)

1681 + audio_lengths = MossAudioEncoder._compute_downsampled_length(

1682 + audio_data_seqlens.to(device=audio_embeds.device, dtype=torch.long)

1683 + ).tolist()

1684 main_embeddings = tuple(audio_embeds.squeeze(0).split(audio_lengths, dim=0))

1685

1686 deepstack_embeddings: list[tuple[torch.Tensor, ...]] = []

vllm/model_executor/models/phi4mm.py

+15 -3

@@ -49,6 +49,7 @@ from vllm.multimodal.processing.processor import (

49 ResolvedPromptUpdate,

50 )

51 from vllm.sequence import IntermediateTensors

52 +from vllm.utils.gpu_sync_debug import gpu_sync_allowed

53 from vllm.utils.tensor_schema import TensorSchema, TensorShape

54

55 from .idefics2_vision_model import Idefics2VisionTransformer

@@ -396,12 +397,23 @@ class Phi4MMImageEncoder(nn.Module):

397 w * base_feat_width // base_feat_height_reduction,

398 )

399 )

399 - useful_height = int(reshaped_image_attention_mask[0, :, 0].sum().item())

400 - useful_width = int(reshaped_image_attention_mask[0, 0, :].sum().item())

400 + # The mask stays on device for the encoder above, so these

401 + # per-image counts have to come back to the host to drive the

402 + # slicing and the Python-level length arithmetic.

403 + with gpu_sync_allowed():

404 + useful_height = int(

405 + reshaped_image_attention_mask[0, :, 0].sum().item()

406 + )

407 + useful_width = int(

408 + reshaped_image_attention_mask[0, 0, :].sum().item()

409 + )

410 + mask_token_count = int(

411 + image_attention_mask[_bs, : B_ + 1, 0::2, 0::2].sum().item()

412 + )

413 sub_img = sub_img[:, :useful_height, :useful_width]

414 temp_sub_GN = self.sub_GN.repeat(1, useful_height, 1, 1)

415 temp_len = (

404 - int(image_attention_mask[_bs, : B_ + 1, 0::2, 0::2].sum().item())

416 + mask_token_count

417 + (useful_height + 1)

418 + base_feat_height // base_feat_height_reduction

419 )

vllm/model_executor/models/qwen2_5_omni_thinker.py

+15 -11

@@ -89,6 +89,7 @@ from vllm.multimodal.processing.processor import (

89 PromptUpdateDetails,

90 )

91 from vllm.sequence import IntermediateTensors

92 +from vllm.utils.gpu_sync_debug import gpu_sync_allowed

93 from vllm.utils.tensor_schema import TensorSchema, TensorShape

94

95 from .interfaces import (

@@ -1036,12 +1037,14 @@ class Qwen2_5OmniConditionalGenerationMixin:

1037 self.audio_tower._get_feat_extract_output_lengths(audio_feature_lengths)

1038 )

1039

1039 - audio_outputs = self.audio_tower(

1040 - input_features.to(self.audio_tower.dtype),

1041 - feature_lens=audio_feature_lengths,

1042 - aftercnn_lens=audio_feat_lengths,

1043 - )

1044 - return audio_outputs.last_hidden_state.split(audio_output_lengths.tolist())

1040 + with gpu_sync_allowed():

1041 + audio_outputs = self.audio_tower(

1042 + input_features.to(self.audio_tower.dtype),

1043 + feature_lens=audio_feature_lengths,

1044 + aftercnn_lens=audio_feat_lengths,

1045 + )

1046 + split_sizes = audio_output_lengths.tolist()

1047 + return audio_outputs.last_hidden_state.split(split_sizes)

1048

1049 def _process_image_input(

1050 self, image_input: Qwen2_5_VLImageInputs

@@ -1526,12 +1529,13 @@ class Qwen2_5OmniThinkerForConditionalGeneration(

1529 video_token_id = self.config.video_token_index

1530 audio_token_id = self.config.audio_token_index

1531

1529 - input_ids_cpu = input_ids.cpu()

1530 - is_video = is_multimodal & (input_ids_cpu == video_token_id)

1531 - is_audio = is_multimodal & (input_ids_cpu == audio_token_id)

1532 + with gpu_sync_allowed():

1533 + input_ids_cpu = input_ids.cpu()

1534 + is_video = is_multimodal & (input_ids_cpu == video_token_id)

1535 + is_audio = is_multimodal & (input_ids_cpu == audio_token_id)

1536

1533 - num_video = is_video.sum().item()

1534 - num_audio = is_audio.sum().item()

1537 + num_video = is_video.sum().item()

1538 + num_audio = is_audio.sum().item()

1539

1540 if check_interleaved_audio_video(is_video, is_audio, num_video, num_audio):

1541 inputs_embeds = self._embed_text_input_ids(

vllm/model_executor/models/qwen2_5_vl.py

+14 -11

@@ -87,6 +87,7 @@ from vllm.utils.torch_utils import PIN_MEMORY, async_tensor_h2d

87 from vllm.v1.attention.backends.registry import AttentionBackendEnum

88 from vllm.v1.worker.encoder_cudagraph_defs import EncoderCudaGraphReplayBuffers

89

90 +from ...utils.gpu_sync_debug import gpu_sync_allowed

91 from .interfaces import (

92 MultiModalEmbeddings,

93 SupportsEagle,

@@ -1569,8 +1570,9 @@ class Qwen2_5_VLForConditionalGeneration(

1570 video_second_per_grid=video_second_per_grid_t.item(),

1571 ).to(emb.device, non_blocking=True)

1572

1572 - emb = emb[retention_mask]

1573 - positions = positions[retention_mask]

1573 + with gpu_sync_allowed():

1574 + emb = emb[retention_mask]

1575 + positions = positions[retention_mask]

1576 emb = torch.cat([emb, positions], dim=1)

1577 video_embeds_out.append(emb)

1578 return tuple(video_embeds_out)

@@ -1624,15 +1626,16 @@ class Qwen2_5_VLForConditionalGeneration(

1626 mm[:, -4:].permute(1, 0).long() for mm in multimodal_embeddings

1627 ]

1628

1627 - positions, mrope_positions_delta = recompute_mrope_positions(

1628 - input_ids_t,

1629 - mm_embeddings_pos,

1630 - mrope_positions,

1631 - num_computed_tokens,

1632 - vision_start_token_id,

1633 - image_token_id,

1634 - video_token_id,

1635 - )

1629 + with gpu_sync_allowed():

1630 + positions, mrope_positions_delta = recompute_mrope_positions(

1631 + input_ids_t,

1632 + mm_embeddings_pos,

1633 + mrope_positions,

1634 + num_computed_tokens,

1635 + vision_start_token_id,

1636 + image_token_id,

1637 + video_token_id,

1638 + )

1639

1640 return mm_embeddings_out, positions, mrope_positions_delta

1641

vllm/model_executor/models/qwen2_audio.py

+9 -5

@@ -60,6 +60,7 @@ from vllm.multimodal.processing import (

60 PromptUpdate,

61 )

62 from vllm.sequence import IntermediateTensors

63 +from vllm.utils.gpu_sync_debug import gpu_sync_allowed

64 from vllm.utils.tensor_schema import TensorSchema, TensorShape

65

66 from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP

@@ -444,12 +445,15 @@ class Qwen2AudioForConditionalGeneration(nn.Module, SupportsMultiModal, Supports

445 )

446 < audio_output_lengths

447 )

447 - masked_audio_features = audio_features[audio_features_mask].view(-1, embed_dim)

448 + with gpu_sync_allowed():

449 + masked_audio_features = audio_features[audio_features_mask].view(

450 + -1, embed_dim

451 + )

452

449 - # Split to tuple of embeddings for individual audio input.

450 - return torch.split(

451 - masked_audio_features, audio_output_lengths.flatten().tolist()

452 - )

453 + # Split to tuple of embeddings for individual audio input.

454 + return torch.split(

455 + masked_audio_features, audio_output_lengths.flatten().tolist()

456 + )

457

458 def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:

459 audio_input = self._parse_and_validate_audio_input(**kwargs)

vllm/model_executor/models/qwen3_asr.py

+4 -1

@@ -92,6 +92,7 @@ from vllm.transformers_utils.processor import cached_processor_from_config

92 from vllm.transformers_utils.processors.qwen3_asr import (

93 Qwen3ASRProcessor,

94 )

95 +from vllm.utils.gpu_sync_debug import gpu_sync_allowed

96

97 logger = init_logger(__name__)

98 _ASR_TEXT_TAG = "<asr_text>"

@@ -467,7 +468,9 @@ class Qwen3ASRForConditionalGeneration(

468 feature_lens=audio_feature_lengths,

469 aftercnn_lens=audio_output_lengths,

470 )

470 - return audio_features.split(audio_output_lengths.tolist())

471 + with gpu_sync_allowed():

472 + split_sizes = audio_output_lengths.tolist()

473 + return audio_features.split(split_sizes)

474

475 def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings | None:

476 mm_input_by_modality = self._parse_and_validate_multimodal_inputs(**kwargs)

vllm/model_executor/models/qwen3_omni_moe_thinker.py

+17 -9

@@ -77,6 +77,7 @@ from vllm.multimodal.processing.processor import (

77 )

78 from vllm.sequence import IntermediateTensors

79 from vllm.transformers_utils.processor import cached_processor_from_config

80 +from vllm.utils.gpu_sync_debug import gpu_sync_allowed

81 from vllm.utils.torch_utils import async_tensor_h2d

82 from vllm.v1.attention.backends.registry import AttentionBackendEnum

83

@@ -434,11 +435,14 @@ class Qwen3OmniMoeAudioEncoder(nn.Module):

435 feature_lens: torch.Tensor,

436 aftercnn_lens: torch.Tensor,

437 ):

437 - # Compute chunk information

438 + # Compute chunk information.

439 chunk_num = torch.ceil(feature_lens / (self.n_window * 2)).long()

440

441 + with gpu_sync_allowed():

442 + total_chunks = int(chunk_num.sum())

443 +

444 chunk_lengths = async_tensor_h2d(

441 - [self.n_window * 2] * chunk_num.sum(),

445 + [self.n_window * 2] * total_chunks,

446 dtype=torch.long,

447 device=feature_lens.device,

448 )

@@ -447,15 +451,18 @@ class Qwen3OmniMoeAudioEncoder(nn.Module):

451 chunk_lengths[chunk_lengths == 0] = self.n_window * 2

452

453 # Split input features into chunks and pad

450 - chunk_list = input_features.T.split(chunk_lengths.tolist(), dim=0)

454 + with gpu_sync_allowed():

455 + chunk_lengths_list = chunk_lengths.tolist()

456 + chunk_list = input_features.T.split(chunk_lengths_list, dim=0)

457 padded_feature = nn.utils.rnn.pad_sequence(

458 chunk_list, batch_first=True

459 ).transpose(1, 2)

460

461 # Compute feature lengths after CNN

462 feature_lens_after_cnn = self._get_cnn_output_lengths(chunk_lengths)

457 - # Vectorized mask creation: avoid creating many small tensors

458 - max_len_after_cnn = feature_lens_after_cnn.max().item()

463 + # Vectorized mask creation: avoid creating many small tensors.

464 + with gpu_sync_allowed():

465 + max_len_after_cnn = feature_lens_after_cnn.max().item()

466 indices = torch.arange(max_len_after_cnn, device=padded_feature.device)

467 padded_mask_after_cnn = indices.unsqueeze(0) < feature_lens_after_cnn.unsqueeze(

468 1

@@ -494,16 +501,17 @@ class Qwen3OmniMoeAudioEncoder(nn.Module):

501 )

502 padded_embed = padded_embed + positional_embedding

503

497 - # Extract valid hidden states and compute cu_seqlens

498 - hidden_states = padded_embed[padded_mask_after_cnn]

504 + with gpu_sync_allowed():

505 + hidden_states = padded_embed[padded_mask_after_cnn]

506 + # Use tolist() for efficient batch conversion from tensor to Python.

507 + aftercnn_lens_list = aftercnn_lens.tolist()

508

509 # Compute cumulative sequence lengths for chunked attention

510 cu_chunk_lens = [0]

511 window_aftercnn = padded_mask_after_cnn.shape[-1] * (

512 self.n_window_infer // (self.n_window * 2)

513 )

505 - # Use tolist() for efficient batch conversion from tensor to Python

506 - for cnn_len in aftercnn_lens.tolist():

514 + for cnn_len in aftercnn_lens_list:

515 num_full_chunks = cnn_len // window_aftercnn

516 remainder = cnn_len % window_aftercnn

517 cu_chunk_lens.extend([window_aftercnn] * num_full_chunks)

vllm/model_executor/models/qwen3_vl.py

+29 -25

@@ -115,6 +115,7 @@ from vllm.utils.math_utils import round_up

115 from vllm.utils.torch_utils import PIN_MEMORY

116 from vllm.v1.worker.encoder_cudagraph_defs import EncoderCudaGraphReplayBuffers

117

118 +from ...utils.gpu_sync_debug import gpu_sync_allowed

119 from ...utils.torch_utils import async_tensor_h2d

120 from .interfaces import (

121 MultiModalEmbeddings,

@@ -2368,19 +2369,17 @@ class Qwen3VLForConditionalGeneration(

2369 spatial_merge_size=self.visual.spatial_merge_size,

2370 q=self.video_pruning_rate,

2371 )

2371 - # Apply retention mask.

2372 - emb = emb[retention_mask]

2372

2374 - # Calculate the actual number of retained tokens per frame.

2375 - num_frames, rows, cols = (

2376 - t,

2377 - h // merge_size,

2378 - w // merge_size,

2379 - )

2380 - retention_mask_thw = retention_mask.reshape(num_frames, rows, cols)

2381 - num_tokens_per_frame = (

2382 - retention_mask_thw.sum(dim=(1, 2)).long().tolist()

2383 - )

2373 + with gpu_sync_allowed():

2374 + # Apply retention mask.

2375 + emb = emb[retention_mask]

2376 +

2377 + # Calculate the actual number of retained tokens per frame.

2378 + num_frames, rows, cols = t, h // merge_size, w // merge_size

2379 + retention_mask_thw = retention_mask.reshape(num_frames, rows, cols)

2380 + num_tokens_per_frame = (

2381 + retention_mask_thw.sum(dim=(1, 2)).long().tolist()

2382 + )

2383 else:

2384 feature_size = emb.shape[0] // num_frames

2385 num_tokens_per_frame = [feature_size] * num_frames

@@ -2548,10 +2547,14 @@ class Qwen3VLForConditionalGeneration(

2547 .permute(1, 0)

2548 )

2549 full_is_video_embed = unpruned_token_ids_tensor == embed_token_id

2551 - expanded_positions[is_video_embed, :3] = original_mrope[full_is_video_embed][

2552 - retention_mask

2553 - ]

2554 - expanded_positions[~is_video_embed, :3] = original_mrope[~full_is_video_embed]

2550 +

2551 + with gpu_sync_allowed():

2552 + expanded_positions[is_video_embed, :3] = original_mrope[

2553 + full_is_video_embed

2554 + ][retention_mask]

2555 + expanded_positions[~is_video_embed, :3] = original_mrope[

2556 + ~full_is_video_embed

2557 + ]

2558 expanded_positions[..., 3] = is_vision_start

2559 expanded_positions[..., 4] = is_video_embed

2560

@@ -2814,15 +2817,16 @@ class Qwen3VLForConditionalGeneration(

2817 torch.empty(5, 0, device=device, dtype=torch.long)

2818 )

2819

2817 - positions, mrope_positions_delta = recompute_mrope_positions(

2818 - input_ids_t,

2819 - mm_embeddings_pos,

2820 - mrope_positions,

2821 - num_computed_tokens,

2822 - vision_start_token_id,

2823 - image_token_id,

2824 - video_token_id,

2825 - )

2820 + with gpu_sync_allowed():

2821 + positions, mrope_positions_delta = recompute_mrope_positions(

2822 + input_ids_t,

2823 + mm_embeddings_pos,

2824 + mrope_positions,

2825 + num_computed_tokens,

2826 + vision_start_token_id,

2827 + image_token_id,

2828 + video_token_id,

2829 + )

2830

2831 return mm_embeddings_out, positions, mrope_positions_delta

2832

vllm/model_executor/models/siglip.py

+20 -15

@@ -54,6 +54,7 @@ from vllm.multimodal.processing import (

54 TimingContext,

55 )

56 from vllm.sequence import IntermediateTensors

57 +from vllm.utils.gpu_sync_debug import gpu_sync_allowed

58 from vllm.utils.tensor_schema import TensorSchema, TensorShape

59

60 from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsQuant

@@ -998,27 +999,31 @@ class SiglipEmbeddingModel(nn.Module, SupportsMultiModal, SupportsQuant):

999 position_diffs = position_ids[1:] - position_ids[:-1]

1000 boundary_mask = position_diffs <= 0

1001

1001 - boundary_indices = torch.cat(

1002 - [

1003 - torch.tensor([0], device=features.device),

1004 - torch.where(boundary_mask)[0] + 1,

1005 - torch.tensor([len(features)], device=features.device),

1006 - ]

1007 - )

1002 + with gpu_sync_allowed():

1003 + boundary_mid_cpu = torch.where(boundary_mask.cpu())[0] + 1

1004 +

1005 + zero = torch.zeros(1, dtype=boundary_mid_cpu.dtype)

1006 + end = torch.full((1,), len(features), dtype=boundary_mid_cpu.dtype)

1007 + boundary_indices_cpu = torch.cat([zero, boundary_mid_cpu, end])

1008

1009 # For each sequence [start, end), position i flips to: start + end - 1 - i

1010 - lengths = boundary_indices[1:] - boundary_indices[:-1]

1011 - starts = boundary_indices[:-1]

1012 - ends = boundary_indices[1:]

1010 + lengths_cpu = boundary_indices_cpu[1:] - boundary_indices_cpu[:-1]

1011 + starts_cpu = boundary_indices_cpu[:-1]

1012 + ends_cpu = boundary_indices_cpu[1:]

1013

1014 # Assign sequence ID to each element

1015 - sequence_ids = torch.arange(

1016 - len(lengths), device=features.device

1017 - ).repeat_interleave(lengths)

1015 + sequence_ids_cpu = torch.arange(

1016 + len(lengths_cpu), dtype=boundary_mid_cpu.dtype

1017 + ).repeat_interleave(lengths_cpu)

1018

1019 # Calculate flipped indices for all positions at once

1020 - current_positions = torch.arange(len(features), device=features.device)

1021 - flip_indices = starts[sequence_ids] + ends[sequence_ids] - 1 - current_positions

1020 + current_positions_cpu = torch.arange(

1021 + len(features), dtype=boundary_mid_cpu.dtype

1022 + )

1023 + flip_indices_cpu = (

1024 + starts_cpu[sequence_ids_cpu] + ends_cpu[sequence_ids_cpu]

1025 + ) - (1 + current_positions_cpu)

1026 + flip_indices = flip_indices_cpu.to(features.device, non_blocking=True)

1027

1028 return features[flip_indices]

1029

vllm/model_executor/models/terratorch.py

+4 -1

@@ -61,6 +61,7 @@ from vllm.multimodal.processing import (

61 TimingContext,

62 )

63 from vllm.sequence import IntermediateTensors

64 +from vllm.utils.gpu_sync_debug import gpu_sync_allowed

65

66 from .interfaces import IsAttentionFree, MultiModalEmbeddings, SupportsMultiModal

67 from .interfaces_base import attn_type

@@ -280,7 +281,9 @@ class Terratorch(nn.Module, IsAttentionFree, SupportsMultiModal):

281 inputs_embeds: torch.Tensor | None = None,

282 **kwargs: object,

283 ):

283 - model_output = self.inference_runner.forward(**kwargs)

284 + # terratorch's forward has internal GPU syncs.

285 + with gpu_sync_allowed():

286 + model_output = self.inference_runner.forward(**kwargs)

287 return model_output.output

288

289 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

vllm/model_executor/models/transformers/multimodal.py

+16 -5

@@ -54,6 +54,7 @@ from vllm.multimodal.processing import (

54 TimingContext,

55 )

56 from vllm.sequence import IntermediateTensors

57 +from vllm.utils.gpu_sync_debug import gpu_sync_allowed

58

59 if TYPE_CHECKING:

60 from transformers import BatchFeature, PreTrainedModel

@@ -726,12 +727,17 @@ class MultiModalMixin(SupportsMultiModal, SupportsMRoPE):

727 kwargs.pop("token_type_ids", None)

728 kwargs.pop("mm_token_type_ids", None)

729

729 - audio_output = self.model.get_audio_features(

730 - input_features, return_dict=True, **kwargs

731 - )

730 + # HuggingFace's `get_audio_features` implementations branch on

731 + # per-sample feature lengths internally.

732 + with gpu_sync_allowed():

733 + audio_output = self.model.get_audio_features(

734 + input_features, return_dict=True, **kwargs

735 + )

736 audio_embeddings = audio_output.pooler_output

737

734 - split_sizes = num_audio_tokens.flatten().tolist()

738 + # Per-audio token counts are needed as Python ints to split.

739 + with gpu_sync_allowed():

740 + split_sizes = num_audio_tokens.flatten().tolist()

741 return self._split_embeddings(audio_embeddings, split_sizes)

742

743 def _process_image_input(self, **kwargs) -> list[torch.Tensor] | None:

@@ -749,7 +755,12 @@ class MultiModalMixin(SupportsMultiModal, SupportsMRoPE):

755

756 num_image_patches = kwargs.pop("num_image_patches")

757

752 - vision_embeddings = self.model.get_image_features(pixel_values, **kwargs)

758 + # The underlying HuggingFace `get_image_features` implementations

759 + # contain model-internal syncs (e.g. Idefics3 filters all-zero

760 + # padding images via boolean-mask indexing, LlavaOnevision

761 + # branches on per-sample batch counts).

762 + with gpu_sync_allowed():

763 + vision_embeddings = self.model.get_image_features(pixel_values, **kwargs)

764

765 # Transformers `v5`, `self.get_image_features` returns a tuple

766 # containing the features and optionally attentions/hidden_states

vllm/model_executor/models/ultravox.py

+13 -8

@@ -43,6 +43,7 @@ from vllm.multimodal.processing import (

43 from vllm.renderers import TokenizeParams

44 from vllm.sequence import IntermediateTensors

45 from vllm.transformers_utils.configs.ultravox import UltravoxConfig

46 +from vllm.utils.gpu_sync_debug import gpu_sync_allowed

47 from vllm.utils.tensor_schema import TensorSchema, TensorShape

48

49 from .interfaces import (

@@ -841,15 +842,19 @@ class UltravoxModel(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA):

842 embeddings.shape[0], -1

843 )

844 mask = indices < audio_token_len[:, None]

844 - # Apply mask and flatten

845 - flattened_embeddings = embeddings[mask]

845

847 - # Return one tensor per input audio

848 - embed_lens = [

849 - chunk_lens.sum().item()

850 - for chunk_lens in audio_token_len.split(audio_input["num_chunks"].tolist())

851 - ]

852 - return flattened_embeddings.split(embed_lens)

846 + with gpu_sync_allowed():

847 + # Apply mask and flatten

848 + flattened_embeddings = embeddings[mask]

849 +

850 + # Return one tensor per input audio

851 + embed_lens = [

852 + chunk_lens.sum().item()

853 + for chunk_lens in audio_token_len.split(

854 + audio_input["num_chunks"].tolist()

855 + )

856 + ]

857 + return flattened_embeddings.split(embed_lens)

858

859 def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:

860 audio_input = self._parse_and_validate_audio_input(**kwargs)

vllm/third_party/flash_linear_attention/ops/index.py

+5 -6

@@ -9,6 +9,7 @@

9 # ruff: noqa: E501

10 import torch

11

12 +from vllm.utils.gpu_sync_debug import gpu_sync_allowed

13 from vllm.triton_utils import triton

14

15 from .utils import tensor_cache

@@ -21,12 +22,10 @@ def prepare_lens(cu_seqlens: torch.Tensor) -> torch.Tensor:

22

23 @tensor_cache

24 def prepare_chunk_indices(cu_seqlens: torch.Tensor, chunk_size: int) -> torch.Tensor:

24 - indices = torch.cat(

25 - [

26 - torch.arange(n)

27 - for n in triton.cdiv(prepare_lens(cu_seqlens), chunk_size).tolist()

28 - ]

29 - )

25 + # This will be fixed by https://github.com/vllm-project/vllm/pull/51540.

26 + with gpu_sync_allowed():

27 + chunk_counts = triton.cdiv(prepare_lens(cu_seqlens), chunk_size).tolist()

28 + indices = torch.cat([torch.arange(n) for n in chunk_counts])

29 chunk_indices = torch.stack([indices.eq(0).cumsum(0) - 1, indices], 1)

30 return chunk_indices.to(

31 device=cu_seqlens.device, dtype=cu_seqlens.dtype, non_blocking=True

vllm/v1/attention/backends/flashinfer.py

+10 -7

@@ -46,6 +46,7 @@ from vllm.utils.flashinfer import (

46 supports_trtllm_attention,

47 use_trtllm_attention,

48 )

49 +from vllm.utils.gpu_sync_debug import gpu_sync_allowed

50 from vllm.utils.math_utils import cdiv

51 from vllm.utils.torch_utils import (

52 PIN_MEMORY,

@@ -1408,13 +1409,15 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

1409 # seq_lens_cpu is not needed since TRTLLM paths use GPU tensors

1410 # (block_tables, seq_lens) directly.

1411 needs_seq_lens_cpu = self.use_dcp or use_cascade or not all_uses_trtllm

1411 - seq_lens_cpu = common_attn_metadata.seq_lens_cpu if needs_seq_lens_cpu else None

1412 - seq_lens_np = seq_lens_cpu.numpy() if seq_lens_cpu is not None else None

1413 - num_blocks_np = (

1414 - (seq_lens_np + (page_size - 1)) // page_size

1415 - if seq_lens_np is not None

1416 - else None

1417 - )

1412 + if needs_seq_lens_cpu:

1413 + with gpu_sync_allowed():

1414 + seq_lens_cpu = common_attn_metadata.seq_lens_cpu

1415 + seq_lens_np = seq_lens_cpu.numpy()

1416 + num_blocks_np = (seq_lens_np + (page_size - 1)) // page_size

1417 + else:

1418 + seq_lens_cpu = None

1419 + seq_lens_np = None

1420 + num_blocks_np = None

1421

1422 # Adjust seq_lens_cpu for DCP

1423 if self.use_dcp:

vllm/v1/attention/ops/vit_attn_wrappers.py

+20 -4

@@ -20,6 +20,7 @@ import torch.nn.functional as F

20

21 from vllm._aiter_ops import rocm_aiter_ops

22 from vllm.platforms import current_platform

23 +from vllm.utils.gpu_sync_debug import gpu_sync_allowed

24 from vllm.utils.torch_utils import direct_register_custom_op

25

26

@@ -50,7 +51,12 @@ def flash_attn_maxseqlen_wrapper(

51 cu_seqlens = torch.arange(

52 0, (batch_size + 1) * q_len, step=q_len, dtype=torch.int32, device=q.device

53 )

53 - max_seqlen = q_len if max_seqlen is None else max_seqlen.item()

54 + if max_seqlen is None:

55 + max_seqlen = q_len

56 + else:

57 + # `flash_attn_varlen_func` needs a Python int for kernel launch bounds.

58 + with gpu_sync_allowed():

59 + max_seqlen = max_seqlen.item()

60

61 q, k, v = (einops.rearrange(x, "b s ... -> (b s) ...") for x in [q, k, v])

62 output = flash_attn_varlen_func(

@@ -159,7 +165,12 @@ def triton_attn_wrapper(

165 cu_seqlens = torch.arange(

166 0, (batch_size + 1) * q_len, step=q_len, dtype=torch.int32, device=q.device

167 )

162 - max_seqlen = q_len if max_seqlen is None else max_seqlen.item()

168 + if max_seqlen is None:

169 + max_seqlen = q_len

170 + else:

171 + # `context_attention_fwd` needs a Python int.

172 + with gpu_sync_allowed():

173 + max_seqlen = max_seqlen.item()

174

175 q, k, v = (einops.rearrange(x, "b s ... -> (b s) ...") for x in [q, k, v])

176 output = torch.empty_like(q)

@@ -261,7 +272,9 @@ def torch_sdpa_wrapper(

272

273 outputs = []

274

264 - lens = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()

275 + # `torch.split` needs Python int sizes.

276 + with gpu_sync_allowed():

277 + lens = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()

278 q_chunks = torch.split(q, lens, dim=1)

279 k_chunks = torch.split(k, lens, dim=1)

280 v_chunks = torch.split(v, lens, dim=1)

@@ -337,7 +350,10 @@ def flashinfer_wrapper(

350 batch_offsets_qko = cu_seqlens[:cu_seqlength].view(-1, 1, 1, 1)

351 batch_offsets_v = cu_seqlens[cu_seqlength:].view(-1, 1, 1, 1)

352 sequence_lengths = sequence_lengths.view(-1, 1, 1, 1)

340 - max_seqlen = max_seqlen.item()

353 + # `cudnn_batch_prefill_with_kv_cache` needs Python ints for the

354 + # max-token-per-seq bounds.

355 + with gpu_sync_allowed():

356 + max_seqlen = max_seqlen.item()

357

358 output, _ = cudnn_batch_prefill_with_kv_cache(

359 q,

vllm/v1/sample/ops/topk_topp_triton.py

+8 -6

@@ -12,6 +12,7 @@ using Pivot-based Truncation and Selection" By Park et al.

12 import torch

13

14 from vllm.triton_utils import tl, triton

15 +from vllm.utils.gpu_sync_debug import gpu_sync_allowed

16 from vllm.utils.math_utils import next_power_of_2

17 from vllm.utils.platform_utils import num_compute_units

18

@@ -919,12 +920,13 @@ def apply_top_k_top_p_triton(

920 # Cache lookup table entries on each device.

921 tables = _TRITON_TABLE_CACHE.get(logits.device)

922 if tables is None:

922 - normal_cdf_to_sigma_table = logits.new_tensor(_NORMAL_CDF_TO_SIGMA_TABLE)

923 - percentile_to_std_table = logits.new_tensor(_PERCENTILE_TO_STD_TABLE)

924 - _TRITON_TABLE_CACHE[logits.device] = (

925 - normal_cdf_to_sigma_table,

926 - percentile_to_std_table,

927 - )

923 + with gpu_sync_allowed():

924 + normal_cdf_to_sigma_table = logits.new_tensor(_NORMAL_CDF_TO_SIGMA_TABLE)

925 + percentile_to_std_table = logits.new_tensor(_PERCENTILE_TO_STD_TABLE)

926 + _TRITON_TABLE_CACHE[logits.device] = (

927 + normal_cdf_to_sigma_table,

928 + percentile_to_std_table,

929 + )

930 else:

931 normal_cdf_to_sigma_table, percentile_to_std_table = tables

932

vllm/v1/sample/sampler.py

+5 -3

@@ -6,6 +6,7 @@ import torch

6 import torch.nn as nn

7

8 from vllm.config.model import LogprobsMode

9 +from vllm.utils.gpu_sync_debug import gpu_sync_allowed

10 from vllm.utils.torch_utils import PIN_MEMORY

11 from vllm.v1.outputs import LogprobsTensors, SamplerOutput

12 from vllm.v1.sample.metadata import SamplingMetadata

@@ -342,9 +343,10 @@ class Sampler(nn.Module):

343 # of the compiled batched_count_greater_than. mark_unbacked makes

344 # the size fully symbolic so dynamo doesn't specialize when

345 # batch_size transitions from 1 to >=2.

345 - torch._dynamo.decorators.mark_unbacked(logprobs, 0)

346 - torch._dynamo.decorators.mark_unbacked(token_logprobs, 0)

347 - token_ranks = batched_count_greater_than(logprobs, token_logprobs)

346 + with gpu_sync_allowed(first_only=True):

347 + torch._dynamo.decorators.mark_unbacked(logprobs, 0)

348 + torch._dynamo.decorators.mark_unbacked(token_logprobs, 0)

349 + token_ranks = batched_count_greater_than(logprobs, token_logprobs)

350

351 # Concatenate together with the topk.

352 indices = torch.cat((token_ids, topk_indices), dim=1)

vllm/v1/worker/dp_utils.py

+30 -21

@@ -1,12 +1,15 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 +from contextlib import nullcontext

5 +

6 import torch

7 import torch.distributed as dist

8

9 from vllm.config import ParallelConfig

10 from vllm.distributed.parallel_state import get_dp_group

11 from vllm.logger import init_logger

12 +from vllm.utils.gpu_sync_debug import gpu_sync_allowed

13 from vllm.v1.worker.ubatch_utils import (

14 check_ubatch_thresholds,

15 is_last_ubatch_empty,

@@ -136,27 +139,33 @@ def _synchronize_dp_ranks(

139 parallel_config=parallel_config,

140 )

141

139 - # Synchronize cudagraph_mode across ranks first (take min).

140 - # This is needed before DP padding decision since we use the synced

141 - # cudagraph mode to determine whether DP padding is needed.

142 - synced_cudagraph_mode = _post_process_cudagraph_mode(tensor)

143 -

144 - # Check conditions for microbatching

145 - should_ubatch = _post_process_ubatch(tensor, parallel_config.num_ubatches)

146 -

147 - # DP padding is needed when cudagraph is enabled (synced across ranks)

148 - # or when ubatching/DBO is active (ubatching requires uniform batch

149 - # sizes across DP ranks currently).

150 - # Use the synced runtime cudagraph mode rather than the compilation config

151 - # so we can avoid padding when cudagraph is not enabled for this step.

152 - should_dp_pad = synced_cudagraph_mode != 0 or should_ubatch

153 -

154 - # Pad all DP ranks up to the maximum token count across ranks if

155 - # should_dp_pad is True

156 - num_tokens_after_padding = _post_process_dp_padding(

157 - tensor,

158 - should_dp_pad,

159 - )

142 + # Only the NCCL path leaves `tensor` on device. With Gloo -- the default

143 + # under async scheduling -- the all-reduce runs on CPU, so the reads below

144 + # are host-side and the check should stay armed.

145 + with (

146 + nullcontext()

147 + if parallel_config.disable_nccl_for_dp_synchronization

148 + else gpu_sync_allowed()

149 + ):

150 + # Synchronize cudagraph_mode across ranks first (take min).

151 + # This is needed before DP padding decision since we use the synced

152 + # cudagraph mode to determine whether DP padding is needed.

153 + synced_cudagraph_mode = _post_process_cudagraph_mode(tensor)

154 +

155 + # Check conditions for microbatching

156 + should_ubatch = _post_process_ubatch(tensor, parallel_config.num_ubatches)

157 +

158 + # DP padding is needed when cudagraph is enabled (synced across ranks)

159 + # or when ubatching/DBO is active (ubatching requires uniform batch

160 + # sizes across DP ranks currently).

161 + # Use the synced runtime cudagraph mode rather than the compilation

162 + # config so we can avoid padding when cudagraph is not enabled for

163 + # this step.

164 + should_dp_pad = synced_cudagraph_mode != 0 or should_ubatch

165 +

166 + # Pad all DP ranks up to the maximum token count across ranks if

167 + # should_dp_pad is True

168 + num_tokens_after_padding = _post_process_dp_padding(tensor, should_dp_pad)

169

170 return should_ubatch, num_tokens_after_padding, synced_cudagraph_mode

171

vllm/v1/worker/encoder_cudagraph.py

+17 -8

@@ -19,6 +19,7 @@ from vllm.model_executor.models.interfaces import (

19 )

20 from vllm.model_executor.models.utils import scatter_output_slices

21 from vllm.model_executor.models.vision import get_load_balance_assignment

22 +from vllm.utils.gpu_sync_debug import gpu_sync_allowed

23 from vllm.v1.worker.encoder_cudagraph_defs import (

24 EncoderCudaGraphConfig,

25 EncoderItemSpec,

@@ -298,7 +299,19 @@ class EncoderCudaGraphManager:

299

300 def _get_item_specs(self, mm_kwargs: dict[str, Any]) -> list[EncoderItemSpec]:

301 """Get item specs from the model."""

301 - return self.model.get_encoder_cudagraph_item_specs(mm_kwargs)

302 + # Implementations read per-item grid/patch counts off device tensors

303 + # to size the cudagraph buffers, so the D2H is inherent here.

304 + with gpu_sync_allowed():

305 + return self.model.get_encoder_cudagraph_item_specs(mm_kwargs)

306 +

307 + def _select_items(

308 + self, mm_kwargs: dict[str, Any], indices: list[int]

309 + ) -> dict[str, Any]:

310 + """Select the mm kwargs for `indices` from the model."""

311 + # Same as `_get_item_specs`: implementations re-read the per-item

312 + # grid/patch counts to slice the batch, so the D2H is inherent.

313 + with gpu_sync_allowed():

314 + return self.model.select_encoder_cudagraph_items(mm_kwargs, indices)

315

316 def _get_per_item_out_tokens(self, mm_kwargs: dict[str, Any]) -> list[int]:

317 """Get per-item output token counts as plain ints."""

@@ -415,9 +428,7 @@ class EncoderCudaGraphManager:

428

429 outputs_by_orig_idx: dict[int, torch.Tensor] = {}

430 for batch_indices, path_budgets in batches:

418 - batch_mm_kwargs = self.model.select_encoder_cudagraph_items(

419 - mm_kwargs, batch_indices

420 - )

431 + batch_mm_kwargs = self._select_items(mm_kwargs, batch_indices)

432 graph_outputs: dict[str, torch.Tensor] = {}

433 all_eager = True

434

@@ -491,11 +502,9 @@ class EncoderCudaGraphManager:

502 ]

503

504 if len(local_indices) > 0:

494 - local_mm_kwargs = self.model.select_encoder_cudagraph_items(

495 - mm_kwargs, local_indices

496 - )

505 + local_mm_kwargs = self._select_items(mm_kwargs, local_indices)

506 else:

498 - local_mm_kwargs = self.model.select_encoder_cudagraph_items(mm_kwargs, [])

507 + local_mm_kwargs = self._select_items(mm_kwargs, [])

508

509 max_output_tokens_per_rank = (

510 max(

vllm/v1/worker/gpu_model_runner.py

+27 -21

@@ -120,6 +120,7 @@ from vllm.sequence import IntermediateTensors

120 from vllm.tasks import GenerationTask, PoolingTask, SupportedTask

121 from vllm.tracing import instrument

122 from vllm.utils import length_from_prompt_token_ids_or_embeds

123 +from vllm.utils.gpu_sync_debug import gpu_sync_allowed

124 from vllm.utils.math_utils import cdiv, round_up

125 from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib

126 from vllm.utils.nvtx_pytorch_hooks import PytHooks

@@ -3848,26 +3849,28 @@ class GPUModelRunner(

3849 self.routed_experts_slot_mapping_device[:total],

3850 non_blocking=True,

3851 )

3851 -

3852 - # Get the valid generated tokens.

3853 - max_gen_len = sampled_token_ids.shape[-1]

3854 - if max_gen_len == 1:

3855 - # No spec decode tokens.

3856 - valid_sampled_token_ids = self._to_list(sampled_token_ids)

3857 - # Mask out the sampled tokens that should not be sampled.

3858 - for i in discard_sampled_tokens_req_indices:

3859 - valid_sampled_token_ids[int(i)].clear()

3860 -

3861 - if logprobs_tensors is not None:

3862 - logprobs_lists = logprobs_tensors.tolists()

3863 - else:

3864 - # Includes spec decode tokens.

3865 - valid_sampled_token_ids, logprobs_lists = RejectionSampler.parse_output(

3866 - sampled_token_ids,

3867 - self.input_batch.vocab_size,

3868 - discard_sampled_tokens_req_indices,

3869 - logprobs_tensors=logprobs_tensors,

3870 - )

3852 + with gpu_sync_allowed():

3853 + # Get the valid generated tokens.

3854 + max_gen_len = sampled_token_ids.shape[-1]

3855 + if max_gen_len == 1:

3856 + # No spec decode tokens.

3857 + valid_sampled_token_ids = self._to_list(sampled_token_ids)

3858 + # Mask out the sampled tokens that should not be sampled.

3859 + for i in discard_sampled_tokens_req_indices:

3860 + valid_sampled_token_ids[int(i)].clear()

3861 +

3862 + if logprobs_tensors is not None:

3863 + logprobs_lists = logprobs_tensors.tolists()

3864 + else:

3865 + # Includes spec decode tokens.

3866 + valid_sampled_token_ids, logprobs_lists = (

3867 + RejectionSampler.parse_output(

3868 + sampled_token_ids,

3869 + self.input_batch.vocab_size,

3870 + discard_sampled_tokens_req_indices,

3871 + logprobs_tensors=logprobs_tensors,

3872 + )

3873 + )

3874 else:

3875 valid_sampled_token_ids = []

3876 invalid_req_indices = discard_sampled_tokens_req_indices.tolist()

@@ -5843,7 +5846,10 @@ class GPUModelRunner(

5846 on device instead; see`AsyncGPUModelRunnerOutput`.

5847 """

5848 try:

5846 - counts = [] if logits is None else count_nans_per_row(logits).tolist()

5849 + # Reporting per-request NaN counts requires them on the host; this

5850 + # path is opt-in diagnostics, so the D2H is intended.

5851 + with gpu_sync_allowed():

5852 + counts = [] if logits is None else count_nans_per_row(logits).tolist()

5853 num_nans_in_logits = nans_to_dict(counts, self.input_batch.req_id_to_index)

5854 if envs.VLLM_RAISE_ON_LOGIT_NANS:

5855 raise_if_nan_logits(num_nans_in_logits)

vllm/v1/worker/mamba_utils.py

+16

@@ -16,6 +16,7 @@ from vllm.model_executor.layers.mamba.mamba_utils import (

16 is_conv_state_dim_first,

17 )

18 from vllm.triton_utils import tl, triton

19 +from vllm.utils.gpu_sync_debug import gpu_sync_allowed

20 from vllm.utils.math_utils import cdiv

21 from vllm.v1.core.sched.output import SchedulerOutput

22 from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec

@@ -754,7 +755,22 @@ class MambaSpecDecodeGPUContext:

755 """

756 if self.is_initialized:

757 return

758 + # This only runs once per worker.

759 + with gpu_sync_allowed():

760 + self._populate_metadata(

761 + kv_cache_config,

762 + forward_context,

763 + mamba_state_copy_funcs,

764 + block_tables,

765 + )

766

767 + def _populate_metadata(

768 + self,

769 + kv_cache_config: KVCacheConfig,

770 + forward_context: dict[str, Any],

771 + mamba_state_copy_funcs: tuple[MambaStateCopyFunc, ...],

772 + block_tables: list[torch.Tensor],

773 + ) -> None:

774 idx = 0

775 for group_local_idx, mamba_group_id in enumerate(self.mamba_group_ids):

776 layer_names = kv_cache_config.kv_cache_groups[mamba_group_id].layer_names

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