From 3d56f70b67f8d37f25b243dc78481cf435f6be30 Mon Sep 17 00:00:00 2001
From: "Chang Liu (Enterprise Products)"
 <9713593+chang-l@users.noreply.github.com>
Date: Wed, 6 Aug 2025 17:22:25 -0700
Subject: [PATCH 01/15] Add multimodal encoder engine path

Signed-off-by: Chang Liu <9713593+chang-l@users.noreply.github.com>
Signed-off-by: Chang Liu (Enterprise Products) <9713593+chang-l@users.noreply.github.com>
---
 tensorrt_llm/_torch/model_config.py           |  9 +-
 tensorrt_llm/_torch/models/modeling_auto.py   | 34 +++++++-
 .../_torch/models/modeling_qwen2vl.py         |  4 +-
 tensorrt_llm/_torch/models/modeling_utils.py  | 15 ++++
 tensorrt_llm/_torch/multimodal_encoder.py     | 85 +++++++++++++++++++
 tensorrt_llm/_torch/pyexecutor/_util.py       |  3 +
 tensorrt_llm/_torch/pyexecutor/config.py      |  5 +-
 tensorrt_llm/_torch/pyexecutor/llm_request.py |  5 ++
 .../_torch/pyexecutor/model_engine.py         |  4 +
 .../_torch/pyexecutor/py_executor_creator.py  |  2 +
 tensorrt_llm/_torch/pyexecutor/sampler.py     | 38 +++++++++
 tensorrt_llm/executor/result.py               | 10 ++-
 tensorrt_llm/llmapi/llm.py                    |  1 +
 13 files changed, 209 insertions(+), 6 deletions(-)
 create mode 100644 tensorrt_llm/_torch/multimodal_encoder.py

diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py
index 232d2ccecd64..209ac3d8b7f5 100644
--- a/tensorrt_llm/_torch/model_config.py
+++ b/tensorrt_llm/_torch/model_config.py
@@ -96,6 +96,9 @@ class ModelConfig(Generic[TConfig]):
 
     _frozen: bool = field(default=False, init=False, repr=False)
 
+    # If true, ONLY the vision encoder part of the full model is loaded/executed.
+    mm_encoder_only: bool = False
+
     def __setattr__(self, key, value):
         """
         Prevent modification of frozen instance attributes.
@@ -115,7 +118,7 @@ def __post_init__(self):
         if self.pretrained_config and hasattr(self.pretrained_config,
                                               "architectures"):
             self.is_generation = self.is_generation_model(
-                self.pretrained_config.architectures)
+                self.pretrained_config.architectures, mm_encoder_only=self.mm_encoder_only)
 
         def get_all_reduce_strategy(strategy: str = "AUTO"):
             maps = {
@@ -164,12 +167,14 @@ def get_quant_config(self, name: Optional[str] = None) -> QuantConfig:
         raise ValueError(f'quant config of {name} is not found')
 
     @staticmethod
-    def is_generation_model(model_architectures: Optional[List[str]]) -> bool:
+    def is_generation_model(model_architectures: Optional[List[str]], mm_encoder_only: bool = False) -> bool:
         if model_architectures is None:
             logger.warning(
                 "Model architectures is None, default to is_generation_model=True"
             )
             return True
+        if mm_encoder_only:
+            return False
         return model_architectures[0] not in [
             "BertForSequenceClassification", "Qwen2ForProcessRewardModel",
             "Qwen2ForRewardModel", "LlamaForTextEmbedding"
diff --git a/tensorrt_llm/_torch/models/modeling_auto.py b/tensorrt_llm/_torch/models/modeling_auto.py
index 17bb6bab55cf..314f92e09268 100644
--- a/tensorrt_llm/_torch/models/modeling_auto.py
+++ b/tensorrt_llm/_torch/models/modeling_auto.py
@@ -3,7 +3,7 @@
 from ..model_config import ModelConfig
 from ..utils import model_extra_attrs
 from .modeling_utils import (MODEL_CLASS_MAPPING, DecoderModelForCausalLM,
-                             TConfig, TModel)
+                             MODEL_CLASS_VISION_ENCODER_MAPPING, TConfig, TModel)
 
 
 class AutoModelForCausalLM(Generic[TModel, TConfig]):
@@ -13,6 +13,15 @@ def from_config(
         config: ModelConfig[TConfig],
     ) -> DecoderModelForCausalLM[TModel, TConfig]:
         model_arch = config.pretrained_config.architectures[0]
+        if config.mm_encoder_only:
+            vision_encoder_info = MODEL_CLASS_VISION_ENCODER_MAPPING.get(model_arch)
+            if vision_encoder_info is None:
+                raise ValueError(
+                    f"Unknown architecture for AutoModelForMultimodalEncoder: {model_arch}"
+                )
+            vision_encoder_cls, vlm_base_model = vision_encoder_info
+            return vision_encoder_cls(config, vlm_base_model)
+
         # Hack to detect eagle3 checkpoints. TODO: should we provide
         # our own checkpoints with the correct arch? It would let us
         # avoid nasty stuff like this.
@@ -35,3 +44,26 @@ def from_config(
             model = cls(config)
         model.extra_attrs = extra_attrs
         return model
+
+class AutoModelForMultimodalEncoder(Generic[TModel, TConfig]):
+
+    @staticmethod
+    def from_config(
+        config: ModelConfig[TConfig],
+    ):
+        model_arch = config.pretrained_config.architectures[0]
+
+        # Direct lookup using architecture name
+        vision_encoder_info = MODEL_CLASS_VISION_ENCODER_MAPPING.get(model_arch)
+        if vision_encoder_info is None:
+            raise ValueError(
+                f"Unknown architecture for AutoModelForMultimodalEncoder: {model_arch}"
+            )
+
+        vision_encoder_cls, vlm_base_model = vision_encoder_info
+        if vlm_base_model is None:
+            model = vision_encoder_cls(config)
+        else:
+            model = vision_encoder_cls(config, vlm_base_model)
+
+        return model
\ No newline at end of file
diff --git a/tensorrt_llm/_torch/models/modeling_qwen2vl.py b/tensorrt_llm/_torch/models/modeling_qwen2vl.py
index 03f15c37b42a..0c43c880bba2 100644
--- a/tensorrt_llm/_torch/models/modeling_qwen2vl.py
+++ b/tensorrt_llm/_torch/models/modeling_qwen2vl.py
@@ -21,7 +21,7 @@
 from .modeling_auto import AutoModelForCausalLM
 from .modeling_multimodal_utils import (find_uncached_mm_embeds,
                                         fuse_input_embeds)
-from .modeling_utils import register_auto_model
+from .modeling_utils import register_auto_model, register_vision_encoder
 
 DISAGG = os.getenv('TLLM_MULTIMODAL_DISAGGREGATED', '0') == '1'
 
@@ -645,6 +645,7 @@ def forward(
 
 
 @register_auto_model("Qwen2VLForConditionalGeneration")
+@register_vision_encoder(Qwen2VisionModelBase, vlm_base_model=Qwen2VLForConditionalGeneration)
 @register_input_processor(Qwen2VLInputProcessorBase, model_type="qwen2_vl")
 class Qwen2VLModel(Qwen2VLModelBase):
 
@@ -657,6 +658,7 @@ def __init__(self, model_config: ModelConfig[PretrainedConfig], *args,
 
 
 @register_auto_model("Qwen2_5_VLForConditionalGeneration")
+@register_vision_encoder(Qwen2VisionModelBase, vlm_base_model=Qwen2_5_VLForConditionalGeneration)
 @register_input_processor(Qwen2VLInputProcessorBase, model_type="qwen2_5_vl")
 class Qwen2_5_VLModel(Qwen2VLModelBase):
 
diff --git a/tensorrt_llm/_torch/models/modeling_utils.py b/tensorrt_llm/_torch/models/modeling_utils.py
index 4252ecd4f533..f00b22e449d5 100755
--- a/tensorrt_llm/_torch/models/modeling_utils.py
+++ b/tensorrt_llm/_torch/models/modeling_utils.py
@@ -573,6 +573,7 @@ def infer_max_seq_len(self) -> int:
 
 
 MODEL_CLASS_MAPPING = {}
+MODEL_CLASS_VISION_ENCODER_MAPPING = {}
 MODEL_CLASS_MAPPER_MAPPING = {}
 MODEL_CLASS_CHECKPOINT_WEIGHT_LOADER_DEFAULT_MAPPING = {}
 MODEL_CLASS_CONFIG_LOADER_DEFAULT_MAPPING = {}
@@ -587,6 +588,20 @@ def decorator(cls):
 
     return decorator
 
+# Need to register the model class before the vision encoder class
+# to ensure that the key (the architecture name) is available here
+def register_vision_encoder(vision_encoder_cls: Type[nn.Module],
+                            vlm_base_model: Optional[Type[nn.Module]] = None):
+
+    def wrapper(model_cls: Type[nn.Module]) -> Type[nn.Module]:
+        for arch_name, registered_cls in MODEL_CLASS_MAPPING.items():
+            if registered_cls == model_cls:
+                MODEL_CLASS_VISION_ENCODER_MAPPING[arch_name] = (vision_encoder_cls, vlm_base_model)
+                break
+
+        return model_cls
+
+    return wrapper
 
 def register_mapper(format: str, name: Optional[str] = None):
 
diff --git a/tensorrt_llm/_torch/multimodal_encoder.py b/tensorrt_llm/_torch/multimodal_encoder.py
new file mode 100644
index 000000000000..3bb219f2c832
--- /dev/null
+++ b/tensorrt_llm/_torch/multimodal_encoder.py
@@ -0,0 +1,85 @@
+from tensorrt_llm.inputs.data import PromptInputs
+from tensorrt_llm.llmapi.llm import BaseLLM, _TorchLLM
+from typing import Any, Union
+from pathlib import Path
+
+
+class MultimodalEncoder(_TorchLLM):
+    """MultimodalEncoder class is the main class for running a multimodal encoder model using PyTorch backend.
+
+    Parameters:
+"""
+
+    def __init__(self,
+                 model: Union[str, Path],
+                 trust_remote_code: bool = False,
+                 tensor_parallel_size: int = 1,
+                 dtype: str = "auto",
+                 **kwargs: Any) -> None:
+
+        # Validate that users don't pass LLM-specific or TRT-specific arguments
+        self._validate_mm_args_for_torch_backend(kwargs)
+
+        super().__init__(model,
+                         trust_remote_code=trust_remote_code,
+                         tensor_parallel_size=tensor_parallel_size,
+                         dtype = dtype,
+                         **kwargs)
+
+    def _build_model(self):
+        super()._build_model()
+        self._executor_config.mm_encoder_only = True
+
+    def _validate_mm_args_for_torch_backend(self, kwargs: dict) -> None:
+        """Validate that users don't pass LLM-specific arguments when using MultimodalEncoder (PyTorch).
+        Placeholder for now.
+        """
+        pass
+
+    def generate(
+        self,
+        inputs: PromptInputs,
+    ) -> RequestOutput:
+        """Generate embeddings (and other multimodal encoder outputs) for multiple multimodal requests in parallel.
+
+        Args:
+            mm_requests: List of multimodal requests to process
+
+        Returns:
+            List of generation results
+        """
+        async def _process_requests():
+            # Submit all requests first
+            futures = []
+            for request in mm_requests:
+                future = await self.generate_async(request)
+                futures.append(future)
+
+            # Then wait for all results
+            results = []
+            for future in futures:
+                result = await future.aresult()
+                results.append(result)
+            return results
+
+        # Run the async operations in an event loop
+        return asyncio.run(_process_requests())
+
+    @nvtx_range_debug("MM_encoder.generate_async", color="green", category="VisionEncoder")
+    def generate_async(
+        self,
+        inputs: PromptInputs,
+        sampling_params: Optional[SamplingParams] = None,
+    ) -> RequestOutput:
+        """Generate output for the given multimodal request in the asynchronous mode.
+        Asynchronous generation accepts single multimodal request only.
+
+        Returns:
+            tensorrt_llm.llmapi.RequestOutput: The output data of the completion request to the LLM.
+        """
+        # TODO: possible preprocess the input
+        result = BaseLLM.generate_async(self, inputs, sampling_params)
+        # TODO: possible postprocess the result
+        return result
+
+
diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py
index b7204afeb433..27be8c643613 100644
--- a/tensorrt_llm/_torch/pyexecutor/_util.py
+++ b/tensorrt_llm/_torch/pyexecutor/_util.py
@@ -594,6 +594,9 @@ def instantiate_sampler(engine: PyTorchModelEngine,
         return TRTLLMSampler(executor_config, engine.model, engine.dtype,
                              mapping, decoding_mode,
                              pytorch_backend_config.disable_overlap_scheduler)
+    if executor_config.mm_encoder_only:
+        # NOTE: handle model outputs specially for mm encoder executor/engine
+        return EarlyStopWithMMResult()
     if not engine.model.model_config.is_generation:
         # NOTE: choose sampler based on model type
         return EarlyStopSampler()
diff --git a/tensorrt_llm/_torch/pyexecutor/config.py b/tensorrt_llm/_torch/pyexecutor/config.py
index 0770643ae35e..cfd3c4947148 100644
--- a/tensorrt_llm/_torch/pyexecutor/config.py
+++ b/tensorrt_llm/_torch/pyexecutor/config.py
@@ -112,6 +112,7 @@ class PyTorchConfig:
     'tokens_per_block',
     'mapping',
     'hf_model_dir',
+    'mm_encoder_only',
 ]
 
 
@@ -126,7 +127,8 @@ def update_executor_config(
         max_input_len: Optional[int] = None,
         max_seq_len: Optional[int] = None,
         checkpoint_format: Optional[str] = None,
-        checkpoint_loader: Optional[BaseCheckpointLoader] = None):
+        checkpoint_loader: Optional[BaseCheckpointLoader] = None,
+        mm_encoder_only: Optional[bool] = False):
     if backend is None:
         return
 
@@ -140,6 +142,7 @@ def update_executor_config(
     executor_config.pytorch_backend_config = pytorch_backend_config
     executor_config.mapping = mapping
     executor_config.speculative_config = speculative_config
+    executor_config.mm_encoder_only = mm_encoder_only
 
     logger.info(f"{executor_config.pytorch_backend_config}")
 
diff --git a/tensorrt_llm/_torch/pyexecutor/llm_request.py b/tensorrt_llm/_torch/pyexecutor/llm_request.py
index 0fb1f06e9640..a5c3cee79793 100644
--- a/tensorrt_llm/_torch/pyexecutor/llm_request.py
+++ b/tensorrt_llm/_torch/pyexecutor/llm_request.py
@@ -172,6 +172,7 @@ def __init__(self,
             max_new_tokens, use_device_memory, exclude_last_generation_logits
         ) if return_generation_logits else None
         self._log_probs = LogProbStorage() if return_log_probs else None
+        self._mm_embeddings = None
 
     def append_context_logits(self, context_logits: torch.Tensor):
         if self._context_logits:
@@ -187,6 +188,10 @@ def append_log_probs(self,
         if self._log_probs:
             self._log_probs.append(log_probs, cum_log_probs)
 
+    def append_mm_embeddings(self, mm_embeddings: torch.Tensor):
+        # TODO: add to_handle to use shared tensor support
+        self._mm_embeddings = mm_embeddings.to(device='cpu', non_blocking=True)
+
     def set_log_probs(self, log_probs: list[TokenLogprobs],
                       cum_log_probs: list[float]):
         """
diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py
index 2d00cee05f09..6af03f4c60f2 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_engine.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -1011,6 +1011,7 @@ def _load_model(self,
             moe_load_balancer=moe_load_balancer,
             lora_config=lora_config,
             allreduce_strategy=self.pytorch_backend_config.allreduce_strategy,
+            mm_encoder_only=self.pytorch_backend_config.mm_encoder_only,
             **kwargs)
 
         validate_and_set_kv_cache_quant(
@@ -2126,12 +2127,15 @@ def forward(
                                                      moe_enable_update)
 
         if kv_cache_manager is None:
+            # here you need to add logic for prepare_tp_inputs_no_cache and forward_step for mm encoder only
             inputs, gather_ids = self._prepare_tp_inputs_no_cache(
                 scheduled_requests, attn_metadata, spec_metadata)
 
             with MoeLoadBalancerIterContext(moe_load_balancer):
+                # need to add key ['logits']=None to the output of forward_step for mm encoder only
                 return self._forward_step(inputs, gather_ids,
                                           gather_context_logits)
+                # lets update requst state here after this 1step finished
         with self._maybe_pad_batch(scheduled_requests, kv_cache_manager,
                                    spec_resource_manager) as scheduled_requests:
             maybe_graph = self._maybe_get_cuda_graph(scheduled_requests)
diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
index bcd006be71e6..c6f99a31ec45 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
@@ -169,6 +169,8 @@ def _mangle_executor_config(executor_config: ExecutorConfig):
         )
         executor_config.enable_chunked_context = False
 
+    if executor_config.mm_encoder_only:
+        pytorch_backend_config.mm_encoder_only = True
 
 def _get_mapping(executor_config: ExecutorConfig) -> Mapping:
     if executor_config.mapping is None:
diff --git a/tensorrt_llm/_torch/pyexecutor/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler.py
index daaac14c5a37..e7edd49f97da 100644
--- a/tensorrt_llm/_torch/pyexecutor/sampler.py
+++ b/tensorrt_llm/_torch/pyexecutor/sampler.py
@@ -87,6 +87,8 @@ def update_requests(self, state: SampleState) -> None:
             request.state = LlmRequestState.GENERATION_COMPLETE
             # NOTE: This is a hack: set finish reason manually and set the beam 0
             request.set_finished_reason(FinishReason.LENGTH, 0)
+            if request.py_return_mm_embeddings:
+                request.py_result.append_mm_embeddings(state.host.mm_embeddings[idx])
             if request.py_return_context_logits:
                 logits = state.host.logits[idx]
                 if logits.ndim == 1:
@@ -97,6 +99,42 @@ def update_requests(self, state: SampleState) -> None:
                 request.py_result.append_context_logits(logits)
 
 
+@dataclass(kw_only=True)
+class MultimodalResult:
+    mm_embeddings: torch.Tensor = None
+
+    def values(self):
+        return vars(self).values()
+
+@dataclass(kw_only=True)
+class SampleStateWithMMResult:
+    scheduled_requests: ScheduledRequests
+
+    data: MultimodalResult = None
+
+
+class EarlyStopWithMMResult(EarlyStopSampler):
+    """
+    Use for skipping decoding step for non generation model, and return the batch_output (such as mm_embeddings)
+    """
+
+    def sample_async(self, scheduled_requests: ScheduledRequests,
+                     model_outputs) -> SampleStateWithMMResult:
+        # from model_outputs to MultimodalResult
+        data = MultimodalResult(mm_embeddings=model_outputs)
+        return SampleStateWithMMResult(scheduled_requests=scheduled_requests, data=data)
+
+    def update_requests(self, state: SampleStateWithMMResult) -> None:
+        assert isinstance(state, SampleStateWithMMResult)
+        scheduled_requests = state.scheduled_requests
+        assert (not scheduled_requests.generation_requests)
+        for idx, request in enumerate(scheduled_requests.context_requests):
+            request.state = LlmRequestState.GENERATION_COMPLETE
+            # NOTE: This is a hack: set finish reason manually and set the beam 0
+            request.set_finished_reason(FinishReason.LENGTH, 0)
+            request.py_result.append_mm_embeddings(state.data.mm_embeddings)
+
+
 def top_k_sampling_batch(logits, top_k=50):
     logits_dim = logits.dim()
     if logits_dim == 1:
diff --git a/tensorrt_llm/executor/result.py b/tensorrt_llm/executor/result.py
index 0408a6c757c0..f1b39ff09e7d 100644
--- a/tensorrt_llm/executor/result.py
+++ b/tensorrt_llm/executor/result.py
@@ -161,6 +161,7 @@ def __init__(self,
             CompletionOutput(i) for i in range(self.sampling_params.best_of)
         ]
         self._context_logits: Optional[torch.Tensor] = None
+        self._mm_embeddings: Optional[torch.Tensor] = None
 
         self._background_error_handler = None
         if background_error_handler is not None:
@@ -197,6 +198,10 @@ def outputs(self) -> List[CompletionOutput]:
     def context_logits(self) -> Optional[torch.Tensor]:
         return self._context_logits
 
