@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 1a014a0a93

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 7, 14:00 UTC

[Model Runner V2] Move MM encoder to Model States [3/N] (#35564)

Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>

Woosuk Kwon committed Feb 27, 2026 at 18:32 UTC 1a014a0a9327ed64a1bdec8e1afa43b9ea70a3c1

7 files changed +135 -111

vllm/v1/worker/gpu/cudagraph_utils.py

-6

@@ -89,7 +89,6 @@ class CudaGraphManager:

89 model: nn.Module,

90 model_state: ModelState,

91 input_buffers: InputBuffers,

92 - inputs_embeds: torch.Tensor | None,

92 block_tables: BlockTables,

93 attn_groups: list[list[AttentionGroup]],

94 kv_cache_config: KVCacheConfig,

@@ -116,9 +115,6 @@ class CudaGraphManager:

115 model_inputs = {

116 "input_ids": input_buffers.input_ids[:num_tokens],

117 "positions": input_buffers.positions[:num_tokens],

119 - "inputs_embeds": (

120 - inputs_embeds[:num_tokens] if inputs_embeds is not None else None

121 - ),

118 # NOTE: Values returned by `prepare_dummy_inputs` will override the

119 # default values above.

120 **model_state.prepare_dummy_inputs(num_reqs, num_tokens),

@@ -255,7 +251,6 @@ class CudaGraphManager:

251 model: nn.Module,

252 model_state: ModelState,

253 input_buffers: InputBuffers,

258 - inputs_embeds: torch.Tensor | None,

254 block_tables: BlockTables,

255 attn_groups: list[list[AttentionGroup]],

256 kv_cache_config: KVCacheConfig,

@@ -267,7 +262,6 @@ class CudaGraphManager:

262 model=model,

263 model_state=model_state,

264 input_buffers=input_buffers,

270 - inputs_embeds=inputs_embeds,

265 block_tables=block_tables,

266 attn_groups=attn_groups,

267 kv_cache_config=kv_cache_config,

vllm/v1/worker/gpu/input_batch.py

-3

@@ -66,8 +66,6 @@ class InputBatch:

66 input_ids: torch.Tensor

67 # [num_tokens_after_padding]

68 positions: torch.Tensor

69 - # [num_tokens_after_padding, hidden_size]

70 - inputs_embeds: torch.Tensor | None

70 # [total_num_logits]

71 logits_indices: torch.Tensor

@@ -138,7 +136,6 @@ class InputBatch:

136 dcp_local_seq_lens=None,

137 input_ids=input_ids,

138 positions=positions,

141 - inputs_embeds=None,

139 logits_indices=logits_indices,

140 cu_num_logits=cu_num_logits,

141 cu_num_logits_np=cu_num_logits_np,

vllm/v1/worker/gpu/mm/encoder_cache.py new

+40

@@ -0,0 +1,40 @@

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

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

3 +import torch

4 +

5 +from vllm.multimodal.inputs import MultiModalFeatureSpec

6 +

7 +

8 +class EncoderCache:

9 + def __init__(self):

10 + # req_id -> MM features

11 + self.mm_features: dict[str, list[MultiModalFeatureSpec]] = {}

12 + # MM hash -> encoder outputs

13 + self.encoder_outputs: dict[str, torch.Tensor] = {}

14 +

15 + def add_request(

16 + self, req_id: str, mm_features: list[MultiModalFeatureSpec]

17 + ) -> None:

18 + self.mm_features[req_id] = mm_features

19 +

20 + def remove_request(self, req_id: str) -> None:

21 + self.mm_features.pop(req_id, None)

22 +

23 + def reset_mm_cache(self) -> None:

24 + """

25 + Clear the multi-modal cache that was used during profiling,

26 + but no longer needed during inference.

27 + """

28 + # TODO: Implement MM budget for encoder dummy run

29 + pass

30 +

31 + def reset_encoder_cache(self) -> None:

32 + """Clear the GPU-side encoder cache storing vision embeddings.

33 +

34 + This should be called when model weights are updated to ensure

35 + stale embeddings computed with old weights are not reused.

36 + """

37 + self.encoder_outputs.clear()

38 +

39 + def free_encoder_cache(self, mm_hash: str) -> None:

40 + self.encoder_outputs.pop(mm_hash, None)

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

+8 -32

@@ -4,8 +4,9 @@ import numpy as np

4 import torch

6 from vllm.model_executor.models.interfaces import SupportsMultiModal

7 -from vllm.multimodal.inputs import MultiModalFeatureSpec, MultiModalKwargsItem

7 +from vllm.multimodal.inputs import MultiModalKwargsItem

8 from vllm.multimodal.utils import group_mm_kwargs_by_modality

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

10 from vllm.v1.worker.utils import sanity_check_mm_encoder_outputs

@@ -14,44 +15,19 @@ class EncoderRunner:

15 self,

16 max_num_tokens: int,

17 hidden_size: int,

18 + encoder_cache: EncoderCache,

19 dtype: torch.dtype,

20 device: torch.device,

21 ):

22 self.max_num_tokens = max_num_tokens

23 self.hidden_size = hidden_size

24 + self.encoder_cache = encoder_cache

25 self.dtype = dtype

26 self.device = device

28 self.inputs_embeds = torch.zeros(

29 max_num_tokens, hidden_size, dtype=dtype, device=device

30 )

28 - self.req_id_to_mm_features: dict[str, list[MultiModalFeatureSpec]] = {}

29 - self.encoder_cache: dict[str, torch.Tensor] = {}

30 -

31 - def reset_mm_cache(self) -> None:

32 - """

33 - Clear the multi-modal cache that was used during profiling,

34 - but no longer needed during inference.

35 - """

36 - # TODO: Implement MM budget for encoder dummy run

37 - pass

38 -

39 - def reset_encoder_cache(self) -> None:

40 - """Clear the GPU-side encoder cache storing vision embeddings.

41 -

42 - This should be called when model weights are updated to ensure

43 - stale embeddings computed with old weights are not reused.

44 - """

45 - self.encoder_cache.clear()

46 -

47 - def add_request(self, req_id: str, mm_features: list[MultiModalFeatureSpec]):

48 - self.req_id_to_mm_features[req_id] = mm_features

49 -

50 - def free_encoder_cache(self, mm_hash: str) -> None:

51 - self.encoder_cache.pop(mm_hash, None)

52 -

53 - def remove_request(self, req_id: str) -> None:

54 - self.req_id_to_mm_features.pop(req_id, None)

32 def prepare_mm_inputs(

33 self, scheduled_encoder_inputs: dict[str, list[int]]

@@ -59,7 +35,7 @@ class EncoderRunner:

35 mm_hashes: list[str] = []

36 mm_kwargs: list[tuple[str, MultiModalKwargsItem]] = []

37 for req_id, encoder_input_ids in scheduled_encoder_inputs.items():

62 - mm_features = self.req_id_to_mm_features[req_id]

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

39 for mm_input_id in encoder_input_ids:

40 mm_feature = mm_features[mm_input_id]

41 if mm_feature.data is None:

@@ -90,7 +66,7 @@ class EncoderRunner:

66 encoder_outputs.extend(curr_group_outputs)

68 # Cache the encoder outputs by mm_hash

93 - self.encoder_cache.update(zip(mm_hashes, encoder_outputs))

69 + self.encoder_cache.encoder_outputs.update(zip(mm_hashes, encoder_outputs))

70 return encoder_outputs

72 def gather_mm_embeddings(

@@ -122,7 +98,7 @@ class EncoderRunner:

98 # OPTIMIZATION: Skip decode requests.

99 continue

100

125 - mm_features = self.req_id_to_mm_features[req_id]

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

102 for mm_feature in mm_features:

103 pos_info = mm_feature.mm_position

104 start_pos = pos_info.offset

@@ -148,7 +124,7 @@ class EncoderRunner:

124 continue

125

126 mm_hash = mm_feature.identifier

151 - encoder_output = self.encoder_cache.get(mm_hash, None)

127 + encoder_output = self.encoder_cache.encoder_outputs.get(mm_hash, None)

128 assert encoder_output is not None, f"Encoder cache miss for {mm_hash}."

129

130 if (is_embed := pos_info.is_embed) is not None:

vllm/v1/worker/gpu/model_runner.py

+34 -64

@@ -77,7 +77,7 @@ from vllm.v1.worker.gpu.kv_connector import (

77 get_kv_connector,

78 )

79 from vllm.v1.worker.gpu.lora_utils import LoraState

80 -from vllm.v1.worker.gpu.mm.encoder_runner import EncoderRunner

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

81 from vllm.v1.worker.gpu.model_states import ModelState

82 from vllm.v1.worker.gpu.pool.pooling_runner import PoolingRunner

83 from vllm.v1.worker.gpu.pp_utils import pp_broadcast, pp_receive

@@ -127,20 +127,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):

127 self.max_model_len = self.model_config.max_model_len

128 self.max_num_tokens = self.scheduler_config.max_num_batched_tokens

129 self.max_num_reqs = self.scheduler_config.max_num_seqs

130 - self.inputs_embeds_size = self.model_config.get_inputs_embeds_size()

131 -

132 - # Multimodal

133 - self.mm_registry = MULTIMODAL_REGISTRY

134 - self.supports_mm_inputs = self.mm_registry.supports_multimodal_inputs(

135 - self.model_config

136 - )

137 - if self.supports_mm_inputs:

138 - self.encoder_runner = EncoderRunner(

139 - max_num_tokens=self.max_num_tokens,

140 - hidden_size=self.inputs_embeds_size,

141 - dtype=self.dtype,

142 - device=self.device,

143 - )

130

131 self.use_async_scheduling = self.scheduler_config.async_scheduling

132 self.output_copy_stream = torch.cuda.Stream(self.device)

@@ -162,6 +148,15 @@ class GPUModelRunner(LoRAModelRunnerMixin):

148 self.dcp_rank = get_dcp_group().rank_in_group if self.use_dcp else 0

149 self.cp_interleave = self.parallel_config.cp_kv_cache_interleave_size

150

151 + # Multimodal

152 + self.mm_registry = MULTIMODAL_REGISTRY

153 + self.supports_mm_inputs = self.mm_registry.supports_multimodal_inputs(

154 + self.model_config

155 + )

156 + self.encoder_cache = None

157 + if self.supports_mm_inputs and self.is_first_pp_rank:

158 + self.encoder_cache = EncoderCache()

159 +

160 self.speculator = None

161 self.num_speculative_steps = 0

162 self.use_aux_hidden_state_outputs = False

@@ -272,7 +267,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):

267 prepare_communication_buffer_for_model(self.speculator)

268

269 # Initialize the components that require the model.

275 - self.model_state = ModelState(self.vllm_config, self.model, self.device)

270 + self.model_state = ModelState(

271 + self.vllm_config, self.model, self.encoder_cache, self.device

272 + )

273 if self.is_pooling_model:

274 self.pooling_runner = PoolingRunner(self.model)

275

@@ -435,12 +432,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):

