move mrope_config and mm_embedding under MultimodalParams

yechank-nvidia · yechank-nvidia · commit a508aee401bb · 2025-07-04T15:19:34.000+09:00
Signed-off-by: yechank &lt;161688079+yechank-nvidia@users.noreply.github.com&gt;
diff --git a/tensorrt_llm/_torch/models/modeling_gemma3vl.py b/tensorrt_llm/_torch/models/modeling_gemma3vl.py
@@ -100,8 +100,10 @@ def __call__(
             "pixel_values"]
         input_ids = preprocess_outputs[0]["mm_processor_kwargs"]["input_ids"]
         mm_features = self._process(pixel_values)
+        multimodal_data = {}
+        multimodal_data["multimodal_embedding"] = mm_features
         return input_ids[0].to(torch.int32).tolist(), {
-            "mm_embedding": mm_features
+            "multimodal_data": multimodal_data
         }
 
 
@@ -163,7 +165,7 @@ def forward(
 
         multimodal_params = kwargs.get("multimodal_params", [])
         mm_embed = [
-            multimodal_param.multimodal_embedding
+            multimodal_param.multimodal_data["multimodal_embedding"]
             for multimodal_param in multimodal_params
         ]
         assert mm_embed == [] or len(
diff --git a/tensorrt_llm/_torch/models/modeling_hyperclovax.py b/tensorrt_llm/_torch/models/modeling_hyperclovax.py
@@ -1,5 +1,6 @@
 import copy
 import math
+import os
 from functools import partial
 from itertools import chain
 from typing import Any, Dict, List, Optional, Tuple, Union
@@ -25,6 +26,8 @@
 from .modeling_siglip import SiglipVisionModel
 from .modeling_utils import register_auto_model
 
+DISAGG = os.getenv('TLLM_MULTIMODAL_DISAGGREGATED', '0') == '1'
+
 
 # Copied from HyperCLOVAX-SEED-Vision-Instruct-3B/modeling_hyperclovax.py
 def select_best_resolution(original_size: tuple,
@@ -969,8 +972,8 @@ def __init__(self, model_config: ModelConfig):
         self.model_config = model_config
         if hasattr(self, "llm"):
             return
-
-        self.mm_encoder = HCXVisionModel(model_config)
+        if not DISAGG:
+            self.mm_encoder = HCXVisionModel(model_config)
         llm_model_config = copy.deepcopy(model_config)
         llm_model_config.pretrained_config = PretrainedConfig.from_dict(
             llm_model_config.pretrained_config.language_config)
@@ -1026,7 +1029,13 @@ def forward(
             assert len(multimodal_params) == num_context_requests == len(
                 multimodal_params
             ), f"Number of multimodal tensors ({len(multimodal_params)}) should be equal to number of context requests ({num_context_requests}) in the batch."
-            mm_embeds = self.mm_encoder.forward(multimodal_params)
+            if not DISAGG:
+                mm_embeds = self.mm_encoder.forward(multimodal_params)
+            else:
+                mm_embeds = [
+                    multimodal_param.multimodal_data["multimodal_embedding"]
+                    for multimodal_param in multimodal_params
+                ]
 
         input_ids, input_embeds = fuse_input_embeds(self.llm.model.embed_tokens,
                                                     input_ids, mm_embeds)
diff --git a/tensorrt_llm/_torch/models/modeling_llama.py b/tensorrt_llm/_torch/models/modeling_llama.py
@@ -851,7 +851,10 @@ def __call__(
             mm_embeds = self.encoder.multi_modal_projector(mm_embeds)
             # for fuse_input_embeds
             token_ids[token_ids == self.image_token_index] = self.vocab_size + 1
-            return token_ids.tolist(), {"mm_embedding": mm_embeds}
+
+            multimodal_data = {}
+            multimodal_data["multimodal_embedding"] = mm_embeds
+            return token_ids.tolist(), {"multimodal_data": multimodal_data}
         else:
             return processed["input_ids"].squeeze().tolist(), {}
 
@@ -882,8 +885,12 @@ def forward(
         spec_metadata: Optional[SpecMetadata] = None,
         **kwargs,
     ) -> torch.Tensor:
-        mm_embed = kwargs.get("multi_modal_data", [])
-        if mm_embed:
+        multimodal_params = kwargs.get("multimodal_params", [])
+        if multimodal_params:
+            mm_embed = [
+                multimodal_param.multimodal_data["multimodal_embedding"]
+                for multimodal_param in multimodal_params
+            ]
             _, inputs_embeds = fuse_input_embeds(self.model.embed_tokens,
                                                  input_ids, mm_embed)
         return super().forward(attn_metadata,
diff --git a/tensorrt_llm/_torch/models/modeling_llava_next.py b/tensorrt_llm/_torch/models/modeling_llava_next.py
@@ -210,8 +210,10 @@ def __call__(
         mm_features = torch.stack(
             [self._process(tensor) for tensor in mm_tensor])
         fused_input_ids, mm_features = self._postprocess(input_ids, mm_features)
+        multimodal_data = {}
+        multimodal_data["multimodal_embedding"] = mm_features
         return fused_input_ids.to(torch.int32).tolist(), {
-            "mm_embedding": mm_features
+            "multimodal_data": multimodal_data
         }
 
 
@@ -273,7 +275,7 @@ def forward(
 
         multimodal_params = kwargs.get("multimodal_params", [])
         mm_embed = [
-            multimodal_param.multimodal_embedding
+            multimodal_param.multimodal_data["multimodal_embedding"]
             for multimodal_param in multimodal_params
         ]
         assert mm_embed == [] or len(
diff --git a/tensorrt_llm/_torch/models/modeling_qwen2vl.py b/tensorrt_llm/_torch/models/modeling_qwen2vl.py
@@ -1,4 +1,5 @@
 import copy
+import os
 from typing import Any, Dict, List, Optional, Tuple, Union
 
 import torch
@@ -20,6 +21,8 @@
 from .modeling_multimodal_utils import fuse_input_embeds
 from .modeling_utils import register_auto_model
 
+DISAGG = os.getenv('TLLM_MULTIMODAL_DISAGGREGATED', '0') == '1'
+
 
 class Qwen2VLInputProcessorBase(InputProcessor):
 
@@ -322,7 +325,8 @@ def get_mrope_config(
         concat_cos_sin = concat_cos_sin.reshape(concat_cos_sin.shape[0], -1)
         mrope_config = {}
         mrope_config['mrope_rotary_cos_sin'] = concat_cos_sin.to('cpu')
-        mrope_config['mrope_position_deltas'] = mrope_position_deltas.to('cpu')
+        mrope_config['mrope_position_deltas'] = mrope_position_deltas.to(
+            'cpu').to(torch.int32)
         return mrope_config
 
     @torch.inference_mode()
@@ -364,11 +368,11 @@ def __call__(
             processed_inputs.get('video_grid_thw', None),
             processed_inputs.get('attention_mask', None),
             processed_inputs.get('second_per_grid_ts', None))
+        multimodal_data["mrope_config"] = mrope_config
 
         fused_input_ids = self._postprocess(input_ids[0])
 
         return fused_input_ids.to(torch.int32).tolist(), {
-            "mrope_config": mrope_config,
             "multimodal_data": multimodal_data,
         }
 
@@ -411,16 +415,14 @@ def _parse_and_batch_multimodal_data(
 
         for multimodal_param in multimodal_params:
             # Process images if present
-            if "image" in multimodal_param.multimodal_data and multimodal_param.multimodal_data[
-                    "image"]:
+            if multimodal_param.multimodal_data.get("image") is not None:
                 pixel_values_list.append(
                     multimodal_param.multimodal_data["image"]["pixel_values"])
                 image_grid_thw_list.append(
                     multimodal_param.multimodal_data["image"]["image_grid_thw"])
 
             # Process videos if present
-            if "video" in multimodal_param.multimodal_data and multimodal_param.multimodal_data[
-                    "video"]:
+            if multimodal_param.multimodal_data.get("video") is not None:
                 pixel_values_videos_list.append(
                     multimodal_param.multimodal_data["video"]
                     ["pixel_values_videos"])
@@ -457,6 +459,8 @@ def forward(self, multimodal_params: List[MultimodalParams]):
 
         mm_content_data, mm_extra_data = self._parse_and_batch_multimodal_data(
             multimodal_params)
+        print(f"mm_content_data: {mm_content_data}")
+        print(f"mm_extra_data: {mm_extra_data}")
         pixel_values = mm_content_data.get("pixel_values", None)
         pixel_values_videos = mm_content_data.get("pixel_values_videos", None)
 
@@ -478,7 +482,6 @@ def forward(self, multimodal_params: List[MultimodalParams]):
             pixel_values_videos = pixel_values_videos.to(self.visual.dtype)
             embeds.append(
                 self.visual(pixel_values_videos, grid_thw=video_grid_thw))
-
         return embeds
 
 
@@ -526,16 +529,19 @@ def _parse_mrope_config(
         mrope_config = {}
         mrope_rotary_cos_sin_list = []
         mrope_position_deltas_list = []
-
         for multimodal_param in multimodal_params:
-            if hasattr(multimodal_param,
-                       'mrope_config') and multimodal_param.mrope_config:
-                if 'mrope_rotary_cos_sin' in multimodal_param.mrope_config:
+            if multimodal_param.multimodal_data and multimodal_param.multimodal_data.get(
+                    'mrope_config'):
+                if multimodal_param.multimodal_data['mrope_config'].get(
+                        'mrope_rotary_cos_sin') is not None:
                     mrope_rotary_cos_sin_list.append(
-                        multimodal_param.mrope_config['mrope_rotary_cos_sin'])
-                if 'mrope_position_deltas' in multimodal_param.mrope_config:
+                        multimodal_param.multimodal_data['mrope_config']
+                        ['mrope_rotary_cos_sin'])
+                if multimodal_param.multimodal_data['mrope_config'].get(
+                        'mrope_position_deltas') is not None:
                     mrope_position_deltas_list.append(
-                        multimodal_param.mrope_config['mrope_position_deltas'])
+                        multimodal_param.multimodal_data['mrope_config']
+                        ['mrope_position_deltas'])
 
         if mrope_rotary_cos_sin_list:
             mrope_config['mrope_rotary_cos_sin'] = torch.cat(
@@ -544,6 +550,8 @@ def _parse_mrope_config(
         if mrope_position_deltas_list:
             mrope_config['mrope_position_deltas'] = torch.cat(
                 mrope_position_deltas_list, dim=0)
+        print(f"mrope_config: {mrope_config}")
+        return mrope_config
 
     @torch.inference_mode()
     def forward(
@@ -568,8 +576,14 @@ def forward(
         mrope_config = {}
 
         if len(multimodal_params) > 0:
-            mm_embeds = self.mm_encoder.forward(
-                multimodal_params[:num_context_requests])
+            if not DISAGG:
+                mm_embeds = self.mm_encoder.forward(
+                    multimodal_params[:num_context_requests])
+            else:
+                mm_embeds = [
+                    multimodal_param.multimodal_data["multimodal_embedding"]
+                    for multimodal_param in multimodal_params
+                ]
             mrope_config = self._parse_mrope_config(multimodal_params)
 
         input_ids, input_embeds = fuse_input_embeds(self.llm.model.embed_tokens,
@@ -592,8 +606,9 @@ class Qwen2VLModel(Qwen2VLModelBase):
 
     def __init__(self, model_config: ModelConfig[PretrainedConfig], *args,
                  **kwargs):
-        self.mm_encoder = Qwen2VisionModelBase(model_config,
-                                               Qwen2VLForConditionalGeneration)
+        if not DISAGG:
+            self.mm_encoder = Qwen2VisionModelBase(
+                model_config, Qwen2VLForConditionalGeneration)
         super().__init__(model_config, *args, **kwargs)
 
 
@@ -603,6 +618,7 @@ class Qwen2_5_VLModel(Qwen2VLModelBase):
 
     def __init__(self, model_config: ModelConfig[PretrainedConfig], *args,
                  **kwargs):
-        self.mm_encoder = Qwen2VisionModelBase(
-            model_config, Qwen2_5_VLForConditionalGeneration)
+        if not DISAGG:
+            self.mm_encoder = Qwen2VisionModelBase(
+                model_config, Qwen2_5_VLForConditionalGeneration)
         super().__init__(model_config, *args, **kwargs)
diff --git a/tensorrt_llm/_torch/models/modeling_vila.py b/tensorrt_llm/_torch/models/modeling_vila.py
@@ -1107,8 +1107,10 @@ def __call__(
         )  # use_fast uses Pytorch GPU preprocessing, otherwise uses PIL CPU preprocessing
         mm_features = self._process(mm_tensor, block_sizes)
         fused_input_ids, mm_features = self._postprocess(input_ids, mm_features)
+        multimodal_data = {}
+        multimodal_data["multimodal_embedding"] = mm_features
         return fused_input_ids.to(torch.int32).tolist(), {
-            "mm_embedding": mm_features
+            "multimodal_data": multimodal_data
         }
 
 
@@ -1163,7 +1165,7 @@ def forward(
         num_context_requests, num_generation_requests = attn_metadata.num_contexts, attn_metadata.num_generations
         multimodal_params = kwargs.get("multimodal_params", [])
         mm_embed = [
-            multimodal_param.multimodal_embedding
+            multimodal_param.multimodal_data["multimodal_embedding"]
             for multimodal_param in multimodal_params
         ]
 
diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -1188,30 +1188,15 @@ def _prepare_tp_inputs(
             prompt_lengths.append(len(prompt_tokens))
             past_seen_token_num = begin_compute
             num_cached_tokens_per_seq.append(past_seen_token_num)
+            request.py_batch_idx = py_batch_idx(request)
 
-            multimodal_embedding = request.multimodal_embedding
-            if multimodal_embedding is not None:
-                # TODO: Visit later once we have the SharedTensor.
-                multimodal_embedding = multimodal_embedding.pin_memory(
-                ) if multimodal_embedding.device == 'cpu' else multimodal_embedding
-                multimodal_embedding = multimodal_embedding.to(
-                    'cuda', non_blocking=True)
-
-            mrope_rotary_cos_sin = request.mrope_rotary_cos_sin
-            if mrope_rotary_cos_sin is not None:
-                # TODO: Visit later once we have the SharedTensor.
-                mrope_rotary_cos_sin = mrope_rotary_cos_sin.pin_memory(
-                ) if mrope_rotary_cos_sin.device == 'cpu' else mrope_rotary_cos_sin
-                mrope_rotary_cos_sin = mrope_rotary_cos_sin.to(
-                    'cuda', non_blocking=True)
-
-            # Create MultimodalParams from request data
+            # Multimodal
             multimodal_params = MultimodalParams(
-                multimodal_embedding=multimodal_embedding,
-                mrope_config={'mrope_rotary_cos_sin': mrope_rotary_cos_sin}
-                if mrope_rotary_cos_sin is not None else {},
-                multimodal_data=request.py_multimodal_data,
-            )
+                multimodal_data=request.py_multimodal_data, )
+            multimodal_params.strip_for_context()
+            multimodal_params.to_device("multimodal_data",
+                                        "cuda",
+                                        pin_memory=True)
 
             if multimodal_params.has_content():
                 multimodal_params_list.append(multimodal_params)
@@ -1243,20 +1228,15 @@ def _prepare_tp_inputs(
                     extend_requests.append(request)
             else:
                 generation_requests.append(request)
-
-            # Handle generation request multimodal params
-            mrope_position_deltas = request.mrope_position_deltas
-            if mrope_position_deltas is not None:
-                mrope_position_deltas_tensor = torch.tensor(
-                    [mrope_position_deltas], dtype=torch.int32, pin_memory=True)
-                multimodal_params = MultimodalParams(
-                    mrope_config={
-                        'mrope_position_deltas':
-                        mrope_position_deltas_tensor.to('cuda',
-                                                        non_blocking=True)
-                    })
-                if multimodal_params.has_content():
-                    multimodal_params_list.append(multimodal_params)
+            # Multimodal
+            multimodal_params = MultimodalParams(
+                multimodal_data=request.py_multimodal_data, )
+            multimodal_params.strip_for_generation()
+            multimodal_params.to_device("multimodal_data",
+                                        "cuda",
+                                        pin_memory=True)
+            if multimodal_params.has_content():
+                multimodal_params_list.append(multimodal_params)
         extend_requests += extend_dummy_requests
 
         if not self._disable_overlap_scheduler and self.is_spec_decode:
diff --git a/tensorrt_llm/executor/worker.py b/tensorrt_llm/executor/worker.py
diff --git a/tensorrt_llm/inputs/multimodal.py b/tensorrt_llm/inputs/multimodal.py
diff --git a/tensorrt_llm/llmapi/llm.py b/tensorrt_llm/llmapi/llm.py