+    @property
+    def mm_embeddings(self) -> Optional[torch.Tensor]:
+        return self._mm_embeddings
+
     def _handle_sequence(self,
                          finish_reasons,
                          response_tensors,
@@ -331,6 +336,9 @@ def _handle_response(self,
             if response_result.context_logits is not None:
                 self._context_logits = response_result.context_logits
 
+            if response_result.mm_embeddings is not None:
+                self._mm_embeddings = response_result.mm_embeddings
+
             # Processing background errors here ASAF during generation.
             if self._background_error_handler and (
                     handler := self._background_error_handler()):
@@ -556,7 +564,7 @@ def _exception(self, timeout: Optional[float] = None):
     def _repr_fields(self):
         return [
             'request_id', 'prompt_token_ids', 'outputs', 'finished',
-            "context_logits"
+            "context_logits", "mm_embeddings"
         ]
 
     def __repr__(self) -> str:
diff --git a/tensorrt_llm/llmapi/llm.py b/tensorrt_llm/llmapi/llm.py
index 321ec11bd759..8ea9e53b2046 100644
--- a/tensorrt_llm/llmapi/llm.py
+++ b/tensorrt_llm/llmapi/llm.py
@@ -53,6 +53,7 @@ class RequestOutput(DetokenizedGenerationResultBase, GenerationResult):
         prompt_token_ids (List[int]): The token ids of the prompt.
         outputs (List[CompletionOutput]): The output sequences of the request.
         context_logits (torch.Tensor, optional): The logits on the prompt token ids.
+        mm_embeddings (torch.Tensor, optional): The multimodal embeddings of the request.
         finished (bool): Whether the whole request is finished.
     """
 

From af28597d3a9d094bbc0b8c22897463b1b644f09c Mon Sep 17 00:00:00 2001
From: "Chang Liu (Enterprise Products)"
 <9713593+chang-l@users.noreply.github.com>
Date: Thu, 7 Aug 2025 14:53:18 -0700
Subject: [PATCH 02/15] Bug fixing and use tensor handle

Signed-off-by: Chang Liu (Enterprise Products) <9713593+chang-l@users.noreply.github.com>
---
 .../_torch/models/modeling_qwen2vl.py         |  20 ++-
 tensorrt_llm/_torch/models/modeling_utils.py  |   2 +-
 tensorrt_llm/_torch/multimodal_encoder.py     | 143 ++++++++++++++----
 tensorrt_llm/_torch/pyexecutor/_util.py       |   2 +-
 tensorrt_llm/_torch/pyexecutor/config.py      |   3 +
 tensorrt_llm/_torch/pyexecutor/llm_request.py |  11 +-
 .../_torch/pyexecutor/model_engine.py         |  66 +++++++-
 .../_torch/pyexecutor/py_executor_creator.py  |   4 +
 tensorrt_llm/_torch/pyexecutor/sampler.py     |  10 +-
 tensorrt_llm/executor/result.py               |  12 +-
 tensorrt_llm/llmapi/llm.py                    |   4 +-
 tensorrt_llm/llmapi/llm_args.py               |   2 +
 12 files changed, 217 insertions(+), 62 deletions(-)

diff --git a/tensorrt_llm/_torch/models/modeling_qwen2vl.py b/tensorrt_llm/_torch/models/modeling_qwen2vl.py
index 0c43c880bba2..c3da24718695 100644
--- a/tensorrt_llm/_torch/models/modeling_qwen2vl.py
+++ b/tensorrt_llm/_torch/models/modeling_qwen2vl.py
@@ -3,6 +3,7 @@
 from typing import Any, Dict, List, Optional, Tuple, Union
 
 import torch
+import torch.nn as nn
 from transformers import (AutoProcessor, AutoTokenizer, PretrainedConfig,
                           PreTrainedModel, Qwen2_5_VLForConditionalGeneration,
                           Qwen2VLForConditionalGeneration)
@@ -341,22 +342,28 @@ def __call__(
         }
 
 
-class Qwen2VisionModelBase:
+class Qwen2VisionModelBase(nn.Module):
 
     def __init__(self, model_config: ModelConfig[PretrainedConfig],
                  model_class: type[PreTrainedModel]):
-        self.pretrained_config = model_config.pretrained_config
+        super().__init__()
+        pretrained_config = model_config.pretrained_config
+        self.model_config = model_config
         self.device = f"cuda:{model_config.mapping.rank}"
 
-        model_path = self.pretrained_config._name_or_path
+        model_path = pretrained_config._name_or_path
         # TODO: Change the model class to TRT-LLM's Qwen2VisionModel
         # Currently, copying vision encoder on all devices.
         # NOTE: Using attn_implementation='flash_attention_2' to avoid the issue of vision model's GPU OOM.
         model = model_class.from_pretrained(
             model_path,
-            torch_dtype=self.pretrained_config.torch_dtype,
+            torch_dtype=pretrained_config.torch_dtype,
             attn_implementation='flash_attention_2').eval()
         self.visual = model.visual.to(self.device)
+        self.post_config()
+
+    def post_config(self):
+        self.config = self.visual.config
 
     def _to_device(
         self, input_tensor: Union[torch.Tensor, List, None]
@@ -644,8 +651,8 @@ def forward(
         return output_prob
 
 
-@register_auto_model("Qwen2VLForConditionalGeneration")
 @register_vision_encoder(Qwen2VisionModelBase, vlm_base_model=Qwen2VLForConditionalGeneration)
+@register_auto_model("Qwen2VLForConditionalGeneration")
 @register_input_processor(Qwen2VLInputProcessorBase, model_type="qwen2_vl")
 class Qwen2VLModel(Qwen2VLModelBase):
 
@@ -656,9 +663,8 @@ def __init__(self, model_config: ModelConfig[PretrainedConfig], *args,
                 model_config, Qwen2VLForConditionalGeneration)
         super().__init__(model_config, *args, **kwargs)
 
-
-@register_auto_model("Qwen2_5_VLForConditionalGeneration")
 @register_vision_encoder(Qwen2VisionModelBase, vlm_base_model=Qwen2_5_VLForConditionalGeneration)
+@register_auto_model("Qwen2_5_VLForConditionalGeneration")
 @register_input_processor(Qwen2VLInputProcessorBase, model_type="qwen2_5_vl")
 class Qwen2_5_VLModel(Qwen2VLModelBase):
 
diff --git a/tensorrt_llm/_torch/models/modeling_utils.py b/tensorrt_llm/_torch/models/modeling_utils.py
index f00b22e449d5..4f0a901802b8 100755
--- a/tensorrt_llm/_torch/models/modeling_utils.py
+++ b/tensorrt_llm/_torch/models/modeling_utils.py
@@ -595,7 +595,7 @@ def register_vision_encoder(vision_encoder_cls: Type[nn.Module],
 
     def wrapper(model_cls: Type[nn.Module]) -> Type[nn.Module]:
         for arch_name, registered_cls in MODEL_CLASS_MAPPING.items():
-            if registered_cls == model_cls:
+            if registered_cls.__name__ == model_cls.__name__:
                 MODEL_CLASS_VISION_ENCODER_MAPPING[arch_name] = (vision_encoder_cls, vlm_base_model)
                 break
 
diff --git a/tensorrt_llm/_torch/multimodal_encoder.py b/tensorrt_llm/_torch/multimodal_encoder.py
index 3bb219f2c832..ba428eae0a9c 100644
--- a/tensorrt_llm/_torch/multimodal_encoder.py
+++ b/tensorrt_llm/_torch/multimodal_encoder.py
@@ -1,8 +1,18 @@
-from tensorrt_llm.inputs.data import PromptInputs
-from tensorrt_llm.llmapi.llm import BaseLLM, _TorchLLM
-from typing import Any, Union
 from pathlib import Path
+from typing import Any, Optional, Union
 
+from tensorrt_llm.inputs.data import PromptInputs
+from tensorrt_llm.llmapi.llm import BaseLLM, _TorchLLM, RequestOutput
+from tensorrt_llm.sampling_params import SamplingParams
+from .._utils import nvtx_range_debug
+from typing import List, Sequence
+from tensorrt_llm.inputs import prompt_inputs
+from tqdm import tqdm
+from tensorrt_llm.inputs import create_input_processor
+from tensorrt_llm.llmapi.llm_args import PybindMirror
+from tensorrt_llm.llmapi.mpi_session import external_mpi_comm_available
+from tensorrt_llm.bindings import executor as tllm
+from tensorrt_llm._torch.pyexecutor.config import update_executor_config
 
 class MultimodalEncoder(_TorchLLM):
     """MultimodalEncoder class is the main class for running a multimodal encoder model using PyTorch backend.
@@ -27,8 +37,62 @@ def __init__(self,
                          **kwargs)
 
     def _build_model(self):
-        super()._build_model()
-        self._executor_config.mm_encoder_only = True
+        BaseLLM._build_model(self)
+        assert self._engine_dir is None
+
+        # Tokenizer loading should be after calling model_loader(), since model_loader() may download the model from HF hub.
+        # It should also be before bindings ExecutorConfig, which may depend on tokenizer info.
+        self._tokenizer = self._try_load_tokenizer()
+
+        # Multimodal special handling:
+        # 1. Default load_tokenizer may fail because MM has different tokenizer configuration. Hence we initialize it inside input processor
+        # 2. May need to modify model weights for MM (e.g., resize vocab embedding). We must do such operation via input processor's __init__
+        self.input_processor = create_input_processor(self._hf_model_dir,
+                                                      self.tokenizer)
+        self.tokenizer = self.input_processor.tokenizer
+
+        max_batch_size = self.args.max_batch_size
+        max_num_tokens = self.args.max_num_tokens
+        max_seq_len = self.args.max_seq_len
+
+        kwargs = {}
+        if self._on_trt_backend:
+            kwargs[
+                "batching_type"] = self.args.batching_type or tllm.BatchingType.INFLIGHT
+
+        self._executor_config = tllm.ExecutorConfig(
+            scheduler_config=PybindMirror.maybe_to_pybind(
+                self.args.scheduler_config),
+            max_batch_size=max_batch_size,
+            max_num_tokens=max_num_tokens,
+            **kwargs)
+
+        max_batch_size = self._executor_config.max_batch_size
+        update_executor_config(
+            self._executor_config,
+            backend=self.args.backend,
+            pytorch_backend_config=self.args.get_pytorch_backend_config()
+            if self.args.backend in ["pytorch", "_autodeploy"] else None,
+            mapping=self.args.parallel_config.to_mapping(),
+            hf_model_dir=self._hf_model_dir,
+            max_input_len=self.args.max_input_len,
+            max_seq_len=max_seq_len,
+            checkpoint_format=None if self.args.backend == "_autodeploy" else
+            self.args.checkpoint_format,
+            checkpoint_loader=None if self.args.backend == "_autodeploy" else
+            self.args.checkpoint_loader,
+            mm_encoder_only=True)
+
+        self._executor = self._executor_cls.create(
+            self._engine_dir,
+            executor_config=self._executor_config,
+            model_world_size=self.args.parallel_config.world_size,
+            mpi_session=self.mpi_session,
+            reuse_mpi_comm=external_mpi_comm_available(
+                self.args.parallel_config.world_size),
+            is_llm_executor=True, # TODO: check if this is correct or needed
+            garbage_collection_gen0_threshold=self.args.
+            garbage_collection_gen0_threshold)
 
     def _validate_mm_args_for_torch_backend(self, kwargs: dict) -> None:
         """Validate that users don't pass LLM-specific arguments when using MultimodalEncoder (PyTorch).
@@ -38,47 +102,66 @@ def _validate_mm_args_for_torch_backend(self, kwargs: dict) -> None:
 
     def generate(
         self,
-        inputs: PromptInputs,
-    ) -> RequestOutput:
-        """Generate embeddings (and other multimodal encoder outputs) for multiple multimodal requests in parallel.
+        inputs: Union[PromptInputs, Sequence[PromptInputs]],
+        use_tqdm: bool = True,
+    ) -> Union[RequestOutput, List[RequestOutput]]:
+        """Generate output for the given prompts in the synchronous mode.
+        Synchronous generation accepts either single prompt or batched prompts.
 
         Args:
-            mm_requests: List of multimodal requests to process
-
+            inputs (tensorrt_llm.inputs.data.PromptInputs, Sequence[tensorrt_llm.inputs.data.PromptInputs]): The prompt text or token ids.
+                It can be single prompt or batched prompts.
         Returns:
-            List of generation results
+            Union[tensorrt_llm.llmapi.RequestOutput, List[tensorrt_llm.llmapi.RequestOutput]]: The output data of the completion request to the LLM.
         """
-        async def _process_requests():
-            # Submit all requests first
-            futures = []
-            for request in mm_requests:
-                future = await self.generate_async(request)
-                futures.append(future)
-
-            # Then wait for all results
-            results = []
-            for future in futures:
-                result = await future.aresult()
-                results.append(result)
-            return results
-
-        # Run the async operations in an event loop
-        return asyncio.run(_process_requests())
+        unbatched = not isinstance(inputs, list)
+        if not unbatched:
+            if isinstance(inputs[0], int):
+                unbatched = True
+
+        if unbatched:
+            inputs = [inputs]
+
+        inputs = [prompt_inputs(i) for i in inputs]
+
+        def _item_at(maybe_batched: Union[Any, Sequence[Any]], pos: int) -> Any:
+            if isinstance(maybe_batched, list):
+                return maybe_batched[pos]
+            else:
+                return maybe_batched
+
+        futures = []
+        for i, request_inputs in enumerate(inputs):
+            future = self.generate_async(
+                request_inputs
+            )
+            futures.append(future)
+
+        for future in tqdm(futures,
+                           desc="Processed requests",
+                           dynamic_ncols=True,
+                           disable=not use_tqdm):
+            future.result()
+
+        if unbatched:
+            futures = futures[0]
+
+        return futures
 
     @nvtx_range_debug("MM_encoder.generate_async", color="green", category="VisionEncoder")
     def generate_async(
         self,
         inputs: PromptInputs,
         sampling_params: Optional[SamplingParams] = None,
-    ) -> RequestOutput:
+    ):
         """Generate output for the given multimodal request in the asynchronous mode.
         Asynchronous generation accepts single multimodal request only.
 
         Returns:
-            tensorrt_llm.llmapi.RequestOutput: The output data of the completion request to the LLM.
+            Future that resolves to tensorrt_llm.llmapi.RequestOutput containing mm_embeddings
         """
         # TODO: possible preprocess the input
-        result = BaseLLM.generate_async(self, inputs, sampling_params)
+        result = super().generate_async(inputs, sampling_params)
         # TODO: possible postprocess the result
         return result
 
diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py
index 27be8c643613..122e3b417a2d 100644
--- a/tensorrt_llm/_torch/pyexecutor/_util.py
+++ b/tensorrt_llm/_torch/pyexecutor/_util.py
@@ -30,7 +30,7 @@
 from .resource_manager import (KVCacheManager, MambaHybridCacheManager,
                                PeftCacheManager, ResourceManager,
                                ResourceManagerType)
-from .sampler import EarlyStopSampler, TorchSampler, TRTLLMSampler
+from .sampler import EarlyStopSampler, TorchSampler, TRTLLMSampler, EarlyStopWithMMResult
 from .scheduler import (BindCapacityScheduler, BindMicroBatchScheduler,
                         SimpleScheduler)
 from .seq_slot_manager import SeqSlotManager
diff --git a/tensorrt_llm/_torch/pyexecutor/config.py b/tensorrt_llm/_torch/pyexecutor/config.py
index cfd3c4947148..f6f3ae8a6c0a 100644
--- a/tensorrt_llm/_torch/pyexecutor/config.py
+++ b/tensorrt_llm/_torch/pyexecutor/config.py
@@ -99,6 +99,9 @@ class PyTorchConfig:
 
     force_dynamic_quantization: bool = False
 
+    # If true, ONLY the vision encoder part of the full model is loaded/executed.
+    mm_encoder_only: bool = False
+
     # If true, adjust PyTorch CUDA memory fraction to correspond to the
     # total GPU memory minus the statically allocated engine memory.
     # If false, set the PyTorch CUDA memory fraction to 1.0.
diff --git a/tensorrt_llm/_torch/pyexecutor/llm_request.py b/tensorrt_llm/_torch/pyexecutor/llm_request.py
index a5c3cee79793..85b1f4c6429b 100644
--- a/tensorrt_llm/_torch/pyexecutor/llm_request.py
+++ b/tensorrt_llm/_torch/pyexecutor/llm_request.py
@@ -1,12 +1,13 @@
 from copy import deepcopy
 from dataclasses import dataclass
-from typing import List, Optional, Union
+from typing import List, Optional, Union, Dict, Any
 
 import torch
 
 import tensorrt_llm.bindings
 from tensorrt_llm.bindings import executor as tllm_executor
 from tensorrt_llm.executor.result import TokenLogprobs
