Add unit test and add back attach_emb for llava model

chang-l · chang-l · commit 810be5459063 · 2025-08-07T23:48:55.000-07:00
Signed-off-by: Chang Liu (Enterprise Products) &lt;9713593+chang-l@users.noreply.github.com&gt;
diff --git 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
@@ -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
@@ -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
+    
