@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 4f6eed3bd4

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 5, 16:00 UTC

[Core] Simplify multimodal masking (#34246)

Signed-off-by: Lukas Geiger <lukas.geiger94@gmail.com>

Lukas Geiger committed Apr 1, 2026 at 09:18 UTC 4f6eed3bd4a92c6bd513460ee85b917d6df88a17

9 files changed +54 -51

tests/models/test_utils.py

+29 -3

@@ -4,9 +4,11 @@

4 import pytest

5 import torch

7 -from vllm.model_executor.models.utils import AutoWeightsLoader

8 -

9 -pytestmark = pytest.mark.cpu_test

7 +from vllm.model_executor.models.utils import (

8 + AutoWeightsLoader,

9 + _merge_multimodal_embeddings,

10 +)

11 +from vllm.platforms import current_platform

14 class ModuleWithBatchNorm(torch.nn.Module):

@@ -27,6 +29,7 @@ class ModuleWithNestedBatchNorm(torch.nn.Module):

29 return self.nested_mod(x)

32 +@pytest.mark.cpu_test

33 def test_module_with_batchnorm_can_load():

34 """Ensure the auto weight loader can load batchnorm stats."""

35 mod = ModuleWithBatchNorm()

@@ -52,6 +55,7 @@ def test_module_with_batchnorm_can_load():

55 assert new_mod.bn.num_batches_tracked.item() == 1

58 +@pytest.mark.cpu_test

59 def test_module_with_child_containing_batchnorm_can_autoload():

60 """Ensure the auto weight loader can load nested modules batchnorm stats."""

61 mod = ModuleWithNestedBatchNorm()

@@ -83,6 +87,7 @@ def test_module_with_child_containing_batchnorm_can_autoload():

87 assert new_mod.nested_mod.bn.num_batches_tracked.item() == 1

90 +@pytest.mark.cpu_test

91 def test_module_skip_prefix():

92 """Ensure the auto weight loader can skip prefix."""

93 mod = ModuleWithNestedBatchNorm()

@@ -119,6 +124,7 @@ def test_module_skip_prefix():

124 assert new_mod.nested_mod.bn.num_batches_tracked.item() == 1

125

126

127 +@pytest.mark.cpu_test

128 def test_module_skip_substr():

129 """Ensure the auto weight loader can skip prefix."""

130 mod = ModuleWithNestedBatchNorm()

@@ -155,3 +161,23 @@ def test_module_skip_substr():

161 )

162 assert torch.all(new_mod.nested_mod.bn.running_var == mod.nested_mod.bn.running_var)

163 assert new_mod.nested_mod.bn.num_batches_tracked.item() == 1

164 +

165 +

166 +class raise_if_cuda_sync:

167 + def __enter__(self):

168 + self.previous_debug_mode = torch.cuda.get_sync_debug_mode()

169 + torch.cuda.set_sync_debug_mode("error")

170 +

171 + def __exit__(self, exception_type, exception_value, traceback):

172 + torch.cuda.set_sync_debug_mode(self.previous_debug_mode)

173 +

174 +

175 +@pytest.mark.skipif(not current_platform.is_cuda(), reason="Skip if not cuda")

176 +def test_merge_multimodal_embeddings_no_sync():

177 + inputs_embeds = torch.zeros([5, 10], dtype=torch.bfloat16, device="cuda:0")

178 + multimodal_embeddings = [torch.ones([3, 10], dtype=torch.bfloat16, device="cuda:0")]

179 + is_multimodal = torch.tensor([True, False, True, True, False], device="cpu")

180 + with raise_if_cuda_sync():

181 + _merge_multimodal_embeddings(

182 + inputs_embeds, multimodal_embeddings, is_multimodal

183 + )

vllm/model_executor/models/interfaces.py

+3 -1

@@ -362,7 +362,9 @@ class SupportsMultiModal(Protocol):

362 # to ensure that any external configuration requiring offset tracking,

363 # e.g., LoRA, are applied correctly regardless of whether or not

364 # we have multimodal tokens.

365 - in_vocab_ids = input_ids.masked_fill(is_multimodal, 0)

365 + in_vocab_ids = input_ids.masked_fill(

366 + is_multimodal.to(device=input_ids.device, non_blocking=True), 0

367 + )

368 return embed_input_ids(in_vocab_ids)

369

370 return embed_input_ids(input_ids)

vllm/model_executor/models/nano_nemotron_vl.py

+2 -3