432 gc.collect()

433

434 def reset_mm_cache(self) -> None:

438 - if self.supports_mm_inputs:

439 - self.encoder_runner.reset_mm_cache()

435 + if self.encoder_cache is not None:

436 + self.encoder_cache.reset_mm_cache()

437

438 def reset_encoder_cache(self) -> None:

442 - if self.supports_mm_inputs:

443 - self.encoder_runner.reset_encoder_cache()

439 + if self.encoder_cache is not None:

440 + self.encoder_cache.reset_encoder_cache()

441

442 def _get_num_input_tokens(self, num_scheduled_tokens: int) -> int:

443 # SP is not supported yet.

@@ -469,14 +466,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):

466 start_free_gpu_memory = torch.cuda.mem_get_info()[0]

467

468 with self.maybe_setup_dummy_loras(self.lora_config):

472 - inputs_embeds = None

473 - if self.supports_mm_inputs:

474 - inputs_embeds = self.encoder_runner.inputs_embeds

469 self.cudagraph_manager.capture(

470 model=self.model,

471 model_state=self.model_state,

472 input_buffers=self.input_buffers,

479 - inputs_embeds=inputs_embeds,

473 block_tables=self.block_tables,

474 attn_groups=self.attn_groups,

475 kv_cache_config=self.kv_cache_config,

@@ -511,15 +504,15 @@ class GPUModelRunner(LoRAModelRunnerMixin):

504 finished_req_ids = finished_req_ids.union(preempted_req_ids)

505 for req_id in finished_req_ids:

506 self.req_states.remove_request(req_id)

514 - if self.supports_mm_inputs:

515 - self.encoder_runner.remove_request(req_id)

507 + if self.encoder_cache is not None:

508 + self.encoder_cache.remove_request(req_id)

509 self.prompt_logprobs_worker.remove_request(req_id)

510 self.lora_state.remove_request(req_id)

511

512 def free_states(self, scheduler_output: SchedulerOutput) -> None:

520 - if self.supports_mm_inputs:

513 + if self.encoder_cache is not None:

514 for mm_hash in scheduler_output.free_encoder_mm_hashes:

522 - self.encoder_runner.free_encoder_cache(mm_hash)

515 + self.encoder_cache.free_encoder_cache(mm_hash)

516

517 def add_requests(self, scheduler_output: SchedulerOutput) -> None:

518 for new_req_data in scheduler_output.scheduled_new_reqs:

@@ -535,8 +528,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):

