@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 84c62e1cbd

[Model Runner V2][MM] Support EVS (#46535)

Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Nick Hill committed Jun 24, 2026 at 07:18 UTC 84c62e1cbdef4250fbfda83782fd250e07ad0256

10 files changed +245 -50

vllm/model_executor/models/diffusion_gemma.py

+9 -9

@@ -49,10 +49,12 @@ from vllm.multimodal import MULTIMODAL_REGISTRY

49 from vllm.v1.outputs import LogprobsTensors

50 from vllm.v1.worker.gpu.attn_utils import build_attn_metadata

51 from vllm.v1.worker.gpu.buffer_utils import UvaBackedTensor, async_copy_to_gpu

52 +from vllm.v1.worker.gpu.input_batch import InputBatch

53 from vllm.v1.worker.gpu.model_states.interface import ModelState

54 from vllm.v1.worker.gpu.sample.logprob import compute_topk_logprobs

55 from vllm.v1.worker.gpu.sample.output import SamplerOutput

56 from vllm.v1.worker.gpu.sample.penalties import use_penalty

57 +from vllm.v1.worker.gpu.states import RequestState

59 from .interfaces import (

60 SupportsMultiModal,

@@ -869,7 +871,12 @@ class DiffusionGemmaModelState(ModelState):

871 if idx is not None:

872 self.diffusion_states.remove_request(idx)

873

872 - def get_mm_embeddings(self, scheduled_encoder_inputs, input_batch):

874 + def get_mm_embeddings(

875 + self,

876 + scheduled_encoder_inputs: dict[str, list[int]],

877 + input_batch: InputBatch,

878 + req_states: RequestState,

879 + ) -> torch.Tensor:

880 if not self.supports_mm_inputs:

881 return None

882

@@ -880,14 +887,7 @@ class DiffusionGemmaModelState(ModelState):

887 encoder_outputs = self.encoder_runner.execute_mm_encoder(mm_kwargs)

888 self.encoder_cache.encoder_outputs.update(zip(mm_hashes, encoder_outputs))

889

883 - mm_embeds, is_mm_embed = self.encoder_runner.gather_mm_embeddings(

884 - input_batch.req_ids,

885 - input_batch.num_tokens,

886 - input_batch.num_scheduled_tokens,

887 - input_batch.query_start_loc_np,

888 - input_batch.prefill_len_np,

889 - input_batch.num_computed_prefill_tokens_np,

890 - )

890 + mm_embeds, is_mm_embed = self.gather_mm_embeddings(input_batch)

891

892 if not mm_embeds:

893 # No MM tokens in this batch (e.g. all-decode step).

vllm/model_executor/models/interfaces.py

+6 -5

@@ -421,11 +421,11 @@ class SupportsMultiModalPruning(Protocol):

421

422 def recompute_mrope_positions(

423 self,

424 - input_ids: list[int],

425 - multimodal_embeddings: MultiModalEmbeddings,

424 + input_ids: list[int] | torch.Tensor,

425 + multimodal_embeddings: Sequence[torch.Tensor],

426 mrope_positions: torch.LongTensor,

427 num_computed_tokens: int,

428 - ) -> tuple[MultiModalEmbeddings, Tensor, int]:

428 + ) -> tuple[Sequence[torch.Tensor], Tensor, int]:

429 """

430 Update part of input mrope positions (starting with

431 num_computed_tokens index). Original mrope_positions are computed

@@ -435,8 +435,9 @@ class SupportsMultiModalPruning(Protocol):

435

436 Args:

437 input_ids: (N,) All input tokens of the prompt containing

438 - entire sequence.

439 - multimodal_embeddings: Tuple of multimodal embeddings that

438 + entire sequence. Either a host-side list or an already

439 + device-resident tensor.

440 + multimodal_embeddings: Sequence of multimodal embeddings that

441 fits into the prefill chunk that is being processed.

442 mrope_positions: Existing mrope positions (3, N) for entire

443 sequence

vllm/model_executor/models/qwen2_5_vl.py

+10 -6

@@ -1583,11 +1583,11 @@ class Qwen2_5_VLForConditionalGeneration(

1583

1584 def recompute_mrope_positions(

1585 self,

1586 - input_ids: list[int],

1587 - multimodal_embeddings: tuple[torch.Tensor, ...],

1586 + input_ids: list[int] | torch.Tensor,

1587 + multimodal_embeddings: Sequence[torch.Tensor],

1588 mrope_positions: torch.LongTensor,

1589 num_computed_tokens: int,

1590 - ) -> tuple[tuple[torch.Tensor, ...], torch.Tensor, int]:

1590 + ) -> tuple[Sequence[torch.Tensor], torch.Tensor, int]:

1591 """

1592 Update part of input mrope positions (starting with

1593 num_computed_tokens index). Original mrope_positions are computed

@@ -1618,8 +1618,12 @@ class Qwen2_5_VLForConditionalGeneration(

1618 else mrope_positions.device

1619 )

1620

1621 - # Tensors.

1622 - input_ids_t = async_tensor_h2d(input_ids, dtype=torch.long, device=device)

1621 + # Tensors. input_ids may already be a (device-side) tensor.

1622 + if isinstance(input_ids, torch.Tensor):

1623 + assert input_ids.device == device

1624 + input_ids_t = input_ids.to(torch.long)

1625 + else:

1626 + input_ids_t = async_tensor_h2d(input_ids, dtype=torch.long, device=device)

1627

1628 mm_embeddings_out = [mm[:, :-4] for mm in multimodal_embeddings]

1629 mm_embeddings_pos = [

@@ -1636,7 +1640,7 @@ class Qwen2_5_VLForConditionalGeneration(

1640 video_token_id,

1641 )

1642

1639 - return tuple(mm_embeddings_out), positions, mrope_positions_delta

1643 + return mm_embeddings_out, positions, mrope_positions_delta

1644

1645 def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict:

1646 mm_input_by_modality = {}

vllm/model_executor/models/qwen3_vl.py

+14 -10

@@ -2657,11 +2657,11 @@ class Qwen3VLForConditionalGeneration(

2657

2658 def recompute_mrope_positions(

2659 self,

2660 - input_ids: list[int],

2661 - multimodal_embeddings: MultiModalEmbeddings,

2660 + input_ids: list[int] | torch.Tensor,

2661 + multimodal_embeddings: Sequence[torch.Tensor],

2662 mrope_positions: torch.LongTensor,

2663 num_computed_tokens: int,

2664 - ) -> tuple[MultiModalEmbeddings, torch.Tensor, int]:

2664 + ) -> tuple[Sequence[torch.Tensor], torch.Tensor, int]:

2665 """

2666 Update part of input mrope positions (starting with

2667 num_computed_tokens index). Original mrope_positions are computed

@@ -2672,7 +2672,7 @@ class Qwen3VLForConditionalGeneration(

2672 Args:

2673 input_ids: (N,) All input tokens of the prompt containing

2674 entire sequence.

2675 - multimodal_embeddings: Tuple of multimodal embeddings that

2675 + multimodal_embeddings: Sequence of multimodal embeddings that

2676 fits into the prefill chunk that is being processed.

2677 mrope_positions: Existing mrope positions (3, N) for entire

2678 sequence

@@ -2694,14 +2694,14 @@ class Qwen3VLForConditionalGeneration(

2694

2695 @staticmethod

2696 def _recompute_mrope_positions(

2697 - input_ids: list[int],

2698 - multimodal_embeddings: MultiModalEmbeddings,

2697 + input_ids: list[int] | torch.Tensor,

2698 + multimodal_embeddings: Sequence[torch.Tensor],

2699 mrope_positions: torch.LongTensor,

2700 num_computed_tokens: int,

2701 vision_start_token_id: int,

2702 image_token_id: int,

2703 video_token_id: int,

2704 - ) -> tuple[MultiModalEmbeddings, torch.Tensor, int]:

2704 + ) -> tuple[Sequence[torch.Tensor], torch.Tensor, int]:

2705 # Device

2706 device = (

2707 multimodal_embeddings[0].device

@@ -2709,8 +2709,12 @@ class Qwen3VLForConditionalGeneration(

2709 else mrope_positions.device

2710 )

2711

2712 - # Tensors

2713 - input_ids_t = async_tensor_h2d(input_ids, device=device, dtype=torch.long)

2712 + # Tensors. input_ids may already be a (device-side) tensor.

2713 + if isinstance(input_ids, torch.Tensor):

2714 + assert input_ids.device == device

2715 + input_ids_t = input_ids.to(dtype=torch.long)

2716 + else:

2717 + input_ids_t = async_tensor_h2d(input_ids, device=device, dtype=torch.long)

2718

2719 mm_embeddings_out = []

2720 mm_embeddings_pos = []

@@ -2738,7 +2742,7 @@ class Qwen3VLForConditionalGeneration(

2742 video_token_id,

2743 )

2744

2741 - return tuple(mm_embeddings_out), positions, mrope_positions_delta

2745 + return mm_embeddings_out, positions, mrope_positions_delta

2746

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

2748 mm_input_by_modality = self._parse_and_validate_multimodal_inputs(**kwargs)

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

+17

@@ -90,6 +90,23 @@ class RopeState:

90 def get_positions(self, num_tokens: int) -> torch.Tensor:

91 return self.positions[:, :num_tokens]

93 + def read_prefill_positions(self, req_idx: int, length: int) -> torch.Tensor:

94 + """Return staged per-request prefill positions as [num_dims, length]."""

95 + base = self.num_dims * req_idx

96 + return self.prefill_positions.gpu[base : base + self.num_dims, :length]

97 +

98 + def update_prefill_positions(

99 + self, req_idx: int, positions: torch.Tensor, delta: int

100 + ) -> None:

101 + """Overwrite a request's staged prefill positions with recomputed values."""

102 + base = self.num_dims * req_idx

103 + length = positions.shape[1]

104 + self.prefill_positions.gpu[base : base + self.num_dims, :length].copy_(

105 + positions

106 + )

107 + if self.has_delta:

108 + self.prefill_delta.np[req_idx] = delta

109 +

110 def prepare_positions(

111 self,

112 idx_mapping: torch.Tensor,

vllm/v1/worker/gpu/model_runner.py

+4 -10

@@ -1229,7 +1229,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1229 scheduled_encoder_inputs=scheduler_output.scheduled_encoder_inputs,

1230 )

1231 inputs_embeds = self.model_state.get_mm_embeddings(

1232 - scheduler_output.scheduled_encoder_inputs, input_batch

1232 + scheduler_output.scheduled_encoder_inputs, input_batch, self.req_states

1233 )

1234 if inputs_embeds is not None and not self.model.requires_raw_input_tokens:

1235 input_ids = None

@@ -1413,15 +1413,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1413 # Get cached multimodal embeddings for draft forward.

1414 # NOTE: This is done here because postprocess updates

1415 # num_computed_prefill_tokens.

1416 - mm_inputs = self.model_state.encoder_runner.gather_mm_embeddings(

1417 - input_batch.req_ids,

1418 - input_batch.num_tokens,

1419 - input_batch.num_scheduled_tokens,

1420 - input_batch.query_start_loc_np,

1421 - input_batch.prefill_len_np,

1422 - input_batch.num_computed_prefill_tokens_np,

1423 - # The EAGLE/MTP drafter reads one position ahead of the target.

1424 - draft_lookahead=1,

1416 + # The EAGLE/MTP drafter reads one position ahead of the target.

1417 + mm_inputs = self.model_state.gather_mm_embeddings(

1418 + input_batch, draft_lookahead=1

1419 )

1420

1421 # Postprocess results and update request states.

vllm/v1/worker/gpu/model_states/default.py

+25 -8

@@ -15,6 +15,7 @@ from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache

15 from vllm.v1.worker.gpu.mm.encoder_runner import EncoderRunner

16 from vllm.v1.worker.gpu.mm.rope import get_rope_state

17 from vllm.v1.worker.gpu.model_states.interface import ModelState

18 +from vllm.v1.worker.gpu.model_states.mm_pruning import maybe_create_mm_pruner

19 from vllm.v1.worker.gpu.states import RequestState

20 from vllm.v1.worker.utils import AttentionGroup

@@ -61,6 +62,11 @@ class DefaultModelState(ModelState):

62 device=self.device,

63 )

65 + # Pruner is used for multimodal embedding pruning (EVS).

66 + self.mm_pruner = maybe_create_mm_pruner(

67 + self.model_config, model, self.rope_state, encoder_cache

68 + )

69 +

70 def add_request(self, req_index: int, new_req_data: NewRequestData) -> None:

71 if self.rope_state is not None:

72 assert new_req_data.prefill_token_ids is not None

@@ -79,6 +85,7 @@ class DefaultModelState(ModelState):

85 self,

86 scheduled_encoder_inputs: dict[str, list[int]],

87 input_batch: InputBatch,

88 + req_states: RequestState,

89 ) -> torch.Tensor:

90 mm_hashes, mm_kwargs = self.encoder_runner.prepare_mm_inputs(

91 scheduled_encoder_inputs

@@ -89,14 +96,13 @@ class DefaultModelState(ModelState):

96 # Cache the encoder outputs by mm_hash

97 self.encoder_cache.encoder_outputs.update(zip(mm_hashes, encoder_outputs))

92 - mm_embeds, is_mm_embed = self.encoder_runner.gather_mm_embeddings(

93 - input_batch.req_ids,

94 - input_batch.num_tokens,

95 - input_batch.num_scheduled_tokens,

96 - input_batch.query_start_loc_np,

97 - input_batch.prefill_len_np,

98 - input_batch.num_computed_prefill_tokens_np,

99 - )

99 + mm_embeds, is_mm_embed = super().gather_mm_embeddings(input_batch)

100 + if self.mm_pruner is not None and mm_embeds:

101 + # EVS: recompute mrope positions for pruned media.

102 + mm_embeds = self.mm_pruner.recompute(mm_embeds, input_batch, req_states)

103 + # We must flush the staged rope updates for prepare_inputs() to pick up.

104 + self.apply_staged_writes()

105 +

106 # Use unpadded input_ids to match is_mm_embed size (num_tokens).

107 # input_batch.input_ids may be padded for CUDA graphs.

108 input_ids_unpadded = input_batch.input_ids[: input_batch.num_tokens]

@@ -105,6 +111,17 @@ class DefaultModelState(ModelState):

111 )

112 return inputs_embeds[: input_batch.num_tokens_after_padding]

113

114 + def gather_mm_embeddings(

115 + self, input_batch: InputBatch, draft_lookahead: int = 0

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

117 + mm_embeds, is_mm_embed = super().gather_mm_embeddings(

118 + input_batch, draft_lookahead

119 + )

120 + if self.mm_pruner is not None:

121 + # EVS: strip the appended mrope-position channels.

122 + mm_embeds = self.mm_pruner.strip(mm_embeds)

123 + return mm_embeds, is_mm_embed

124 +

125 def prepare_inputs(

126 self, input_batch: InputBatch, req_states: RequestState

127 ) -> dict[str, torch.Tensor | None]:

vllm/v1/worker/gpu/model_states/encoder_decoder.py

+4 -1

@@ -85,7 +85,10 @@ class EncoderDecoderModelState(ModelState):

85 self.encoder_outputs: list[torch.Tensor] = []

87 def get_mm_embeddings(

88 - self, scheduled_encoder_inputs: dict[str, list[int]], input_batch: InputBatch

88 + self,

89 + scheduled_encoder_inputs: dict[str, list[int]],

90 + input_batch: InputBatch,

91 + req_states: RequestState,

92 ) -> None:

93 # Ensure encoder inputs are ordered consistently with input_batch.req_ids.

94 encoder_inputs: dict[str, list[int]] = {}

vllm/v1/worker/gpu/model_states/interface.py

+21 -1

@@ -13,6 +13,7 @@ from vllm.v1.core.sched.output import NewRequestData

13 from vllm.v1.kv_cache_interface import KVCacheConfig

14 from vllm.v1.worker.gpu.input_batch import InputBatch

15 from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache

16 +from vllm.v1.worker.gpu.mm.encoder_runner import EncoderRunner

17 from vllm.v1.worker.gpu.states import RequestState

18 from vllm.v1.worker.utils import AttentionGroup

@@ -47,6 +48,8 @@ class ModelState(ABC):

48 raise NotImplementedError

50 model: nn.Module

51 + # Set by mm-capable states; used by the default gather_mm_embeddings().

52 + encoder_runner: EncoderRunner

54 def get_supported_generation_tasks(self) -> tuple[GenerationTask, ...]:

55 from vllm.model_executor.models.interfaces import (

@@ -82,10 +85,27 @@ class ModelState(ABC):

86 @abstractmethod

87 def get_mm_embeddings(

85 - self, scheduled_encoder_inputs: dict[str, list[int]], input_batch: InputBatch

88 + self,

89 + scheduled_encoder_inputs: dict[str, list[int]],

90 + input_batch: InputBatch,

91 + req_states: RequestState,

92 ) -> torch.Tensor | None:

93 raise NotImplementedError

95 + def gather_mm_embeddings(

96 + self, input_batch: InputBatch, draft_lookahead: int = 0

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

98 + """Gather cached multimodal embeddings for a speculator's draft forward."""

99 + return self.encoder_runner.gather_mm_embeddings(

100 + input_batch.req_ids,

101 + input_batch.num_tokens,

102 + input_batch.num_scheduled_tokens,

103 + input_batch.query_start_loc_np,

104 + input_batch.prefill_len_np,

105 + input_batch.num_computed_prefill_tokens_np,

106 + draft_lookahead=draft_lookahead,

107 + )

108 +

109 @abstractmethod

110 def prepare_inputs(

111 self, input_batch: InputBatch, req_states: RequestState

vllm/v1/worker/gpu/model_states/mm_pruning.py new

+135

@@ -0,0 +1,135 @@

1 +# SPDX-License-Identifier: Apache-2.0

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

3 +import torch

4 +import torch.nn as nn

5 +

6 +from vllm.config import ModelConfig

7 +from vllm.model_executor.models.interfaces import supports_multimodal_pruning

8 +from vllm.multimodal.utils import get_mm_features_in_window

9 +from vllm.v1.worker.gpu.input_batch import InputBatch

10 +from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache

11 +from vllm.v1.worker.gpu.mm.rope import RopeState

12 +from vllm.v1.worker.gpu.states import RequestState

13 +

14 +

15 +class MultiModalPruner:

16 + """Recomputes M-RoPE positions for multimodal models that prune embeddings

17 + (e.g. Qwen2.5-VL / Qwen3-VL / Nemotron-Nano-VL Efficient Video Sampling).

18 +

19 + Pruning models append their mrope-position channels to the (variable-count)

20 + media embeddings from `embed_multimodal`. Those channels must be split off and

21 + used to recompute mrope positions before the embeddings are merged.

22 + """

23 +

24 + def __init__(

25 + self,

26 + model: nn.Module,

27 + rope_state: RopeState,

28 + encoder_cache: EncoderCache,

29 + inputs_embeds_size: int,

30 + ) -> None:

31 + self.model = model

32 + self.rope_state = rope_state

33 + self.encoder_cache = encoder_cache

34 + # The cleaned embedding width: pruning models append their mrope-position

35 + # channels as trailing columns, so embeds[:, :inputs_embeds_size] strips them.

36 + self.inputs_embeds_size = inputs_embeds_size

37 +

38 + def strip(self, mm_embeds: list[torch.Tensor]) -> list[torch.Tensor]:

39 + """Draft forward: strip the appended position channels only.

40 +

41 + Stripping is per-embedding, so no per-request segmentation is needed. The

42 + speculator reuses the target's already-recomputed positions, hence there is

43 + no position write-back here.

44 + """

45 + return [mm[:, : self.inputs_embeds_size] for mm in mm_embeds]

46 +

47 + def recompute(

48 + self,

49 + mm_embeds: list[torch.Tensor],

50 + input_batch: InputBatch,

51 + req_states: RequestState,

52 + ) -> list[torch.Tensor]:

53 + """Target forward: split the appended mrope-position channels off each

54 + request's media embeddings, recompute the corrected mrope positions, and

55 + stage them back into RopeState. Returns the cleaned, flattened embeddings.

56 + """

57 + cleaned: list[torch.Tensor] = []

58 + pos = 0

59 + req_idx_list = input_batch.idx_mapping_np.tolist()

60 + prefill_lens_list = input_batch.prefill_len_np.tolist()

61 + num_computed_list = input_batch.num_computed_prefill_tokens_np.tolist()

62 + num_scheduled_list = input_batch.num_scheduled_tokens.tolist()

63 + for batch_idx, req_id in enumerate(input_batch.req_ids):

64 + num_computed = num_computed_list[batch_idx]

65 + query_end = num_computed + num_scheduled_list[batch_idx]

66 + num_req_embeds = self._num_window_embeds(req_id, num_computed, query_end)

67 + if num_req_embeds == 0:

68 + continue

69 + req_embeds = mm_embeds[pos : pos + num_req_embeds]

70 + pos += num_req_embeds

71 +

72 + req_idx = req_idx_list[batch_idx]

73 + prefill_len = prefill_lens_list[batch_idx]

74 + input_ids = req_states.all_token_ids.gpu[req_idx, :prefill_len]

75 + mrope_positions = self.rope_state.read_prefill_positions(

76 + req_idx, prefill_len

77 + ).long()

78 + req_cleaned, new_positions, delta = self.model.recompute_mrope_positions(

79 + input_ids=input_ids,

80 + multimodal_embeddings=req_embeds,

81 + mrope_positions=mrope_positions,

82 + num_computed_tokens=num_computed,

83 + )

84 + self.rope_state.update_prefill_positions(req_idx, new_positions, delta)

85 + cleaned.extend(req_cleaned)

86 +

87 + assert pos == len(mm_embeds)

88 + return cleaned

89 +

90 + def _num_window_embeds(self, req_id: str, query_start: int, query_end: int) -> int:

91 + """Count the media items contributing embeddings to [query_start,

92 + query_end), mirroring EncoderRunner.gather_mm_embeddings' per-request

93 + windowing so the flat mm_embeds list can be re-segmented per request.

94 +

95 + Note: This logic is intentionally duplicated here rather than being emitted

96 + from gather_mm_embeddings, to keep the main path cleaner, since this is a niche

97 + feature.

98 + """

99 + mm_features = self.encoder_cache.mm_features[req_id]

100 + lo, hi = get_mm_features_in_window(

101 + mm_features, start=query_start, end=query_end

102 + )

103 + count = 0

104 + for mm_feature in mm_features[lo:hi]:

105 + pos_info = mm_feature.mm_position

106 + start_idx = max(query_start - pos_info.offset, 0)

107 + end_idx = min(query_end - pos_info.offset, pos_info.length)

108 + embeds_start, embeds_end = pos_info.get_embeds_indices_in_range(

109 + start_idx, end_idx

110 + )

111 + if embeds_start != embeds_end:

112 + count += 1

113 + return count

114 +

115 +

116 +def maybe_create_mm_pruner(

117 + model_config: ModelConfig,

118 + model: nn.Module,

119 + rope_state: RopeState | None,

120 + encoder_cache: EncoderCache | None,

121 +) -> MultiModalPruner | None:

122 + """Create a MultiModalPruner if the model prunes embeddings and uses M-RoPE."""

123 + if (

124 + not rope_state

125 + or not rope_state.has_delta

126 + or not encoder_cache

127 + or not model_config.multimodal_config

128 + or not model_config.multimodal_config.is_multimodal_pruning_enabled()

129 + or not supports_multimodal_pruning(model)

130 + ):

131 + return None

132 +

133 + return MultiModalPruner(

134 + model, rope_state, encoder_cache, model_config.get_inputs_embeds_size()

135 + )