+from tensorrt_llm._torch.shared_tensor import SharedTensorContainer
 
 SamplingConfig = tensorrt_llm.bindings.SamplingConfig
 '''
@@ -189,8 +190,7 @@ def append_log_probs(self,
             self._log_probs.append(log_probs, cum_log_probs)
 
     def append_mm_embeddings(self, mm_embeddings: torch.Tensor):
-        # TODO: add to_handle to use shared tensor support
-        self._mm_embeddings = mm_embeddings.to(device='cpu', non_blocking=True)
+        self._mm_embeddings = SharedTensorContainer.from_tensor(mm_embeddings).dump_to_dict()
 
     def set_log_probs(self, log_probs: list[TokenLogprobs],
                       cum_log_probs: list[float]):
@@ -229,11 +229,14 @@ def log_probs(self) -> list[TokenLogprobs] | None:
     def cum_log_probs(self) -> list[float] | None:
         return self._log_probs and self._log_probs.cum_log_probs
 
+    @property
+    def mm_embedding_handle(self) -> Dict[str, Any] | None:
+        return self._mm_embeddings
 
 class LlmResult:
     """LlmResult wraps `bindings.executor.Result` but detour some features to Python implementation"""
     py_result_properties = frozenset(
-        ('context_logits', 'generation_logits', 'log_probs', 'cum_log_probs'))
+        ('context_logits', 'generation_logits', 'log_probs', 'cum_log_probs', 'mm_embedding_handle'))
 
     def __init__(self,
                  result: Union[bytes, tensorrt_llm.bindings.executor.Result],
diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py
index 6af03f4c60f2..962874b2cb17 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_engine.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -1074,6 +1074,10 @@ def init_meta_tensor(t: torch.Tensor):
             elif load_format == LoadFormat.DUMMY:
                 initialize_dummy_weights(model)
 
+            elif load_format == LoadFormat.VISION_ONLY:
+                # Vision weights are already loaded within model
+                pass
+
             else:
                 raise NotImplementedError(
                     f"No load support for load format: {load_format}")
@@ -1099,7 +1103,11 @@ def _call_load_weights(self, load_method, weights, weight_mapper):
             load_method(weights)
 
     def _init_max_seq_len(self):
-        inferred_max_seq_len = self.model.infer_max_seq_len()
+        if self.pytorch_backend_config.mm_encoder_only:
+            # TODO: hardcoded for now, need to infer as: max_num_token_per_image * max_num_images (can be given as a parameter or defined same as max_seq_len)
+            inferred_max_seq_len = 32768
+        else:
+            inferred_max_seq_len = self.model.infer_max_seq_len()
         if self.max_seq_len is None:
             logger.info(
                 f"max_seq_len is not specified, using inferred value {inferred_max_seq_len}"
@@ -1618,6 +1626,7 @@ def _prepare_tp_inputs_no_cache(
         multi_modal_data = []
         draft_lens = []
         request_ids = []
+        multimodal_params_list = []
 
         for request in scheduled_requests.context_requests:
             prompt_tokens = request.get_tokens(0)
@@ -1634,6 +1643,19 @@ def _prepare_tp_inputs_no_cache(
             if multimodal_embedding is not None:
                 multi_modal_data.append(multimodal_embedding)
 
+            # Multimodal
+            multimodal_params = MultimodalParams(
+                multimodal_data=request.py_multimodal_data,
+            )
+            multimodal_params.to_device("multimodal_data",
+                                        "cuda",
+                                        pin_memory=True)
+
+            if multimodal_params.has_content():
+                multimodal_params_list.append(multimodal_params)
+
+            request.py_batch_idx = request.py_seq_slot
+
         num_tokens = len(input_ids)
         assert num_tokens <= self.max_num_tokens, (
             "num_tokens should be less than or equal to max_num_tokens")
@@ -1685,7 +1707,7 @@ def _prepare_tp_inputs_no_cache(
             'input_ids': self.input_ids_cuda[:num_tokens],
             'position_ids': self.position_ids_cuda[:num_tokens].unsqueeze(0),
             'inputs_embeds': None,
-            'multi_modal_data': multi_modal_data
+            "multimodal_params": multimodal_params_list
         }
 
         if bool(lora_params):
@@ -2127,15 +2149,15 @@ def forward(
                                                      moe_enable_update)
 
         if kv_cache_manager is None:
-            # here you need to add logic for prepare_tp_inputs_no_cache and forward_step for mm encoder only
             inputs, gather_ids = self._prepare_tp_inputs_no_cache(
                 scheduled_requests, attn_metadata, spec_metadata)
 
             with MoeLoadBalancerIterContext(moe_load_balancer):
-                # need to add key ['logits']=None to the output of forward_step for mm encoder only
-                return self._forward_step(inputs, gather_ids,
-                                          gather_context_logits)
-                # lets update requst state here after this 1step finished
+                # Special handling for multimodal encoder only mode
+                if self.pytorch_backend_config.mm_encoder_only:
+                    return self._forward_step_mm_encoder_only(inputs, scheduled_requests)
+                else:
+                    return self._forward_step(inputs, gather_ids, gather_context_logits)
         with self._maybe_pad_batch(scheduled_requests, kv_cache_manager,
                                    spec_resource_manager) as scheduled_requests:
             maybe_graph = self._maybe_get_cuda_graph(scheduled_requests)
@@ -2229,6 +2251,36 @@ def _forward_step(self,
         else:
             return {'logits': logits}
 
+    @nvtx_range("_forward_step_mm_encoder_only")
+    def _forward_step_mm_encoder_only(self,
+                                      inputs: Dict[str, Any],
+                                      scheduled_requests: ScheduledRequests) -> Dict[str, Any]:
+        """Forward step for multimodal encoder only mode - returns mm_embeddings instead of logits."""
+        # Get multimodal parameters from inputs
+        multimodal_params = inputs.get("multimodal_params", [])
+        if not multimodal_params or len(multimodal_params) == 0:
+            # Return empty embeddings if no multimodal data
+            return {'mm_embeddings': []}
+        if getattr(scheduled_requests.context_requests[0], 'multimodal_lengths', None) is None:
+            multimodal_chunks = None
+        else:
+            multimodal_chunks = [
+                sum(request.multimodal_lengths)
+                for request in scheduled_requests.context_requests
+                if request.multimodal_lengths is not None
+            ]
+        # For mm_encoder_only mode, we only run the vision encoder part
+        # The model should be a vision encoder (e.g., Qwen2VisionModelBase)
+        mm_embeddings = self.model.forward(multimodal_params)
+        assert len(mm_embeddings) == 1, "mm_embeddings should be a 1-element list, mix modality (video+image) is not supported"
+
+        if multimodal_chunks is None or len(multimodal_chunks) != len(multimodal_params):
+            mm_embeddings = list(torch.chunk(mm_embeddings[0], len(scheduled_requests.context_requests), dim=0))
+        else:
+            mm_embeddings = list(torch.split(mm_embeddings[0], multimodal_chunks, dim=0))
+
+        return {'mm_embeddings': mm_embeddings, 'logits': None}
+
     def _init_userbuffers(self, hidden_size):
         if self.mapping.tp_size <= 1:
             return False
diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
index c6f99a31ec45..22fcf982d26d 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
@@ -170,7 +170,11 @@ def _mangle_executor_config(executor_config: ExecutorConfig):
         executor_config.enable_chunked_context = False
 
     if executor_config.mm_encoder_only:
+        from tensorrt_llm.llmapi.llm_args import LoadFormat
         pytorch_backend_config.mm_encoder_only = True
+        pytorch_backend_config.load_format = LoadFormat.VISION_ONLY
+        # TODO: add comment and print warning here
+        pytorch_backend_config.disable_overlap_scheduler = True
 
 def _get_mapping(executor_config: ExecutorConfig) -> Mapping:
     if executor_config.mapping is None:
diff --git a/tensorrt_llm/_torch/pyexecutor/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler.py
index e7edd49f97da..1fbbcfc898e6 100644
--- a/tensorrt_llm/_torch/pyexecutor/sampler.py
+++ b/tensorrt_llm/_torch/pyexecutor/sampler.py
@@ -1,7 +1,7 @@
 from abc import ABC, abstractmethod
 from collections.abc import Iterable
 from dataclasses import dataclass
-from typing import Literal
+from typing import Literal, List
 
 import torch
 
@@ -101,7 +101,7 @@ def update_requests(self, state: SampleState) -> None:
 
 @dataclass(kw_only=True)
 class MultimodalResult:
-    mm_embeddings: torch.Tensor = None
+    mm_embeddings: List[torch.Tensor] = None
 
     def values(self):
         return vars(self).values()
@@ -121,18 +121,20 @@ class EarlyStopWithMMResult(EarlyStopSampler):
     def sample_async(self, scheduled_requests: ScheduledRequests,
                      model_outputs) -> SampleStateWithMMResult:
         # from model_outputs to MultimodalResult
-        data = MultimodalResult(mm_embeddings=model_outputs)
+        data = MultimodalResult(mm_embeddings=model_outputs['mm_embeddings'])
         return SampleStateWithMMResult(scheduled_requests=scheduled_requests, data=data)
 
     def update_requests(self, state: SampleStateWithMMResult) -> None:
         assert isinstance(state, SampleStateWithMMResult)
         scheduled_requests = state.scheduled_requests
         assert (not scheduled_requests.generation_requests)
+        mm_embeddings = state.data.mm_embeddings
         for idx, request in enumerate(scheduled_requests.context_requests):
             request.state = LlmRequestState.GENERATION_COMPLETE
             # NOTE: This is a hack: set finish reason manually and set the beam 0
             request.set_finished_reason(FinishReason.LENGTH, 0)
-            request.py_result.append_mm_embeddings(state.data.mm_embeddings)
+            if idx < len(mm_embeddings):
+                request.py_result.append_mm_embeddings(mm_embeddings[idx])
 
 
 def top_k_sampling_batch(logits, top_k=50):
diff --git a/tensorrt_llm/executor/result.py b/tensorrt_llm/executor/result.py
index f1b39ff09e7d..d2b8ae79e537 100644
--- a/tensorrt_llm/executor/result.py
+++ b/tensorrt_llm/executor/result.py
@@ -3,7 +3,7 @@
 import weakref
 from dataclasses import dataclass, field
 from queue import Empty, Queue
-from typing import (TYPE_CHECKING, Any, Callable, List, Literal, NamedTuple,
+from typing import (TYPE_CHECKING, Any, Callable, Dict, List, Literal, NamedTuple,
                     Optional, TypeAlias, Union)
 from weakref import WeakMethod
 
@@ -161,7 +161,7 @@ def __init__(self,
             CompletionOutput(i) for i in range(self.sampling_params.best_of)
         ]
         self._context_logits: Optional[torch.Tensor] = None
-        self._mm_embeddings: Optional[torch.Tensor] = None
+        self._mm_embedding_handle: Optional[Dict[str, Any]] = None
 
         self._background_error_handler = None
         if background_error_handler is not None:
@@ -199,8 +199,8 @@ def context_logits(self) -> Optional[torch.Tensor]:
         return self._context_logits
 
     @property
-    def mm_embeddings(self) -> Optional[torch.Tensor]:
-        return self._mm_embeddings
+    def mm_embedding_handle(self) -> Optional[Dict[str, Any]]:
+        return self._mm_embedding_handle
 
     def _handle_sequence(self,
                          finish_reasons,
@@ -336,8 +336,8 @@ def _handle_response(self,
             if response_result.context_logits is not None:
                 self._context_logits = response_result.context_logits
 
-            if response_result.mm_embeddings is not None:
-                self._mm_embeddings = response_result.mm_embeddings
+            if response_result.mm_embedding_handle is not None:
+                self._mm_embedding_handle = response_result.mm_embedding_handle
 
             # Processing background errors here ASAF during generation.
             if self._background_error_handler and (
diff --git a/tensorrt_llm/llmapi/llm.py b/tensorrt_llm/llmapi/llm.py
index 8ea9e53b2046..dc67b00c8c55 100644
--- a/tensorrt_llm/llmapi/llm.py
+++ b/tensorrt_llm/llmapi/llm.py
@@ -53,7 +53,7 @@ class RequestOutput(DetokenizedGenerationResultBase, GenerationResult):
         prompt_token_ids (List[int]): The token ids of the prompt.
         outputs (List[CompletionOutput]): The output sequences of the request.
         context_logits (torch.Tensor, optional): The logits on the prompt token ids.
-        mm_embeddings (torch.Tensor, optional): The multimodal embeddings of the request.
+        mm_embedding_handle (Dict[str, Any], optional): The multimodal embedding handle of the request.
         finished (bool): Whether the whole request is finished.
     """
 
@@ -82,7 +82,7 @@ def prompt(self) -> Optional[str]:
 
     def _repr_fields(self):
         return [
-            "request_id", "prompt", "prompt_token_ids", "outputs", "finished"
+            "request_id", "prompt", "prompt_token_ids", "outputs", "finished", "mm_embedding_handle"
         ]
 
 
diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py
index b7d46ed6fa25..62f96d01030c 100644
--- a/tensorrt_llm/llmapi/llm_args.py
+++ b/tensorrt_llm/llmapi/llm_args.py
@@ -1940,6 +1940,8 @@ class LoadFormat(Enum):
     AUTO = 0
     # Initialize all weights randomly.
     DUMMY = 1
+    # Only load the multimodal(vision) encoder weights
+    VISION_ONLY = 2
 
 
 class TorchCompileConfig(StrictBaseModel):

From 3c7c788a46e5d9b04484eb91a516fbb0817fd42a Mon Sep 17 00:00:00 2001
From: "Chang Liu (Enterprise Products)"
 <9713593+chang-l@users.noreply.github.com>
Date: Thu, 7 Aug 2025 19:48:55 -0700
Subject: [PATCH 03/15] Enable trtllm-serve endpoint

Signed-off-by: Chang Liu (Enterprise Products) <9713593+chang-l@users.noreply.github.com>
---
 tensorrt_llm/__init__.py                      |  2 +
 tensorrt_llm/commands/serve.py                | 91 ++++++++++++++++++-
 tensorrt_llm/executor/result.py               |  2 +-
 tensorrt_llm/llmapi/__init__.py               |  2 +
 tensorrt_llm/llmapi/disagg_utils.py           |  1 +
 .../mm_encoder.py}                            | 14 +--
 tensorrt_llm/serve/openai_protocol.py         |  3 +
 tensorrt_llm/serve/openai_server.py           | 86 +++++++++++++++++-
 8 files changed, 186 insertions(+), 15 deletions(-)
 rename tensorrt_llm/{_torch/multimodal_encoder.py => llmapi/mm_encoder.py} (95%)

diff --git a/tensorrt_llm/__init__.py b/tensorrt_llm/__init__.py
index f54026a8cbc0..ee5fa8c967bd 100644
--- a/tensorrt_llm/__init__.py
+++ b/tensorrt_llm/__init__.py
@@ -51,6 +51,7 @@ def _add_trt_llm_dll_directory():
 from .disaggregated_params import DisaggregatedParams
 from .functional import Tensor, constant
 from .llmapi import LLM, LlmArgs
+from .llmapi import MultimodalEncoder
 from .llmapi.llm_args import LlmArgs, TorchLlmArgs, TrtLlmArgs
 from .logger import logger
 from .mapping import Mapping
@@ -103,6 +104,7 @@ def _add_trt_llm_dll_directory():
     'quantization',
     'tools',
     'LLM',
+    'MultimodalEncoder',
     'LlmArgs',
     'TorchLlmArgs',
     'TrtLlmArgs',
diff --git a/tensorrt_llm/commands/serve.py b/tensorrt_llm/commands/serve.py
index 4f26be6579b3..521d9eb62e6c 100644
--- a/tensorrt_llm/commands/serve.py
+++ b/tensorrt_llm/commands/serve.py
@@ -12,6 +12,7 @@
 from torch.cuda import device_count
 
 from tensorrt_llm import LLM as PyTorchLLM
+from tensorrt_llm import MultimodalEncoder
 from tensorrt_llm._tensorrt_engine import LLM
 from tensorrt_llm._utils import mpi_rank
 from tensorrt_llm.executor.utils import LlmLauncherEnvs
@@ -172,6 +173,19 @@ def launch_server(host: str,
 
     asyncio.run(server(host, port))
 
+def launch_mm_encoder_server(host: str,
+                  port: int,
+                  encoder_args: dict,
+                  metadata_server_cfg: Optional[MetadataServerConfig] = None,
+                  ):
+    model = encoder_args["model"]
+    mm_encoder = MultimodalEncoder(**encoder_args)
+
+    server = OpenAIServer(llm=mm_encoder,
+                          model=model,
+                          server_role=ServerRole.MM_ENCODER,
+                          metadata_server_cfg=metadata_server_cfg)
+    asyncio.run(server(host, port))
 
 @click.command("serve")
 @click.argument("model", type=str)
@@ -332,6 +346,80 @@ def serve(
     launch_server(host, port, llm_args, metadata_server_cfg, server_role)
 
 
+@click.command("mmemb_serve")
+@click.argument("model", type=str)
+@click.option("--host",
+              type=str,
+              default="localhost",
+              help="Hostname of the server.")
+@click.option("--port", type=int, default=8000, help="Port of the server.")
+@click.option('--log_level',
+              type=click.Choice(severity_map.keys()),
+              default='info',
+              help="The logging level.")
+@click.option("--max_batch_size",
+              type=int,
+              default=BuildConfig.max_batch_size,
+              help="Maximum number of requests that the engine can schedule.")
+@click.option("--gpus_per_node",
+              type=int,
+              default=None,
+              help="Number of GPUs per node. Default to None, and it will be "
+              "detected automatically.")
+@click.option("--trust_remote_code",
+              is_flag=True,
+              default=False,
+              help="Flag for HF transformers.")
+@click.option(
+    "--extra_encoder_options",
+    type=str,
+    default=None,
+    help=
+    "Path to a YAML file that overwrites the parameters specified by trtllm-serve."
+)
+@click.option("--metadata_server_config_file",
+              type=str,
+              default=None,
+              help="Path to metadata server config file")
+def serve_encoder(model: str, host: str, port: int,
+          log_level: str, max_batch_size: int,
+          gpus_per_node: Optional[int],
+          trust_remote_code: bool,
+          extra_encoder_options: Optional[str],
+          metadata_server_config_file: Optional[str]):
+    """Running an OpenAI API compatible server
+
+    MODEL: model name | HF checkpoint path | TensorRT engine path
+    """
+    logger.set_level(log_level)
+
+    # TODO: expose more argument progressivly
+    llm_args, _ = get_llm_args(
+        model=model,
+        max_batch_size=max_batch_size,
+        gpus_per_node=gpus_per_node,
+        trust_remote_code=trust_remote_code)
+
+    encoder_args_extra_dict = {}
+    if extra_encoder_options is not None:
+        with open(extra_encoder_options, 'r') as f:
+            encoder_args_extra_dict = yaml.safe_load(f)
+    encoder_args = update_llm_args_with_extra_dict(llm_args, encoder_args_extra_dict)
+
+    metadata_server_cfg = parse_metadata_server_config_file(
+        metadata_server_config_file)
+
+    if metadata_server_cfg is not None:
+        assert server_role is not None, "server_role is required when metadata_server_cfg is provided"
+        try:
+            server_role = ServerRole[server_role.upper()]
+            assert server_role == ServerRole.MM_ENCODER, "server_role must be MM_ENCODER for multimodal encoder"
+        except ValueError:
+            raise ValueError(f"Invalid server role: {server_role}. " \
+                             f"Must be one of: {', '.join([role.name for role in ServerRole])}")
+    launch_mm_encoder_server(host, port, encoder_args, metadata_server_cfg)
+
+
 def get_ctx_gen_server_urls(
         server_configs: List[CtxGenServerConfig]) -> List[str]:
     ctx_server_urls = []
@@ -637,7 +725,8 @@ def resolve_command(self, ctx, args):
     commands={
         "serve": serve,
         "disaggregated": disaggregated,
-        "disaggregated_mpi_worker": disaggregated_mpi_worker
+        "disaggregated_mpi_worker": disaggregated_mpi_worker,
+        "mmemb_serve": serve_encoder
     })
 
 if __name__ == "__main__":
diff --git a/tensorrt_llm/executor/result.py b/tensorrt_llm/executor/result.py
index d2b8ae79e537..7374f0ceb7b3 100644
--- a/tensorrt_llm/executor/result.py
+++ b/tensorrt_llm/executor/result.py
@@ -564,7 +564,7 @@ def _exception(self, timeout: Optional[float] = None):
     def _repr_fields(self):
         return [
             'request_id', 'prompt_token_ids', 'outputs', 'finished',
-            "context_logits", "mm_embeddings"
+            "context_logits", "mm_embedding_handle"
         ]
 
     def __repr__(self) -> str:
diff --git a/tensorrt_llm/llmapi/__init__.py b/tensorrt_llm/llmapi/__init__.py
index a5f49917886c..4e6c906ad3eb 100644
--- a/tensorrt_llm/llmapi/__init__.py
+++ b/tensorrt_llm/llmapi/__init__.py
@@ -3,6 +3,7 @@
 from ..sampling_params import GuidedDecodingParams, SamplingParams
 from .build_cache import BuildCacheConfig
 from .llm import LLM, RequestOutput
+from .mm_encoder import MultimodalEncoder
 # yapf: disable
 from .llm_args import (AttentionDpConfig, AutoDecodingConfig, BatchingType,
                        CacheTransceiverConfig, CalibConfig,
@@ -20,6 +21,7 @@
 
 __all__ = [
     'LLM',
+    'MultimodalEncoder',
     'CompletionOutput',
     'RequestOutput',
     'GuidedDecodingParams',
diff --git a/tensorrt_llm/llmapi/disagg_utils.py b/tensorrt_llm/llmapi/disagg_utils.py
index f929c701fe4c..faa9bd08b34a 100644
--- a/tensorrt_llm/llmapi/disagg_utils.py
+++ b/tensorrt_llm/llmapi/disagg_utils.py
@@ -19,6 +19,7 @@
 class ServerRole(Enum):
     CONTEXT = 0
     GENERATION = 1
+    MM_ENCODER = 2
 
 
 @dataclass
diff --git a/tensorrt_llm/_torch/multimodal_encoder.py b/tensorrt_llm/llmapi/mm_encoder.py
similarity index 95%
rename from tensorrt_llm/_torch/multimodal_encoder.py
rename to tensorrt_llm/llmapi/mm_encoder.py
index ba428eae0a9c..5d25c6e190ef 100644
--- a/tensorrt_llm/_torch/multimodal_encoder.py
+++ b/tensorrt_llm/llmapi/mm_encoder.py
@@ -2,22 +2,19 @@
 from typing import Any, Optional, Union
 
 from tensorrt_llm.inputs.data import PromptInputs
-from tensorrt_llm.llmapi.llm import BaseLLM, _TorchLLM, RequestOutput
+from .llm import BaseLLM, _TorchLLM, RequestOutput
 from tensorrt_llm.sampling_params import SamplingParams
-from .._utils import nvtx_range_debug
+from tensorrt_llm._utils import nvtx_range_debug
 from typing import List, Sequence
 from tensorrt_llm.inputs import prompt_inputs
 from tqdm import tqdm
 from tensorrt_llm.inputs import create_input_processor
-from tensorrt_llm.llmapi.llm_args import PybindMirror
-from tensorrt_llm.llmapi.mpi_session import external_mpi_comm_available
+from .llm_args import PybindMirror
+from .mpi_session import external_mpi_comm_available
 from tensorrt_llm.bindings import executor as tllm
-from tensorrt_llm._torch.pyexecutor.config import update_executor_config
 
 class MultimodalEncoder(_TorchLLM):
     """MultimodalEncoder class is the main class for running a multimodal encoder model using PyTorch backend.
-
-    Parameters:
 """
 
     def __init__(self,
@@ -66,7 +63,7 @@ def _build_model(self):
             max_batch_size=max_batch_size,
             max_num_tokens=max_num_tokens,
             **kwargs)
-
+        from tensorrt_llm._torch.pyexecutor.config import update_executor_config
         max_batch_size = self._executor_config.max_batch_size
         update_executor_config(
             self._executor_config,
@@ -165,4 +162,3 @@ def generate_async(
         # TODO: possible postprocess the result
         return result
 
-
diff --git a/tensorrt_llm/serve/openai_protocol.py b/tensorrt_llm/serve/openai_protocol.py
index 59be57ffd907..b50ef0085136 100644
--- a/tensorrt_llm/serve/openai_protocol.py
+++ b/tensorrt_llm/serve/openai_protocol.py
@@ -390,6 +390,9 @@ class ChatCompletionResponseChoice(OpenAIBaseModel):
     logprobs: Optional[ChatCompletionLogProbs] = None
     finish_reason: Optional[str] = None
     stop_reason: Optional[Union[int, str]] = None
+    # TODO: progressivly add more info like input_ids, specific_token_ids, mrope, mm_hashes, etc
+    # TODO: refer to ChatCompletionLogProbs
+    mm_embedding_handle: Optional[Dict[str, Any]] = None
 
     disaggregated_params: Optional[DisaggregatedParams] = Field(default=None)
 
diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py
index d90578ce36b3..3cc498bfcce4 100644
--- a/tensorrt_llm/serve/openai_server.py
+++ b/tensorrt_llm/serve/openai_server.py
@@ -7,7 +7,7 @@
 from datetime import datetime
 from http import HTTPStatus
 from pathlib import Path
-from typing import Any, AsyncGenerator, AsyncIterator, List, Optional
+from typing import Any, AsyncGenerator, AsyncIterator, List, Optional, Union
 
 import uvicorn
 from fastapi import FastAPI, Request
@@ -16,6 +16,7 @@
 from transformers import AutoConfig, AutoProcessor
 
 from tensorrt_llm._tensorrt_engine import LLM
+from tensorrt_llm.llmapi import MultimodalEncoder
 # yapf: disable
 from tensorrt_llm.executor import CppExecutorError
 from tensorrt_llm.executor.postproc_worker import PostprocParams
@@ -41,7 +42,7 @@
     chat_stream_post_processor, completion_response_post_processor,
     completion_stream_post_processor)
 from tensorrt_llm.version import __version__ as VERSION
-
+from tensorrt_llm.serve.openai_protocol import ChatMessage, ChatCompletionResponseChoice, UsageInfo
 from .._utils import nvtx_mark
 
 # yapf: enale
@@ -51,7 +52,7 @@
 class OpenAIServer:
 
     def __init__(self,
-                 llm: LLM,
+                 llm: Union[LLM, MultimodalEncoder],
                  model: str,
                  server_role: Optional[ServerRole],
                  metadata_server_cfg: MetadataServerConfig):
@@ -108,7 +109,11 @@ async def lifespan(app: FastAPI):
         async def validation_exception_handler(_, exc):
             return self.create_error_response(message=str(exc))
 
-        self.register_routes()
+        if self.server_role is not ServerRole.MM_ENCODER:
+            self.register_routes()
+        else:
+            assert isinstance(self.llm, MultimodalEncoder), "llm must be a MultimodalEncoder for multimodal encoder"
+            self.register_mm_encoder_routes()
 
     async def await_disconnected(self, raw_request: Request, promise):
         if raw_request is None:
@@ -152,6 +157,17 @@ def register_routes(self):
                                self.openai_chat,
                                methods=["POST"])
 
+    def register_mm_encoder_routes(self):
+        self.app.add_api_route("/health", self.health, methods=["GET"])
+        self.app.add_api_route("/health_generate", self.health_generate, methods=["GET"])
+        self.app.add_api_route("/version", self.version, methods=["GET"])
+        self.app.add_api_route("/v1/models", self.get_model, methods=["GET"])
+        # TODO: the metrics endpoint only reports iteration stats, not the runtime stats for now
+        self.app.add_api_route("/metrics", self.get_iteration_stats, methods=["GET"])
+        self.app.add_api_route("/v1/chat/completions",
+                               self.openai_mm_encoder,
+                               methods=["POST"])
+
     async def health(self) -> Response:
         return Response(status_code=200)
 
@@ -324,6 +340,68 @@ async def create_chat_response(
             logger.error(traceback.format_exc())
             return self.create_error_response(str(e))
 
+    async def openai_mm_encoder(self, request: ChatCompletionRequest, raw_request: Request) -> Response:
+
+        async def create_mm_embedding_response(
+                promise: RequestOutput):
+            await promise.aresult()
+            mm_embedding_handle = promise.mm_embedding_handle
+            num_tokens = mm_embedding_handle["tensor_size"][0]
+            # TODO: need to add more info like input_ids, specific_token_ids, mrope, mm_hashes, etc
+            return ChatCompletionResponse(
+                id=str(promise.request_id),
+                model=self.model,
+                choices=[ChatCompletionResponseChoice(index=0, message=ChatMessage(role="assistant", content="dummy"), mm_embedding_handle=mm_embedding_handle, finish_reason="length")],
+                usage=UsageInfo(prompt_tokens=num_tokens, completion_tokens=1, total_tokens=num_tokens+1))
+
+        try:
+            check_multiple_response(request.n, self.llm.args.backend)
+            conversation: List[ConversationMessage] = []
+            tool_dicts = None if request.tools is None else [
+                tool.model_dump() for tool in request.tools
+            ]
+            # TODO: placeholder for multimodal disagg e2e
+            disaggregated_params = to_llm_disaggregated_params(request.disaggregated_params)
+
+            conversation, mm_coroutines, mm_placeholder_counts = parse_chat_messages_coroutines(request.messages, self.model_config)
+
+            if request.prompt_token_ids is not None:
+                prompt = request.prompt_token_ids
+            else:
+                prompt: str = apply_chat_template(
+                    model_type=self.model_config.model_type,
+                    tokenizer=self.tokenizer,
+                    processor=self.processor,
+                    conversation=conversation,
+                    add_generation_prompt=request.add_generation_prompt,
+                    mm_placeholder_counts=mm_placeholder_counts,
+                    tools=tool_dicts,
+                    documents=request.documents,
+                    chat_template=request.chat_template,
+                    chat_template_kwargs=request.chat_template_kwargs or {},
+                )
+            prompt = prompt_inputs(prompt)
+
+            mm_data = await mm_coroutines
+            if mm_data is not None:
+                prompt["multi_modal_data"] = mm_data
+
+            promise = self.llm.generate_async(
+                inputs=prompt,
+            )
+            asyncio.create_task(self.await_disconnected(raw_request, promise))
+
+            response = await create_mm_embedding_response(promise)
+            return JSONResponse(content=response.model_dump())
+
+        except CppExecutorError:
+            logger.error(traceback.format_exc())
+            # If internal executor error is raised, shutdown the server
+            signal.raise_signal(signal.SIGINT)
+        except Exception as e:
+            logger.error(traceback.format_exc())
+            return self.create_error_response(str(e))
+
     async def openai_completion(self, request: CompletionRequest, raw_request: Request) -> Response:
 
         async def completion_response(promise: RequestOutput,

From 29a23e2338e191aa765318aa3a8fb4561639e4e4 Mon Sep 17 00:00:00 2001
From: "Chang Liu (Enterprise Products)"
 <9713593+chang-l@users.noreply.github.com>
Date: Thu, 7 Aug 2025 21:44:08 -0700
Subject: [PATCH 04/15] Enable llava-next

Signed-off-by: Chang Liu (Enterprise Products) <9713593+chang-l@users.noreply.github.com>
---
 .../_torch/models/modeling_llava_next.py      | 91 +++++++++++++++++--
 tensorrt_llm/_torch/pyexecutor/sampler.py     |  3 +
 tensorrt_llm/llmapi/mm_encoder.py             |  3 +-
 3 files changed, 86 insertions(+), 11 deletions(-)

diff --git a/tensorrt_llm/_torch/models/modeling_llava_next.py b/tensorrt_llm/_torch/models/modeling_llava_next.py
index ac26cb647366..16604f439fe4 100644
--- a/tensorrt_llm/_torch/models/modeling_llava_next.py
+++ b/tensorrt_llm/_torch/models/modeling_llava_next.py
@@ -24,7 +24,7 @@
 from .modeling_auto import AutoModelForCausalLM
 from .modeling_clip import CLIPVisionModel
 from .modeling_multimodal_utils import fuse_input_embeds
-from .modeling_utils import ModelConfig, filter_weights, register_auto_model
+from .modeling_utils import ModelConfig, filter_weights, register_auto_model, register_vision_encoder
 
 DISAGG = os.getenv('TLLM_MULTIMODAL_DISAGGREGATED', '0') == '1'
 
@@ -51,6 +51,73 @@ def __init__(self,
 
         self.image_token_index = model_config.image_token_index
         self.vocab_size = model_config.vocab_size
+        self.vision_feature_select_strategy = getattr(
+            model_config, "vision_feature_select_strategy",
+            "default")
+        self.config = model_config.vision_config
+
+    def _get_num_unpadded_features(
+        self,
+        *,
+        original_height: int,
+        original_width: int,
+        npatches: int,
+        num_patch_height: int,
+        num_patch_width: int,
+    ) -> tuple[int, int]:
+        current_height = npatches * num_patch_height
+        current_width = npatches * num_patch_width
+
+        aspect_ratio = original_width / original_height
+        current_aspect_ratio = current_width / current_height
+
+        if aspect_ratio > current_aspect_ratio:
+            new_height = int(
+                round(original_height * (current_width / original_width), 7))
+            padding = (current_height - new_height) // 2
+            current_height = current_height - (2 * padding)
+        else:
+            new_width = int(
+                round(original_width * (current_height / original_height), 7))
+            padding = (current_width - new_width) // 2
+            current_width = current_width - (2 * padding)
+
+        unpadded_features = current_height * current_width
+        newline_features = current_height
+
+        return (unpadded_features, newline_features)
+
+    # Based on: https://github.com/huggingface/text-generation-inference/blob/v3.0.1/server/text_generation_server/models/vlm_causal_lm.py#L113
+    def get_num_tokens_per_image(
+        self,
+        *,
+        image_width: int,
+        image_height: int,
+    ) -> int:
+        patch_grid_length = self.config.image_size // self.config.patch_size
+        base_feature_size = patch_grid_length**2 + 1
+
+        if self.vision_feature_select_strategy == "default":
+            base_feature_size = base_feature_size - 1
+
+        num_patch_height, num_patch_width = get_anyres_image_grid_shape(
+            image_size=(image_width, image_height),
+            grid_pinpoints=self.model_config.image_grid_pinpoints,
+            patch_size=self.config.image_size,
+        )
+
+        (
+            unpadded_feature_size,
+            newline_feature_size,
+        ) = self._get_num_unpadded_features(
+            original_height=image_height,
+            original_width=image_width,
+            npatches=patch_grid_length,
+            num_patch_height=num_patch_height,
+            num_patch_width=num_patch_width,
+        )
+        return unpadded_feature_size + newline_feature_size + base_feature_size
+
 
     @torch.inference_mode()
     def __call__(
@@ -90,9 +157,10 @@ class LlavaNextVisionModel(nn.Module):
     def __init__(self, model_config: ModelConfig[PretrainedConfig], *args,
                  **kwargs) -> None:
         super().__init__()
-        self.pretrained_config = model_config.pretrained_config
+        self.model_config = model_config
+        pretrained_config = model_config.pretrained_config
         self.device = f"cuda:{model_config.mapping.rank}"
-        model_path = self.pretrained_config._name_or_path
+        model_path = pretrained_config._name_or_path
 
         # Determine the actual local path for model files
         if os.path.isdir(model_path):
@@ -145,6 +213,11 @@ def __init__(self, model_config: ModelConfig[PretrainedConfig], *args,
             model_config.pretrained_config, "vision_feature_select_strategy",
             "default")
 
+        self.post_config()
+
+    def post_config(self):
+        self.config = self.model_config.pretrained_config.vision_config
+
     # Copied from https://github.com/huggingface/transformers/blob/main/src/transformers/models/llava_next/modeling_llava_next.py#L284
     def pack_image_features(self,
                             image_features,
@@ -157,12 +230,12 @@ def pack_image_features(self,
             if image_feature.shape[0] > 1:
                 base_image_feature = image_feature[0]
                 image_feature = image_feature[1:]
-                height = width = self.pretrained_config.vision_config.image_size // self.pretrained_config.vision_config.patch_size
+                height = width = self.config.image_size // self.config.patch_size
 
                 num_patch_height, num_patch_width = get_anyres_image_grid_shape(
                     image_sizes[image_idx],
-                    self.pretrained_config.image_grid_pinpoints,
-                    self.pretrained_config.vision_config.image_size,
+                    self.model_config.pretrained_config.image_grid_pinpoints,
+                    self.config.image_size,
                 )
 
                 if (np.prod(image_feature.shape) %
@@ -223,8 +296,8 @@ def forward(self, multimodal_params: List[MultimodalParams]):
         image_num_patches = [
             image_size_to_num_patches(
                 image_size=imsize,
-                grid_pinpoints=self.pretrained_config.image_grid_pinpoints,
-                patch_size=self.pretrained_config.vision_config.image_size,
+                grid_pinpoints=self.model_config.pretrained_config.image_grid_pinpoints,
+                patch_size=self.config.image_size,
             ) for imsize in image_sizes
         ]
 
@@ -261,7 +334,7 @@ def forward(self, multimodal_params: List[MultimodalParams]):
         image_features = torch.cat(image_features, dim=0)
         return [image_features]
 
-
+@register_vision_encoder(LlavaNextVisionModel)
 @register_auto_model("LlavaNextForConditionalGeneration")
 @register_input_processor(LlavaNextInputProcessor, model_type="llava_next")
 class LlavaNextModel(PreTrainedModel):
diff --git a/tensorrt_llm/_torch/pyexecutor/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler.py
index 1fbbcfc898e6..edee7dd01d63 100644
--- a/tensorrt_llm/_torch/pyexecutor/sampler.py
+++ b/tensorrt_llm/_torch/pyexecutor/sampler.py
@@ -134,6 +134,9 @@ def update_requests(self, state: SampleStateWithMMResult) -> None:
             # NOTE: This is a hack: set finish reason manually and set the beam 0
             request.set_finished_reason(FinishReason.LENGTH, 0)
             if idx < len(mm_embeddings):
+                if len(mm_embeddings[idx]) != sum(request.multimodal_lengths):
+                    raise ValueError(f"mm_embedding shape mismatch: {len(mm_embeddings[idx])} != {sum(request.multimodal_lengths)}")
+
                 request.py_result.append_mm_embeddings(mm_embeddings[idx])
 
 
diff --git a/tensorrt_llm/llmapi/mm_encoder.py b/tensorrt_llm/llmapi/mm_encoder.py
index 5d25c6e190ef..0f599431b856 100644
--- a/tensorrt_llm/llmapi/mm_encoder.py
+++ b/tensorrt_llm/llmapi/mm_encoder.py
@@ -157,8 +157,7 @@ def generate_async(
         Returns:
             Future that resolves to tensorrt_llm.llmapi.RequestOutput containing mm_embeddings
         """
-        # TODO: possible preprocess the input
         result = super().generate_async(inputs, sampling_params)
-        # TODO: possible postprocess the result
+        # TODO: possible postprocess the result for disaggregated serving
         return result
 

From e7bc7951a38c7e102e1cc5aac9668d4f41d1035d Mon Sep 17 00:00:00 2001
From: "Chang Liu (Enterprise Products)"
 <9713593+chang-l@users.noreply.github.com>
Date: Thu, 7 Aug 2025 23:48:55 -0700
Subject: [PATCH 05/15] Add unit test and add back attach_emb for llava model

Signed-off-by: Chang Liu (Enterprise Products) <9713593+chang-l@users.noreply.github.com>
---
 .../_torch/models/modeling_llava_next.py      | 146 +++++++++++++-
 .../_torch/pyexecutor/py_executor_creator.py  |   2 +
 .../multimodal/test_mm_encoder_standalone.py  | 183 ++++++++++++++++++
 3 files changed, 322 insertions(+), 9 deletions(-)
 create mode 100644 tests/unittest/_torch/multimodal/test_mm_encoder_standalone.py

diff --git a/tensorrt_llm/_torch/models/modeling_llava_next.py b/tensorrt_llm/_torch/models/modeling_llava_next.py
index 16604f439fe4..88119bb28225 100644
--- a/tensorrt_llm/_torch/models/modeling_llava_next.py
+++ b/tensorrt_llm/_torch/models/modeling_llava_next.py
@@ -1,6 +1,6 @@
 import copy
 import os
-from typing import List, Optional, Tuple
+from typing import List, Optional, Tuple, Dict
 
 import numpy as np
 import torch
@@ -118,6 +118,128 @@ def get_num_tokens_per_image(
         )
         return unpadded_feature_size + newline_feature_size + base_feature_size
 
+    def _postprocess(self, input_ids, mm_features):
+        # Define model specific variables here before shared logic
+        mm_tokens = torch.tensor([self.model_config.image_token_index
+                                  ]).to(input_ids.device)
+        model_hidden_size = self.model_config.text_config.hidden_size
+        vocab_size = self.model_config.text_config.vocab_size
+        start_len = end_len = 0  # for llava, need not append start/end token around each image token
+        # End model specific variables
+
+        ## find mm token positions in input_ids
+        mm_token_positions = torch.where(torch.isin(input_ids, mm_tokens))[0]
+        num_medias = num_mm_tokens = len(mm_token_positions)
+        if num_medias > 1 and isinstance(mm_features, torch.Tensor):
+            mm_features = list(
+                mm_features.split(mm_features.shape[0] // num_medias))
+
+        if isinstance(mm_features, torch.Tensor):
+            # 1 prompt + 1 media
+            # "split" means what a single mm_token in the input_ids should represent
+            # image: one split --> one frame
+            # video: one split --> N frames
+            num_frames, mm_feature_length, mm_hidden_dim = mm_features.shape
+            mm_lengths_per_split = [mm_feature_length * num_frames]
+            mm_lengths_per_frame = [mm_feature_length]
+        elif isinstance(mm_features, list):
+            # 1 prompt + N media
+            num_frames = len(mm_features) if mm_features[0].dim() == 2 else sum(
+                [f.shape[0] for f in mm_features])
+            mm_lengths_per_split = [
+                f.shape[0] if f.dim() == 2 else f.shape[0] * f.shape[1]
+                for f in mm_features
+            ]
+            mm_lengths_per_frame = [
+                f.shape[0] if f.dim() == 2 else f.shape[1] for f in mm_features
+            ]
+            mm_hidden_dim = mm_features[0].shape[-1]
+            mm_features = torch.cat(mm_features, dim=0)
+        else:
+            raise ValueError(
+                f"Invalid multimodal features type: {type(mm_features)}")
+        mm_total_length = sum(mm_lengths_per_split)
+        assert mm_hidden_dim == model_hidden_size, "Multimodal embedding_dim must match model hidden_size"
+
+        ## split input_ids into segments by isolating mm tokens
+        mm_split_positions = torch.cat(
+            [mm_token_positions, mm_token_positions + 1]).unique()
+        input_ids_splits = list(input_ids.tensor_split(mm_split_positions.cpu(
+        )))  # len(input_ids_splits) = num_segments after mm tokens are isolated
+        mm_ids_splits = list(
+            torch.arange(vocab_size,
+                         vocab_size + mm_total_length,
+                         device=input_ids.device).split(mm_lengths_per_split)
+        )  # len(mm_ids_splits) = num_mm_segments
+
+        for i, mm_ids in enumerate(mm_ids_splits):
+            mm_ids = mm_ids.reshape(-1, mm_lengths_per_frame[i])
+            mm_ids_splits[i] = mm_ids.flatten()
+
+        ## replace mm token ids with the expanded out-of-vocab ids
+        mm_split_idx = 0
+        for i, split in enumerate(input_ids_splits):
+            if torch.isin(split, mm_tokens).any().item():
+                input_ids_splits[i] = mm_ids_splits[mm_split_idx]
+                mm_split_idx += 1
+        assert mm_split_idx == len(
+            mm_ids_splits), "All mm_ids_splits should be consumed"
+
+        ## concat text & mm input_ids, wrap mm feature in prompt tuning config
+        fused_input_ids = torch.cat(input_ids_splits).to(
+            device=input_ids.device)
+        fused_length = len(input_ids) + mm_total_length + num_frames * (
+            start_len + end_len) - num_medias
+        assert len(
+            fused_input_ids
+        ) == fused_length, f"Fused input_ids length {len(fused_input_ids)} should match the sum of text and multimodal embedding lengths {fused_length}"
+
+        # [num_frames, feature_length, hidden_dim] -> [num_frames * feature_length, hidden_dim]
+        mm_features = mm_features.view(-1, mm_features.shape[-1])
+        return fused_input_ids, mm_features
+
+
+    def attach_multimodal_embeddings(
+        self, inputs: TextPrompt,
+        multimodal_embedding: Dict[str, List[torch.Tensor]],
+        sampling_params: SamplingParams
+    ) -> Tuple[List[int], Optional[ExtraProcessedInputs]]:
+        """
+        Attach pre-processed multimodal embeddings into text token stream for LlavaNext model.
+        This method skips vision processing and works with externally provided embeddings.
+        It replaces/expands image placeholders in the text with appropriate tokens and prepares
+        the embeddings for model forward pass.
+        Args:
+            inputs: Text prompt containing image placeholders
+            multimodal_embedding: Dictionary containing pre-processed image embedding data
+        Returns:
+            Tuple of (token_ids, extra_processed_inputs) where:
+            - token_ids: List of processed token IDs with image placeholders
+            - extra_processed_inputs: Optional dictionary containing multimodal embeddings
+        """
+        text_prompt = inputs.get("prompt")
+        if not text_prompt:
+            raise ValueError("Text prompt is required but not provided")
+
+
+
+        if not isinstance(multimodal_embedding, dict):
+            raise ValueError("multimodal_embedding must be a dictionary")
+
+        if 'image' not in multimodal_embedding:
+            raise ValueError(
+                "Only image modality is supported for external multimodal embedding"
+            )
+
+        input_ids = self.tokenizer(
+            text_prompt, return_tensors="pt").input_ids[0]
+        mm_features = torch.stack(multimodal_embedding['image'])
+        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(), {
+            "multimodal_data": multimodal_data
+        }
 
     @torch.inference_mode()
     def __call__(
@@ -158,9 +280,9 @@ def __init__(self, model_config: ModelConfig[PretrainedConfig], *args,
                  **kwargs) -> None:
         super().__init__()
         self.model_config = model_config
-        pretrained_config = model_config.pretrained_config
+        self.pretrained_config = model_config.pretrained_config
         self.device = f"cuda:{model_config.mapping.rank}"
-        model_path = pretrained_config._name_or_path
+        model_path = self.pretrained_config._name_or_path
 
         # Determine the actual local path for model files
         if os.path.isdir(model_path):
@@ -200,7 +322,7 @@ def __init__(self, model_config: ModelConfig[PretrainedConfig], *args,
             self.vision_tower = hf_vision_tower.to(self.device)
         else:
             vision_model_config = ModelConfig(
-                pretrained_config=model_config.pretrained_config.vision_config,
+                pretrained_config=self.pretrained_config.vision_config,
                 attn_backend="TRTLLM")
             self.vision_tower = CLIPVisionModel(vision_model_config).to(
                 self.device).to(self.dtype)
@@ -210,13 +332,13 @@ def __init__(self, model_config: ModelConfig[PretrainedConfig], *args,
         self.mm_projector = hf_mm_projector
         self.image_newline = hf_image_newline
         self.vision_feature_select_strategy = getattr(
-            model_config.pretrained_config, "vision_feature_select_strategy",
+            self.pretrained_config, "vision_feature_select_strategy",
             "default")
 
         self.post_config()
 
     def post_config(self):
-        self.config = self.model_config.pretrained_config.vision_config
+        self.config = self.pretrained_config.vision_config
 
     # Copied from https://github.com/huggingface/transformers/blob/main/src/transformers/models/llava_next/modeling_llava_next.py#L284
     def pack_image_features(self,
@@ -234,7 +356,7 @@ def pack_image_features(self,
 
                 num_patch_height, num_patch_width = get_anyres_image_grid_shape(
                     image_sizes[image_idx],
-                    self.model_config.pretrained_config.image_grid_pinpoints,
+                    self.pretrained_config.image_grid_pinpoints,
                     self.config.image_size,
                 )
 
@@ -296,7 +418,7 @@ def forward(self, multimodal_params: List[MultimodalParams]):
         image_num_patches = [
             image_size_to_num_patches(
                 image_size=imsize,
-                grid_pinpoints=self.model_config.pretrained_config.image_grid_pinpoints,
+                grid_pinpoints=self.pretrained_config.image_grid_pinpoints,
                 patch_size=self.config.image_size,
             ) for imsize in image_sizes
         ]
@@ -396,7 +518,13 @@ def forward(
         mm_embeds = []
         if len(multimodal_params) > 0:
             if not DISAGG:
-                mm_embeds = self.mm_encoder.forward(multimodal_params)
+                if  multimodal_params[0].multimodal_data.get("multimodal_embedding", None) is not None:
+                    mm_embeds = [
+                        multimodal_param.multimodal_data["multimodal_embedding"]
+                        for multimodal_param in multimodal_params
+                    ]
+                else:
+                    mm_embeds = self.mm_encoder.forward(multimodal_params)
             else:
                 mm_embeds = [
                     multimodal_param.multimodal_data["multimodal_embedding"]
diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
index 22fcf982d26d..f3b72100de03 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
@@ -175,6 +175,8 @@ def _mangle_executor_config(executor_config: ExecutorConfig):
         pytorch_backend_config.load_format = LoadFormat.VISION_ONLY
         # TODO: add comment and print warning here
         pytorch_backend_config.disable_overlap_scheduler = True
+        # TODO: add comment here to infer it by max_num_images and image_token_sizen
+        executor_config.max_num_tokens = 16384
 
 def _get_mapping(executor_config: ExecutorConfig) -> Mapping:
     if executor_config.mapping is None:
diff --git a/tests/unittest/_torch/multimodal/test_mm_encoder_standalone.py b/tests/unittest/_torch/multimodal/test_mm_encoder_standalone.py
new file mode 100644
index 000000000000..fa1694ff304c
--- /dev/null
+++ b/tests/unittest/_torch/multimodal/test_mm_encoder_standalone.py
@@ -0,0 +1,183 @@
+import os
+import pytest
+import copy
+import json
+
+from tensorrt_llm import MultimodalEncoder
+from tensorrt_llm._torch.shared_tensor import SharedTensorContainer
+from tensorrt_llm.llmapi.llm import LLM, SamplingParams
+from tensorrt_llm.llmapi import KvCacheConfig
+from tensorrt_llm.inputs import default_multimodal_input_loader
+
+example_images = [
+    "https://huggingface.co/datasets/YiYiXu/testing-images/resolve/main/seashore.png",
+    "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/inpaint.png",
+    "https://huggingface.co/datasets/Sayali9141/traffic_signal_images/resolve/main/61.jpg",
+]
+
+
+@pytest.fixture(scope="function")
+def multimodal_model_config():
+    """Get multimodal model configuration similar to integration tests"""
+    # You can extend this to support multiple models or get from environment
+    model_configs = {
+        'llava-v1.6-mistral-7b-hf': {
+            'model_name': 'llava-v1.6-mistral-7b-hf',
+            'hf_model_dir': 'llava-hf/llava-v1.6-mistral-7b-hf',  # HuggingFace model ID
+        }
+    }
+    
+    return model_configs['llava-v1.6-mistral-7b-hf']
+
+
+@pytest.mark.parametrize("model_key", [
+    "llava-v1.6-mistral-7b-hf",
+])
+def test_single_image_chat(model_key, multimodal_model_config):
+    """Test processing single image using disaggregated encoder + LLM API.
+    
+    This test verifies that disaggregated multimodal generation produces identical
+    results to standard multimodal generation by comparing outputs.
+    """
+    # Get model configuration
+    if model_key != "llava-v1.6-mistral-7b-hf":
+        pytest.skip(f"Skipping test for {model_key} - only testing llava-v1.6-mistral-7b-hf for now")
+    
+    # Extract model information from config
+    model_name = multimodal_model_config['model_name']
+    encoder_model_dir = multimodal_model_config['hf_model_dir']
+    
+    # Test configuration
+    max_tokens = 64
+    free_gpu_memory_fraction = 0.6
+    max_batch_size = 1
+    
+    # Test data - OpenAI chat completion format
+    prompts = ["Describe the natural environment in the image."]
+    media = [example_images[0]]
+    
+    # Create OpenAI chat messages format
+    messages_list = []
+    for prompt, image_url in zip(prompts, media):
+        messages = [{
+            "role": "user",
+            "content": [{
+                "type": "text",
+                "text": prompt
+            }, {
+                "type": "image_url",
+                "image_url": {
+                    "url": image_url
+                }
+            }]
+        }]
+        messages_list.append(messages)
+    
+    # Sampling configuration
+    sampling_params = SamplingParams(max_tokens=max_tokens)
+    kv_cache_config = KvCacheConfig(
+        enable_block_reuse=False,
+        free_gpu_memory_fraction=free_gpu_memory_fraction,
+    )
+    
+    # Step 1: Process multimodal data using disaggregated encoder
+    encoder = None
+    llm = None
+    
+    try:
+        # Step 1: Initialize encoder
+        encoder = MultimodalEncoder(model=encoder_model_dir, max_batch_size=max_batch_size)
+        
+        # Step 2: Initialize LLM and prepare inputs
+        llm = LLM(
+            model=encoder_model_dir, 
+            backend='pytorch', 
+            kv_cache_config=kv_cache_config, 
+            trust_remote_code=True
+        )
+        
+        # Load model configuration
+        config_path = os.path.join(llm._hf_model_dir, 'config.json')
+        assert os.path.exists(config_path), f"Model config not found at {config_path}"
+        
+        with open(config_path, 'r') as f:
+            model_config = json.load(f)
+        model_type = model_config['model_type']
+        
+        # Prepare multimodal inputs
+        inputs = default_multimodal_input_loader(
+            tokenizer=llm.tokenizer,
+            model_dir=llm._hf_model_dir,
+            model_type=model_type,
+            modality="image",
+            prompts=prompts,  
+            media=media,      
+            image_data_format="pt"
+        )
+        
+        # Validate inputs structure
+        assert len(inputs) == len(prompts), f"Expected {len(prompts)} inputs, got {len(inputs)}"
+        # Step 3: Generate reference output with raw multimodal inputs
+        outputs_ref = llm.generate(inputs, sampling_params=sampling_params)
+        
+        # Validate reference outputs
+        assert outputs_ref is not None, "Reference generation returned None"
+        assert len(outputs_ref) == len(prompts), f"Expected {len(prompts)} reference outputs, got {len(outputs_ref)}"
+        for i, output in enumerate(outputs_ref):
+            assert len(output.outputs) > 0, f"Reference generation has no output text for input {i}"
+        
+        # Step 4: Prepare inputs for disaggregated multimodal generation
+        encoder_outputs = encoder.generate(inputs)
+        inputs = default_multimodal_input_loader(
+            tokenizer=llm.tokenizer,
+            model_dir=llm._hf_model_dir,
+            model_type=model_type,
+            modality="image",
+            prompts=prompts,  
+            mm_embeddings=[SharedTensorContainer.from_dict(output.mm_embedding_handle).get_local_view() for output in encoder_outputs],
+            image_data_format="pt"
+        )
+
+        # Step 5: Generate output using disaggregated multimodal parameters
+        # Note: For batch processing, we need to match mm_params with inputs
+        outputs = llm.generate(inputs, sampling_params=sampling_params)
+        
+        # Validate disaggregated outputs
+        assert len(outputs) == len(prompts), f"Expected {len(prompts)} disaggregated outputs, got {len(outputs)}"
+        for i, output in enumerate(outputs):
+            assert len(output.outputs) > 0, f"Disaggregated generation has no output text for input {i}"
+        
+        # Step 6: Compare outputs - they should match exactly
+        assert len(outputs_ref) == len(outputs), f"Number of outputs don't match: {len(outputs_ref)} vs {len(outputs)}"
+        
+        for i, (ref_output, test_output) in enumerate(zip(outputs_ref, outputs)):
+            # Compare prompts
+            assert ref_output.prompt == test_output.prompt, \
+                f"Prompts don't match for output {i}:\nReference: {ref_output.prompt!r}\nTest: {test_output.prompt!r}"
+            
+            # Compare number of generated outputs
+            assert len(ref_output.outputs) == len(test_output.outputs), \
+                f"Number of generated outputs don't match for output {i}: {len(ref_output.outputs)} vs {len(test_output.outputs)}"
+            
+            # Compare generated text and other attributes
+            for j, (ref_gen, test_gen) in enumerate(zip(ref_output.outputs, test_output.outputs)):
+                assert ref_gen.text == test_gen.text, \
+                    f"Generated text doesn't match for output {i}, generation {j}:\nReference: {ref_gen.text!r}\nTest: {test_gen.text!r}"
+                
+                # Compare token IDs if available
+                if hasattr(ref_gen, 'token_ids') and hasattr(test_gen, 'token_ids'):
+                    assert ref_gen.token_ids == test_gen.token_ids, \
+                        f"Token IDs don't match for output {i}, generation {j}"
+                
+                # Compare log probabilities if available
+                if hasattr(ref_gen, 'logprobs') and hasattr(test_gen, 'logprobs'):
+                    assert ref_gen.logprobs == test_gen.logprobs, \
+                        f"Log probabilities don't match for output {i}, generation {j}"
+                
+    finally:
+        # Cleanup resources
+        if encoder is not None:
+            del encoder
+        if llm is not None:
+            del llm
+    

From 2c8c31e9b6deb85bb204522baad91cf421c5721d Mon Sep 17 00:00:00 2001
From: "Chang Liu (Enterprise Products)"
 <9713593+chang-l@users.noreply.github.com>
Date: Fri, 8 Aug 2025 00:43:15 -0700
Subject: [PATCH 06/15] Minor update/cleanup

Signed-off-by: Chang Liu (Enterprise Products) <9713593+chang-l@users.noreply.github.com>
---
 tensorrt_llm/_torch/models/modeling_auto.py   | 23 -------------------
 tensorrt_llm/_torch/pyexecutor/sampler.py     |  2 --
 .../multimodal/test_mm_encoder_standalone.py  | 19 ++++++++-------
 3 files changed, 9 insertions(+), 35 deletions(-)

diff --git a/tensorrt_llm/_torch/models/modeling_auto.py b/tensorrt_llm/_torch/models/modeling_auto.py
index 314f92e09268..edba367210d0 100644
--- a/tensorrt_llm/_torch/models/modeling_auto.py
+++ b/tensorrt_llm/_torch/models/modeling_auto.py
@@ -44,26 +44,3 @@ def from_config(
             model = cls(config)
         model.extra_attrs = extra_attrs
         return model
-
-class AutoModelForMultimodalEncoder(Generic[TModel, TConfig]):
-
-    @staticmethod
-    def from_config(
-        config: ModelConfig[TConfig],
-    ):
-        model_arch = config.pretrained_config.architectures[0]
-
-        # Direct lookup using architecture name
-        vision_encoder_info = MODEL_CLASS_VISION_ENCODER_MAPPING.get(model_arch)
-        if vision_encoder_info is None:
-            raise ValueError(
-                f"Unknown architecture for AutoModelForMultimodalEncoder: {model_arch}"
-            )
-
-        vision_encoder_cls, vlm_base_model = vision_encoder_info
-        if vlm_base_model is None:
-            model = vision_encoder_cls(config)
-        else:
-            model = vision_encoder_cls(config, vlm_base_model)
-
-        return model
\ No newline at end of file
diff --git a/tensorrt_llm/_torch/pyexecutor/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler.py
index edee7dd01d63..10adbe8eca30 100644
--- a/tensorrt_llm/_torch/pyexecutor/sampler.py
+++ b/tensorrt_llm/_torch/pyexecutor/sampler.py
@@ -87,8 +87,6 @@ def update_requests(self, state: SampleState) -> None:
             request.state = LlmRequestState.GENERATION_COMPLETE
             # NOTE: This is a hack: set finish reason manually and set the beam 0
             request.set_finished_reason(FinishReason.LENGTH, 0)
-            if request.py_return_mm_embeddings:
-                request.py_result.append_mm_embeddings(state.host.mm_embeddings[idx])
             if request.py_return_context_logits:
                 logits = state.host.logits[idx]
                 if logits.ndim == 1:
diff --git a/tests/unittest/_torch/multimodal/test_mm_encoder_standalone.py b/tests/unittest/_torch/multimodal/test_mm_encoder_standalone.py
index fa1694ff304c..7d493789bbd3 100644
--- a/tests/unittest/_torch/multimodal/test_mm_encoder_standalone.py
+++ b/tests/unittest/_torch/multimodal/test_mm_encoder_standalone.py
@@ -34,17 +34,16 @@ def multimodal_model_config():
     "llava-v1.6-mistral-7b-hf",
 ])
 def test_single_image_chat(model_key, multimodal_model_config):
-    """Test processing single image using disaggregated encoder + LLM API.
+    """Test processing single image using encoder (pass mm_embeddings) + LLM API.
     
-    This test verifies that disaggregated multimodal generation produces identical
-    results to standard multimodal generation by comparing outputs.
+    This test verifies that encoder (pass mm_embeddings) + LLM API produces identical
+    results to standard llm generation (pass raw image) by comparing outputs.
     """
     # Get model configuration
     if model_key != "llava-v1.6-mistral-7b-hf":
         pytest.skip(f"Skipping test for {model_key} - only testing llava-v1.6-mistral-7b-hf for now")
     
     # Extract model information from config
-    model_name = multimodal_model_config['model_name']
     encoder_model_dir = multimodal_model_config['hf_model_dir']
     
     # Test configuration
@@ -80,7 +79,7 @@ def test_single_image_chat(model_key, multimodal_model_config):
         free_gpu_memory_fraction=free_gpu_memory_fraction,
     )
     
-    # Step 1: Process multimodal data using disaggregated encoder
+    # Step 1: Process multimodal data using encoder (pass mm_embeddings)
     encoder = None
     llm = None
     
@@ -126,7 +125,7 @@ def test_single_image_chat(model_key, multimodal_model_config):
         for i, output in enumerate(outputs_ref):
             assert len(output.outputs) > 0, f"Reference generation has no output text for input {i}"
         
-        # Step 4: Prepare inputs for disaggregated multimodal generation
+        # Step 4: Prepare inputs for llm (pass mm_embeddings)
         encoder_outputs = encoder.generate(inputs)
         inputs = default_multimodal_input_loader(
             tokenizer=llm.tokenizer,
@@ -138,14 +137,14 @@ def test_single_image_chat(model_key, multimodal_model_config):
             image_data_format="pt"
         )
 
-        # Step 5: Generate output using disaggregated multimodal parameters
+        # Step 5: Generate output using llm (pass mm_embeddings)
         # Note: For batch processing, we need to match mm_params with inputs
         outputs = llm.generate(inputs, sampling_params=sampling_params)
         
-        # Validate disaggregated outputs
-        assert len(outputs) == len(prompts), f"Expected {len(prompts)} disaggregated outputs, got {len(outputs)}"
+        # Validate outputs
+        assert len(outputs) == len(prompts), f"Expected {len(prompts)} outputs, got {len(outputs)}"
         for i, output in enumerate(outputs):
-            assert len(output.outputs) > 0, f"Disaggregated generation has no output text for input {i}"
+            assert len(output.outputs) > 0, f"generation has no output text for input {i}"
         
         # Step 6: Compare outputs - they should match exactly
         assert len(outputs_ref) == len(outputs), f"Number of outputs don't match: {len(outputs_ref)} vs {len(outputs)}"

From 262ebcb74393789b3648d680acb3e732b3edb79b Mon Sep 17 00:00:00 2001
From: Chang Liu <9713593+chang-l@users.noreply.github.com>
Date: Fri, 8 Aug 2025 21:46:51 -0700
Subject: [PATCH 07/15] Address comments and add more UTw

Signed-off-by: Chang Liu <9713593+chang-l@users.noreply.github.com>
---
 tensorrt_llm/__init__.py                      |   3 +-
 tensorrt_llm/_torch/model_config.py           |   6 +-
 tensorrt_llm/_torch/models/modeling_auto.py   |   8 +-
 .../_torch/models/modeling_llava_next.py      |  83 ++----
 .../_torch/models/modeling_qwen2vl.py         |   7 +-
 tensorrt_llm/_torch/models/modeling_utils.py  |  28 +-
 tensorrt_llm/_torch/pyexecutor/_util.py       |   3 +-
 tensorrt_llm/_torch/pyexecutor/config.py      |   2 +-
 tensorrt_llm/_torch/pyexecutor/llm_request.py |  11 +-
 .../_torch/pyexecutor/model_engine.py         |  52 ++--
 .../_torch/pyexecutor/py_executor_creator.py  |   9 +-
 tensorrt_llm/_torch/pyexecutor/sampler.py     |  20 +-
 tensorrt_llm/commands/serve.py                |  37 +--
 tensorrt_llm/executor/result.py               |   4 +-
 tensorrt_llm/llmapi/__init__.py               |   2 +-
 tensorrt_llm/llmapi/llm.py                    |   3 +-
 tensorrt_llm/llmapi/mm_encoder.py             |  42 +--
 tensorrt_llm/serve/openai_server.py           |   9 +-
 .../multimodal/test_external_embedding.py     | 181 +++++++++++++
 .../multimodal/test_find_num_image_tokens.py  | 159 ++++++++++++
 .../multimodal/test_mm_encoder_standalone.py  | 241 ++++++++----------
 21 files changed, 616 insertions(+), 294 deletions(-)
 create mode 100644 tests/unittest/_torch/multimodal/test_external_embedding.py
 create mode 100644 tests/unittest/_torch/multimodal/test_find_num_image_tokens.py

diff --git a/tensorrt_llm/__init__.py b/tensorrt_llm/__init__.py
index ee5fa8c967bd..f0c8a1941db5 100644
--- a/tensorrt_llm/__init__.py
+++ b/tensorrt_llm/__init__.py
@@ -50,8 +50,7 @@ def _add_trt_llm_dll_directory():
 from .builder import BuildConfig, Builder, BuilderConfig, build
 from .disaggregated_params import DisaggregatedParams
 from .functional import Tensor, constant
-from .llmapi import LLM, LlmArgs
-from .llmapi import MultimodalEncoder
+from .llmapi import LLM, LlmArgs, MultimodalEncoder
 from .llmapi.llm_args import LlmArgs, TorchLlmArgs, TrtLlmArgs
 from .logger import logger
 from .mapping import Mapping
diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py
index 209ac3d8b7f5..8e6ba15c33a1 100644
--- a/tensorrt_llm/_torch/model_config.py
+++ b/tensorrt_llm/_torch/model_config.py
@@ -118,7 +118,8 @@ def __post_init__(self):
         if self.pretrained_config and hasattr(self.pretrained_config,
                                               "architectures"):
             self.is_generation = self.is_generation_model(
-                self.pretrained_config.architectures, mm_encoder_only=self.mm_encoder_only)
+                self.pretrained_config.architectures,
+                mm_encoder_only=self.mm_encoder_only)
 
         def get_all_reduce_strategy(strategy: str = "AUTO"):
             maps = {
@@ -167,7 +168,8 @@ def get_quant_config(self, name: Optional[str] = None) -> QuantConfig:
         raise ValueError(f'quant config of {name} is not found')
 
     @staticmethod
-    def is_generation_model(model_architectures: Optional[List[str]], mm_encoder_only: bool = False) -> bool:
+    def is_generation_model(model_architectures: Optional[List[str]],
+                            mm_encoder_only: bool = False) -> bool:
         if model_architectures is None:
             logger.warning(
                 "Model architectures is None, default to is_generation_model=True"
diff --git a/tensorrt_llm/_torch/models/modeling_auto.py b/tensorrt_llm/_torch/models/modeling_auto.py
index edba367210d0..cd73919ca2cf 100644
--- a/tensorrt_llm/_torch/models/modeling_auto.py
+++ b/tensorrt_llm/_torch/models/modeling_auto.py
@@ -2,8 +2,9 @@
 
 from ..model_config import ModelConfig
 from ..utils import model_extra_attrs
-from .modeling_utils import (MODEL_CLASS_MAPPING, DecoderModelForCausalLM,
-                             MODEL_CLASS_VISION_ENCODER_MAPPING, TConfig, TModel)
+from .modeling_utils import (MODEL_CLASS_MAPPING,
+                             MODEL_CLASS_VISION_ENCODER_MAPPING,
+                             DecoderModelForCausalLM, TConfig, TModel)
 
 
 class AutoModelForCausalLM(Generic[TModel, TConfig]):
@@ -14,7 +15,8 @@ def from_config(
     ) -> DecoderModelForCausalLM[TModel, TConfig]:
         model_arch = config.pretrained_config.architectures[0]
         if config.mm_encoder_only:
-            vision_encoder_info = MODEL_CLASS_VISION_ENCODER_MAPPING.get(model_arch)
+            vision_encoder_info = MODEL_CLASS_VISION_ENCODER_MAPPING.get(
+                model_arch)
             if vision_encoder_info is None:
                 raise ValueError(
                     f"Unknown architecture for AutoModelForMultimodalEncoder: {model_arch}"
diff --git a/tensorrt_llm/_torch/models/modeling_llava_next.py b/tensorrt_llm/_torch/models/modeling_llava_next.py
index 88119bb28225..4e23f8634c95 100644
--- a/tensorrt_llm/_torch/models/modeling_llava_next.py
+++ b/tensorrt_llm/_torch/models/modeling_llava_next.py
@@ -1,6 +1,6 @@
 import copy
 import os
-from typing import List, Optional, Tuple, Dict
+from typing import Dict, List, Optional, Tuple
 
 import numpy as np
 import torch
@@ -24,7 +24,8 @@
 from .modeling_auto import AutoModelForCausalLM
 from .modeling_clip import CLIPVisionModel
 from .modeling_multimodal_utils import fuse_input_embeds
-from .modeling_utils import ModelConfig, filter_weights, register_auto_model, register_vision_encoder
+from .modeling_utils import (filter_weights, register_auto_model,
+                             register_vision_encoder)
 
 DISAGG = os.getenv('TLLM_MULTIMODAL_DISAGGREGATED', '0') == '1'
 
@@ -52,71 +53,19 @@ def __init__(self,
         self.image_token_index = model_config.image_token_index
         self.vocab_size = model_config.vocab_size
         self.vision_feature_select_strategy = getattr(
-            model_config, "vision_feature_select_strategy",
-            "default")
+            model_config, "vision_feature_select_strategy", "default")
         self.config = model_config.vision_config
 
-    def _get_num_unpadded_features(
-        self,
-        *,
-        original_height: int,
-        original_width: int,
-        npatches: int,
-        num_patch_height: int,
-        num_patch_width: int,
-    ) -> tuple[int, int]:
-        current_height = npatches * num_patch_height
-        current_width = npatches * num_patch_width
-
-        aspect_ratio = original_width / original_height
-        current_aspect_ratio = current_width / current_height
-
-        if aspect_ratio > current_aspect_ratio:
-            new_height = int(
-                round(original_height * (current_width / original_width), 7))
-            padding = (current_height - new_height) // 2
-            current_height = current_height - (2 * padding)
-        else:
-            new_width = int(
-                round(original_width * (current_height / original_height), 7))
-            padding = (current_width - new_width) // 2
-            current_width = current_width - (2 * padding)
-
-        unpadded_features = current_height * current_width
-        newline_features = current_height
-
-        return (unpadded_features, newline_features)
-
-    # Based on: https://github.com/huggingface/text-generation-inference/blob/v3.0.1/server/text_generation_server/models/vlm_causal_lm.py#L113
     def get_num_tokens_per_image(
         self,
         *,
         image_width: int,
         image_height: int,
     ) -> int:
-        patch_grid_length = self.config.image_size // self.config.patch_size
-        base_feature_size = patch_grid_length**2 + 1
-
-        if self.vision_feature_select_strategy == "default":
-            base_feature_size = base_feature_size - 1
-
-        num_patch_height, num_patch_width = get_anyres_image_grid_shape(
-            image_size=(image_width, image_height),
-            grid_pinpoints=self.model_config.image_grid_pinpoints,
-            patch_size=self.config.image_size,
-        )
-
-        (
-            unpadded_feature_size,
-            newline_feature_size,
-        ) = self._get_num_unpadded_features(
-            original_height=image_height,
-            original_width=image_width,
-            npatches=patch_grid_length,
-            num_patch_height=num_patch_height,
-            num_patch_width=num_patch_width,
-        )
-        return unpadded_feature_size + newline_feature_size + base_feature_size
+        image_size = (image_height, image_width)
+        num_image_tokens = self.processor._get_num_multimodal_tokens(
+            [image_size])["num_image_tokens"][0]
+        return num_image_tokens
 
     def _postprocess(self, input_ids, mm_features):
         # Define model specific variables here before shared logic
@@ -198,7 +147,6 @@ def _postprocess(self, input_ids, mm_features):
         mm_features = mm_features.view(-1, mm_features.shape[-1])
         return fused_input_ids, mm_features
 
-
     def attach_multimodal_embeddings(
         self, inputs: TextPrompt,
         multimodal_embedding: Dict[str, List[torch.Tensor]],
@@ -221,8 +169,6 @@ def attach_multimodal_embeddings(
         if not text_prompt:
             raise ValueError("Text prompt is required but not provided")
 
-
-
         if not isinstance(multimodal_embedding, dict):
             raise ValueError("multimodal_embedding must be a dictionary")
 
@@ -231,9 +177,9 @@ def attach_multimodal_embeddings(
                 "Only image modality is supported for external multimodal embedding"
             )
 
-        input_ids = self.tokenizer(
-            text_prompt, return_tensors="pt").input_ids[0]
-        mm_features = torch.stack(multimodal_embedding['image'])
+        input_ids = self.tokenizer(text_prompt,
+                                   return_tensors="pt").input_ids[0]
+        mm_features = multimodal_embedding['image']
         fused_input_ids, mm_features = self._postprocess(input_ids, mm_features)
         multimodal_data = {}
         multimodal_data["multimodal_embedding"] = mm_features
@@ -332,8 +278,7 @@ def __init__(self, model_config: ModelConfig[PretrainedConfig], *args,
         self.mm_projector = hf_mm_projector
         self.image_newline = hf_image_newline
         self.vision_feature_select_strategy = getattr(
-            self.pretrained_config, "vision_feature_select_strategy",
-            "default")
+            self.pretrained_config, "vision_feature_select_strategy", "default")
 
         self.post_config()
 
@@ -456,6 +401,7 @@ def forward(self, multimodal_params: List[MultimodalParams]):
         image_features = torch.cat(image_features, dim=0)
         return [image_features]
 
+
 @register_vision_encoder(LlavaNextVisionModel)
 @register_auto_model("LlavaNextForConditionalGeneration")
 @register_input_processor(LlavaNextInputProcessor, model_type="llava_next")
@@ -518,7 +464,8 @@ def forward(
         mm_embeds = []
         if len(multimodal_params) > 0:
             if not DISAGG:
-                if  multimodal_params[0].multimodal_data.get("multimodal_embedding", None) is not None:
+                if multimodal_params[0].multimodal_data.get(
+                        "multimodal_embedding", None) is not None:
                     mm_embeds = [
                         multimodal_param.multimodal_data["multimodal_embedding"]
                         for multimodal_param in multimodal_params
diff --git a/tensorrt_llm/_torch/models/modeling_qwen2vl.py b/tensorrt_llm/_torch/models/modeling_qwen2vl.py
index c3da24718695..e7533a4b4ab2 100644
--- a/tensorrt_llm/_torch/models/modeling_qwen2vl.py
+++ b/tensorrt_llm/_torch/models/modeling_qwen2vl.py
@@ -651,7 +651,8 @@ def forward(
         return output_prob
 
 
-@register_vision_encoder(Qwen2VisionModelBase, vlm_base_model=Qwen2VLForConditionalGeneration)
+@register_vision_encoder(Qwen2VisionModelBase,
+                         vlm_base_model=Qwen2VLForConditionalGeneration)
 @register_auto_model("Qwen2VLForConditionalGeneration")
 @register_input_processor(Qwen2VLInputProcessorBase, model_type="qwen2_vl")
 class Qwen2VLModel(Qwen2VLModelBase):
@@ -663,7 +664,9 @@ def __init__(self, model_config: ModelConfig[PretrainedConfig], *args,
                 model_config, Qwen2VLForConditionalGeneration)
         super().__init__(model_config, *args, **kwargs)
 
-@register_vision_encoder(Qwen2VisionModelBase, vlm_base_model=Qwen2_5_VLForConditionalGeneration)
+
+@register_vision_encoder(Qwen2VisionModelBase,
+                         vlm_base_model=Qwen2_5_VLForConditionalGeneration)
 @register_auto_model("Qwen2_5_VLForConditionalGeneration")
 @register_input_processor(Qwen2VLInputProcessorBase, model_type="qwen2_5_vl")
 class Qwen2_5_VLModel(Qwen2VLModelBase):
diff --git a/tensorrt_llm/_torch/models/modeling_utils.py b/tensorrt_llm/_torch/models/modeling_utils.py
index 4f0a901802b8..87f23c34db9b 100755
--- a/tensorrt_llm/_torch/models/modeling_utils.py
+++ b/tensorrt_llm/_torch/models/modeling_utils.py
@@ -588,21 +588,39 @@ def decorator(cls):
 
     return decorator
 
-# Need to register the model class before the vision encoder class
-# to ensure that the key (the architecture name) is available here
-def register_vision_encoder(vision_encoder_cls: Type[nn.Module],
-                            vlm_base_model: Optional[Type[nn.Module]] = None):
+
+def register_vision_encoder(
+    vision_encoder_cls: Type[nn.Module],
+    vlm_base_model: Optional[Type[nn.Module]] = None,
+):
+    """Decorator to register a vision encoder implementation for a pre-registered model architecture.
+
+    Usage:
+        @register_vision_encoder(MyVisionEncoder, MyVLMBaseModel)
+        @register_auto_model("SomeVLModel")
+        class SomeVLModel(...):
+            ...
+    The register_auto_model decorator must be applied (executed) before this one (i.e., placed lower)
+    so that the architecture name is present in MODEL_CLASS_MAPPING.
+    """
 
     def wrapper(model_cls: Type[nn.Module]) -> Type[nn.Module]:
         for arch_name, registered_cls in MODEL_CLASS_MAPPING.items():
             if registered_cls.__name__ == model_cls.__name__:
-                MODEL_CLASS_VISION_ENCODER_MAPPING[arch_name] = (vision_encoder_cls, vlm_base_model)
+                MODEL_CLASS_VISION_ENCODER_MAPPING[arch_name] = (
+                    vision_encoder_cls, vlm_base_model)
                 break
+        else:
+            raise ValueError(
+                f"register_vision_encoder: model class {model_cls.__name__} is not registered "
+                f"via register_auto_model; decorator order must ensure registration occurs first."
+            )
 
         return model_cls
 
     return wrapper
 
+
 def register_mapper(format: str, name: Optional[str] = None):
 
     def decorator(cls):
diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py
index 122e3b417a2d..72a43bea61ec 100644
--- a/tensorrt_llm/_torch/pyexecutor/_util.py
+++ b/tensorrt_llm/_torch/pyexecutor/_util.py
@@ -30,7 +30,8 @@
 from .resource_manager import (KVCacheManager, MambaHybridCacheManager,
                                PeftCacheManager, ResourceManager,
                                ResourceManagerType)
-from .sampler import EarlyStopSampler, TorchSampler, TRTLLMSampler, EarlyStopWithMMResult
+from .sampler import (EarlyStopSampler, EarlyStopWithMMResult, TorchSampler,
+                      TRTLLMSampler)
 from .scheduler import (BindCapacityScheduler, BindMicroBatchScheduler,
                         SimpleScheduler)
 from .seq_slot_manager import SeqSlotManager
diff --git a/tensorrt_llm/_torch/pyexecutor/config.py b/tensorrt_llm/_torch/pyexecutor/config.py
index f6f3ae8a6c0a..4911fdb9e94b 100644
--- a/tensorrt_llm/_torch/pyexecutor/config.py
+++ b/tensorrt_llm/_torch/pyexecutor/config.py
@@ -131,7 +131,7 @@ def update_executor_config(
         max_seq_len: Optional[int] = None,
         checkpoint_format: Optional[str] = None,
         checkpoint_loader: Optional[BaseCheckpointLoader] = None,
-        mm_encoder_only: Optional[bool] = False):
+        mm_encoder_only: bool = False):
     if backend is None:
         return
 
diff --git a/tensorrt_llm/_torch/pyexecutor/llm_request.py b/tensorrt_llm/_torch/pyexecutor/llm_request.py
index 85b1f4c6429b..cc0bf5f6ceed 100644
--- a/tensorrt_llm/_torch/pyexecutor/llm_request.py
+++ b/tensorrt_llm/_torch/pyexecutor/llm_request.py
@@ -1,13 +1,13 @@
 from copy import deepcopy
 from dataclasses import dataclass
-from typing import List, Optional, Union, Dict, Any
+from typing import Any, Dict, List, Optional, Union
 
 import torch
 
 import tensorrt_llm.bindings
+from tensorrt_llm._torch.shared_tensor import SharedTensorContainer
 from tensorrt_llm.bindings import executor as tllm_executor
 from tensorrt_llm.executor.result import TokenLogprobs
-from tensorrt_llm._torch.shared_tensor import SharedTensorContainer
 
 SamplingConfig = tensorrt_llm.bindings.SamplingConfig
 '''