528 )

529 req_index = self.req_states.req_id_to_index[req_id]

530

538 - if self.supports_mm_inputs:

539 - self.encoder_runner.add_request(req_id, new_req_data.mm_features)

531 + if self.encoder_cache is not None:

532 + self.encoder_cache.add_request(req_id, new_req_data.mm_features)

533

534 self.model_state.add_request(req_index, new_req_data)

535 self.block_tables.append_block_ids(

@@ -695,7 +688,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):

688 dcp_local_seq_lens=dcp_local_seq_lens,

689 input_ids=self.input_buffers.input_ids[:num_tokens_after_padding],

690 positions=self.input_buffers.positions[:num_tokens_after_padding],

698 - inputs_embeds=None,

691 logits_indices=logits_indices,

692 cu_num_logits=cu_num_logits,

693 cu_num_logits_np=cu_num_logits_np,

@@ -724,26 +716,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):

716 )

717 return block_tables, slot_mappings

718

727 - @torch.inference_mode()

728 - def get_mm_embeddings(

729 - self,

730 - scheduled_encoder_inputs: dict[str, list[int]],

731 - input_batch: InputBatch,

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

733 - mm_hashes, mm_kwargs = self.encoder_runner.prepare_mm_inputs(

734 - scheduled_encoder_inputs

735 - )

736 - self.encoder_runner.execute_mm_encoder(self.model, mm_hashes, mm_kwargs)

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

738 - input_batch.req_ids,

739 - input_batch.num_tokens,

740 - input_batch.num_scheduled_tokens,

741 - input_batch.query_start_loc_np,

742 - self.req_states.prefill_len.np[input_batch.idx_mapping_np],

743 - self.req_states.num_computed_prefill_tokens[input_batch.idx_mapping_np],

744 - )