@@ -1215,7 +1215,6 @@ class NemotronH_Nano_VL_V2(

1215 These embeddings will replace the placeholder embeddings to create

1216 input_embeds for the LLM.

1217 """

1218 - device = video_embeddings.device

1218 tokenizer = cached_tokenizer_from_config(self.model_config)

1219

1220 # Generate video replacement token IDs using get_video_repl

@@ -1234,10 +1233,10 @@ class NemotronH_Nano_VL_V2(

1233 )

1234

1235 # video_repl.full is a list of token IDs

1237 - repl_token_ids = torch.tensor(video_repl.full, device=device)

1236 + repl_token_ids = torch.tensor(video_repl.full)

1237

1238 # Get embedding token IDs for image context (use pre-tokenized version)

1240 - embed_token_ids = torch.tensor(self._img_context_token_ids, device=device)

1239 + embed_token_ids = torch.tensor(self._img_context_token_ids)

1240

1241 # Create mask for video embedding positions

1242 is_video_embed = torch.isin(repl_token_ids, embed_token_ids)

vllm/model_executor/models/qwen2_5_omni_thinker.py

+6 -8

@@ -211,15 +211,12 @@ def merge_interleaved_embeddings(

211

212 # Scatter each modality to its positions

213 if video_embeds:

214 - video_positions = is_video.nonzero(as_tuple=True)[0]

215 - inputs_embeds[video_positions] = torch.cat(video_embeds, dim=0)

214 + inputs_embeds[is_video] = torch.cat(video_embeds, dim=0)

215 if audio_embeds:

217 - audio_positions = is_audio.nonzero(as_tuple=True)[0]

218 - inputs_embeds[audio_positions] = torch.cat(audio_embeds, dim=0)

216 + inputs_embeds[is_audio] = torch.cat(audio_embeds, dim=0)

217 if other_embeds:

218 other_mask = is_multimodal & ~is_video & ~is_audio

221 - other_positions = other_mask.nonzero(as_tuple=True)[0]

222 - inputs_embeds[other_positions] = torch.cat(other_embeds, dim=0)

219 + inputs_embeds[other_mask] = torch.cat(other_embeds, dim=0)

220

221 return inputs_embeds

222

@@ -1457,8 +1454,9 @@ class Qwen2_5OmniThinkerForConditionalGeneration(

1454 video_token_id = self.config.video_token_index

1455 audio_token_id = self.config.audio_token_index

1456

1460 - is_video = is_multimodal & (input_ids == video_token_id)

1461 - is_audio = is_multimodal & (input_ids == audio_token_id)

1457 + input_ids_cpu = input_ids.cpu()

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

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

1460

1461 num_video = is_video.sum().item()

1462 num_audio = is_audio.sum().item()

vllm/model_executor/models/qwen3_omni_moe_thinker.py

+3 -2

@@ -1869,8 +1869,9 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(

1869 # both the deepstack path and the final embedding merge.

1870 video_token_id = self.config.video_token_id

1871 audio_token_id = self.config.audio_token_id

1872 - is_video = is_multimodal & (input_ids == video_token_id)

1873 - is_audio = is_multimodal & (input_ids == audio_token_id)

1872 + input_ids_cpu = input_ids.cpu()

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

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

1875 num_video = is_video.sum().item()

1876 num_audio = is_audio.sum().item()

1877

vllm/model_executor/models/qwen3_vl.py

+4 -3

@@ -1977,7 +1977,6 @@ class Qwen3VLForConditionalGeneration(

1977 These embeddings will replace the placeholder embeddings to create

1978 input_embeds for the LLM.

1979 """

1980 - device = video_embeddings.device

1980

1981 # Generate video replacement token IDs using get_video_repl

1982 # This tokenizes each frame separator independently, then uses pre-tokenized

@@ -1993,8 +1992,10 @@ class Qwen3VLForConditionalGeneration(

1992 select_token_id=self.is_multimodal_pruning_enabled,

1993 )

1994

1996 - repl_token_ids = torch.tensor(video_repl.full, device=device)

1997 - embed_token_id = _cached_tensor(self.config.video_token_id, device=device)

1995 + repl_token_ids = torch.tensor(video_repl.full)

1996 + embed_token_id = _cached_tensor(

1997 + self.config.video_token_id, repl_token_ids.device

1998 + )

1999 is_video_embed = torch.isin(repl_token_ids, embed_token_id)

2000

2001 # Get text embeddings for indicator tokens (has only `visual_dim``).

vllm/model_executor/models/utils.py

+3 -9

@@ -468,14 +468,8 @@ def _merge_multimodal_embeddings(

468 input_dtype = inputs_embeds.dtype

469

470 try:

471 - # For debugging

472 - # inputs_embeds[is_multimodal] = mm_embeds_flat.to(dtype=input_dtype)

473 -

474 - # NOTE: This can avoid D2H sync (#22105), but fails to

475 - # raise an error if is_multimodal.sum() < len(mm_embeds_flat)

476 - inputs_embeds.masked_scatter_(

477 - is_multimodal.unsqueeze(-1), mm_embeds_flat.to(dtype=input_dtype)

478 - )

471 + # If is_multimodal is on CPU this avoids a D2H sync

472 + inputs_embeds[is_multimodal] = mm_embeds_flat.to(dtype=input_dtype)

473 except RuntimeError as e:

474 num_actual_tokens = len(mm_embeds_flat)

475 num_expected_tokens = is_multimodal.sum().item()

@@ -488,7 +482,7 @@ def _merge_multimodal_embeddings(

482 f"multimodal tokens to {num_expected_tokens} placeholders"

483 ) from e

484

491 - raise ValueError("Error during masked scatter operation") from e

485 + raise ValueError("Error during index put operation") from e

486

487 return inputs_embeds

488

vllm/v1/worker/gpu/mm/encoder_runner.py

+1 -3

@@ -83,7 +83,7 @@ class EncoderRunner:

84 mm_embeds: list[torch.Tensor] = []

85 is_mm_embed = torch.zeros(

86 - total_num_scheduled_tokens, dtype=torch.bool, device="cpu", pin_memory=True

86 + total_num_scheduled_tokens, dtype=torch.bool, device="cpu"

87 )

88 for i, req_id in enumerate(req_ids):

89 if not is_prefilling[i]:

@@ -131,8 +131,6 @@ class EncoderRunner:

131 )

132 mm_embeds.append(mm_embeds_item)

133

134 - # Copy the is_mm_embed tensor to the GPU.

135 - is_mm_embed = is_mm_embed.to(device=self.device, non_blocking=True)

134 return mm_embeds, is_mm_embed

135

136 @torch.inference_mode()

vllm/v1/worker/gpu_model_runner.py

+3 -19

@@ -719,16 +719,6 @@ class GPUModelRunner(

719 self.max_num_reqs, dtype=torch.int32

720 )

721

722 - # Only relevant for multimodal models

723 - if self.supports_mm_inputs:

724 - # Double buffer to avoid race condition: previous iteration's async

725 - # copy may still be reading from CPU while current iteration writes.

726 - self.is_mm_embed_buffers = [

727 - self._make_buffer(self.max_num_tokens, dtype=torch.bool),

728 - self._make_buffer(self.max_num_tokens, dtype=torch.bool),

729 - ]

730 - self.is_mm_embed_idx = 0

731 -

722 # Only relevant for models using M-RoPE (e.g, Qwen2-VL)

723 if self.uses_mrope:

724 # NOTE: `mrope_positions` is implemented with one additional dummy

@@ -2910,14 +2900,10 @@ class GPUModelRunner(

2900 ) -> tuple[list[torch.Tensor], torch.Tensor]:

2901 total_num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens

2902

2913 - # Swap to the other buffer to avoid race condition with previous

2914 - # iteration's async copy that may still be reading from CPU.

2915 - self.is_mm_embed_idx = 1 - self.is_mm_embed_idx

2916 - is_mm_embed_buf = self.is_mm_embed_buffers[self.is_mm_embed_idx]

2917 -

2903 mm_embeds = list[torch.Tensor]()

2919 - is_mm_embed = is_mm_embed_buf.cpu

2920 - is_mm_embed[:total_num_scheduled_tokens] = False

2904 + is_mm_embed = torch.zeros(

2905 + total_num_scheduled_tokens, dtype=torch.bool, device="cpu"

2906 + )

2907

2908 req_start_idx = 0

2909 should_sync_mrope_positions = False

@@ -3000,8 +2986,6 @@ class GPUModelRunner(

2986 mm_embeds.extend(mm_embeds_req)

2987 req_start_idx += num_scheduled_tokens

2988

3003 - is_mm_embed = is_mm_embed_buf.copy_to_gpu(total_num_scheduled_tokens)

3004 -

2989 if should_sync_mrope_positions:

2990 self._calc_mrope_positions(scheduler_output)

2991 self.mrope_positions.copy_to_gpu(total_num_scheduled_tokens)