@@ -190,7 +190,8 @@ def append_log_probs(self,
             self._log_probs.append(log_probs, cum_log_probs)
 
     def append_mm_embeddings(self, mm_embeddings: torch.Tensor):
-        self._mm_embeddings = SharedTensorContainer.from_tensor(mm_embeddings).dump_to_dict()
+        self._mm_embeddings = SharedTensorContainer.from_tensor(
+            mm_embeddings).dump_to_dict()
 
     def set_log_probs(self, log_probs: list[TokenLogprobs],
                       cum_log_probs: list[float]):
@@ -233,10 +234,12 @@ def cum_log_probs(self) -> list[float] | None:
     def mm_embedding_handle(self) -> Dict[str, Any] | None:
         return self._mm_embeddings
 
+
 class LlmResult:
     """LlmResult wraps `bindings.executor.Result` but detour some features to Python implementation"""
     py_result_properties = frozenset(
-        ('context_logits', 'generation_logits', 'log_probs', 'cum_log_probs', 'mm_embedding_handle'))
+        ('context_logits', 'generation_logits', 'log_probs', 'cum_log_probs',
+         'mm_embedding_handle'))
 
     def __init__(self,
                  result: Union[bytes, tensorrt_llm.bindings.executor.Result],
diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py
index 962874b2cb17..44bc0462fba7 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_engine.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -1075,8 +1075,10 @@ def init_meta_tensor(t: torch.Tensor):
                 initialize_dummy_weights(model)
 
             elif load_format == LoadFormat.VISION_ONLY:
-                # Vision weights are already loaded within model
-                pass
+                # Vision weights are already loaded within the model.
+                logger.info(
+                    "LoadFormat.VISION_ONLY: skipping weight loading; using preloaded vision weights."
+                )
 
             else:
                 raise NotImplementedError(
@@ -1644,14 +1646,12 @@ def _prepare_tp_inputs_no_cache(
                 multi_modal_data.append(multimodal_embedding)
 
             # Multimodal
-            multimodal_params = MultimodalParams(
-                multimodal_data=request.py_multimodal_data,
-            )
-            multimodal_params.to_device("multimodal_data",
-                                        "cuda",
-                                        pin_memory=True)
-
-            if multimodal_params.has_content():
+            if request.py_multimodal_data is not None:
+                multimodal_params = MultimodalParams(
+                    multimodal_data=request.py_multimodal_data, )
+                multimodal_params.to_device("multimodal_data",
+                                            "cuda",
+                                            pin_memory=True)
                 multimodal_params_list.append(multimodal_params)
 
             request.py_batch_idx = request.py_seq_slot
@@ -2155,9 +2155,11 @@ def forward(
             with MoeLoadBalancerIterContext(moe_load_balancer):
                 # Special handling for multimodal encoder only mode
                 if self.pytorch_backend_config.mm_encoder_only:
-                    return self._forward_step_mm_encoder_only(inputs, scheduled_requests)
+                    return self._forward_step_mm_encoder_only(
+                        inputs, scheduled_requests)
                 else:
-                    return self._forward_step(inputs, gather_ids, gather_context_logits)
+                    return self._forward_step(inputs, gather_ids,
+                                              gather_context_logits)
         with self._maybe_pad_batch(scheduled_requests, kv_cache_manager,
                                    spec_resource_manager) as scheduled_requests:
             maybe_graph = self._maybe_get_cuda_graph(scheduled_requests)
@@ -2252,16 +2254,17 @@ def _forward_step(self,
             return {'logits': logits}
 
     @nvtx_range("_forward_step_mm_encoder_only")
-    def _forward_step_mm_encoder_only(self,
-                                      inputs: Dict[str, Any],
-                                      scheduled_requests: ScheduledRequests) -> Dict[str, Any]:
+    def _forward_step_mm_encoder_only(
+            self, inputs: Dict[str, Any],
+            scheduled_requests: ScheduledRequests) -> Dict[str, Any]:
         """Forward step for multimodal encoder only mode - returns mm_embeddings instead of logits."""
         # Get multimodal parameters from inputs
         multimodal_params = inputs.get("multimodal_params", [])
         if not multimodal_params or len(multimodal_params) == 0:
             # Return empty embeddings if no multimodal data
             return {'mm_embeddings': []}
-        if getattr(scheduled_requests.context_requests[0], 'multimodal_lengths', None) is None:
+        if getattr(scheduled_requests.context_requests[0], 'multimodal_lengths',
+                   None) is None:
             multimodal_chunks = None
         else:
             multimodal_chunks = [
@@ -2272,12 +2275,19 @@ def _forward_step_mm_encoder_only(self,
         # For mm_encoder_only mode, we only run the vision encoder part
         # The model should be a vision encoder (e.g., Qwen2VisionModelBase)
         mm_embeddings = self.model.forward(multimodal_params)
-        assert len(mm_embeddings) == 1, "mm_embeddings should be a 1-element list, mix modality (video+image) is not supported"
-
-        if multimodal_chunks is None or len(multimodal_chunks) != len(multimodal_params):
-            mm_embeddings = list(torch.chunk(mm_embeddings[0], len(scheduled_requests.context_requests), dim=0))
+        assert len(
+            mm_embeddings
+        ) == 1, "mm_embeddings should be a 1-element list, mix modality (video+image) is not supported"
+
+        if multimodal_chunks is None or len(multimodal_chunks) != len(
+                multimodal_params):
+            mm_embeddings = list(
+                torch.chunk(mm_embeddings[0],
+                            len(scheduled_requests.context_requests),
+                            dim=0))
         else:
-            mm_embeddings = list(torch.split(mm_embeddings[0], multimodal_chunks, dim=0))
+            mm_embeddings = list(
+                torch.split(mm_embeddings[0], multimodal_chunks, dim=0))
 
         return {'mm_embeddings': mm_embeddings, 'logits': None}
 
diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
index f3b72100de03..a3e8319355d1 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
@@ -173,10 +173,13 @@ def _mangle_executor_config(executor_config: ExecutorConfig):
         from tensorrt_llm.llmapi.llm_args import LoadFormat
         pytorch_backend_config.mm_encoder_only = True
         pytorch_backend_config.load_format = LoadFormat.VISION_ONLY
-        # TODO: add comment and print warning here
+        # Disable overlap scheduler for multimodal encoder-only mode
+        logger.warning(
+            "Disabling overlap scheduler for multimodal encoder-only mode. "
+            "The overlap scheduler is designed for generation models and is not needed "
+            "when only processing vision encoder inputs.")
         pytorch_backend_config.disable_overlap_scheduler = True
-        # TODO: add comment here to infer it by max_num_images and image_token_sizen
-        executor_config.max_num_tokens = 16384
+
 
 def _get_mapping(executor_config: ExecutorConfig) -> Mapping:
     if executor_config.mapping is None:
diff --git a/tensorrt_llm/_torch/pyexecutor/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler.py
index 10adbe8eca30..a9ce22da58b6 100644
--- a/tensorrt_llm/_torch/pyexecutor/sampler.py
+++ b/tensorrt_llm/_torch/pyexecutor/sampler.py
@@ -1,7 +1,7 @@
 from abc import ABC, abstractmethod
 from collections.abc import Iterable
 from dataclasses import dataclass
-from typing import Literal, List
+from typing import List, Literal
 
 import torch
 
@@ -104,6 +104,7 @@ class MultimodalResult:
     def values(self):
         return vars(self).values()
 
+
 @dataclass(kw_only=True)
 class SampleStateWithMMResult:
     scheduled_requests: ScheduledRequests
@@ -111,7 +112,7 @@ class SampleStateWithMMResult:
     data: MultimodalResult = None
 
 
-class EarlyStopWithMMResult(EarlyStopSampler):
+class EarlyStopWithMMResult(Sampler):
     """
     Use for skipping decoding step for non generation model, and return the batch_output (such as mm_embeddings)
     """
@@ -120,22 +121,25 @@ def sample_async(self, scheduled_requests: ScheduledRequests,
                      model_outputs) -> SampleStateWithMMResult:
         # from model_outputs to MultimodalResult
         data = MultimodalResult(mm_embeddings=model_outputs['mm_embeddings'])
-        return SampleStateWithMMResult(scheduled_requests=scheduled_requests, data=data)
+        return SampleStateWithMMResult(scheduled_requests=scheduled_requests,
+                                       data=data)
 
     def update_requests(self, state: SampleStateWithMMResult) -> None:
         assert isinstance(state, SampleStateWithMMResult)
         scheduled_requests = state.scheduled_requests
         assert (not scheduled_requests.generation_requests)
         mm_embeddings = state.data.mm_embeddings
-        for idx, request in enumerate(scheduled_requests.context_requests):
+        for request, mm_embedding in zip(scheduled_requests.context_requests,
+                                         mm_embeddings):
             request.state = LlmRequestState.GENERATION_COMPLETE
             # NOTE: This is a hack: set finish reason manually and set the beam 0
             request.set_finished_reason(FinishReason.LENGTH, 0)
-            if idx < len(mm_embeddings):
-                if len(mm_embeddings[idx]) != sum(request.multimodal_lengths):
-                    raise ValueError(f"mm_embedding shape mismatch: {len(mm_embeddings[idx])} != {sum(request.multimodal_lengths)}")
+            if len(mm_embedding) != sum(request.multimodal_lengths):
+                raise ValueError(
+                    f"mm_embedding shape mismatch: {len(mm_embedding)} != {sum(request.multimodal_lengths)}"
+                )
 
-                request.py_result.append_mm_embeddings(mm_embeddings[idx])
+            request.py_result.append_mm_embeddings(mm_embedding)
 
 
 def top_k_sampling_batch(logits, top_k=50):
diff --git a/tensorrt_llm/commands/serve.py b/tensorrt_llm/commands/serve.py
index 521d9eb62e6c..aac70039f3cf 100644
--- a/tensorrt_llm/commands/serve.py
+++ b/tensorrt_llm/commands/serve.py
@@ -173,11 +173,13 @@ def launch_server(host: str,
 
     asyncio.run(server(host, port))
 
-def launch_mm_encoder_server(host: str,
-                  port: int,
-                  encoder_args: dict,
-                  metadata_server_cfg: Optional[MetadataServerConfig] = None,
-                  ):
+
+def launch_mm_encoder_server(
+    host: str,
+    port: int,
+    encoder_args: dict,
+    metadata_server_cfg: Optional[MetadataServerConfig] = None,
+):
     model = encoder_args["model"]
     mm_encoder = MultimodalEncoder(**encoder_args)
 
@@ -187,6 +189,7 @@ def launch_mm_encoder_server(host: str,
                           metadata_server_cfg=metadata_server_cfg)
     asyncio.run(server(host, port))
 
+
 @click.command("serve")
 @click.argument("model", type=str)
 @click.option("--tokenizer",
@@ -381,12 +384,10 @@ def serve(
               type=str,
               default=None,
               help="Path to metadata server config file")
-def serve_encoder(model: str, host: str, port: int,
-          log_level: str, max_batch_size: int,
-          gpus_per_node: Optional[int],
-          trust_remote_code: bool,
-          extra_encoder_options: Optional[str],
-          metadata_server_config_file: Optional[str]):
+def serve_encoder(model: str, host: str, port: int, log_level: str,
+                  max_batch_size: int, gpus_per_node: Optional[int],
+                  trust_remote_code: bool, extra_encoder_options: Optional[str],
+                  metadata_server_config_file: Optional[str]):
     """Running an OpenAI API compatible server
 
     MODEL: model name | HF checkpoint path | TensorRT engine path
@@ -394,17 +395,17 @@ def serve_encoder(model: str, host: str, port: int,
     logger.set_level(log_level)
 
     # TODO: expose more argument progressivly
-    llm_args, _ = get_llm_args(
-        model=model,
-        max_batch_size=max_batch_size,
-        gpus_per_node=gpus_per_node,
-        trust_remote_code=trust_remote_code)
+    llm_args, _ = get_llm_args(model=model,
+                               max_batch_size=max_batch_size,
+                               gpus_per_node=gpus_per_node,
+                               trust_remote_code=trust_remote_code)
 
     encoder_args_extra_dict = {}
     if extra_encoder_options is not None:
         with open(extra_encoder_options, 'r') as f:
             encoder_args_extra_dict = yaml.safe_load(f)
-    encoder_args = update_llm_args_with_extra_dict(llm_args, encoder_args_extra_dict)
+    encoder_args = update_llm_args_with_extra_dict(llm_args,
+                                                   encoder_args_extra_dict)
 
     metadata_server_cfg = parse_metadata_server_config_file(
         metadata_server_config_file)
@@ -726,7 +727,7 @@ def resolve_command(self, ctx, args):
         "serve": serve,
         "disaggregated": disaggregated,
         "disaggregated_mpi_worker": disaggregated_mpi_worker,
-        "mmemb_serve": serve_encoder
+        "mm_embedding_serve": serve_encoder
     })
 
 if __name__ == "__main__":
diff --git a/tensorrt_llm/executor/result.py b/tensorrt_llm/executor/result.py
index 7374f0ceb7b3..dada83a16a03 100644
--- a/tensorrt_llm/executor/result.py
+++ b/tensorrt_llm/executor/result.py
@@ -3,8 +3,8 @@
 import weakref
 from dataclasses import dataclass, field
 from queue import Empty, Queue
-from typing import (TYPE_CHECKING, Any, Callable, Dict, List, Literal, NamedTuple,
-                    Optional, TypeAlias, Union)
+from typing import (TYPE_CHECKING, Any, Callable, Dict, List, Literal,
+                    NamedTuple, Optional, TypeAlias, Union)
 from weakref import WeakMethod
 
 import torch
diff --git a/tensorrt_llm/llmapi/__init__.py b/tensorrt_llm/llmapi/__init__.py
index 4e6c906ad3eb..4981b1639170 100644
--- a/tensorrt_llm/llmapi/__init__.py
+++ b/tensorrt_llm/llmapi/__init__.py
@@ -3,7 +3,6 @@
 from ..sampling_params import GuidedDecodingParams, SamplingParams
 from .build_cache import BuildCacheConfig
 from .llm import LLM, RequestOutput
-from .mm_encoder import MultimodalEncoder
 # yapf: disable
 from .llm_args import (AttentionDpConfig, AutoDecodingConfig, BatchingType,
                        CacheTransceiverConfig, CalibConfig,
@@ -17,6 +16,7 @@
                        UserProvidedDecodingConfig)
 from .llm_utils import (BuildConfig, KvCacheRetentionConfig, QuantAlgo,
                         QuantConfig)
+from .mm_encoder import MultimodalEncoder
 from .mpi_session import MpiCommSession
 
 __all__ = [
diff --git a/tensorrt_llm/llmapi/llm.py b/tensorrt_llm/llmapi/llm.py
index dc67b00c8c55..eb17e9cbd946 100644
--- a/tensorrt_llm/llmapi/llm.py
+++ b/tensorrt_llm/llmapi/llm.py
@@ -82,7 +82,8 @@ def prompt(self) -> Optional[str]:
 
     def _repr_fields(self):
         return [
-            "request_id", "prompt", "prompt_token_ids", "outputs", "finished", "mm_embedding_handle"
+            "request_id", "prompt", "prompt_token_ids", "outputs", "finished",
+            "mm_embedding_handle"
         ]
 
 
diff --git a/tensorrt_llm/llmapi/mm_encoder.py b/tensorrt_llm/llmapi/mm_encoder.py
index 0f599431b856..a827ced96737 100644
--- a/tensorrt_llm/llmapi/mm_encoder.py
+++ b/tensorrt_llm/llmapi/mm_encoder.py
@@ -1,17 +1,18 @@
 from pathlib import Path
-from typing import Any, Optional, Union
+from typing import Any, List, Literal, Optional, Sequence, Union
 
+from tqdm import tqdm
+
+from tensorrt_llm._utils import nvtx_range_debug
+from tensorrt_llm.bindings import executor as tllm
+from tensorrt_llm.inputs import create_input_processor, prompt_inputs
 from tensorrt_llm.inputs.data import PromptInputs
-from .llm import BaseLLM, _TorchLLM, RequestOutput
 from tensorrt_llm.sampling_params import SamplingParams
-from tensorrt_llm._utils import nvtx_range_debug
-from typing import List, Sequence
-from tensorrt_llm.inputs import prompt_inputs
-from tqdm import tqdm
-from tensorrt_llm.inputs import create_input_processor
+
+from .llm import BaseLLM, RequestOutput, _TorchLLM
 from .llm_args import PybindMirror
 from .mpi_session import external_mpi_comm_available
-from tensorrt_llm.bindings import executor as tllm
+
 
 class MultimodalEncoder(_TorchLLM):
     """MultimodalEncoder class is the main class for running a multimodal encoder model using PyTorch backend.
@@ -21,16 +22,23 @@ def __init__(self,
                  model: Union[str, Path],
                  trust_remote_code: bool = False,
                  tensor_parallel_size: int = 1,
-                 dtype: str = "auto",
+                 dtype: Literal["auto", "float16", "float32",
+                                "bfloat16"] = "auto",
                  **kwargs: Any) -> None:
 
         # Validate that users don't pass LLM-specific or TRT-specific arguments
         self._validate_mm_args_for_torch_backend(kwargs)
 
+        # Set higher default max_num_tokens for multimodal encoder (16384 vs 8192 default)
+        # Vision encoders can handle more tokens than text-only models
+        # TODO: Make this adaptive based on model-specific max_mm_token_length (see _test_llm_multimodal_general)
+        if 'max_num_tokens' not in kwargs or kwargs['max_num_tokens'] is None:
+            kwargs['max_num_tokens'] = 16384
+
         super().__init__(model,
                          trust_remote_code=trust_remote_code,
                          tensor_parallel_size=tensor_parallel_size,
-                         dtype = dtype,
+                         dtype=dtype,
                          **kwargs)
 
     def _build_model(self):
@@ -87,7 +95,7 @@ def _build_model(self):
             mpi_session=self.mpi_session,
             reuse_mpi_comm=external_mpi_comm_available(
                 self.args.parallel_config.world_size),
-            is_llm_executor=True, # TODO: check if this is correct or needed
+            is_llm_executor=True,  # TODO: check if this is correct or needed
             garbage_collection_gen0_threshold=self.args.
             garbage_collection_gen0_threshold)
 
@@ -95,7 +103,6 @@ def _validate_mm_args_for_torch_backend(self, kwargs: dict) -> None:
         """Validate that users don't pass LLM-specific arguments when using MultimodalEncoder (PyTorch).
         Placeholder for now.
         """
-        pass
 
     def generate(
         self,
@@ -129,9 +136,7 @@ def _item_at(maybe_batched: Union[Any, Sequence[Any]], pos: int) -> Any:
 
         futures = []
         for i, request_inputs in enumerate(inputs):
-            future = self.generate_async(
-                request_inputs
-            )
+            future = self.generate_async(request_inputs)
             futures.append(future)
 
         for future in tqdm(futures,
@@ -145,12 +150,14 @@ def _item_at(maybe_batched: Union[Any, Sequence[Any]], pos: int) -> Any:
 
         return futures
 
-    @nvtx_range_debug("MM_encoder.generate_async", color="green", category="VisionEncoder")
+    @nvtx_range_debug("MM_encoder.generate_async",
+                      color="green",
+                      category="VisionEncoder")
     def generate_async(
         self,
         inputs: PromptInputs,
         sampling_params: Optional[SamplingParams] = None,
-    ):
+    ) -> RequestOutput:
         """Generate output for the given multimodal request in the asynchronous mode.
         Asynchronous generation accepts single multimodal request only.
 
@@ -160,4 +167,3 @@ def generate_async(
         result = super().generate_async(inputs, sampling_params)
         # TODO: possible postprocess the result for disaggregated serving
         return result
-
diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py
index 3cc498bfcce4..3d9e31d48413 100644
--- a/tensorrt_llm/serve/openai_server.py
+++ b/tensorrt_llm/serve/openai_server.py
@@ -16,13 +16,13 @@
 from transformers import AutoConfig, AutoProcessor
 
 from tensorrt_llm._tensorrt_engine import LLM
-from tensorrt_llm.llmapi import MultimodalEncoder
 # yapf: disable
 from tensorrt_llm.executor import CppExecutorError
 from tensorrt_llm.executor.postproc_worker import PostprocParams
 from tensorrt_llm.inputs import prompt_inputs
 from tensorrt_llm.inputs.utils import ConversationMessage, apply_chat_template
 from tensorrt_llm.llmapi import DisaggregatedParams as LlmDisaggregatedParams
+from tensorrt_llm.llmapi import MultimodalEncoder
 from tensorrt_llm.llmapi.disagg_utils import MetadataServerConfig, ServerRole
 from tensorrt_llm.llmapi.llm import RequestOutput
 from tensorrt_llm.logger import logger
@@ -31,7 +31,8 @@
 from tensorrt_llm.serve.metadata_server import create_metadata_server
 from tensorrt_llm.serve.openai_protocol import (ChatCompletionRequest,
                                                 ChatCompletionResponse,
-                                                CompletionRequest,
+                                                ChatCompletionResponseChoice,
+                                                ChatMessage, CompletionRequest,
                                                 CompletionResponse,
                                                 CompletionResponseChoice,
                                                 ErrorResponse, ModelCard,
@@ -42,7 +43,7 @@
     chat_stream_post_processor, completion_response_post_processor,
     completion_stream_post_processor)
 from tensorrt_llm.version import __version__ as VERSION
-from tensorrt_llm.serve.openai_protocol import ChatMessage, ChatCompletionResponseChoice, UsageInfo
+
 from .._utils import nvtx_mark
 
 # yapf: enale
@@ -361,7 +362,7 @@ async def create_mm_embedding_response(
                 tool.model_dump() for tool in request.tools
             ]
             # TODO: placeholder for multimodal disagg e2e
-            disaggregated_params = to_llm_disaggregated_params(request.disaggregated_params)
+            to_llm_disaggregated_params(request.disaggregated_params)
 
             conversation, mm_coroutines, mm_placeholder_counts = parse_chat_messages_coroutines(request.messages, self.model_config)
 
diff --git a/tests/unittest/_torch/multimodal/test_external_embedding.py b/tests/unittest/_torch/multimodal/test_external_embedding.py
new file mode 100644
index 000000000000..24d670b79fcc
--- /dev/null
+++ b/tests/unittest/_torch/multimodal/test_external_embedding.py
@@ -0,0 +1,181 @@
+from unittest.mock import Mock, patch
+import pytest
+
+from tensorrt_llm.inputs.data import TextPrompt
+from tensorrt_llm.sampling_params import SamplingParams
+from tensorrt_llm._torch.models.modeling_llava_next import LlavaNextInputProcessor
+
+
+# Model configurations for different processors
+MODEL_CONFIGS = {
+    "LlavaNextInputProcessor": {
+        "processor_class": LlavaNextInputProcessor,
+        "module_path": "tensorrt_llm._torch.models.modeling_llava_next",
+        "config_setup": lambda mock_config: setattr_multiple(mock_config, {
+            "image_token_index": 32001,
+            "vocab_size": 32001,
+            "text_config.hidden_size": 4096,
+            "text_config.vocab_size": 32000,
+            "vision_config": Mock(),
+            "vision_feature_select_strategy": "default"
+        })
+    },
+    # TODO: Add test for more VLM models
+    # Future model configurations can be added here
+    # "LlamaInputProcessor": {
+    #     "processor_class": LlamaInputProcessor,
+    #     "module_path": "tensorrt_llm._torch.models.modeling_llama",
+    # },
+    # "QwenInputProcessor": {
+    #     "processor_class": QwenInputProcessor,
+    #     "module_path": "tensorrt_llm._torch.models.modeling_qwen2vl",
+    # },
+}
+
+def setattr_multiple(obj, attr_dict):
+    """Helper function to set multiple nested attributes."""
+    for attr_path, value in attr_dict.items():
+        if '.' in attr_path:
+            # Handle nested attributes like 'text_config.hidden_size'
+            attrs = attr_path.split('.')
+            current_obj = obj
+            for attr in attrs[:-1]:
+                if not hasattr(current_obj, attr):
+                    setattr(current_obj, attr, Mock())
+                current_obj = getattr(current_obj, attr)
+            setattr(current_obj, attrs[-1], value)
+        else:
+            setattr(obj, attr_path, value)
+
+
+@pytest.fixture(params=["LlavaNextInputProcessor"])
+def processor_setup(request):
+    """Fixture to set up different input processors based on the parameter."""
+    processor_name = request.param
+    config = MODEL_CONFIGS[processor_name]
+
+    # Mock model configuration
+    mock_config = Mock()
+    config["config_setup"](mock_config)
+
+    # Mock tokenizer and processor
+    mock_tokenizer = Mock()
+    mock_processor = Mock()
+
+    # Create processor instance with mocks
+    with patch(f'{config["module_path"]}.AutoTokenizer') as mock_auto_tokenizer, \
+         patch(f'{config["module_path"]}.AutoProcessor') as mock_auto_processor:
+
+        mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer
+        mock_auto_processor.from_pretrained.return_value = mock_processor
+
+        processor = config["processor_class"](
+            model_path="dummy_path",
+            model_config=mock_config,
+            tokenizer=mock_tokenizer,
+            trust_remote_code=True
+        )
+
+        # Return processor along with mocks for test usage
+        return {
+            "processor": processor,
+            "mock_tokenizer": mock_tokenizer,
+            "mock_processor": mock_processor,
+            "mock_config": mock_config,
+            "processor_name": processor_name
+        }
+
+
+class TestExternalEmbedding:
+    # TODO: Add test for more modalities (audio, video, etc.)
+    def test_attach_multimodal_embeddings_image_basic(self, processor_setup):
+        """Test basic functionality of attach_multimodal_embeddings."""
+        import torch
+
+        # Prepare test inputs
+        text_prompt: TextPrompt = {
+            "prompt": "What is in this image? <image>"
+        }
+
+        # Create mock embedding tensors (batch_size=1, seq_len=144, hidden_dim=768)
+        image_embeddings = [
+            torch.randn(144, 4096),  # Single image embedding so num_frames=1 always
+        ]
+
+        multimodal_embedding = {
+            "image": image_embeddings
+        }
+
+        sampling_params = SamplingParams()
+
+        # Mock tokenizer output - simulate tokenizing the text with image token
+        mock_token_ids = torch.tensor([1, 2, 3, 32001, 4, 5])  # image token at index 3
+        processor_setup["mock_tokenizer"].return_value.input_ids = [mock_token_ids]
+
+        result_token_ids, extra_processed_inputs = processor_setup["processor"].attach_multimodal_embeddings(
+            text_prompt, multimodal_embedding, sampling_params
+        )
+
+        # Verify outputs
+        assert isinstance(result_token_ids, list)
+        assert len(result_token_ids) == len(mock_token_ids) + len(image_embeddings[0]) - 1
+        assert isinstance(extra_processed_inputs, dict)
+        assert "multimodal_data" in extra_processed_inputs
+        assert "multimodal_embedding" in extra_processed_inputs["multimodal_data"]
+
+        # Check that multimodal embedding is properly formatted
+        mm_embedding = extra_processed_inputs["multimodal_data"]["multimodal_embedding"]
+        assert mm_embedding.shape[-1] == 4096  # Hidden dimension should match
+        assert mm_embedding.shape[0] == 144
+
+        # Verify tokenizer was called with correct prompt
+        processor_setup["mock_tokenizer"].assert_called_once_with(text_prompt["prompt"], return_tensors="pt")
+
+    def test_attach_multimodal_embeddings_multiple_images(self, processor_setup):
+        """Test with multiple image embeddings."""
+        import torch
+
+        text_prompt: TextPrompt = {
+            "prompt": "Compare these images: <image> and <image>"
+        }
+
+        # Create multiple image embeddings
+        image_embeddings = [
+            torch.randn(144, 4096),  # First image
+            torch.randn(288, 4096),  # Second image
+        ]
+
+        # multiple images are concatenated along the first dimension
+        multimodal_embedding = {
+            "image": image_embeddings
+        }
+
+        sampling_params = SamplingParams()
+
+        # Mock tokenizer output with two image tokens
+        mock_token_ids = torch.tensor([1, 2, 32001, 3, 4, 32001, 5])
+        processor_setup["mock_tokenizer"].return_value.input_ids = [mock_token_ids]
+
+        # Call the function
+        result_token_ids, extra_processed_inputs = processor_setup["processor"].attach_multimodal_embeddings(
+            text_prompt, multimodal_embedding, sampling_params
+        )
+
+        # Verify outputs
+        assert isinstance(result_token_ids, list)
+        assert len(result_token_ids) == len(mock_token_ids) + torch.cat(image_embeddings, dim=0).shape[0] - 2
+        assert isinstance(extra_processed_inputs, dict)
+        assert "multimodal_data" in extra_processed_inputs
+        assert "multimodal_embedding" in extra_processed_inputs["multimodal_data"]
+
+        # Check that multimodal embedding is properly formatted
+        mm_embedding = extra_processed_inputs["multimodal_data"]["multimodal_embedding"]
+        assert mm_embedding.shape[-1] == 4096  # Hidden dimension should match
+        assert mm_embedding.shape[0] == 144 + 288
+
+        # Verify tokenizer was called with correct prompt
+        processor_setup["mock_tokenizer"].assert_called_once_with(text_prompt["prompt"], return_tensors="pt")
+
+
+if __name__ == "__main__":
+    pytest.main([__file__])
diff --git a/tests/unittest/_torch/multimodal/test_find_num_image_tokens.py b/tests/unittest/_torch/multimodal/test_find_num_image_tokens.py
new file mode 100644
index 000000000000..d102f8d4e009
--- /dev/null
+++ b/tests/unittest/_torch/multimodal/test_find_num_image_tokens.py
@@ -0,0 +1,159 @@
+import json
+import os
+from typing import Dict, Any
+
+import pytest
+import torch
+from PIL import Image
+import requests
+import io
+
+from tensorrt_llm import MultimodalEncoder
+from tensorrt_llm._torch.shared_tensor import SharedTensorContainer
+from tensorrt_llm.inputs import default_multimodal_input_loader
+from transformers import AutoTokenizer, AutoConfig, AutoProcessor
+from tensorrt_llm._torch.models.modeling_llava_next import LlavaNextInputProcessor
+from tensorrt_llm._torch.models.modeling_qwen2vl import Qwen2VLInputProcessorBase
+
+example_images = [
+    "https://huggingface.co/datasets/YiYiXu/testing-images/resolve/main/seashore.png",
+    "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/inpaint.png",
+    "https://huggingface.co/datasets/Sayali9141/traffic_signal_images/resolve/main/61.jpg",
+]
+
+
+def download_image(url: str) -> Image.Image:
+    """Download image from URL and return as PIL Image."""
+    response = requests.get(url)
+    response.raise_for_status()
+    return Image.open(io.BytesIO(response.content))
+
+
+@pytest.fixture(scope="function")
+def multimodal_model_configs():
+    """Get multimodal model configurations for testing."""
+    model_configs = {
+        'llava-v1.6-mistral-7b-hf': {
+            'hf_model_dir': 'llava-hf/llava-v1.6-mistral-7b-hf',
+            'model_type': 'llava_next',
+        },
+        'qwen2.5-vl': {
+            'hf_model_dir': 'Qwen/Qwen2.5-VL-3B-Instruct',
+            'model_type': 'qwen2_5_vl',
+        }
+    }
+    return model_configs
+
+
+@pytest.mark.parametrize("model_key", [
+    "llava-v1.6-mistral-7b-hf",
+    "qwen2.5-vl",
+])
+def test_get_num_tokens_per_image(model_key, multimodal_model_configs):
+    """Test that get_num_tokens_per_image predicts the correct number of tokens.
+
+    This test verifies that the get_num_tokens_per_image method correctly predicts
+    the number of image tokens by comparing against actual encoder output shapes.
+    Tests all example images in a single test run.
+    """
+    # Get model configuration
+    if model_key not in multimodal_model_configs:
+        pytest.skip(f"Skipping test for {model_key} - model not available")
+
+    model_config = multimodal_model_configs[model_key]
+    encoder_model_dir = model_config['hf_model_dir']
+    model_type = model_config['model_type']
+
+    # Test configuration
+    max_batch_size = len(example_images)  # Batch size to handle all test images
+
+    encoder = None
+
+    try:
+        encoder = MultimodalEncoder(model=encoder_model_dir,
+                                    max_batch_size=max_batch_size)
+
+        # Load model configuration and create input processor once
+        model_config_dict = AutoConfig.from_pretrained(encoder_model_dir, trust_remote_code=True)
+        tokenizer = AutoTokenizer.from_pretrained(encoder_model_dir, trust_remote_code=True)
+
+        # Create input processor once
+        if model_type == 'llava_next':
+            input_processor = LlavaNextInputProcessor(
+                model_path=encoder_model_dir,
+                model_config=model_config_dict,
+                tokenizer=tokenizer,
+                trust_remote_code=True
+            )
+        elif model_type == 'qwen2_5_vl':
+            input_processor = Qwen2VLInputProcessorBase(
+                model_path=encoder_model_dir,
+                model_config=model_config_dict,
+                tokenizer=tokenizer,
+                trust_remote_code=True
+            )
+        else:
+            pytest.fail(f"Unsupported model type: {model_type}")
+
+        # Prepare batch inputs with all 3 images
+        prompts = ["dummy"] * len(example_images)  # One prompt per image
+        media = example_images  # All image URLs
+
+        # Prepare inputs and get actual embeddings for all images
+        inputs = default_multimodal_input_loader(
+            tokenizer=tokenizer,
+            model_dir=encoder_model_dir,
+            model_type=model_type,
+            modality="image",
+            prompts=prompts,
+            media=media,
+            image_data_format="pt"
+        )
+
+        # Get actual embeddings from encoder (batch processing)
+        encoder_outputs = encoder.generate(inputs)
+        assert len(encoder_outputs) == len(example_images), f"Expected {len(example_images)} encoder outputs, got {len(encoder_outputs)}"
+
+        for image_idx, test_image_url in enumerate(example_images):
+
+            # Get test image dimensions
+            test_image = download_image(test_image_url)
+            image_width, image_height = test_image.size
+
+            # Get actual embedding tensor for this image
+            actual_embedding = SharedTensorContainer.from_dict(
+                encoder_outputs[image_idx].mm_embedding_handle
+            ).get_local_view()
+
+            # The first dimension should be the number of image tokens
+            actual_num_tokens = actual_embedding.shape[0]
+
+            # Get predicted number of tokens using get_num_tokens_per_image
+            if model_type == 'llava_next':
+                predicted_num_tokens = input_processor.get_num_tokens_per_image(
+                    image_width=image_width,
+                    image_height=image_height
+                )
+            elif model_type == 'qwen2_5_vl':
+                predicted_num_tokens = input_processor.get_num_tokens_per_image(
+                    image_width=image_width,
+                    image_height=image_height,
+                    num_frames=1,
+                    do_resize=True
+                )
+
+            # The key assertion: predicted should match actual
+            assert predicted_num_tokens == actual_num_tokens, \
+                f"Token count mismatch for {model_key} with image {image_idx} ({image_width}x{image_height}): " \
+                f"predicted={predicted_num_tokens}, actual={actual_num_tokens}"
+
+            # Additional validation: ensure we got reasonable token counts
+            assert actual_num_tokens > 0, f"Got zero image tokens for {model_key} image {image_idx}"
+            assert actual_embedding.ndim >= 2, f"Expected at least 2D embedding, got {actual_embedding.ndim}D"
+
+    finally:
+        # Cleanup resources
+        if encoder is not None:
+            del encoder
+
+
diff --git a/tests/unittest/_torch/multimodal/test_mm_encoder_standalone.py b/tests/unittest/_torch/multimodal/test_mm_encoder_standalone.py
index 7d493789bbd3..e42e07654fbf 100644
--- a/tests/unittest/_torch/multimodal/test_mm_encoder_standalone.py
+++ b/tests/unittest/_torch/multimodal/test_mm_encoder_standalone.py
@@ -1,13 +1,13 @@
+import json
 import os
+
 import pytest
-import copy
-import json
 
 from tensorrt_llm import MultimodalEncoder
 from tensorrt_llm._torch.shared_tensor import SharedTensorContainer
-from tensorrt_llm.llmapi.llm import LLM, SamplingParams
-from tensorrt_llm.llmapi import KvCacheConfig
 from tensorrt_llm.inputs import default_multimodal_input_loader
+from tensorrt_llm.llmapi import KvCacheConfig
+from tensorrt_llm.llmapi.llm import LLM, SamplingParams
 
 example_images = [
     "https://huggingface.co/datasets/YiYiXu/testing-images/resolve/main/seashore.png",
@@ -23,10 +23,10 @@ def multimodal_model_config():
     model_configs = {
         'llava-v1.6-mistral-7b-hf': {
             'model_name': 'llava-v1.6-mistral-7b-hf',
-            'hf_model_dir': 'llava-hf/llava-v1.6-mistral-7b-hf',  # HuggingFace model ID
+            'hf_model_dir': 'llava-hf/llava-v1.6-mistral-7b-hf',
         }
     }
-    
+
     return model_configs['llava-v1.6-mistral-7b-hf']
 
 
@@ -35,148 +35,129 @@ def multimodal_model_config():
 ])
 def test_single_image_chat(model_key, multimodal_model_config):
     """Test processing single image using encoder (pass mm_embeddings) + LLM API.
-    
+
     This test verifies that encoder (pass mm_embeddings) + LLM API produces identical
     results to standard llm generation (pass raw image) by comparing outputs.
     """
     # Get model configuration
     if model_key != "llava-v1.6-mistral-7b-hf":
-        pytest.skip(f"Skipping test for {model_key} - only testing llava-v1.6-mistral-7b-hf for now")
-    
+        #TODO: add more model tests progressively here
+        pytest.skip(
+            f"Skipping test for {model_key} - only testing llava-v1.6-mistral-7b-hf for now"
+        )
+
     # Extract model information from config
     encoder_model_dir = multimodal_model_config['hf_model_dir']
-    
+
     # Test configuration
     max_tokens = 64
     free_gpu_memory_fraction = 0.6
     max_batch_size = 1
-    
+
     # Test data - OpenAI chat completion format
     prompts = ["Describe the natural environment in the image."]
     media = [example_images[0]]
-    
-    # Create OpenAI chat messages format
-    messages_list = []
-    for prompt, image_url in zip(prompts, media):
-        messages = [{
-            "role": "user",
-            "content": [{
-                "type": "text",
-                "text": prompt
-            }, {
-                "type": "image_url",
-                "image_url": {
-                    "url": image_url
-                }
-            }]
-        }]
-        messages_list.append(messages)
-    
+
     # Sampling configuration
     sampling_params = SamplingParams(max_tokens=max_tokens)
     kv_cache_config = KvCacheConfig(
         enable_block_reuse=False,
         free_gpu_memory_fraction=free_gpu_memory_fraction,
     )
-    
-    # Step 1: Process multimodal data using encoder (pass mm_embeddings)
-    encoder = None
-    llm = None
-    
-    try:
-        # Step 1: Initialize encoder
-        encoder = MultimodalEncoder(model=encoder_model_dir, max_batch_size=max_batch_size)
-        
-        # Step 2: Initialize LLM and prepare inputs
-        llm = LLM(
-            model=encoder_model_dir, 
-            backend='pytorch', 
-            kv_cache_config=kv_cache_config, 
-            trust_remote_code=True
-        )
-        
-        # Load model configuration
-        config_path = os.path.join(llm._hf_model_dir, 'config.json')
-        assert os.path.exists(config_path), f"Model config not found at {config_path}"
-        
-        with open(config_path, 'r') as f:
-            model_config = json.load(f)
-        model_type = model_config['model_type']
-        
-        # Prepare multimodal inputs
-        inputs = default_multimodal_input_loader(
-            tokenizer=llm.tokenizer,
-            model_dir=llm._hf_model_dir,
-            model_type=model_type,
-            modality="image",
-            prompts=prompts,  
-            media=media,      
-            image_data_format="pt"
-        )
-        
-        # Validate inputs structure
-        assert len(inputs) == len(prompts), f"Expected {len(prompts)} inputs, got {len(inputs)}"
-        # Step 3: Generate reference output with raw multimodal inputs
-        outputs_ref = llm.generate(inputs, sampling_params=sampling_params)
-        
-        # Validate reference outputs
-        assert outputs_ref is not None, "Reference generation returned None"
-        assert len(outputs_ref) == len(prompts), f"Expected {len(prompts)} reference outputs, got {len(outputs_ref)}"
-        for i, output in enumerate(outputs_ref):
-            assert len(output.outputs) > 0, f"Reference generation has no output text for input {i}"
-        
-        # Step 4: Prepare inputs for llm (pass mm_embeddings)
-        encoder_outputs = encoder.generate(inputs)
-        inputs = default_multimodal_input_loader(
-            tokenizer=llm.tokenizer,
-            model_dir=llm._hf_model_dir,
-            model_type=model_type,
-            modality="image",
-            prompts=prompts,  
-            mm_embeddings=[SharedTensorContainer.from_dict(output.mm_embedding_handle).get_local_view() for output in encoder_outputs],
-            image_data_format="pt"
-        )
 
-        # Step 5: Generate output using llm (pass mm_embeddings)
-        # Note: For batch processing, we need to match mm_params with inputs
-        outputs = llm.generate(inputs, sampling_params=sampling_params)
-        
-        # Validate outputs
-        assert len(outputs) == len(prompts), f"Expected {len(prompts)} outputs, got {len(outputs)}"
-        for i, output in enumerate(outputs):
-            assert len(output.outputs) > 0, f"generation has no output text for input {i}"
-        
-        # Step 6: Compare outputs - they should match exactly
-        assert len(outputs_ref) == len(outputs), f"Number of outputs don't match: {len(outputs_ref)} vs {len(outputs)}"
-        
-        for i, (ref_output, test_output) in enumerate(zip(outputs_ref, outputs)):
-            # Compare prompts
-            assert ref_output.prompt == test_output.prompt, \
-                f"Prompts don't match for output {i}:\nReference: {ref_output.prompt!r}\nTest: {test_output.prompt!r}"
-            
-            # Compare number of generated outputs
-            assert len(ref_output.outputs) == len(test_output.outputs), \
-                f"Number of generated outputs don't match for output {i}: {len(ref_output.outputs)} vs {len(test_output.outputs)}"
-            
-            # Compare generated text and other attributes
-            for j, (ref_gen, test_gen) in enumerate(zip(ref_output.outputs, test_output.outputs)):
-                assert ref_gen.text == test_gen.text, \
-                    f"Generated text doesn't match for output {i}, generation {j}:\nReference: {ref_gen.text!r}\nTest: {test_gen.text!r}"
-                
-                # Compare token IDs if available
-                if hasattr(ref_gen, 'token_ids') and hasattr(test_gen, 'token_ids'):
-                    assert ref_gen.token_ids == test_gen.token_ids, \
-                        f"Token IDs don't match for output {i}, generation {j}"
-                
-                # Compare log probabilities if available
-                if hasattr(ref_gen, 'logprobs') and hasattr(test_gen, 'logprobs'):
-                    assert ref_gen.logprobs == test_gen.logprobs, \
-                        f"Log probabilities don't match for output {i}, generation {j}"
-                
-    finally:
-        # Cleanup resources
-        if encoder is not None:
-            del encoder
-        if llm is not None:
-            del llm
-    
+    # Process multimodal data using encoder (pass mm_embeddings)
+    encoder = MultimodalEncoder(model=encoder_model_dir,
+                                max_batch_size=max_batch_size)
+    llm = LLM(model=encoder_model_dir,
+              backend='pytorch',
+              kv_cache_config=kv_cache_config,
+              trust_remote_code=True)
+
+    # Load model configuration
+    config_path = os.path.join(llm._hf_model_dir, 'config.json')
+    assert os.path.exists(
+        config_path), f"Model config not found at {config_path}"
+
+    with open(config_path, 'r') as f:
+        model_config = json.load(f)
+    model_type = model_config['model_type']
+
+    # Prepare multimodal inputs
+    inputs = default_multimodal_input_loader(tokenizer=llm.tokenizer,
+                                             model_dir=llm._hf_model_dir,
+                                             model_type=model_type,
+                                             modality="image",
+                                             prompts=prompts,
+                                             media=media,
+                                             image_data_format="pt")
+
+    # Validate inputs structure
+    assert len(inputs) == len(
+        prompts), f"Expected {len(prompts)} inputs, got {len(inputs)}"
+    # Generate reference output with raw multimodal inputs
+    outputs_ref = llm.generate(inputs, sampling_params=sampling_params)
+
+    # Validate reference outputs
+    assert outputs_ref is not None, "Reference generation returned None"
+    assert len(outputs_ref) == len(
+        prompts
+    ), f"Expected {len(prompts)} reference outputs, got {len(outputs_ref)}"
+    for i, output in enumerate(outputs_ref):
+        assert len(
+            output.outputs
+        ) > 0, f"Reference generation has no output text for input {i}"
+
+    # Prepare inputs for llm (pass mm_embeddings)
+    encoder_outputs = encoder.generate(inputs)
+    inputs = default_multimodal_input_loader(
+        tokenizer=llm.tokenizer,
+        model_dir=llm._hf_model_dir,
+        model_type=model_type,
+        modality="image",
+        prompts=prompts,
+        mm_embeddings=[
+            SharedTensorContainer.from_dict(
+                output.mm_embedding_handle).get_local_view()
+            for output in encoder_outputs
+        ],
+        image_data_format="pt")
+
+    # Generate output using llm (pass mm_embeddings)
+    outputs = llm.generate(inputs, sampling_params=sampling_params)
+
+    # Validate outputs
+    assert len(outputs) == len(
+        prompts), f"Expected {len(prompts)} outputs, got {len(outputs)}"
+    for i, output in enumerate(outputs):
+        assert len(
+            output.outputs) > 0, f"generation has no output text for input {i}"
+
+    # Compare outputs - they should match exactly
+    assert len(outputs_ref) == len(
+        outputs
+    ), f"Number of outputs don't match: {len(outputs_ref)} vs {len(outputs)}"
+
+    for i, (ref_output, test_output) in enumerate(zip(outputs_ref, outputs)):
+        # Compare prompts
+        assert ref_output.prompt == test_output.prompt, \
+            f"Prompts don't match for output {i}:\nReference: {ref_output.prompt!r}\nTest: {test_output.prompt!r}"
+
+        # Compare number of generated outputs
+        assert len(ref_output.outputs) == len(test_output.outputs), \
+            f"Number of generated outputs don't match for output {i}: {len(ref_output.outputs)} vs {len(test_output.outputs)}"
+
+        # Compare generated text and other attributes
+        for j, (ref_gen, test_gen) in enumerate(
+                zip(ref_output.outputs, test_output.outputs)):
+            assert ref_gen.text == test_gen.text, \
+                f"Generated text doesn't match for output {i}, generation {j}:\nReference: {ref_gen.text!r}\nTest: {test_gen.text!r}"
+
+            # Compare token IDs if available
+            if hasattr(ref_gen, 'token_ids') and hasattr(test_gen, 'token_ids'):
+                assert ref_gen.token_ids == test_gen.token_ids, \
+                    f"Token IDs don't match for output {i}, generation {j}"
+
+            # Compare log probabilities if available
+            if hasattr(ref_gen, 'logprobs') and hasattr(test_gen, 'logprobs'):
+                assert ref_gen.logprobs == test_gen.logprobs, \
+                    f"Log probabilities don't match for output {i}, generation {j}"

From b2cd0bbe08608dc27ac2a0a9ed927690c714b955 Mon Sep 17 00:00:00 2001
From: Chang Liu <9713593+chang-l@users.noreply.github.com>
Date: Fri, 8 Aug 2025 21:58:30 -0700
Subject: [PATCH 08/15] add type hints

Signed-off-by: Chang Liu <9713593+chang-l@users.noreply.github.com>
---
 tensorrt_llm/_torch/models/modeling_llava_next.py | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/tensorrt_llm/_torch/models/modeling_llava_next.py b/tensorrt_llm/_torch/models/modeling_llava_next.py
index 4e23f8634c95..f91b021a61ab 100644
--- a/tensorrt_llm/_torch/models/modeling_llava_next.py
+++ b/tensorrt_llm/_torch/models/modeling_llava_next.py
@@ -1,6 +1,6 @@
 import copy
 import os
-from typing import Dict, List, Optional, Tuple
+from typing import Dict, List, Optional, Tuple, Union
 
 import numpy as np
 import torch
@@ -67,7 +67,10 @@ def get_num_tokens_per_image(
             [image_size])["num_image_tokens"][0]
         return num_image_tokens
 
-    def _postprocess(self, input_ids, mm_features):
+    def _postprocess(
+        self, input_ids: torch.Tensor, mm_features: Union[torch.Tensor,
+                                                          List[torch.Tensor]]
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
         # Define model specific variables here before shared logic
         mm_tokens = torch.tensor([self.model_config.image_token_index
                                   ]).to(input_ids.device)

From ca90390fead3ca682908ff91d6fc08cd9094c3ce Mon Sep 17 00:00:00 2001
From: Chang Liu <9713593+chang-l@users.noreply.github.com>
Date: Sat, 9 Aug 2025 00:12:52 -0700
Subject: [PATCH 09/15] Add integration (endpoint server) test

Signed-off-by: Chang Liu <9713593+chang-l@users.noreply.github.com>
---
 .../_torch/models/modeling_llava_next.py      |   8 +-
 .../_torch/models/modeling_qwen2vl.py         |   3 +-
 tensorrt_llm/serve/openai_server.py           |   1 -
 tests/integration/defs/test_e2e.py            |   7 +
 .../multimodal/test_external_embedding.py     | 107 ++++++------
 .../multimodal/test_find_num_image_tokens.py  |  61 +++----
 .../llmapi/apps/_test_openai_mmencoder.py     | 154 ++++++++++++++++++
 7 files changed, 253 insertions(+), 88 deletions(-)
 create mode 100644 tests/unittest/llmapi/apps/_test_openai_mmencoder.py

diff --git a/tensorrt_llm/_torch/models/modeling_llava_next.py b/tensorrt_llm/_torch/models/modeling_llava_next.py
index f91b021a61ab..b55a2713f4be 100644
--- a/tensorrt_llm/_torch/models/modeling_llava_next.py
+++ b/tensorrt_llm/_torch/models/modeling_llava_next.py
@@ -52,8 +52,6 @@ def __init__(self,
 
         self.image_token_index = model_config.image_token_index
         self.vocab_size = model_config.vocab_size
-        self.vision_feature_select_strategy = getattr(
-            model_config, "vision_feature_select_strategy", "default")
         self.config = model_config.vision_config
 
     def get_num_tokens_per_image(
@@ -75,7 +73,6 @@ def _postprocess(
         mm_tokens = torch.tensor([self.model_config.image_token_index
                                   ]).to(input_ids.device)
         model_hidden_size = self.model_config.text_config.hidden_size
-        vocab_size = self.model_config.text_config.vocab_size
         start_len = end_len = 0  # for llava, need not append start/end token around each image token
         # End model specific variables
 
@@ -119,8 +116,8 @@ def _postprocess(
         input_ids_splits = list(input_ids.tensor_split(mm_split_positions.cpu(
         )))  # len(input_ids_splits) = num_segments after mm tokens are isolated
         mm_ids_splits = list(
-            torch.arange(vocab_size,
-                         vocab_size + mm_total_length,
+            torch.arange(self.vocab_size,
+                         self.vocab_size + mm_total_length,
                          device=input_ids.device).split(mm_lengths_per_split)
         )  # len(mm_ids_splits) = num_mm_segments
 
@@ -431,7 +428,6 @@ def __init__(self, model_config: ModelConfig[PretrainedConfig], *args,
         self.llm = AutoModelForCausalLM.from_config(llm_model_config)
 
         self.model_config = model_config
-        self.vocab_size = config.vocab_size
         self.model_dtype = getattr(config.text_config, "torch_dtype",
                                    torch.float16)
         logger.info(f"{self.dtype=} {self.model_dtype=}")
diff --git a/tensorrt_llm/_torch/models/modeling_qwen2vl.py b/tensorrt_llm/_torch/models/modeling_qwen2vl.py
index e7533a4b4ab2..8fbec053ad9d 100644
--- a/tensorrt_llm/_torch/models/modeling_qwen2vl.py
+++ b/tensorrt_llm/_torch/models/modeling_qwen2vl.py
@@ -308,7 +308,8 @@ def __call__(
                                                 mm_processor_kwargs)
         if not mm_data:
             fused_input_ids = processed_inputs['input_ids']
-            return fused_input_ids.to(torch.int32).tolist(), {}
+            # Flatten the tensor to get a simple list of integers
+            return fused_input_ids.flatten().to(torch.int32).tolist(), {}
 
         pixel_values = processed_inputs.get('pixel_values', None)
         pixel_values_videos = processed_inputs.get('pixel_values_videos', None)
diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py
index 3d9e31d48413..8307afb1dc0c 100644
--- a/tensorrt_llm/serve/openai_server.py
+++ b/tensorrt_llm/serve/openai_server.py
@@ -160,7 +160,6 @@ def register_routes(self):
 
     def register_mm_encoder_routes(self):
         self.app.add_api_route("/health", self.health, methods=["GET"])
-        self.app.add_api_route("/health_generate", self.health_generate, methods=["GET"])
         self.app.add_api_route("/version", self.version, methods=["GET"])
         self.app.add_api_route("/v1/models", self.get_model, methods=["GET"])
         # TODO: the metrics endpoint only reports iteration stats, not the runtime stats for now
diff --git a/tests/integration/defs/test_e2e.py b/tests/integration/defs/test_e2e.py
index a7d45fb37f94..b6031236cb76 100644
--- a/tests/integration/defs/test_e2e.py
+++ b/tests/integration/defs/test_e2e.py
@@ -1509,6 +1509,13 @@ def test_openai_chat_multimodal_example(llm_root, llm_venv):
          str(test_root / "_test_openai_chat_multimodal.py")])
 
 
+def test_openai_mmencoder_example(llm_root, llm_venv):
+    test_root = unittest_path() / "llmapi" / "apps"
+    llm_venv.run_cmd(
+        ["-m", "pytest",
+         str(test_root / "_test_openai_mmencoder.py")])
+
+
 def test_openai_chat_structural_tag_example(llm_venv):
     test_root = unittest_path() / "llmapi" / "apps"
     llm_venv.run_cmd([
diff --git a/tests/unittest/_torch/multimodal/test_external_embedding.py b/tests/unittest/_torch/multimodal/test_external_embedding.py
index 24d670b79fcc..c2734b67765f 100644
--- a/tests/unittest/_torch/multimodal/test_external_embedding.py
+++ b/tests/unittest/_torch/multimodal/test_external_embedding.py
@@ -1,24 +1,29 @@
 from unittest.mock import Mock, patch
+
 import pytest
 
+from tensorrt_llm._torch.models.modeling_llava_next import \
+    LlavaNextInputProcessor
 from tensorrt_llm.inputs.data import TextPrompt
 from tensorrt_llm.sampling_params import SamplingParams
-from tensorrt_llm._torch.models.modeling_llava_next import LlavaNextInputProcessor
-
 
 # Model configurations for different processors
 MODEL_CONFIGS = {
     "LlavaNextInputProcessor": {
-        "processor_class": LlavaNextInputProcessor,
-        "module_path": "tensorrt_llm._torch.models.modeling_llava_next",
-        "config_setup": lambda mock_config: setattr_multiple(mock_config, {
-            "image_token_index": 32001,
-            "vocab_size": 32001,
-            "text_config.hidden_size": 4096,
-            "text_config.vocab_size": 32000,
-            "vision_config": Mock(),
-            "vision_feature_select_strategy": "default"
-        })
+        "processor_class":
+        LlavaNextInputProcessor,
+        "module_path":
+        "tensorrt_llm._torch.models.modeling_llava_next",
+        "config_setup":
+        lambda mock_config: setattr_multiple(
+            mock_config, {
+                "image_token_index": 32001,
+                "vocab_size": 32001,
+                "text_config.hidden_size": 4096,
+                "text_config.vocab_size": 32000,
+                "vision_config": Mock(),
+                "vision_feature_select_strategy": "default"
+            })
     },
     # TODO: Add test for more VLM models
     # Future model configurations can be added here
@@ -32,6 +37,7 @@
     # },
 }
 
+
 def setattr_multiple(obj, attr_dict):
     """Helper function to set multiple nested attributes."""
     for attr_path, value in attr_dict.items():
@@ -69,12 +75,10 @@ def processor_setup(request):
         mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer
         mock_auto_processor.from_pretrained.return_value = mock_processor
 
-        processor = config["processor_class"](
-            model_path="dummy_path",
-            model_config=mock_config,
-            tokenizer=mock_tokenizer,
-            trust_remote_code=True
-        )
+        processor = config["processor_class"](model_path="dummy_path",
+                                              model_config=mock_config,
+                                              tokenizer=mock_tokenizer,
+                                              trust_remote_code=True)
 
         # Return processor along with mocks for test usage
         return {
@@ -93,45 +97,51 @@ def test_attach_multimodal_embeddings_image_basic(self, processor_setup):
         import torch
 
         # Prepare test inputs
-        text_prompt: TextPrompt = {
-            "prompt": "What is in this image? <image>"
-        }
+        text_prompt: TextPrompt = {"prompt": "What is in this image? <image>"}
 
         # Create mock embedding tensors (batch_size=1, seq_len=144, hidden_dim=768)
         image_embeddings = [
-            torch.randn(144, 4096),  # Single image embedding so num_frames=1 always
+            torch.randn(144,
+                        4096),  # Single image embedding so num_frames=1 always
         ]
 
-        multimodal_embedding = {
-            "image": image_embeddings
-        }
+        multimodal_embedding = {"image": image_embeddings}
 
         sampling_params = SamplingParams()
 
         # Mock tokenizer output - simulate tokenizing the text with image token
-        mock_token_ids = torch.tensor([1, 2, 3, 32001, 4, 5])  # image token at index 3
-        processor_setup["mock_tokenizer"].return_value.input_ids = [mock_token_ids]
+        mock_token_ids = torch.tensor([1, 2, 3, 32001, 4,
+                                       5])  # image token at index 3
+        processor_setup["mock_tokenizer"].return_value.input_ids = [
+            mock_token_ids
+        ]
 
-        result_token_ids, extra_processed_inputs = processor_setup["processor"].attach_multimodal_embeddings(
-            text_prompt, multimodal_embedding, sampling_params
-        )
+        result_token_ids, extra_processed_inputs = processor_setup[
+            "processor"].attach_multimodal_embeddings(text_prompt,
+                                                      multimodal_embedding,
+                                                      sampling_params)
 
         # Verify outputs
         assert isinstance(result_token_ids, list)
-        assert len(result_token_ids) == len(mock_token_ids) + len(image_embeddings[0]) - 1
+        assert len(result_token_ids) == len(mock_token_ids) + len(
+            image_embeddings[0]) - 1
         assert isinstance(extra_processed_inputs, dict)
         assert "multimodal_data" in extra_processed_inputs
-        assert "multimodal_embedding" in extra_processed_inputs["multimodal_data"]
+        assert "multimodal_embedding" in extra_processed_inputs[
+            "multimodal_data"]
 
         # Check that multimodal embedding is properly formatted
-        mm_embedding = extra_processed_inputs["multimodal_data"]["multimodal_embedding"]
+        mm_embedding = extra_processed_inputs["multimodal_data"][
+            "multimodal_embedding"]
         assert mm_embedding.shape[-1] == 4096  # Hidden dimension should match
         assert mm_embedding.shape[0] == 144
 
         # Verify tokenizer was called with correct prompt
-        processor_setup["mock_tokenizer"].assert_called_once_with(text_prompt["prompt"], return_tensors="pt")
+        processor_setup["mock_tokenizer"].assert_called_once_with(
+            text_prompt["prompt"], return_tensors="pt")
 
-    def test_attach_multimodal_embeddings_multiple_images(self, processor_setup):
+    def test_attach_multimodal_embeddings_multiple_images(
+            self, processor_setup):
         """Test with multiple image embeddings."""
         import torch
 
@@ -146,35 +156,40 @@ def test_attach_multimodal_embeddings_multiple_images(self, processor_setup):
         ]
 
         # multiple images are concatenated along the first dimension
-        multimodal_embedding = {
-            "image": image_embeddings
-        }
+        multimodal_embedding = {"image": image_embeddings}
 
         sampling_params = SamplingParams()
 
         # Mock tokenizer output with two image tokens
         mock_token_ids = torch.tensor([1, 2, 32001, 3, 4, 32001, 5])
-        processor_setup["mock_tokenizer"].return_value.input_ids = [mock_token_ids]
+        processor_setup["mock_tokenizer"].return_value.input_ids = [
+            mock_token_ids
+        ]
 
         # Call the function
-        result_token_ids, extra_processed_inputs = processor_setup["processor"].attach_multimodal_embeddings(
-            text_prompt, multimodal_embedding, sampling_params
-        )
+        result_token_ids, extra_processed_inputs = processor_setup[
+            "processor"].attach_multimodal_embeddings(text_prompt,
+                                                      multimodal_embedding,
+                                                      sampling_params)
 
         # Verify outputs
         assert isinstance(result_token_ids, list)
-        assert len(result_token_ids) == len(mock_token_ids) + torch.cat(image_embeddings, dim=0).shape[0] - 2
+        assert len(result_token_ids) == len(mock_token_ids) + torch.cat(
+            image_embeddings, dim=0).shape[0] - 2
         assert isinstance(extra_processed_inputs, dict)
         assert "multimodal_data" in extra_processed_inputs
-        assert "multimodal_embedding" in extra_processed_inputs["multimodal_data"]
+        assert "multimodal_embedding" in extra_processed_inputs[
+            "multimodal_data"]
 
         # Check that multimodal embedding is properly formatted
-        mm_embedding = extra_processed_inputs["multimodal_data"]["multimodal_embedding"]
+        mm_embedding = extra_processed_inputs["multimodal_data"][
+            "multimodal_embedding"]
         assert mm_embedding.shape[-1] == 4096  # Hidden dimension should match
         assert mm_embedding.shape[0] == 144 + 288
 
         # Verify tokenizer was called with correct prompt
-        processor_setup["mock_tokenizer"].assert_called_once_with(text_prompt["prompt"], return_tensors="pt")
+        processor_setup["mock_tokenizer"].assert_called_once_with(
+            text_prompt["prompt"], return_tensors="pt")
 
 
 if __name__ == "__main__":
diff --git a/tests/unittest/_torch/multimodal/test_find_num_image_tokens.py b/tests/unittest/_torch/multimodal/test_find_num_image_tokens.py
index d102f8d4e009..3fb463665bb9 100644
--- a/tests/unittest/_torch/multimodal/test_find_num_image_tokens.py
+++ b/tests/unittest/_torch/multimodal/test_find_num_image_tokens.py
@@ -1,19 +1,17 @@
-import json
-import os
-from typing import Dict, Any
+import io
 
 import pytest
-import torch
-from PIL import Image
 import requests
-import io
+from PIL import Image
+from transformers import AutoConfig, AutoTokenizer
 
 from tensorrt_llm import MultimodalEncoder
+from tensorrt_llm._torch.models.modeling_llava_next import \
+    LlavaNextInputProcessor
+from tensorrt_llm._torch.models.modeling_qwen2vl import \
+    Qwen2VLInputProcessorBase
 from tensorrt_llm._torch.shared_tensor import SharedTensorContainer
 from tensorrt_llm.inputs import default_multimodal_input_loader
-from transformers import AutoTokenizer, AutoConfig, AutoProcessor
-from tensorrt_llm._torch.models.modeling_llava_next import LlavaNextInputProcessor
-from tensorrt_llm._torch.models.modeling_qwen2vl import Qwen2VLInputProcessorBase
 
 example_images = [
     "https://huggingface.co/datasets/YiYiXu/testing-images/resolve/main/seashore.png",
@@ -74,8 +72,10 @@ def test_get_num_tokens_per_image(model_key, multimodal_model_configs):
                                     max_batch_size=max_batch_size)
 
         # Load model configuration and create input processor once
-        model_config_dict = AutoConfig.from_pretrained(encoder_model_dir, trust_remote_code=True)
-        tokenizer = AutoTokenizer.from_pretrained(encoder_model_dir, trust_remote_code=True)
+        model_config_dict = AutoConfig.from_pretrained(encoder_model_dir,
+                                                       trust_remote_code=True)
+        tokenizer = AutoTokenizer.from_pretrained(encoder_model_dir,
+                                                  trust_remote_code=True)
 
         # Create input processor once
         if model_type == 'llava_next':
@@ -83,15 +83,13 @@ def test_get_num_tokens_per_image(model_key, multimodal_model_configs):
                 model_path=encoder_model_dir,
                 model_config=model_config_dict,
                 tokenizer=tokenizer,
-                trust_remote_code=True
-            )
+                trust_remote_code=True)
         elif model_type == 'qwen2_5_vl':
             input_processor = Qwen2VLInputProcessorBase(
                 model_path=encoder_model_dir,
                 model_config=model_config_dict,
                 tokenizer=tokenizer,
-                trust_remote_code=True
-            )
+                trust_remote_code=True)
         else:
             pytest.fail(f"Unsupported model type: {model_type}")
 
@@ -100,19 +98,19 @@ def test_get_num_tokens_per_image(model_key, multimodal_model_configs):
         media = example_images  # All image URLs
 
         # Prepare inputs and get actual embeddings for all images
-        inputs = default_multimodal_input_loader(
-            tokenizer=tokenizer,
-            model_dir=encoder_model_dir,
-            model_type=model_type,
-            modality="image",
-            prompts=prompts,
-            media=media,
-            image_data_format="pt"
-        )
+        inputs = default_multimodal_input_loader(tokenizer=tokenizer,
+                                                 model_dir=encoder_model_dir,
+                                                 model_type=model_type,
+                                                 modality="image",
+                                                 prompts=prompts,
+                                                 media=media,
+                                                 image_data_format="pt")
 
         # Get actual embeddings from encoder (batch processing)
         encoder_outputs = encoder.generate(inputs)
-        assert len(encoder_outputs) == len(example_images), f"Expected {len(example_images)} encoder outputs, got {len(encoder_outputs)}"
+        assert len(encoder_outputs) == len(
+            example_images
+        ), f"Expected {len(example_images)} encoder outputs, got {len(encoder_outputs)}"
 
         for image_idx, test_image_url in enumerate(example_images):
 
@@ -122,8 +120,8 @@ def test_get_num_tokens_per_image(model_key, multimodal_model_configs):
 
             # Get actual embedding tensor for this image
             actual_embedding = SharedTensorContainer.from_dict(
-                encoder_outputs[image_idx].mm_embedding_handle
-            ).get_local_view()
+                encoder_outputs[image_idx].mm_embedding_handle).get_local_view(
+                )
 
             # The first dimension should be the number of image tokens
             actual_num_tokens = actual_embedding.shape[0]
@@ -131,16 +129,13 @@ def test_get_num_tokens_per_image(model_key, multimodal_model_configs):
             # Get predicted number of tokens using get_num_tokens_per_image
             if model_type == 'llava_next':
                 predicted_num_tokens = input_processor.get_num_tokens_per_image(
-                    image_width=image_width,
-                    image_height=image_height
-                )
+                    image_width=image_width, image_height=image_height)
             elif model_type == 'qwen2_5_vl':
                 predicted_num_tokens = input_processor.get_num_tokens_per_image(
                     image_width=image_width,
                     image_height=image_height,
                     num_frames=1,
-                    do_resize=True
-                )
+                    do_resize=True)
 
             # The key assertion: predicted should match actual
             assert predicted_num_tokens == actual_num_tokens, \
@@ -155,5 +150,3 @@ def test_get_num_tokens_per_image(model_key, multimodal_model_configs):
         # Cleanup resources
         if encoder is not None:
             del encoder
-
-
diff --git a/tests/unittest/llmapi/apps/_test_openai_mmencoder.py b/tests/unittest/llmapi/apps/_test_openai_mmencoder.py
new file mode 100644
index 000000000000..8daa9c38f74d
--- /dev/null
+++ b/tests/unittest/llmapi/apps/_test_openai_mmencoder.py
@@ -0,0 +1,154 @@
+# Adapted from
+# https://github.com/vllm-project/vllm/blob/aae6927be06dedbda39c6b0c30f6aa3242b84388/tests/entrypoints/openai/test_chat.py
+import os
+import tempfile
+from typing import List
+
+import openai
+import pytest
+import requests
+import yaml
+
+from ..test_llm import get_model_path
+from .openai_server import RemoteOpenAIServer
+
+pytestmark = pytest.mark.threadleak(enabled=False)
+
+
+class RemoteMMEncoderServer(RemoteOpenAIServer):
+    """Remote server for testing multimodal encoder endpoints."""
+
+    def __init__(self,
+                 model: str,
+                 cli_args: List[str] = None,
+                 port: int = None) -> None:
+        # Reuse parent initialization but change the command
+        import subprocess
+        import sys
+
+        from tensorrt_llm.llmapi.mpi_session import find_free_port
+
+        self.host = "localhost"
+        self.port = port if port is not None else find_free_port()
+        self.rank = os.environ.get("SLURM_PROCID", 0)
+
+        args = ["--host", f"{self.host}", "--port", f"{self.port}"]
+        if cli_args:
+            args += cli_args
+
+        # Use mm_embedding_serve command instead of regular serve
+        launch_cmd = ["trtllm-serve", "mm_embedding_serve"] + [model] + args
+
+        self.proc = subprocess.Popen(launch_cmd,
+                                     stdout=sys.stdout,
+                                     stderr=sys.stderr)
+        self._wait_for_server(url=self.url_for("health"),
+                              timeout=self.MAX_SERVER_START_WAIT_S)
+
+
+@pytest.fixture(scope="module", ids=["Qwen2.5-VL-3B-Instruct"])
+def model_name():
+    return "Qwen2.5-VL-3B-Instruct"
+
+
+@pytest.fixture(scope="module",
+                params=[True, False],
+                ids=["extra_options", "no_extra_options"])
+def extra_encoder_options(request):
+    return request.param
+
+
+@pytest.fixture(scope="module")
+def temp_extra_encoder_options_file(request):
+    temp_dir = tempfile.gettempdir()
+    temp_file_path = os.path.join(temp_dir, "extra_encoder_options.yaml")
+    try:
+        extra_encoder_options_dict = {
+            "max_batch_size": 8,
+            "max_num_tokens": 16384
+        }
+
+        with open(temp_file_path, 'w') as f:
+            yaml.dump(extra_encoder_options_dict, f)
+
+        yield temp_file_path
+    finally:
+        if os.path.exists(temp_file_path):
+            os.remove(temp_file_path)
+
+
+@pytest.fixture(scope="module")
+def server(model_name: str, extra_encoder_options: bool,
+           temp_extra_encoder_options_file: str):
+    model_path = get_model_path(model_name)
+    args = ["--max_batch_size", "8"]
+    if extra_encoder_options:
+        args.extend(
+            ["--extra_encoder_options", temp_extra_encoder_options_file])
+
+    with RemoteMMEncoderServer(model_path, args) as remote_server:
+        yield remote_server
+
+
+@pytest.fixture(scope="module")
+def client(server: RemoteMMEncoderServer):
+    return server.get_client()
+
+
+@pytest.fixture(scope="module")
+def async_client(server: RemoteMMEncoderServer):
+    return server.get_async_client()
+
+
+def test_multimodal_content_mm_encoder(client: openai.OpenAI, model_name: str):
+
+    content_text = "Describe the natural environment in the image."
+    image_url = "https://huggingface.co/datasets/YiYiXu/testing-images/resolve/main/seashore.png"
+    messages = [{
+        "role":
+        "user",
+        "content": [{
+            "type": "text",
+            "text": content_text
+        }, {
+            "type": "image_url",
+            "image_url": {
+                "url": image_url
+            }
+        }],
+    }]
+
+    chat_completion = client.chat.completions.create(
+        model=model_name,
+        messages=messages,
+        temperature=0.0,
+    )
+    assert chat_completion.id is not None
+    assert len(chat_completion.choices) == 1
+    choice = chat_completion.choices[0]
+    # Verify mm_embedding_handle is present
+    assert hasattr(choice, 'mm_embedding_handle')
+    assert choice.mm_embedding_handle is not None
+    # Verify the handle contains tensor information
+    mm_handle = choice.mm_embedding_handle
+    assert "tensor_size" in mm_handle
+    assert mm_handle["tensor_size"][
+        0] == 324  # qwen2.5-vl: 324 tokens for the same image
+    assert mm_handle["tensor_size"][
+        1] == 2048  # qwen2.5-vl: hidden_size of the vision encoder
+
+
+def test_health(server: RemoteMMEncoderServer):
+    health_url = server.url_for("health")
+    response = requests.get(health_url)
+    assert response.status_code == 200
+
+
+def test_models_endpoint(client: openai.OpenAI, model_name: str):
+    models = client.models.list()
+    assert len(models.data) >= 1
+
+    model_names = [model.id for model in models.data]
+    # The model name might be transformed, so check if any model contains our base name
+    expected_name = model_name.split('/')[-1]
+    assert any(expected_name in name for name in model_names)

From 672d476040384a150a99b79bc9aa6e590ecd28a3 Mon Sep 17 00:00:00 2001
From: Chang Liu <9713593+chang-l@users.noreply.github.com>
Date: Sat, 9 Aug 2025 19:22:09 -0700
Subject: [PATCH 10/15] Format

Signed-off-by: Chang Liu <9713593+chang-l@users.noreply.github.com>
---
 tensorrt_llm/_torch/pyexecutor/py_executor_creator.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
index 7db887acb41c..b05888daa5a1 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
@@ -177,7 +177,7 @@ def _mangle_executor_config(executor_config: ExecutorConfig):
                 f"Disable overlap scheduler for speculation mode {spec_config.spec_dec_mode.name}"
             )
             executor_config.pytorch_backend_config.disable_overlap_scheduler = True
-    
+
     if executor_config.mm_encoder_only:
         from tensorrt_llm.llmapi.llm_args import LoadFormat
         pytorch_backend_config.mm_encoder_only = True

From 3ce172e9f6ad27c564afa9f5bab1058523fe764d Mon Sep 17 00:00:00 2001
From: Chang Liu <9713593+chang-l@users.noreply.github.com>
Date: Sat, 9 Aug 2025 21:30:44 -0700
Subject: [PATCH 11/15] bug fix for adjust init order slightly

Signed-off-by: Chang Liu <9713593+chang-l@users.noreply.github.com>
---
 tensorrt_llm/_torch/models/modeling_qwen2vl.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensorrt_llm/_torch/models/modeling_qwen2vl.py b/tensorrt_llm/_torch/models/modeling_qwen2vl.py
index 8fbec053ad9d..0efefd48c4a5 100644
--- a/tensorrt_llm/_torch/models/modeling_qwen2vl.py
+++ b/tensorrt_llm/_torch/models/modeling_qwen2vl.py
@@ -660,10 +660,10 @@ class Qwen2VLModel(Qwen2VLModelBase):
 
     def __init__(self, model_config: ModelConfig[PretrainedConfig], *args,
                  **kwargs):
+        super().__init__(model_config, *args, **kwargs)
         if not DISAGG:
             self.mm_encoder = Qwen2VisionModelBase(
                 model_config, Qwen2VLForConditionalGeneration)
-        super().__init__(model_config, *args, **kwargs)
 
 
 @register_vision_encoder(Qwen2VisionModelBase,
@@ -674,7 +674,7 @@ class Qwen2_5_VLModel(Qwen2VLModelBase):
 
     def __init__(self, model_config: ModelConfig[PretrainedConfig], *args,
                  **kwargs):
+        super().__init__(model_config, *args, **kwargs)
         if not DISAGG:
             self.mm_encoder = Qwen2VisionModelBase(
                 model_config, Qwen2_5_VLForConditionalGeneration)
-        super().__init__(model_config, *args, **kwargs)

From 8b5b64b83e53bb5fd154f33b33859ca7d0b29393 Mon Sep 17 00:00:00 2001
From: Chang Liu <9713593+chang-l@users.noreply.github.com>
Date: Sat, 9 Aug 2025 23:26:07 -0700
Subject: [PATCH 12/15] address some comments from coderabbit

Signed-off-by: Chang Liu <9713593+chang-l@users.noreply.github.com>
---
 tensorrt_llm/__init__.py                             |  2 +-
 tensorrt_llm/_torch/models/modeling_utils.py         | 10 ++++++++--
 tensorrt_llm/_torch/pyexecutor/model_engine.py       |  6 +-----
 tensorrt_llm/commands/serve.py                       | 10 +---------
 tests/unittest/llmapi/apps/_test_openai_mmencoder.py |  2 --
 5 files changed, 11 insertions(+), 19 deletions(-)

diff --git a/tensorrt_llm/__init__.py b/tensorrt_llm/__init__.py
index f0c8a1941db5..1fcc4be9b127 100644
--- a/tensorrt_llm/__init__.py
+++ b/tensorrt_llm/__init__.py
@@ -50,7 +50,7 @@ def _add_trt_llm_dll_directory():
 from .builder import BuildConfig, Builder, BuilderConfig, build
 from .disaggregated_params import DisaggregatedParams
 from .functional import Tensor, constant
-from .llmapi import LLM, LlmArgs, MultimodalEncoder
+from .llmapi import LLM, MultimodalEncoder
 from .llmapi.llm_args import LlmArgs, TorchLlmArgs, TrtLlmArgs
 from .logger import logger
 from .mapping import Mapping
diff --git a/tensorrt_llm/_torch/models/modeling_utils.py b/tensorrt_llm/_torch/models/modeling_utils.py
index 6730a8595481..ad200f1a3e66 100755
--- a/tensorrt_llm/_torch/models/modeling_utils.py
+++ b/tensorrt_llm/_torch/models/modeling_utils.py
@@ -563,8 +563,14 @@ def infer_max_seq_len(self) -> int:
 
         # Step 1: Find the upper bound of max_seq_len
         inferred_max_seq_len = 2048
-        if getattr(self.config, 'max_position_embeddings', None) is not None:
-            inferred_max_seq_len = self.config.max_position_embeddings
+        max_position_embeddings = getattr(self.config,
+                                          'max_position_embeddings', None)
+        if max_position_embeddings is None and hasattr(self.config,
+                                                       'text_config'):
+            max_position_embeddings = getattr(self.config.text_config,
+                                              'max_position_embeddings', None)
+        if max_position_embeddings is not None:
+            inferred_max_seq_len = max_position_embeddings
 
         # Step 2: Scale max_seq_len with rotary scaling
         if rope_factor != 1:
diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py
index c72463778051..19e71f53d6d4 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_engine.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -1112,11 +1112,7 @@ def _call_load_weights(self, load_method, weights, weight_mapper):
             load_method(weights)
 
     def _init_max_seq_len(self):
-        if self.pytorch_backend_config.mm_encoder_only:
-            # TODO: hardcoded for now, need to infer as: max_num_token_per_image * max_num_images (can be given as a parameter or defined same as max_seq_len)
-            inferred_max_seq_len = 32768
-        else:
-            inferred_max_seq_len = self.model.infer_max_seq_len()
+        inferred_max_seq_len = self.model.infer_max_seq_len()
         if self.max_seq_len is None:
             logger.info(
                 f"max_seq_len is not specified, using inferred value {inferred_max_seq_len}"
diff --git a/tensorrt_llm/commands/serve.py b/tensorrt_llm/commands/serve.py
index aac70039f3cf..3bd6dcd86ae4 100644
--- a/tensorrt_llm/commands/serve.py
+++ b/tensorrt_llm/commands/serve.py
@@ -349,7 +349,7 @@ def serve(
     launch_server(host, port, llm_args, metadata_server_cfg, server_role)
 
 
-@click.command("mmemb_serve")
+@click.command("mm_embedding_serve")
 @click.argument("model", type=str)
 @click.option("--host",
               type=str,
@@ -410,14 +410,6 @@ def serve_encoder(model: str, host: str, port: int, log_level: str,
     metadata_server_cfg = parse_metadata_server_config_file(
         metadata_server_config_file)
 
-    if metadata_server_cfg is not None:
-        assert server_role is not None, "server_role is required when metadata_server_cfg is provided"
-        try:
-            server_role = ServerRole[server_role.upper()]
-            assert server_role == ServerRole.MM_ENCODER, "server_role must be MM_ENCODER for multimodal encoder"
-        except ValueError:
-            raise ValueError(f"Invalid server role: {server_role}. " \
-                             f"Must be one of: {', '.join([role.name for role in ServerRole])}")
     launch_mm_encoder_server(host, port, encoder_args, metadata_server_cfg)
 
 
diff --git a/tests/unittest/llmapi/apps/_test_openai_mmencoder.py b/tests/unittest/llmapi/apps/_test_openai_mmencoder.py
index 8daa9c38f74d..15a1f66cd501 100644
--- a/tests/unittest/llmapi/apps/_test_openai_mmencoder.py
+++ b/tests/unittest/llmapi/apps/_test_openai_mmencoder.py
@@ -1,5 +1,3 @@
-# Adapted from
-# https://github.com/vllm-project/vllm/blob/aae6927be06dedbda39c6b0c30f6aa3242b84388/tests/entrypoints/openai/test_chat.py
 import os
 import tempfile
 from typing import List

From cb7ddec0f9b2932020142afbb153c57fef27c7b7 Mon Sep 17 00:00:00 2001
From: Chang Liu <9713593+chang-l@users.noreply.github.com>
Date: Sun, 10 Aug 2025 13:15:25 -0700
Subject: [PATCH 13/15] fix ci and address suggestions from ai

Signed-off-by: Chang Liu <9713593+chang-l@users.noreply.github.com>
---
 .../_torch/pyexecutor/model_engine.py         | 47 ++++++++++++++++++-
 tensorrt_llm/commands/serve.py                | 13 ++++-
 tensorrt_llm/executor/result.py               |  3 +-
 tensorrt_llm/llmapi/llm.py                    |  4 +-
 tensorrt_llm/llmapi/mm_encoder.py             |  2 +-
 tensorrt_llm/serve/openai_protocol.py         |  2 +-
 tensorrt_llm/serve/openai_server.py           | 32 +++++++++----
 .../multimodal/test_find_num_image_tokens.py  |  5 +-
 8 files changed, 89 insertions(+), 19 deletions(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py
index 19e71f53d6d4..1575ad261ef6 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_engine.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -1112,7 +1112,12 @@ def _call_load_weights(self, load_method, weights, weight_mapper):
             load_method(weights)
 
     def _init_max_seq_len(self):
-        inferred_max_seq_len = self.model.infer_max_seq_len()
+        # For mm_encoder_only mode, infer_max_seq_len() is for LLM decoder models
+        if hasattr(self.model, 'infer_max_seq_len'):
+            inferred_max_seq_len = self.model.infer_max_seq_len()
+        else:
+            inferred_max_seq_len = self._infer_max_seq_len_from_config()
+
         if self.max_seq_len is None:
             logger.info(
                 f"max_seq_len is not specified, using inferred value {inferred_max_seq_len}"
@@ -1128,6 +1133,46 @@ def _init_max_seq_len(self):
             )
             self.max_seq_len = inferred_max_seq_len
 
+    def _infer_max_seq_len_from_config(self) -> int:
+
+        if hasattr(self.model, 'model_config') and self.model.model_config:
+            model_config = self.model.model_config.pretrained_config
+            rope_scaling = getattr(model_config, 'rope_scaling', None)
+            rope_factor = 1
+            if rope_scaling is not None:
+                rope_type = rope_scaling.get('type',
+                                             rope_scaling.get('rope_type'))
+                if rope_type not in ("su", "longrope", "llama3", "yarn"):
+                    rope_factor = rope_scaling.get('factor', 1.0)
+
+            # Step 1: Find the upper bound of max_seq_len
+            inferred_max_seq_len = 2048
+            max_position_embeddings = getattr(model_config,
+                                              'max_position_embeddings', None)
+            if max_position_embeddings is None and hasattr(
+                    model_config, 'text_config'):
+                max_position_embeddings = getattr(model_config.text_config,
+                                                  'max_position_embeddings',
+                                                  None)
+            if max_position_embeddings is not None:
+                inferred_max_seq_len = max_position_embeddings
+
+            # Step 2: Scale max_seq_len with rotary scaling
+            if rope_factor != 1:
+                inferred_max_seq_len = int(
+                    math.ceil(inferred_max_seq_len * rope_factor))
+                logger.warning(
+                    f'max_seq_len is scaled to {inferred_max_seq_len} by rope scaling {rope_factor}'
+                )
+
+            return inferred_max_seq_len
+
+        default_max_seq_len = 8192
+        logger.warning(
+            f"Could not infer max_seq_len from model config, using default value: {default_max_seq_len}"
+        )
+        return default_max_seq_len
+
     def _init_max_num_tokens(self):
         # Modified from tensorrt_llm/_common.py check_max_num_tokens
         if self.max_num_tokens is None:
diff --git a/tensorrt_llm/commands/serve.py b/tensorrt_llm/commands/serve.py
index 3bd6dcd86ae4..26a30b1b6ee3 100644
--- a/tensorrt_llm/commands/serve.py
+++ b/tensorrt_llm/commands/serve.py
@@ -364,6 +364,13 @@ def serve(
               type=int,
               default=BuildConfig.max_batch_size,
               help="Maximum number of requests that the engine can schedule.")
+@click.option(
+    "--max_num_tokens",
+    type=int,
+    default=16384,  # set higher default max_num_tokens for multimodal encoder
+    help=
+    "Maximum number of batched input tokens after padding is removed in each batch."
+)
 @click.option("--gpus_per_node",
               type=int,
               default=None,
@@ -385,8 +392,9 @@ def serve(
               default=None,
               help="Path to metadata server config file")
 def serve_encoder(model: str, host: str, port: int, log_level: str,
-                  max_batch_size: int, gpus_per_node: Optional[int],
-                  trust_remote_code: bool, extra_encoder_options: Optional[str],
+                  max_batch_size: int, max_num_tokens: int,
+                  gpus_per_node: Optional[int], trust_remote_code: bool,
+                  extra_encoder_options: Optional[str],
                   metadata_server_config_file: Optional[str]):
     """Running an OpenAI API compatible server
 
@@ -397,6 +405,7 @@ def serve_encoder(model: str, host: str, port: int, log_level: str,
     # TODO: expose more argument progressivly
     llm_args, _ = get_llm_args(model=model,
                                max_batch_size=max_batch_size,
+                               max_num_tokens=max_num_tokens,
                                gpus_per_node=gpus_per_node,
                                trust_remote_code=trust_remote_code)
 
diff --git a/tensorrt_llm/executor/result.py b/tensorrt_llm/executor/result.py
index c5114ac696c2..fd1c362f7bd1 100644
--- a/tensorrt_llm/executor/result.py
+++ b/tensorrt_llm/executor/result.py
@@ -359,7 +359,8 @@ def _handle_response(self,
             if response_result.context_logits is not None:
                 self._context_logits = response_result.context_logits
 
-            if response_result.mm_embedding_handle is not None:
+            if hasattr(response_result, 'mm_embedding_handle'
+                       ) and response_result.mm_embedding_handle is not None:
                 self._mm_embedding_handle = response_result.mm_embedding_handle
 
             # Processing background errors here ASAF during generation.
diff --git a/tensorrt_llm/llmapi/llm.py b/tensorrt_llm/llmapi/llm.py
index da10be88bbc0..071b58638fa2 100644
--- a/tensorrt_llm/llmapi/llm.py
+++ b/tensorrt_llm/llmapi/llm.py
@@ -792,7 +792,7 @@ def _build_model(self):
         # 2. May need to modify model weights for MM (e.g., resize vocab embedding). We must do such operation via input processor's __init__
         self.input_processor = create_input_processor(self._hf_model_dir,
                                                       self.tokenizer)
-        self.tokenizer = self.input_processor.tokenizer
+        self._tokenizer = self.input_processor.tokenizer
 
         max_batch_size = self.args.max_batch_size
         max_num_tokens = self.args.max_num_tokens
@@ -957,7 +957,7 @@ def _build_model(self):
         # 2. May need to modify model weights for MM (e.g., resize vocab embedding). We must do such operation via input processor's __init__
         self.input_processor = create_input_processor(self._hf_model_dir,
                                                       self.tokenizer)
-        self.tokenizer = self.input_processor.tokenizer
+        self._tokenizer = self.input_processor.tokenizer
 
         max_batch_size = self.args.max_batch_size
         max_num_tokens = self.args.max_num_tokens
diff --git a/tensorrt_llm/llmapi/mm_encoder.py b/tensorrt_llm/llmapi/mm_encoder.py
index a827ced96737..541068a9a6df 100644
--- a/tensorrt_llm/llmapi/mm_encoder.py
+++ b/tensorrt_llm/llmapi/mm_encoder.py
@@ -54,7 +54,7 @@ def _build_model(self):
         # 2. May need to modify model weights for MM (e.g., resize vocab embedding). We must do such operation via input processor's __init__
         self.input_processor = create_input_processor(self._hf_model_dir,
                                                       self.tokenizer)
-        self.tokenizer = self.input_processor.tokenizer
+        self._tokenizer = self.input_processor.tokenizer
 
         max_batch_size = self.args.max_batch_size
         max_num_tokens = self.args.max_num_tokens
diff --git a/tensorrt_llm/serve/openai_protocol.py b/tensorrt_llm/serve/openai_protocol.py
index b50ef0085136..4aba90562203 100644
--- a/tensorrt_llm/serve/openai_protocol.py
+++ b/tensorrt_llm/serve/openai_protocol.py
@@ -391,7 +391,7 @@ class ChatCompletionResponseChoice(OpenAIBaseModel):
     finish_reason: Optional[str] = None
     stop_reason: Optional[Union[int, str]] = None
     # TODO: progressivly add more info like input_ids, specific_token_ids, mrope, mm_hashes, etc
-    # TODO: refer to ChatCompletionLogProbs
+    # TODO: and use a JSON-safe handle to refer to the server-side output
     mm_embedding_handle: Optional[Dict[str, Any]] = None
 
     disaggregated_params: Optional[DisaggregatedParams] = Field(default=None)
diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py
index 6afb743d2737..cecbdb5d1d4c 100644
--- a/tensorrt_llm/serve/openai_server.py
+++ b/tensorrt_llm/serve/openai_server.py
@@ -382,17 +382,33 @@ async def create_chat_response(
 
     async def openai_mm_encoder(self, request: ChatCompletionRequest, raw_request: Request) -> Response:
 
-        async def create_mm_embedding_response(
-                promise: RequestOutput):
+        async def create_mm_embedding_response(promise: RequestOutput):
             await promise.aresult()
-            mm_embedding_handle = promise.mm_embedding_handle
-            num_tokens = mm_embedding_handle["tensor_size"][0]
-            # TODO: need to add more info like input_ids, specific_token_ids, mrope, mm_hashes, etc
+            # TODO: Replace mm_embedding_handle with a dedicated OpenAIBaseModel(JSON-safe), when enable multimodal disagg E2E
+            mm_embedding_handle = getattr(promise, "mm_embedding_handle", None)
+            if not mm_embedding_handle or "tensor_size" not in mm_embedding_handle:
+                return self.create_error_response(
+                    message="Multimodal embedding handle missing in response",
+                    err_type="InternalServerError",
+                    status_code=HTTPStatus.INTERNAL_SERVER_ERROR)
+            num_tokens = int(mm_embedding_handle["tensor_size"][0])
             return ChatCompletionResponse(
                 id=str(promise.request_id),
                 model=self.model,
-                choices=[ChatCompletionResponseChoice(index=0, message=ChatMessage(role="assistant", content="dummy"), mm_embedding_handle=mm_embedding_handle, finish_reason="length")],
-                usage=UsageInfo(prompt_tokens=num_tokens, completion_tokens=1, total_tokens=num_tokens+1))
+                choices=[
+                    ChatCompletionResponseChoice(
+                        index=0,
+                        message=ChatMessage(role="assistant", content="dummy"),
+                        mm_embedding_handle=mm_embedding_handle,
+                        finish_reason="length",
+                    )
+                ],
+                usage=UsageInfo(
+                    prompt_tokens=num_tokens,
+                    completion_tokens=1,
+                    total_tokens=num_tokens + 1,
+                ),
+            )
 
         try:
             check_multiple_response(request.n, self.llm.args.backend)
@@ -400,8 +416,6 @@ async def create_mm_embedding_response(
             tool_dicts = None if request.tools is None else [
                 tool.model_dump() for tool in request.tools
             ]
-            # TODO: placeholder for multimodal disagg e2e
-            to_llm_disaggregated_params(request.disaggregated_params)
 
             conversation, mm_coroutines, mm_placeholder_counts = parse_chat_messages_coroutines(request.messages, self.model_config)
 
diff --git a/tests/unittest/_torch/multimodal/test_find_num_image_tokens.py b/tests/unittest/_torch/multimodal/test_find_num_image_tokens.py
index 3fb463665bb9..915970d1cea0 100644
--- a/tests/unittest/_torch/multimodal/test_find_num_image_tokens.py
+++ b/tests/unittest/_torch/multimodal/test_find_num_image_tokens.py
@@ -22,9 +22,10 @@
 
 def download_image(url: str) -> Image.Image:
     """Download image from URL and return as PIL Image."""
-    response = requests.get(url)
+    response = requests.get(url, timeout=30)
     response.raise_for_status()
-    return Image.open(io.BytesIO(response.content))
+    img = Image.open(io.BytesIO(response.content))
+    return img.convert("RGB")
 
 
 @pytest.fixture(scope="function")

From 37b24bd03a85190d712be4feb3c6b889285a5584 Mon Sep 17 00:00:00 2001
From: Chang Liu <9713593+chang-l@users.noreply.github.com>
Date: Sun, 10 Aug 2025 23:13:16 -0700
Subject: [PATCH 14/15] fix ci

Signed-off-by: Chang Liu <9713593+chang-l@users.noreply.github.com>
---
 tensorrt_llm/_torch/models/modeling_qwen2vl.py | 1 +
 tensorrt_llm/_torch/pyexecutor/model_engine.py | 7 ++++++-
 2 files changed, 7 insertions(+), 1 deletion(-)

diff --git a/tensorrt_llm/_torch/models/modeling_qwen2vl.py b/tensorrt_llm/_torch/models/modeling_qwen2vl.py
index 0efefd48c4a5..f009034f6051 100644
--- a/tensorrt_llm/_torch/models/modeling_qwen2vl.py
+++ b/tensorrt_llm/_torch/models/modeling_qwen2vl.py
@@ -360,6 +360,7 @@ def __init__(self, model_config: ModelConfig[PretrainedConfig],
             model_path,
             torch_dtype=pretrained_config.torch_dtype,
             attn_implementation='flash_attention_2').eval()
+        # TODO: Make vision model compatible with meta init mode and load_weights at the same place
         self.visual = model.visual.to(self.device)
         self.post_config()
 
diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py
index 1575ad261ef6..8b4015cf4c75 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_engine.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -1,5 +1,6 @@
 import bisect
 import contextlib
+import copy
 import functools
 import gc
 import inspect
@@ -1033,9 +1034,12 @@ def _load_model(self,
 
         with timing("Model init total"), maybe_create_moe_load_balancer(
                 config, self.mapping) as moe_load_balancer:
+
             try:
+                # config will be modified in-place for some models, like Qwen2
+                config_copy = copy.deepcopy(config)
                 with MetaInitMode():
-                    model = AutoModelForCausalLM.from_config(config)
+                    model = AutoModelForCausalLM.from_config(config_copy)
 
                 memo = dict()
 
@@ -1047,6 +1051,7 @@ def init_meta_tensor(t: torch.Tensor):
                     return memo[t]
 
                 model._apply(init_meta_tensor)
+                config = config_copy
 
             except Exception:
                 logger.info(

From 34b500ea6208782c32156ca3907c29b102326aef Mon Sep 17 00:00:00 2001
From: Chang Liu <9713593+chang-l@users.noreply.github.com>
Date: Mon, 11 Aug 2025 22:49:07 -0700
Subject: [PATCH 15/15] update llmapi reference

---
 .../api_stability/references_committed/request_output.yaml     | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/tests/unittest/api_stability/references_committed/request_output.yaml b/tests/unittest/api_stability/references_committed/request_output.yaml
index c8f745129cb4..62e8ec234710 100644
--- a/tests/unittest/api_stability/references_committed/request_output.yaml
+++ b/tests/unittest/api_stability/references_committed/request_output.yaml
@@ -27,3 +27,6 @@ properties:
   finished:
     annotation: bool
     default: inspect._empty
+  mm_embedding_handle:
+    annotation: Optional[Dict[str, Any]]
+    default: inspect._empty