745 - return mm_embeds, is_mm_embed

746 -

719 def sample(

720 self,

721 hidden_states: torch.Tensor,

@@ -890,18 +862,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):

862 input_batch.num_scheduled_tokens,

863 )

864 self._set_active_loras(*lora_inputs)

893 -

894 - # Only first PP rank prepares multimodal embeddings.

895 - if self.supports_mm_inputs and self.is_first_pp_rank:

896 - mm_embeds, is_mm_embed = self.get_mm_embeddings(

897 - scheduler_output.scheduled_encoder_inputs, input_batch

898 - )

899 - inputs_embeds = self.encoder_runner.get_inputs_embeds(

900 - self.model, input_batch.input_ids, mm_embeds, is_mm_embed

901 - )

902 - input_batch.inputs_embeds = inputs_embeds[

903 - : input_batch.num_tokens_after_padding

904 - ]

865 else:

866 # No actual tokens to run. A dummy run for DP or memory profiling.

867 num_reqs = min(num_tokens_after_padding, self.max_num_reqs)

@@ -934,10 +894,20 @@ class GPUModelRunner(LoRAModelRunnerMixin):

894 self.kv_cache_config,

895 )

896

897 + inputs_embeds = None

898 + if self.supports_mm_inputs and self.is_first_pp_rank and not dummy_run:

899 + # Run MM encoder (if needed) and get multimodal embeddings.

900 + # Only first PP rank prepares multimodal embeddings.

901 + inputs_embeds = self.model_state.get_mm_embeddings(

902 + scheduler_output.scheduled_encoder_inputs,

903 + input_batch,

904 + self.req_states,

905 + )

906 +

907 model_inputs = {

908 "input_ids": input_batch.input_ids,

909 "positions": input_batch.positions,

940 - "inputs_embeds": input_batch.inputs_embeds,

910 + "inputs_embeds": inputs_embeds,

911 # NOTE: Values returned by `prepare_inputs` will override the default

912 # values above.

913 **self.model_state.prepare_inputs(input_batch, self.req_states),

vllm/v1/worker/gpu/model_states.py

+53 -5

@@ -10,22 +10,43 @@ from vllm.v1.core.sched.output import NewRequestData

10 from vllm.v1.kv_cache_interface import KVCacheConfig

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

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

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

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

15 from vllm.v1.worker.gpu.mm.mrope_utils import MRopeState

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

17 from vllm.v1.worker.utils import AttentionGroup

20 class ModelState:

19 - def __init__(self, vllm_config: VllmConfig, model: nn.Module, device: torch.device):

21 + def __init__(

22 + self,

23 + vllm_config: VllmConfig,

24 + model: nn.Module,

25 + encoder_cache: EncoderCache | None,

26 + device: torch.device,

27 + ):

28 self.vllm_config = vllm_config

29 self.model_config = vllm_config.model_config

30 self.scheduler_config = vllm_config.scheduler_config

31 self.model = model

32 self.device = device

34 + self.supports_mm_inputs = encoder_cache is not None

35 self.max_model_len = self.model_config.max_model_len

36 self.max_num_reqs = self.scheduler_config.max_num_seqs

37 self.max_num_tokens = self.scheduler_config.max_num_batched_tokens

38 + self.inputs_embeds_size = self.model_config.get_inputs_embeds_size()

39 + self.dtype = self.model_config.dtype

40 +

41 + if self.supports_mm_inputs:

42 + assert encoder_cache is not None

43 + self.encoder_runner = EncoderRunner(

44 + max_num_tokens=self.max_num_tokens,

45 + hidden_size=self.inputs_embeds_size,

46 + encoder_cache=encoder_cache,

47 + dtype=self.dtype,

48 + device=self.device,

49 + )

51 self.uses_mrope = self.model_config.uses_mrope

52 if self.uses_mrope:

@@ -51,6 +72,29 @@ class ModelState:

72 if self.uses_mrope:

73 self.mrope_state.apply_staged_writes()

75 + def get_mm_embeddings(

76 + self,

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

78 + input_batch: InputBatch,

79 + req_states: RequestState,

80 + ) -> torch.Tensor:

81 + mm_hashes, mm_kwargs = self.encoder_runner.prepare_mm_inputs(

82 + scheduled_encoder_inputs

83 + )

84 + self.encoder_runner.execute_mm_encoder(self.model, mm_hashes, mm_kwargs)

85 + mm_embeds, is_mm_embed = self.encoder_runner.gather_mm_embeddings(

86 + input_batch.req_ids,

87 + input_batch.num_tokens,

88 + input_batch.num_scheduled_tokens,

89 + input_batch.query_start_loc_np,

90 + req_states.prefill_len.np[input_batch.idx_mapping_np],

91 + req_states.num_computed_prefill_tokens[input_batch.idx_mapping_np],

92 + )

93 + inputs_embeds = self.encoder_runner.get_inputs_embeds(

94 + self.model, input_batch.input_ids, mm_embeds, is_mm_embed

95 + )

96 + return inputs_embeds[: input_batch.num_tokens_after_padding]

97 +

98 def prepare_inputs(

99 self, input_batch: InputBatch, req_states: RequestState

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

@@ -73,10 +117,14 @@ class ModelState:

117 def prepare_dummy_inputs(

118 self, num_reqs: int, num_tokens: int

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

76 - if not self.uses_mrope:

77 - return {}

78 - mrope_positions = self.mrope_state.mrope_positions[:, :num_tokens]

79 - return {"positions": mrope_positions}

120 + model_inputs = {}

121 + if self.supports_mm_inputs:

122 + inputs_embeds = self.encoder_runner.inputs_embeds[:num_tokens]

123 + model_inputs["inputs_embeds"] = inputs_embeds

124 + if self.uses_mrope:

125 + mrope_positions = self.mrope_state.mrope_positions[:, :num_tokens]

126 + model_inputs["positions"] = mrope_positions

127 + return model_inputs

128

129 def prepare_attn(

130 self,

vllm/v1/worker/gpu/spec_decode/eagle/speculator.py

-1

@@ -44,7 +44,6 @@ class EagleSpeculator:

44 # the draft model's hidden size can be different from the target model's

45 # hidden size (e.g., Llama 3.3 70B).

46 self.hidden_size = self.draft_model_config.get_hidden_size()

47 - self.inputs_embeds_size = self.draft_model_config.get_inputs_embeds_size()

47 self.vocab_size = self.draft_model_config.get_vocab_size()

48 self.dtype = vllm_config.model_config.dtype