From 8ad3a36c3558302202afa1feb039c506c3a4cb40 Mon Sep 17 00:00:00 2001
From: john calderon <jcalderon@nvidia.com>
Date: Thu, 4 Sep 2025 16:01:35 +0000
Subject: [PATCH 01/24] add multimodal dummy request to profiling - drafting
 stage

Signed-off-by: john calderon <jcalderon@nvidia.com>

rebase to main

Signed-off-by: John Calderon <johncalesp@gmail.com>
Signed-off-by: John Calderon <jcalderon@nvidia.com>
---
 tensorrt_llm/_torch/pyexecutor/_util.py | 53 +++++++++++++++++++++++++
 1 file changed, 53 insertions(+)

diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py
index 2fa2eeb47684..205ada675310 100644
--- a/tensorrt_llm/_torch/pyexecutor/_util.py
+++ b/tensorrt_llm/_torch/pyexecutor/_util.py
@@ -3,6 +3,7 @@
 from typing import Dict, List, Optional
 
 import torch
+from transformers import AutoTokenizer
 
 import tensorrt_llm
 import tensorrt_llm.bindings.executor as trtllm
@@ -133,6 +134,58 @@ def _cal_max_memory(self, peak_memory, total_gpu_memory, fraction,
             f", tmp kv_mem { (allocated_bytes) / (GB):.2f} GiB")
         return int(available_kv_mem)
 
+    def _create_dummy_mm_context_request(
+            self, input_seq_len: int) -> List[trtllm.Request]:
+        self._model_name_or_path = getattr(self._model_engine.model,
+                                           "name_or_path", None)
+        self._tokenizer = AutoTokenizer.from_pretrained(
+            self._model_name_or_path)
+        input_processor = create_input_processor(self._model_name_or_path,
+                                                 self._tokenizer)
+        if not (hasattr(input_processor, "get_prompt_for_profiling")):
+            logger.warning("The input processor of the model does not have the method [get_prompt_for_profiling] implemented." \
+            "Profiling with the default input dummy context request. This may not take into account the memory consumption of " \
+            "ViT's encoder")
+            return self._create_dummy_context_requests(input_seq_len)
+        text_prompt = input_processor.get_prompt_for_profiling()
+        max_beam_width = self._executor_config.max_beam_width
+        input_processor_with_hash = create_input_processor_with_hash(
+            input_processor)
+        prompt_token_ids, extra_processed_inputs = input_processor_with_hash(
+            text_prompt, None)
+        multimodal_input = extra_processed_inputs.get('multimodal_input')
+        multimodal_data = extra_processed_inputs.get('multimodal_data')
+
+        requests = []
+        max_num_tokens = len(prompt_token_ids)
+        remaining_tokens = max(max_num_tokens, input_seq_len)
+        # add +1 to max_num_tokens to avoid assert in line 772 of tensorrt_llm/_torch/attention_backend/trtllm.py
+        self._executor_config.max_seq_len = remaining_tokens + 1
+        if remaining_tokens > input_seq_len:
+            logger.warning(f"Profiling with multimedia prompt which contains more tokens than the allowed input_seq_len." \
+                           f"Multimedia prompt has {remaining_tokens} while the input_seq_len is: {input_seq_len}")
+        while remaining_tokens > 0:
+            req_mm_input = trtllm.MultimodalInput(
+                multimodal_hashes=multimodal_input.multimodal_hashes,
+                multimodal_positions=multimodal_input.multimodal_positions,
+                multimodal_lengths=multimodal_input.multimodal_lengths)
+            request = trtllm.Request(prompt_token_ids,
+                                     max_tokens=1,
+                                     streaming=False,
+                                     sampling_config=trtllm.SamplingConfig(
+                                         beam_width=max_beam_width, ),
+                                     output_config=trtllm.OutputConfig(),
+                                     end_id=-1,
+                                     multimodal_input=req_mm_input)
+            request.py_multimodal_data = multimodal_data
+            remaining_tokens -= max_num_tokens
+            requests.append(request)
+
+        if self._mapping.enable_attention_dp:
+            requests = requests * self._mapping.tp_size
+
+        return requests
+
     def _create_dummy_context_requests(
             self, input_seq_len: int) -> List[trtllm.Request]:
         vocab_size = self._model_engine.model.model_config.pretrained_config.vocab_size

From 7f8e1692d634a6ac9b22ca5c5f5477de90a022f1 Mon Sep 17 00:00:00 2001
From: john calderon <jcalderon@nvidia.com>
Date: Thu, 4 Sep 2025 16:07:28 +0000
Subject: [PATCH 02/24] guard against multimodal_input being None

Signed-off-by: john calderon <jcalderon@nvidia.com>
Signed-off-by: John Calderon <johncalesp@gmail.com>
Signed-off-by: John Calderon <jcalderon@nvidia.com>
---
 tensorrt_llm/_torch/pyexecutor/_util.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py
index 205ada675310..531708af8f0d 100644
--- a/tensorrt_llm/_torch/pyexecutor/_util.py
+++ b/tensorrt_llm/_torch/pyexecutor/_util.py
@@ -168,7 +168,8 @@ def _create_dummy_mm_context_request(
             req_mm_input = trtllm.MultimodalInput(
                 multimodal_hashes=multimodal_input.multimodal_hashes,
                 multimodal_positions=multimodal_input.multimodal_positions,
-                multimodal_lengths=multimodal_input.multimodal_lengths)
+                multimodal_lengths=multimodal_input.multimodal_lengths
+            ) if multimodal_input else None
             request = trtllm.Request(prompt_token_ids,
                                      max_tokens=1,
                                      streaming=False,

From d59fb38685b7de78c6c9aad206ee761f2a4ee728 Mon Sep 17 00:00:00 2001
From: John Calderon <jcalderon@nvidia.com>
Date: Tue, 9 Sep 2025 15:54:51 +0000
Subject: [PATCH 03/24] rebase and change logic

Signed-off-by: John Calderon <jcalderon@nvidia.com>
Signed-off-by: John Calderon <johncalesp@gmail.com>
Signed-off-by: John Calderon <jcalderon@nvidia.com>
---
 .../_torch/models/modeling_qwen2vl.py         | 21 +++++++++++++++++++
 tensorrt_llm/_torch/pyexecutor/_util.py       | 15 ++++++++++---
 2 files changed, 33 insertions(+), 3 deletions(-)

diff --git a/tensorrt_llm/_torch/models/modeling_qwen2vl.py b/tensorrt_llm/_torch/models/modeling_qwen2vl.py
index 1db70fdbfb45..daa6f3d20553 100644
--- a/tensorrt_llm/_torch/models/modeling_qwen2vl.py
+++ b/tensorrt_llm/_torch/models/modeling_qwen2vl.py
@@ -277,6 +277,27 @@ def get_rope_index(
             mrope_position_deltas, device=input_ids.device).unsqueeze(1)
         return position_ids, mrope_position_deltas
 
+    def get_prompt_for_profiling(self):
+        "Send prompt with largest image resolution for profiling the worst case"
+        max_width = 9999999
+        max_height = 9999999
+        resized_height, resized_width = smart_resize(
+            height=max_height,
+            width=max_width,
+            factor=self.model_config.vision_config.patch_size *
+            self.model_config.vision_config.spatial_merge_size,
+            min_pixels=self.processor.image_processor.min_pixels,
+            max_pixels=self.processor.image_processor.max_pixels,
+        )
+        img_tensor = torch.rand([3, resized_width, resized_height],
+                                device="cpu")
+        mm_data = {"image": [img_tensor]}
+
+        text_prompt = TextPrompt(
+            prompt="<|vision_start|><|image_pad|><|vision_end|>",
+            multi_modal_data=mm_data)
+        return text_prompt
+
     def _preprocess(self, text: dict[str, any], mm_data: dict[str, any],
                     mm_processor_kwargs: Dict[str, Any]):
         images = mm_data.get("image")
diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py
index 531708af8f0d..ba7cbd16e7c1 100644
--- a/tensorrt_llm/_torch/pyexecutor/_util.py
+++ b/tensorrt_llm/_torch/pyexecutor/_util.py
@@ -8,6 +8,8 @@
 import tensorrt_llm
 import tensorrt_llm.bindings.executor as trtllm
 from tensorrt_llm._torch.model_config import ModelConfig
+from tensorrt_llm._torch.models.modeling_utils import \
+    MODEL_CLASS_VISION_ENCODER_MAPPING
 from tensorrt_llm._utils import str_dtype_to_binding, torch_dtype_to_str
 from tensorrt_llm.bindings.executor import DecodingMode
 from tensorrt_llm.llmapi.llm_args import (EagleDecodingConfig, KvCacheConfig,
@@ -136,6 +138,7 @@ def _cal_max_memory(self, peak_memory, total_gpu_memory, fraction,
 
     def _create_dummy_mm_context_request(
             self, input_seq_len: int) -> List[trtllm.Request]:
+        requests = []
         self._model_name_or_path = getattr(self._model_engine.model,
                                            "name_or_path", None)
         self._tokenizer = AutoTokenizer.from_pretrained(
@@ -146,7 +149,7 @@ def _create_dummy_mm_context_request(
             logger.warning("The input processor of the model does not have the method [get_prompt_for_profiling] implemented." \
             "Profiling with the default input dummy context request. This may not take into account the memory consumption of " \
             "ViT's encoder")
-            return self._create_dummy_context_requests(input_seq_len)
+            return requests
         text_prompt = input_processor.get_prompt_for_profiling()
         max_beam_width = self._executor_config.max_beam_width
         input_processor_with_hash = create_input_processor_with_hash(
@@ -156,7 +159,6 @@ def _create_dummy_mm_context_request(
         multimodal_input = extra_processed_inputs.get('multimodal_input')
         multimodal_data = extra_processed_inputs.get('multimodal_data')
 
-        requests = []
         max_num_tokens = len(prompt_token_ids)
         remaining_tokens = max(max_num_tokens, input_seq_len)
         # add +1 to max_num_tokens to avoid assert in line 772 of tensorrt_llm/_torch/attention_backend/trtllm.py
@@ -189,11 +191,18 @@ def _create_dummy_mm_context_request(
 
     def _create_dummy_context_requests(
             self, input_seq_len: int) -> List[trtllm.Request]:
+        requests = []
+        if MODEL_CLASS_VISION_ENCODER_MAPPING.get(
+                self._model_engine.model.original_arch, None):
+            requests = self._create_dummy_mm_context_request(input_seq_len)
+        # if succeed profiling with multimodal requests then return, otherwise profile
+        # with default case
+        if requests:
+            return requests
         vocab_size = self._model_engine.model.model_config.pretrained_config.vocab_size
         max_num_tokens = self._max_num_tokens
         max_beam_width = self._max_beam_width
 
-        requests = []
         input_seq_len = min(max_num_tokens, input_seq_len)
         remaining_tokens = max_num_tokens
         while remaining_tokens > 0:

From 22b8d3733c3e97f77a8fbbecc506747641df540c Mon Sep 17 00:00:00 2001
From: John Calderon <jcalderon@nvidia.com>
Date: Tue, 9 Sep 2025 20:41:57 +0000
Subject: [PATCH 04/24] check attribute original_arch is present

Signed-off-by: John Calderon <jcalderon@nvidia.com>
Signed-off-by: John Calderon <johncalesp@gmail.com>
Signed-off-by: John Calderon <jcalderon@nvidia.com>
---
 tensorrt_llm/_torch/pyexecutor/_util.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py
index ba7cbd16e7c1..aa0ac4b3d5d4 100644
--- a/tensorrt_llm/_torch/pyexecutor/_util.py
+++ b/tensorrt_llm/_torch/pyexecutor/_util.py
@@ -192,8 +192,9 @@ def _create_dummy_mm_context_request(
     def _create_dummy_context_requests(
             self, input_seq_len: int) -> List[trtllm.Request]:
         requests = []
-        if MODEL_CLASS_VISION_ENCODER_MAPPING.get(
-                self._model_engine.model.original_arch, None):
+        if hasattr(self._model_engine.model,
+                   "original_arch") and MODEL_CLASS_VISION_ENCODER_MAPPING.get(
+                       self._model_engine.model.original_arch, None):
             requests = self._create_dummy_mm_context_request(input_seq_len)
         # if succeed profiling with multimodal requests then return, otherwise profile
         # with default case

From 11018622705edfcadf09be26b13dec382eaeb4f6 Mon Sep 17 00:00:00 2001
From: John Calderon <jcalderon@nvidia.com>
Date: Mon, 15 Sep 2025 19:36:37 +0000
Subject: [PATCH 05/24] rebase from main and add fix for only-text request for
 qwen2.5 and add unit test

Signed-off-by: John Calderon <jcalderon@nvidia.com>
Signed-off-by: John Calderon <johncalesp@gmail.com>
Signed-off-by: John Calderon <jcalderon@nvidia.com>

rebase from upstream
---
 tensorrt_llm/_torch/pyexecutor/_util.py       |  7 +-
 .../_torch/pyexecutor/py_executor_creator.py  |  1 +
 tensorrt_llm/commands/serve.py                | 66 +++++++--------
 .../unittest/llmapi/test_memory_profiling.py  | 81 +++++++++++++++++++
 4 files changed, 114 insertions(+), 41 deletions(-)
 create mode 100644 tests/unittest/llmapi/test_memory_profiling.py

diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py
index aa0ac4b3d5d4..5dfef892865f 100644
--- a/tensorrt_llm/_torch/pyexecutor/_util.py
+++ b/tensorrt_llm/_torch/pyexecutor/_util.py
@@ -151,7 +151,7 @@ def _create_dummy_mm_context_request(
             "ViT's encoder")
             return requests
         text_prompt = input_processor.get_prompt_for_profiling()
-        max_beam_width = self._executor_config.max_beam_width
+        max_beam_width = self._max_beam_width
         input_processor_with_hash = create_input_processor_with_hash(
             input_processor)
         prompt_token_ids, extra_processed_inputs = input_processor_with_hash(
@@ -161,8 +161,6 @@ def _create_dummy_mm_context_request(
 
         max_num_tokens = len(prompt_token_ids)
         remaining_tokens = max(max_num_tokens, input_seq_len)
-        # add +1 to max_num_tokens to avoid assert in line 772 of tensorrt_llm/_torch/attention_backend/trtllm.py
-        self._executor_config.max_seq_len = remaining_tokens + 1
         if remaining_tokens > input_seq_len:
             logger.warning(f"Profiling with multimedia prompt which contains more tokens than the allowed input_seq_len." \
                            f"Multimedia prompt has {remaining_tokens} while the input_seq_len is: {input_seq_len}")
@@ -413,6 +411,9 @@ def configure_kv_cache_capacity(self, py_executor: PyExecutor) -> None:
         )
         # set max_gpu_total_bytes
         self._kv_cache_config.max_gpu_total_bytes = kv_cache_max_memory
+        if isinstance(self._profiling_stage_data, dict):
+            self._profiling_stage_data[
+                "max_gpu_total_bytes"] = kv_cache_max_memory
         # ---------------------------handle max_gpu_total_bytes---------------------------------
 
     def _create_kv_cache_manager(
diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
index ff5807519768..14ab81138463 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
@@ -207,6 +207,7 @@ def create_py_executor(
     tokenizer: Optional[TokenizerBase] = None,
     lora_config: Optional[LoraConfig] = None,
     kv_connector_config: Optional[KvCacheConnectorConfig] = None,
+    profiling_stage_data: Optional[dict] = None,
 ) -> PyExecutor:
 
     garbage_collection_gen0_threshold = llm_args.garbage_collection_gen0_threshold
diff --git a/tensorrt_llm/commands/serve.py b/tensorrt_llm/commands/serve.py
index 7da0930264be..9842c6d6c986 100644
--- a/tensorrt_llm/commands/serve.py
+++ b/tensorrt_llm/commands/serve.py
@@ -87,6 +87,7 @@ def get_llm_args(model: str,
                  trust_remote_code: bool = False,
                  reasoning_parser: Optional[str] = None,
                  fail_fast_on_attention_window_too_large: bool = False,
+                 enable_chunked_prefill: bool = False,
                  **llm_args_extra_dict: Any):
 
     if gpus_per_node is None:
@@ -109,44 +110,27 @@ def get_llm_args(model: str,
         dynamic_batch_config=dynamic_batch_config,
     )
     llm_args = {
-        "model":
-        model,
-        "scheduler_config":
-        scheduler_config,
-        "tokenizer":
-        tokenizer,
-        "tensor_parallel_size":
-        tensor_parallel_size,
-        "pipeline_parallel_size":
-        pipeline_parallel_size,
-        "moe_expert_parallel_size":
-        moe_expert_parallel_size,
-        "gpus_per_node":
-        gpus_per_node,
-        "trust_remote_code":
-        trust_remote_code,
-        "build_config":
-        build_config,
-        "max_batch_size":
-        max_batch_size,
-        "max_num_tokens":
-        max_num_tokens,
-        "max_beam_width":
-        max_beam_width,
-        "max_seq_len":
-        max_seq_len,
-        "kv_cache_config":
-        kv_cache_config,
-        "backend":
-        backend,
-        "num_postprocess_workers":
-        num_postprocess_workers,
-        "postprocess_tokenizer_dir":
-        tokenizer or model,
-        "reasoning_parser":
-        reasoning_parser,
+        "model": model,
+        "scheduler_config": scheduler_config,
+        "tokenizer": tokenizer,
+        "tensor_parallel_size": tensor_parallel_size,
+        "pipeline_parallel_size": pipeline_parallel_size,
+        "moe_expert_parallel_size": moe_expert_parallel_size,
+        "gpus_per_node": gpus_per_node,
+        "trust_remote_code": trust_remote_code,
+        "build_config": build_config,
+        "max_batch_size": max_batch_size,
+        "max_num_tokens": max_num_tokens,
+        "max_beam_width": max_beam_width,
+        "max_seq_len": max_seq_len,
+        "kv_cache_config": kv_cache_config,
+        "backend": backend,
+        "num_postprocess_workers": num_postprocess_workers,
+        "postprocess_tokenizer_dir": tokenizer or model,
+        "reasoning_parser": reasoning_parser,
         "fail_fast_on_attention_window_too_large":
         fail_fast_on_attention_window_too_large,
+        "enable_chunked_prefill": enable_chunked_prefill,
     }
 
     return llm_args, llm_args_extra_dict
@@ -329,6 +313,10 @@ def convert(self, value: Any, param: Optional["click.Parameter"],
     help=
     "Exit with runtime error when attention window is too large to fit even a single sequence in the KV cache."
 )
+@click.option("--enable_chunked_prefill",
+              is_flag=True,
+              default=False,
+              help="Enable chunked prefill")
 def serve(
         model: str, tokenizer: Optional[str], host: str, port: int,
         log_level: str, backend: str, max_beam_width: int, max_batch_size: int,
@@ -338,7 +326,8 @@ def serve(
         num_postprocess_workers: int, trust_remote_code: bool,
         extra_llm_api_options: Optional[str], reasoning_parser: Optional[str],
         metadata_server_config_file: Optional[str], server_role: Optional[str],
-        fail_fast_on_attention_window_too_large: bool):
+        fail_fast_on_attention_window_too_large: bool,
+        enable_chunked_prefill: bool):
     """Running an OpenAI API compatible server
 
     MODEL: model name | HF checkpoint path | TensorRT engine path
@@ -363,7 +352,8 @@ def serve(
         trust_remote_code=trust_remote_code,
         reasoning_parser=reasoning_parser,
         fail_fast_on_attention_window_too_large=
-        fail_fast_on_attention_window_too_large)
+        fail_fast_on_attention_window_too_large,
+        enable_chunked_prefill=enable_chunked_prefill)
 
     llm_args_extra_dict = {}
     if extra_llm_api_options is not None:
diff --git a/tests/unittest/llmapi/test_memory_profiling.py b/tests/unittest/llmapi/test_memory_profiling.py
new file mode 100644
index 000000000000..b5a812d2fb04
--- /dev/null
+++ b/tests/unittest/llmapi/test_memory_profiling.py
@@ -0,0 +1,81 @@
+import pytest
+import torch
+
+from tensorrt_llm._torch.pyexecutor.py_executor_creator import \
+    create_py_executor
+from tensorrt_llm.llmapi import (BuildConfig, CapacitySchedulerPolicy,
+                                 DynamicBatchConfig, SchedulerConfig)
+from tensorrt_llm.llmapi.llm_args import (CudaGraphConfig, KvCacheConfig,
+                                          TorchLlmArgs)
+
+pytestmark = pytest.mark.threadleak(enabled=False)
+
+
+def test_profile_kvcache():
+    kv_cache_config = KvCacheConfig(enable_block_reuse=False,
+                                    free_gpu_memory_fraction=0.9)
+    cuda_graph_config = CudaGraphConfig(max_batch_size=512)
+
+    VLM_MODEL = "Qwen/Qwen2.5-VL-3B-Instruct"
+    VLM_MODEL_PATH = "/workspace/.cache/huggingface/hub/models--Qwen--Qwen2.5-VL-3B-Instruct/snapshots/66285546d2b821cf421d4f5eb2576359d3770cd3"
+    LLM_MODEL = "Qwen/Qwen2.5-3B-Instruct"
+    LLM_MODEL_PATH = "/workspace/.cache/huggingface/hub/models--Qwen--Qwen2.5-3B-Instruct/snapshots/aa8e72537993ba99e69dfaafa59ed015b17504d1"
+    build_config = BuildConfig(max_batch_size=2048,
+                               max_num_tokens=8192,
+                               max_beam_width=1,
+                               max_seq_len=None)
+    kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.9, )
+
+    dynamic_batch_config = DynamicBatchConfig(
+        enable_batch_size_tuning=True,
+        enable_max_num_tokens_tuning=False,
+        dynamic_batch_moving_average_window=128)
+    scheduler_config = SchedulerConfig(
+        capacity_scheduler_policy=CapacitySchedulerPolicy.GUARANTEED_NO_EVICT,
+        dynamic_batch_config=dynamic_batch_config,
+    )
+    backend = "pytorch"
+    llm_args = {
+        "model": VLM_MODEL,
+        "scheduler_config": scheduler_config,
+        "tokenizer": None,
+        "tensor_parallel_size": 1,
+        "pipeline_parallel_size": 1,
+        "moe_expert_parallel_size": None,
+        "gpus_per_node": 1,
+        "trust_remote_code": False,
+        "build_config": build_config,
+        "max_batch_size": build_config.max_batch_size,
+        "max_num_tokens": build_config.max_num_tokens,
+        "max_beam_width": build_config.max_beam_width,
+        "max_seq_len": build_config.max_seq_len,
+        "kv_cache_config": kv_cache_config,
+        "backend": backend,
+        "num_postprocess_workers": 0,
+        "postprocess_tokenizer_dir": VLM_MODEL,
+        "reasoning_parser": None,
+        "fail_fast_on_attention_window_too_large": False,
+        "enable_chunked_prefill": True,
+        "cuda_graph_config": cuda_graph_config,
+    }
+
+    torchllm_args = TorchLlmArgs(**llm_args)
+
+    profiling_data = dict()
+    py_executor = create_py_executor(llm_args=torchllm_args,
+                                     checkpoint_dir=VLM_MODEL_PATH,
+                                     profiling_stage_data=profiling_data)
+    vlm_max_gpu_total_bytes = profiling_data["max_gpu_total_bytes"]
+    py_executor.shutdown()
+    torch.cuda.empty_cache()
+
+    profiling_data = dict()
+    llm_args["model"] = LLM_MODEL
+    llm_args["postprocess_tokenizer_dir"] = LLM_MODEL
+    torchllm_args = TorchLlmArgs(**llm_args)
+    create_py_executor(llm_args=torchllm_args,
+                       checkpoint_dir=LLM_MODEL_PATH,
+                       profiling_stage_data=profiling_data)
+    llm_max_gpu_total_bytes = profiling_data["max_gpu_total_bytes"]
+
+    assert vlm_max_gpu_total_bytes < llm_max_gpu_total_bytes, f"available KVCache for VLMs is expected to be less than LLMs, but got {vlm_max_gpu_total_bytes} for VLM and {llm_max_gpu_total_bytes} for LLM"

From 5449a63f886d733d26735b9d8bf64e22d3737cbb Mon Sep 17 00:00:00 2001
From: John Calderon <jcalderon@nvidia.com>
Date: Tue, 16 Sep 2025 17:29:29 +0000
Subject: [PATCH 06/24] fix max_seq_len is less than max_num_tokens during
 profiling

Signed-off-by: John Calderon <jcalderon@nvidia.com>
Signed-off-by: John Calderon <johncalesp@gmail.com>
Signed-off-by: John Calderon <jcalderon@nvidia.com>
---
 tensorrt_llm/_torch/pyexecutor/_util.py        |  8 +++++---
 tests/unittest/llmapi/test_memory_profiling.py | 13 +++++++++----
 2 files changed, 14 insertions(+), 7 deletions(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py
index 5dfef892865f..eda098f663f7 100644
--- a/tensorrt_llm/_torch/pyexecutor/_util.py
+++ b/tensorrt_llm/_torch/pyexecutor/_util.py
@@ -92,7 +92,6 @@ def __init__(
         self._speculative_config = speculative_config
         self._sparse_attention_config = sparse_attention_config
         self._tokens_per_block = tokens_per_block
-        self._max_seq_len = max_seq_len
         self._max_batch_size = max_batch_size
         self._net_max_seq_len = net_max_seq_len
         self._dummy_reqs = None
@@ -162,8 +161,11 @@ def _create_dummy_mm_context_request(
         max_num_tokens = len(prompt_token_ids)
         remaining_tokens = max(max_num_tokens, input_seq_len)
         if remaining_tokens > input_seq_len:
-            logger.warning(f"Profiling with multimedia prompt which contains more tokens than the allowed input_seq_len." \
-                           f"Multimedia prompt has {remaining_tokens} while the input_seq_len is: {input_seq_len}")
+            logger.warning(f"Profiling with multimedia prompt which contains more tokens than the allowed input_seq_len. " \
+                           f"Multimodal prompt has {remaining_tokens} while the input_seq_len is: {input_seq_len}")
+        ## add + 1 to avoid error: RuntimeError: The max KV cache length of input sequences (X + 1) exceeds the KV cache manager's maximum supported length X.
+        ## at line "/code/tensorrt_llm/tensorrt_llm/_torch/attention_backend/trtllm.py", line 837
+        self._max_seq_len = remaining_tokens + 1
         while remaining_tokens > 0:
             req_mm_input = trtllm.MultimodalInput(
                 multimodal_hashes=multimodal_input.multimodal_hashes,
diff --git a/tests/unittest/llmapi/test_memory_profiling.py b/tests/unittest/llmapi/test_memory_profiling.py
index b5a812d2fb04..0ae2bf687858 100644
--- a/tests/unittest/llmapi/test_memory_profiling.py
+++ b/tests/unittest/llmapi/test_memory_profiling.py
@@ -8,6 +8,10 @@
 from tensorrt_llm.llmapi.llm_args import (CudaGraphConfig, KvCacheConfig,
                                           TorchLlmArgs)
 
+# isort: off
+from .test_llm import get_model_path
+# isort: on
+
 pytestmark = pytest.mark.threadleak(enabled=False)
 
 
@@ -16,10 +20,11 @@ def test_profile_kvcache():
                                     free_gpu_memory_fraction=0.9)
     cuda_graph_config = CudaGraphConfig(max_batch_size=512)
 
-    VLM_MODEL = "Qwen/Qwen2.5-VL-3B-Instruct"
-    VLM_MODEL_PATH = "/workspace/.cache/huggingface/hub/models--Qwen--Qwen2.5-VL-3B-Instruct/snapshots/66285546d2b821cf421d4f5eb2576359d3770cd3"
-    LLM_MODEL = "Qwen/Qwen2.5-3B-Instruct"
-    LLM_MODEL_PATH = "/workspace/.cache/huggingface/hub/models--Qwen--Qwen2.5-3B-Instruct/snapshots/aa8e72537993ba99e69dfaafa59ed015b17504d1"
+    VLM_MODEL = "Qwen2.5-VL-7B-Instruct"
+    VLM_MODEL_PATH = get_model_path(VLM_MODEL)
+    LLM_MODEL = "Qwen2.5-7B-Instruct"
+    LLM_MODEL_PATH = get_model_path(LLM_MODEL)
+
     build_config = BuildConfig(max_batch_size=2048,
                                max_num_tokens=8192,
                                max_beam_width=1,

From f20c8fbec31fe0597b11039e8c3cd23c6f08fb6a Mon Sep 17 00:00:00 2001
From: John Calderon <jcalderon@nvidia.com>
Date: Thu, 18 Sep 2025 18:00:35 +0000
Subject: [PATCH 07/24] address yechang comments - p1

Signed-off-by: John Calderon <jcalderon@nvidia>
Signed-off-by: John Calderon <johncalesp@gmail.com>
Signed-off-by: John Calderon <jcalderon@nvidia.com>
---
 .../_torch/models/modeling_qwen2vl.py         | 21 ++++++++++++-------
 tensorrt_llm/_torch/pyexecutor/_util.py       |  2 +-
 tensorrt_llm/inputs/__init__.py               |  6 ++++--
 tensorrt_llm/inputs/registry.py               | 10 +++++++++
 4 files changed, 29 insertions(+), 10 deletions(-)

diff --git a/tensorrt_llm/_torch/models/modeling_qwen2vl.py b/tensorrt_llm/_torch/models/modeling_qwen2vl.py
index daa6f3d20553..7ecbc61cb131 100644
--- a/tensorrt_llm/_torch/models/modeling_qwen2vl.py
+++ b/tensorrt_llm/_torch/models/modeling_qwen2vl.py
@@ -277,8 +277,14 @@ def get_rope_index(
             mrope_position_deltas, device=input_ids.device).unsqueeze(1)
         return position_ids, mrope_position_deltas
 
-    def get_prompt_for_profiling(self):
-        "Send prompt with largest image resolution for profiling the worst case"
+    def get_dummy_text(self, mm_data: dict[str, int]):
+        num_images = mm_data.get("image", 0)
+        img_token = self.processor.image_token
+
+        return img_token * num_images
+
+    def get_dummy_images(self, mm_data: dict[str, int]):
+        num_images = mm_data.get("image", 0)
         max_width = 9999999
         max_height = 9999999
         resized_height, resized_width = smart_resize(
@@ -291,12 +297,13 @@ def get_prompt_for_profiling(self):
         )
         img_tensor = torch.rand([3, resized_width, resized_height],
                                 device="cpu")
-        mm_data = {"image": [img_tensor]}
+        return num_images * [img_tensor]
 
-        text_prompt = TextPrompt(
-            prompt="<|vision_start|><|image_pad|><|vision_end|>",
-            multi_modal_data=mm_data)
-        return text_prompt
+    def get_dummy_prompt(self):
+        mm_reqs = {"image": 1}
+        text = self.get_dummy_text(mm_reqs)
+        images = self.get_dummy_images(mm_reqs)
+        return TextPrompt(text=text, multi_modal_data={"image": images})
 
     def _preprocess(self, text: dict[str, any], mm_data: dict[str, any],
                     mm_processor_kwargs: Dict[str, Any]):
diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py
index eda098f663f7..e103b950e0d8 100644
--- a/tensorrt_llm/_torch/pyexecutor/_util.py
+++ b/tensorrt_llm/_torch/pyexecutor/_util.py
@@ -149,7 +149,7 @@ def _create_dummy_mm_context_request(
             "Profiling with the default input dummy context request. This may not take into account the memory consumption of " \
             "ViT's encoder")
             return requests
-        text_prompt = input_processor.get_prompt_for_profiling()
+        text_prompt = input_processor.get_dummy_prompt()
         max_beam_width = self._max_beam_width
         input_processor_with_hash = create_input_processor_with_hash(
             input_processor)
diff --git a/tensorrt_llm/inputs/__init__.py b/tensorrt_llm/inputs/__init__.py
index d411731e5198..1bf38ed5a7b7 100644
--- a/tensorrt_llm/inputs/__init__.py
+++ b/tensorrt_llm/inputs/__init__.py
@@ -1,7 +1,8 @@
 from .data import PromptInputs, TextPrompt, TokensPrompt, prompt_inputs
 from .multimodal import MultimodalInput
-from .registry import (BaseMultimodalInputProcessor, ExtraProcessedInputs,
-                       InputProcessor, MultimodalPlaceholderMetadata,
+from .registry import (BaseDummyInputsBuilder, BaseMultimodalInputProcessor,
+                       ExtraProcessedInputs, InputProcessor,
+                       MultimodalPlaceholderMetadata,
                        MultimodalPlaceholderPlacement, create_input_processor,
                        create_input_processor_with_hash,
                        register_input_processor,
@@ -30,6 +31,7 @@
     "register_input_processor",
     "support_multimodal_disaggregated",
     "ExtraProcessedInputs",
+    "BaseDummyInputsBuilder",
     "BaseMultimodalInputProcessor",
     "MultimodalPlaceholderMetadata",
     "MultimodalPlaceholderPlacement",
diff --git a/tensorrt_llm/inputs/registry.py b/tensorrt_llm/inputs/registry.py
index 99c984a77fbf..31226dd74d78 100644
--- a/tensorrt_llm/inputs/registry.py
+++ b/tensorrt_llm/inputs/registry.py
@@ -42,6 +42,16 @@ def __call__(
         ...
 
 
+class BaseDummyInputsBuilder:
+    """
+    Base class for generating dummy inputs. Specially for profiling
+    """
+
+    def get_dummy_prompt(self):
+        raise NotImplementedError(
+            "Please ensure this method is implemented in your inherited class")
+
+
 class BaseMultimodalInputProcessor:
     """
     Base class for multimodal input processors with default implementations

From 0a1fd4acee4c00b533266f450458f8a97dced6cc Mon Sep 17 00:00:00 2001
From: John Calderon <jcalderon@nvidia.com>
Date: Thu, 18 Sep 2025 18:12:25 +0000
Subject: [PATCH 08/24] fix rebase to 80dd8fe1973323eb8f01060788c0d5485a0ce0f8

Signed-off-by: John Calderon <jcalderon@nvidia.com>
Signed-off-by: John Calderon <johncalesp@gmail.com>
Signed-off-by: John Calderon <jcalderon@nvidia.com>
---
 tensorrt_llm/_torch/pyexecutor/py_executor_creator.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
index 14ab81138463..c8aafeff429c 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
@@ -571,6 +571,7 @@ def drafting_loop_wrapper(model):
             kv_cache_config=kv_cache_config,
             pytorch_backend_config=pytorch_backend_config,
             speculative_config=spec_config,
+            profiling_stage_data=profiling_stage_data,
             sparse_attention_config=sparse_attention_config,
         )
         estimating_kv_cache = kv_cache_creator.try_prepare_estimation()

From 39e78ff49dbd632776faeb79646a9b5dfc48e10e Mon Sep 17 00:00:00 2001
From: John Calderon <jcalderon@nvidia.com>
Date: Thu, 18 Sep 2025 20:09:20 +0000
Subject: [PATCH 09/24] check for new function name and fix TextPrompt
 attribute for qwen

Signed-off-by: John Calderon <jcalderon@nvidia.com>
---
 tensorrt_llm/_torch/models/modeling_qwen2vl.py | 2 +-
 tensorrt_llm/_torch/pyexecutor/_util.py        | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensorrt_llm/_torch/models/modeling_qwen2vl.py b/tensorrt_llm/_torch/models/modeling_qwen2vl.py
index 7ecbc61cb131..2d7aefe25527 100644
--- a/tensorrt_llm/_torch/models/modeling_qwen2vl.py
+++ b/tensorrt_llm/_torch/models/modeling_qwen2vl.py
@@ -303,7 +303,7 @@ def get_dummy_prompt(self):
         mm_reqs = {"image": 1}
         text = self.get_dummy_text(mm_reqs)
         images = self.get_dummy_images(mm_reqs)
-        return TextPrompt(text=text, multi_modal_data={"image": images})
+        return TextPrompt(prompt=text, multi_modal_data={"image": images})
 
     def _preprocess(self, text: dict[str, any], mm_data: dict[str, any],
                     mm_processor_kwargs: Dict[str, Any]):
diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py
index e103b950e0d8..0651ddbf8413 100644
--- a/tensorrt_llm/_torch/pyexecutor/_util.py
+++ b/tensorrt_llm/_torch/pyexecutor/_util.py
@@ -144,7 +144,7 @@ def _create_dummy_mm_context_request(
             self._model_name_or_path)
         input_processor = create_input_processor(self._model_name_or_path,
                                                  self._tokenizer)
-        if not (hasattr(input_processor, "get_prompt_for_profiling")):
+        if not (hasattr(input_processor, "get_dummy_prompt")):
             logger.warning("The input processor of the model does not have the method [get_prompt_for_profiling] implemented." \
             "Profiling with the default input dummy context request. This may not take into account the memory consumption of " \
             "ViT's encoder")

From 38bb1d24863e3c3f16d70a9826d991990941ac12 Mon Sep 17 00:00:00 2001
From: John Calderon <jcalderon@nvidia.com>
Date: Fri, 19 Sep 2025 18:07:17 +0000
Subject: [PATCH 10/24] address comments - change design to use
 default_multimodal_input_loader

Signed-off-by: John Calderon <jcalderon@nvidia.com>
---
 .../_torch/models/modeling_qwen2vl.py         | 56 ++++++++++---------
 tensorrt_llm/_torch/pyexecutor/_util.py       |  2 +-
 2 files changed, 30 insertions(+), 28 deletions(-)

diff --git a/tensorrt_llm/_torch/models/modeling_qwen2vl.py b/tensorrt_llm/_torch/models/modeling_qwen2vl.py
index 2d7aefe25527..ee099c669901 100644
--- a/tensorrt_llm/_torch/models/modeling_qwen2vl.py
+++ b/tensorrt_llm/_torch/models/modeling_qwen2vl.py
@@ -1,5 +1,6 @@
 import copy
 import os
+import random
 from typing import Any, Dict, List, Optional, Tuple, Union
 
 import torch
@@ -28,6 +29,7 @@
 from ...inputs import (BaseMultimodalInputProcessor, ExtraProcessedInputs,
                        InputProcessor, MultimodalPlaceholderMetadata,
                        MultimodalPlaceholderPlacement, TextPrompt,
+                       default_multimodal_input_loader,
                        register_input_processor)
 from ...logger import logger
 from ...sampling_params import SamplingParams
@@ -93,6 +95,7 @@ def __init__(self,
         self.model_config = model_config
         self.tokenizer = tokenizer
         self.use_fast = True
+        self.model_path = model_path
         self.processor = AutoProcessor.from_pretrained(
             model_path,
             use_fast=self.use_fast,
@@ -277,33 +280,32 @@ def get_rope_index(
             mrope_position_deltas, device=input_ids.device).unsqueeze(1)
         return position_ids, mrope_position_deltas
 
-    def get_dummy_text(self, mm_data: dict[str, int]):
-        num_images = mm_data.get("image", 0)
-        img_token = self.processor.image_token
-
-        return img_token * num_images
-
-    def get_dummy_images(self, mm_data: dict[str, int]):
-        num_images = mm_data.get("image", 0)
-        max_width = 9999999
-        max_height = 9999999
-        resized_height, resized_width = smart_resize(
-            height=max_height,
-            width=max_width,
-            factor=self.model_config.vision_config.patch_size *
-            self.model_config.vision_config.spatial_merge_size,
-            min_pixels=self.processor.image_processor.min_pixels,
-            max_pixels=self.processor.image_processor.max_pixels,
-        )
-        img_tensor = torch.rand([3, resized_width, resized_height],
-                                device="cpu")
-        return num_images * [img_tensor]
-
-    def get_dummy_prompt(self):
-        mm_reqs = {"image": 1}
-        text = self.get_dummy_text(mm_reqs)
-        images = self.get_dummy_images(mm_reqs)
-        return TextPrompt(prompt=text, multi_modal_data={"image": images})
+    def get_dummy_text(self, input_seq_len: int):
+        return self.tokenizer.decode([
+            random.randint(0, self.model_config.vocab_size - 1)
+            for _ in range(input_seq_len)
+        ])
+
+    def get_dummy_images(self,
+                         max_width: int,
+                         max_height: int,
+                         num_images: int = 1):
+        image = Image.new("RGB", (max_width, max_height), color=255)
+        return [image] * num_images
+
+    def get_dummy_prompt(self, input_seq_len: int):
+        text = self.get_dummy_text(input_seq_len)
+        images = self.get_dummy_images(
+            max_width=3584,
+            max_height=3584)  #sqrt of max_pixels value (12845056)
+        return default_multimodal_input_loader(
+            tokenizer=self.tokenizer,
+            model_dir=self.model_path,
+            model_type=self.model_config.model_type,
+            modality="image",
+            prompts=[text],
+            media=[images],
+            image_data_format="pt")[0]
 
     def _preprocess(self, text: dict[str, any], mm_data: dict[str, any],
                     mm_processor_kwargs: Dict[str, Any]):
diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py
index 0651ddbf8413..6c55e8f1c729 100644
--- a/tensorrt_llm/_torch/pyexecutor/_util.py
+++ b/tensorrt_llm/_torch/pyexecutor/_util.py
@@ -149,7 +149,7 @@ def _create_dummy_mm_context_request(
             "Profiling with the default input dummy context request. This may not take into account the memory consumption of " \
             "ViT's encoder")
             return requests
-        text_prompt = input_processor.get_dummy_prompt()
+        text_prompt = input_processor.get_dummy_prompt(input_seq_len)
         max_beam_width = self._max_beam_width
         input_processor_with_hash = create_input_processor_with_hash(
             input_processor)

From 62679a3b31b3489e26945ebcdfc975ff0827daad Mon Sep 17 00:00:00 2001
From: John Calderon <jcalderon@nvidia.com>
Date: Fri, 19 Sep 2025 18:37:28 +0000
Subject: [PATCH 11/24] add additional arguments for mm data

Signed-off-by: John Calderon <jcalderon@nvidia.com>
---
 tensorrt_llm/_torch/models/modeling_qwen2vl.py | 13 ++++++-------
 tensorrt_llm/_torch/pyexecutor/_util.py        |  3 ++-
 tensorrt_llm/inputs/registry.py                |  2 +-
 3 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/tensorrt_llm/_torch/models/modeling_qwen2vl.py b/tensorrt_llm/_torch/models/modeling_qwen2vl.py
index ee099c669901..e64188c01f4b 100644
--- a/tensorrt_llm/_torch/models/modeling_qwen2vl.py
+++ b/tensorrt_llm/_torch/models/modeling_qwen2vl.py
@@ -286,18 +286,17 @@ def get_dummy_text(self, input_seq_len: int):
             for _ in range(input_seq_len)
         ])
 
-    def get_dummy_images(self,
-                         max_width: int,
-                         max_height: int,
-                         num_images: int = 1):
+    def get_dummy_images(self, max_width: int, max_height: int,
+                         num_images: int):
         image = Image.new("RGB", (max_width, max_height), color=255)
         return [image] * num_images
 
-    def get_dummy_prompt(self, input_seq_len: int):
+    def get_dummy_prompt(self, input_seq_len: int, mm_data: dict):
+        num_images = mm_data.get("image", 0)
         text = self.get_dummy_text(input_seq_len)
         images = self.get_dummy_images(
-            max_width=3584,
-            max_height=3584)  #sqrt of max_pixels value (12845056)
+            max_width=3584, max_height=3584,
+            num_images=num_images)  #w, h is sqrt of max_pixels value (12845056)
         return default_multimodal_input_loader(
             tokenizer=self.tokenizer,
             model_dir=self.model_path,
diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py
index 6c55e8f1c729..4046f707d2ea 100644
--- a/tensorrt_llm/_torch/pyexecutor/_util.py
+++ b/tensorrt_llm/_torch/pyexecutor/_util.py
@@ -149,7 +149,8 @@ def _create_dummy_mm_context_request(
             "Profiling with the default input dummy context request. This may not take into account the memory consumption of " \
             "ViT's encoder")
             return requests
-        text_prompt = input_processor.get_dummy_prompt(input_seq_len)
+        text_prompt = input_processor.get_dummy_prompt(input_seq_len,
+                                                       {'image': 1})
         max_beam_width = self._max_beam_width
         input_processor_with_hash = create_input_processor_with_hash(
             input_processor)
diff --git a/tensorrt_llm/inputs/registry.py b/tensorrt_llm/inputs/registry.py
index 31226dd74d78..f86f2ecac64b 100644
--- a/tensorrt_llm/inputs/registry.py
+++ b/tensorrt_llm/inputs/registry.py
@@ -47,7 +47,7 @@ class BaseDummyInputsBuilder:
     Base class for generating dummy inputs. Specially for profiling
     """
 
-    def get_dummy_prompt(self):
+    def get_dummy_prompt(self, input_seq_len: int, mm_data: dict):
         raise NotImplementedError(
             "Please ensure this method is implemented in your inherited class")
 

From 8b80b3213fe2dece905d5a3c7dfd21e03277384b Mon Sep 17 00:00:00 2001
From: John Calderon <jcalderon@nvidia.com>
Date: Mon, 22 Sep 2025 15:43:35 +0000
Subject: [PATCH 12/24] address comments: change unit test, and add more
 asserts

Signed-off-by: John Calderon <jcalderon@nvidia.com>
---
 .../_torch/models/modeling_qwen2vl.py         | 12 ++++++-----
 tensorrt_llm/_torch/pyexecutor/_util.py       | 21 ++++++++++++-------
 .../unittest/llmapi/test_memory_profiling.py  | 19 +++++++----------
 3 files changed, 27 insertions(+), 25 deletions(-)

diff --git a/tensorrt_llm/_torch/models/modeling_qwen2vl.py b/tensorrt_llm/_torch/models/modeling_qwen2vl.py
index e64188c01f4b..03df9fdaa4d8 100644
--- a/tensorrt_llm/_torch/models/modeling_qwen2vl.py
+++ b/tensorrt_llm/_torch/models/modeling_qwen2vl.py
@@ -1,8 +1,8 @@
 import copy
 import os
-import random
 from typing import Any, Dict, List, Optional, Tuple, Union
 
+import numpy as np
 import torch
 import torch.nn as nn
 from torch.nn import functional as F
@@ -281,10 +281,12 @@ def get_rope_index(
         return position_ids, mrope_position_deltas
 
     def get_dummy_text(self, input_seq_len: int):
-        return self.tokenizer.decode([
-            random.randint(0, self.model_config.vocab_size - 1)
-            for _ in range(input_seq_len)
-        ])
+        return self.tokenizer.decode(
+            np.random.randint(
+                low=0,
+                high=self.model_config.
+                vocab_size,  # Note: high is exclusive in NumPy
+                size=input_seq_len))
 
     def get_dummy_images(self, max_width: int, max_height: int,
                          num_images: int):
diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py
index 4046f707d2ea..791312d149ce 100644
--- a/tensorrt_llm/_torch/pyexecutor/_util.py
+++ b/tensorrt_llm/_torch/pyexecutor/_util.py
@@ -138,16 +138,20 @@ def _cal_max_memory(self, peak_memory, total_gpu_memory, fraction,
     def _create_dummy_mm_context_request(
             self, input_seq_len: int) -> List[trtllm.Request]:
         requests = []
-        self._model_name_or_path = getattr(self._model_engine.model,
-                                           "name_or_path", None)
-        self._tokenizer = AutoTokenizer.from_pretrained(
-            self._model_name_or_path)
-        input_processor = create_input_processor(self._model_name_or_path,
-                                                 self._tokenizer)
+        if isinstance(
+                self._profiling_stage_data,
+                dict) and not self._profiling_stage_data.get("enable_mm_reqs"):
+            return requests
+
+        model_name_or_path = getattr(self._model_engine.model, "name_or_path",
+                                     None)
+        assert model_name_or_path is not None, "Could not determine model name or path"
+        tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
+        input_processor = create_input_processor(model_name_or_path, tokenizer)
         if not (hasattr(input_processor, "get_dummy_prompt")):
-            logger.warning("The input processor of the model does not have the method [get_prompt_for_profiling] implemented." \
+            logger.warning("The input processor of the model does not have the method [get_dummy_prompt] implemented." \
             "Profiling with the default input dummy context request. This may not take into account the memory consumption of " \
-            "ViT's encoder")
+            "the image encoder")
             return requests
         text_prompt = input_processor.get_dummy_prompt(input_seq_len,
                                                        {'image': 1})
@@ -160,6 +164,7 @@ def _create_dummy_mm_context_request(
         multimodal_data = extra_processed_inputs.get('multimodal_data')
 
         max_num_tokens = len(prompt_token_ids)
+        assert max_num_tokens > 0, "the length of the prompt of the dummy mm req is less than or equal to 0"
         remaining_tokens = max(max_num_tokens, input_seq_len)
         if remaining_tokens > input_seq_len:
             logger.warning(f"Profiling with multimedia prompt which contains more tokens than the allowed input_seq_len. " \
diff --git a/tests/unittest/llmapi/test_memory_profiling.py b/tests/unittest/llmapi/test_memory_profiling.py
index 0ae2bf687858..9089b20cee7a 100644
--- a/tests/unittest/llmapi/test_memory_profiling.py
+++ b/tests/unittest/llmapi/test_memory_profiling.py
@@ -22,13 +22,10 @@ def test_profile_kvcache():
 
     VLM_MODEL = "Qwen2.5-VL-7B-Instruct"
     VLM_MODEL_PATH = get_model_path(VLM_MODEL)
-    LLM_MODEL = "Qwen2.5-7B-Instruct"
-    LLM_MODEL_PATH = get_model_path(LLM_MODEL)
 
     build_config = BuildConfig(max_batch_size=2048,
-                               max_num_tokens=8192,
                                max_beam_width=1,
-                               max_seq_len=None)
+                               max_seq_len=8192)
     kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.9, )
 
     dynamic_batch_config = DynamicBatchConfig(
@@ -66,21 +63,19 @@ def test_profile_kvcache():
 
     torchllm_args = TorchLlmArgs(**llm_args)
 
-    profiling_data = dict()
+    profiling_data = {"enable_mm_reqs": True}
     py_executor = create_py_executor(llm_args=torchllm_args,
                                      checkpoint_dir=VLM_MODEL_PATH,
                                      profiling_stage_data=profiling_data)
-    vlm_max_gpu_total_bytes = profiling_data["max_gpu_total_bytes"]
+    vlm_max_gpu_total_bytes_with_mm_reqs = profiling_data["max_gpu_total_bytes"]
     py_executor.shutdown()
     torch.cuda.empty_cache()
 
-    profiling_data = dict()
-    llm_args["model"] = LLM_MODEL
-    llm_args["postprocess_tokenizer_dir"] = LLM_MODEL
+    profiling_data = {"enable_mm_reqs": False}
     torchllm_args = TorchLlmArgs(**llm_args)
     create_py_executor(llm_args=torchllm_args,
-                       checkpoint_dir=LLM_MODEL_PATH,
+                       checkpoint_dir=VLM_MODEL_PATH,
                        profiling_stage_data=profiling_data)
-    llm_max_gpu_total_bytes = profiling_data["max_gpu_total_bytes"]
+    vlm_max_gpu_total_bytes_no_mm_reqs = profiling_data["max_gpu_total_bytes"]
 
-    assert vlm_max_gpu_total_bytes < llm_max_gpu_total_bytes, f"available KVCache for VLMs is expected to be less than LLMs, but got {vlm_max_gpu_total_bytes} for VLM and {llm_max_gpu_total_bytes} for LLM"
+    assert vlm_max_gpu_total_bytes_with_mm_reqs < vlm_max_gpu_total_bytes_no_mm_reqs, f"available KVCache for VLMs is expected to be less when profiling with mm reqs, but got {vlm_max_gpu_total_bytes_with_mm_reqs} for mm reqs and {vlm_max_gpu_total_bytes_no_mm_reqs} without mm reqs"

From 1512044ce51dc824b69b83fc0cc38767d9199586 Mon Sep 17 00:00:00 2001
From: John Calderon <jcalderon@nvidia.com>
Date: Mon, 22 Sep 2025 21:16:03 +0000
Subject: [PATCH 13/24] fix rebase to commit
 b1738c3f189560a857ea1adcfdfb8e68c571c81d

Signed-off-by: John Calderon <jcalderon@nvidia.com>
---
 tensorrt_llm/_torch/models/modeling_qwen2vl.py | 11 +++++++----
 tensorrt_llm/_torch/pyexecutor/_util.py        |  2 ++
 2 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/tensorrt_llm/_torch/models/modeling_qwen2vl.py b/tensorrt_llm/_torch/models/modeling_qwen2vl.py
index 03df9fdaa4d8..c94cb518feb4 100644
--- a/tensorrt_llm/_torch/models/modeling_qwen2vl.py
+++ b/tensorrt_llm/_torch/models/modeling_qwen2vl.py
@@ -5,6 +5,7 @@
 import numpy as np
 import torch
 import torch.nn as nn
+from PIL import Image
 from torch.nn import functional as F
 from transformers import (AutoProcessor, AutoTokenizer, PretrainedConfig,
                           PreTrainedModel)
@@ -26,8 +27,9 @@
 from tensorrt_llm.inputs.multimodal import MultimodalParams
 
 from ..._utils import nvtx_range, nvtx_range_debug
-from ...inputs import (BaseMultimodalInputProcessor, ExtraProcessedInputs,
-                       InputProcessor, MultimodalPlaceholderMetadata,
+from ...inputs import (BaseDummyInputsBuilder, BaseMultimodalInputProcessor,
+                       ExtraProcessedInputs, InputProcessor,
+                       MultimodalPlaceholderMetadata,
                        MultimodalPlaceholderPlacement, TextPrompt,
                        default_multimodal_input_loader,
                        register_input_processor)
@@ -85,7 +87,8 @@ def process_weights(weights: Dict,
     return filtered_weights
 
 
-class Qwen2VLInputProcessorBase(BaseMultimodalInputProcessor, InputProcessor):
+class Qwen2VLInputProcessorBase(BaseDummyInputsBuilder,
+                                BaseMultimodalInputProcessor, InputProcessor):
 
     def __init__(self,
                  model_path: str,
@@ -821,7 +824,7 @@ def __init__(
         **kwargs,
     ) -> None:
         model_config.pretrained_config.rope_scaling['type'] = 'mrope'
-
+        self.original_arch = model_config.pretrained_config.architectures[0]
         # NOTE: Setting disable_fuse_rope to True to do mrope fusion in the model engine by pre-computing rotary_cos_sin in the model engine
         disabble_fuse_rope = kwargs.get('disable_fuse_rope', False)
         model_config.pretrained_config.text_config.disable_fuse_rope = disabble_fuse_rope
diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py
index 791312d149ce..13b3287fb218 100644
--- a/tensorrt_llm/_torch/pyexecutor/_util.py
+++ b/tensorrt_llm/_torch/pyexecutor/_util.py
@@ -92,9 +92,11 @@ def __init__(
         self._speculative_config = speculative_config
         self._sparse_attention_config = sparse_attention_config
         self._tokens_per_block = tokens_per_block
+        self._max_seq_len = max_seq_len
         self._max_batch_size = max_batch_size
         self._net_max_seq_len = net_max_seq_len
         self._dummy_reqs = None
+        self._profiling_stage_data = profiling_stage_data
         self._kv_cache_manager_cls = get_kv_cache_manager_cls(
             model_engine.model.model_config)
 

From 3658a44e448f0b2de44666071579c955794e3333 Mon Sep 17 00:00:00 2001
From: John Calderon <jcalderon@nvidia.com>
Date: Wed, 24 Sep 2025 17:19:00 +0000
Subject: [PATCH 14/24] address code rabbit comments && remove
 mrope_config.mrope_position_ids from modeling qwen

Signed-off-by: John Calderon <jcalderon@nvidia.com>
---
 .../_torch/models/modeling_qwen2vl.py         | 21 ++++++++++---------
 .../unittest/llmapi/test_memory_profiling.py  | 11 +++++-----
 2 files changed, 16 insertions(+), 16 deletions(-)

diff --git a/tensorrt_llm/_torch/models/modeling_qwen2vl.py b/tensorrt_llm/_torch/models/modeling_qwen2vl.py
index c94cb518feb4..ed42f462acdb 100644
--- a/tensorrt_llm/_torch/models/modeling_qwen2vl.py
+++ b/tensorrt_llm/_torch/models/modeling_qwen2vl.py
@@ -283,13 +283,14 @@ def get_rope_index(
             mrope_position_deltas, device=input_ids.device).unsqueeze(1)
         return position_ids, mrope_position_deltas
 
-    def get_dummy_text(self, input_seq_len: int):
-        return self.tokenizer.decode(
-            np.random.randint(
-                low=0,
-                high=self.model_config.
-                vocab_size,  # Note: high is exclusive in NumPy
-                size=input_seq_len))
+    def get_dummy_text(self, input_seq_len: int) -> str:
+        ids = np.random.randint(
+            low=0,
+            high=int(
+                self.model_config.vocab_size),  # high is exclusive in NumPy
+            size=input_seq_len,
+        ).tolist()
+        return self.tokenizer.decode(ids, skip_special_tokens=True)
 
     def get_dummy_images(self, max_width: int, max_height: int,
                          num_images: int):
@@ -1013,7 +1014,7 @@ def multimodal_data_device_paths(self) -> List[str]:
         return [
             "image.pixel_values", "image.image_grid_thw",
             "video.pixel_values_videos", "video.video_grid_thw",
-            "multimodal_embedding", "mrope_config.mrope_position_ids"
+            "multimodal_embedding"
         ]
 
     def load_weights(self, weights, weight_mapper: BaseWeightMapper):
@@ -1066,12 +1067,12 @@ def multimodal_data_device_paths(self) -> List[str]:
             return [
                 "image.pixel_values", "video.pixel_values_videos",
                 "image.image_grid_thw", "video.video_grid_thw",
-                "multimodal_embedding", "mrope_config.mrope_position_ids"
+                "multimodal_embedding"
             ]
         else:
             return [
                 "image.pixel_values", "video.pixel_values_videos",
-                "multimodal_embedding", "mrope_config.mrope_position_ids"
+                "multimodal_embedding"
             ]
 
     def load_weights(self, weights, weight_mapper: BaseWeightMapper):
diff --git a/tests/unittest/llmapi/test_memory_profiling.py b/tests/unittest/llmapi/test_memory_profiling.py
index 9089b20cee7a..b62a161a1402 100644
--- a/tests/unittest/llmapi/test_memory_profiling.py
+++ b/tests/unittest/llmapi/test_memory_profiling.py
@@ -19,15 +19,12 @@ def test_profile_kvcache():
     kv_cache_config = KvCacheConfig(enable_block_reuse=False,
                                     free_gpu_memory_fraction=0.9)
     cuda_graph_config = CudaGraphConfig(max_batch_size=512)
-
     VLM_MODEL = "Qwen2.5-VL-7B-Instruct"
     VLM_MODEL_PATH = get_model_path(VLM_MODEL)
 
     build_config = BuildConfig(max_batch_size=2048,
                                max_beam_width=1,
                                max_seq_len=8192)
-    kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.9, )
-
     dynamic_batch_config = DynamicBatchConfig(
         enable_batch_size_tuning=True,
         enable_max_num_tokens_tuning=False,
@@ -73,9 +70,11 @@ def test_profile_kvcache():
 
     profiling_data = {"enable_mm_reqs": False}
     torchllm_args = TorchLlmArgs(**llm_args)
-    create_py_executor(llm_args=torchllm_args,
-                       checkpoint_dir=VLM_MODEL_PATH,
-                       profiling_stage_data=profiling_data)
+    py_executor_2 = create_py_executor(llm_args=torchllm_args,
+                                       checkpoint_dir=VLM_MODEL_PATH,
+                                       profiling_stage_data=profiling_data)
     vlm_max_gpu_total_bytes_no_mm_reqs = profiling_data["max_gpu_total_bytes"]
+    py_executor_2.shutdown()
+    torch.cuda.empty_cache()
 
     assert vlm_max_gpu_total_bytes_with_mm_reqs < vlm_max_gpu_total_bytes_no_mm_reqs, f"available KVCache for VLMs is expected to be less when profiling with mm reqs, but got {vlm_max_gpu_total_bytes_with_mm_reqs} for mm reqs and {vlm_max_gpu_total_bytes_no_mm_reqs} without mm reqs"

From b7f0f7e3a33280c028791dfb786c8aa5ee7f30ec Mon Sep 17 00:00:00 2001
From: John Calderon <jcalderon@nvidia.com>
Date: Wed, 1 Oct 2025 19:47:45 +0000
Subject: [PATCH 15/24] integrate latest feedback

Signed-off-by: John Calderon <jcalderon@nvidia.com>
---
 .../_torch/models/modeling_qwen2vl.py         | 69 ++++++++++++++++---
 tensorrt_llm/_torch/pyexecutor/_util.py       | 21 ++----
 .../_torch/pyexecutor/model_engine.py         |  8 ++-
 tensorrt_llm/inputs/registry.py               |  2 +-
 4 files changed, 73 insertions(+), 27 deletions(-)

diff --git a/tensorrt_llm/_torch/models/modeling_qwen2vl.py b/tensorrt_llm/_torch/models/modeling_qwen2vl.py
index ed42f462acdb..ad3a72f3f55b 100644
--- a/tensorrt_llm/_torch/models/modeling_qwen2vl.py
+++ b/tensorrt_llm/_torch/models/modeling_qwen2vl.py
@@ -292,24 +292,71 @@ def get_dummy_text(self, input_seq_len: int) -> str:
         ).tolist()
         return self.tokenizer.decode(ids, skip_special_tokens=True)
 
-    def get_dummy_images(self, max_width: int, max_height: int,
-                         num_images: int):
+    def get_dummy_image(self, max_width: int, max_height: int):
         image = Image.new("RGB", (max_width, max_height), color=255)
-        return [image] * num_images
-
-    def get_dummy_prompt(self, input_seq_len: int, mm_data: dict):
-        num_images = mm_data.get("image", 0)
-        text = self.get_dummy_text(input_seq_len)
-        images = self.get_dummy_images(
-            max_width=3584, max_height=3584,
-            num_images=num_images)  #w, h is sqrt of max_pixels value (12845056)
+        return image
+
+    def get_dummy_prompt(self, input_seq_len: int):
+        text = ""
+        # we use the max resolution as starting point
+        img_max_dim = 3584
+        image = self.get_dummy_image(max_width=img_max_dim,
+                                     max_height=img_max_dim)
+
+        test_mm_prompt = default_multimodal_input_loader(
+            tokenizer=self.tokenizer,
+            model_dir=self.model_path,
+            model_type=self.model_config.model_type,
+            modality="image",
+            prompts=[text],
+            media=[[image]],
+            image_data_format="pt")[0]
+
+        prompt_token_ids_single_img, _ = self(test_mm_prompt, None)
+
+        # if the max img resolution results in a number of tokens greater then
+        # input_seq_len, we keep lowering the resolution such as to find the
+        # max resolution such as it does not exceed the input_seq_len
+        while len(prompt_token_ids_single_img) > input_seq_len:
+            # reduce img resolution
+            img_max_dim = img_max_dim >> 1
+
+            image = self.get_dummy_image(
+                max_width=img_max_dim, max_height=img_max_dim
+            )  #w, h is sqrt of min_pixels value (3136)
+
+            test_mm_prompt = default_multimodal_input_loader(
+                tokenizer=self.tokenizer,
+                model_dir=self.model_path,
+                model_type=self.model_config.model_type,
+                modality="image",
+                prompts=[text],
+                media=[[image]],
+                image_data_format="pt")[0]
+
+            prompt_token_ids_single_img, _ = self(test_mm_prompt, None)
+
+        len_prompt_tokens_ids = len(prompt_token_ids_single_img)
+        # There are corner cases where if we strictly try to generate a text based
+        # on how many tokens we need to complete the input_seq_len, the output of
+        # default_multimodal_input_loader may give more tokens then the input_seq_len and this
+        # can lead to errors.
+        # That is why we try to clipped the variable text_token_left to a lower threshold
+        # but close enough to the actual input_seq_len
+        text_generation_perc_threshold = 0.95
+        text_token_left = int((input_seq_len - len_prompt_tokens_ids) *
+                              text_generation_perc_threshold)
+
+        if text_token_left > 0:
+            text = self.get_dummy_text(text_token_left)
+
         return default_multimodal_input_loader(
             tokenizer=self.tokenizer,
             model_dir=self.model_path,
             model_type=self.model_config.model_type,
             modality="image",
             prompts=[text],
-            media=[images],
+            media=[[image]],
             image_data_format="pt")[0]
 
     def _preprocess(self, text: dict[str, any], mm_data: dict[str, any],
diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py
index 13b3287fb218..e92dfa027b2f 100644
--- a/tensorrt_llm/_torch/pyexecutor/_util.py
+++ b/tensorrt_llm/_torch/pyexecutor/_util.py
@@ -3,7 +3,6 @@
 from typing import Dict, List, Optional
 
 import torch
-from transformers import AutoTokenizer
 
 import tensorrt_llm
 import tensorrt_llm.bindings.executor as trtllm
@@ -145,23 +144,17 @@ def _create_dummy_mm_context_request(
                 dict) and not self._profiling_stage_data.get("enable_mm_reqs"):
             return requests
 
-        model_name_or_path = getattr(self._model_engine.model, "name_or_path",
-                                     None)
-        assert model_name_or_path is not None, "Could not determine model name or path"
-        tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
-        input_processor = create_input_processor(model_name_or_path, tokenizer)
+        input_processor = self._model_engine.input_processor
         if not (hasattr(input_processor, "get_dummy_prompt")):
             logger.warning("The input processor of the model does not have the method [get_dummy_prompt] implemented." \
             "Profiling with the default input dummy context request. This may not take into account the memory consumption of " \
             "the image encoder")
             return requests
-        text_prompt = input_processor.get_dummy_prompt(input_seq_len,
-                                                       {'image': 1})
-        max_beam_width = self._max_beam_width
-        input_processor_with_hash = create_input_processor_with_hash(
-            input_processor)
-        prompt_token_ids, extra_processed_inputs = input_processor_with_hash(
-            text_prompt, None)
+        prompt = input_processor.get_dummy_prompt(input_seq_len)
+
+        prompt_token_ids, extra_processed_inputs = self._model_engine.input_processor_with_hash(
+            prompt, None)
+
         multimodal_input = extra_processed_inputs.get('multimodal_input')
         multimodal_data = extra_processed_inputs.get('multimodal_data')
 
@@ -184,7 +177,7 @@ def _create_dummy_mm_context_request(
                                      max_tokens=1,
                                      streaming=False,
                                      sampling_config=trtllm.SamplingConfig(
-                                         beam_width=max_beam_width, ),
+                                         beam_width=self._max_beam_width, ),
                                      output_config=trtllm.OutputConfig(),
                                      end_id=-1,
                                      multimodal_input=req_mm_input)
diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py
index 53d98fc83d84..49eb4e8e5222 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_engine.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -12,12 +12,15 @@
 
 import torch
 import torch._dynamo.config
+from transformers import AutoTokenizer
 
 import tensorrt_llm.bindings.internal.userbuffers as ub
 from tensorrt_llm._utils import (is_trace_enabled, nvtx_range, release_gc,
                                  torch_dtype_to_str, trace_func)
 from tensorrt_llm.inputs.multimodal import (MultimodalParams,
                                             MultimodalRuntimeData)
+from tensorrt_llm.inputs.registry import (create_input_processor,
+                                          create_input_processor_with_hash)
 from tensorrt_llm.logger import logger
 from tensorrt_llm.lora_helper import LoraConfig
 from tensorrt_llm.lora_manager import LoraModelConfig
@@ -171,7 +174,10 @@ def __init__(
 
         self.attn_runtime_features = attn_runtime_features or AttentionRuntimeFeatures(
         )
-
+        tokenizer = AutoTokenizer.from_pretrained(model_path)
+        self.input_processor = create_input_processor(model_path, tokenizer)
+        self.input_processor_with_hash = create_input_processor_with_hash(
+            self.input_processor)
         if model is None:
             loader = ModelLoader(
                 pytorch_backend_config=pytorch_backend_config,
diff --git a/tensorrt_llm/inputs/registry.py b/tensorrt_llm/inputs/registry.py
index f86f2ecac64b..b1c93ca58981 100644
--- a/tensorrt_llm/inputs/registry.py
+++ b/tensorrt_llm/inputs/registry.py
@@ -47,7 +47,7 @@ class BaseDummyInputsBuilder:
     Base class for generating dummy inputs. Specially for profiling
     """
 
-    def get_dummy_prompt(self, input_seq_len: int, mm_data: dict):
+    def get_dummy_prompt(self, input_seq_len: int):
         raise NotImplementedError(
             "Please ensure this method is implemented in your inherited class")
 

From 52afab22ed6b7e7736c3ff08b6b1715ebb9590ef Mon Sep 17 00:00:00 2001
From: John Calderon <jcalderon@nvidia.com>
Date: Wed, 1 Oct 2025 21:23:13 +0000
Subject: [PATCH 16/24] add unit test to test-db

Signed-off-by: John Calderon <jcalderon@nvidia.com>
---
 tests/integration/test_lists/test-db/l0_a100.yml | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tests/integration/test_lists/test-db/l0_a100.yml b/tests/integration/test_lists/test-db/l0_a100.yml
index 5d8d5162c36a..619bb57b2754 100644
--- a/tests/integration/test_lists/test-db/l0_a100.yml
+++ b/tests/integration/test_lists/test-db/l0_a100.yml
@@ -15,6 +15,7 @@ l0_a100:
   tests:
     - unittest/llmapi/test_llm_pytorch.py
     - unittest/llmapi/test_mpi_session.py ISOLATION
+    - unittest/llmapi/test_memory_profiling.py # profile kvcache for vision encoder
     - unittest/trt/model_api/test_model_quantization.py
     # executor
     - unittest/executor/test_base_worker.py

From 187b80fb4596ab93a888793455655d02f2d5287f Mon Sep 17 00:00:00 2001
From: John Calderon <jcalderon@nvidia.com>
Date: Fri, 3 Oct 2025 17:48:09 +0000
Subject: [PATCH 17/24] fix unit test by adding chunked prefill parameter

Signed-off-by: John Calderon <jcalderon@nvidia.com>
---
 tests/integration/defs/accuracy/test_llm_api_pytorch.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py
index 213d0597ad8b..acbb8a27dca7 100644
--- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py
+++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py
@@ -3550,6 +3550,7 @@ class TestQwen2_VL_7B(LlmapiAccuracyTestHarness):
     def test_auto_dtype(self):
         with LLM(self.MODEL_PATH,
                  max_num_tokens=16384,
+                 enable_chunked_prefill=True,
                  kv_cache_config=self.kv_cache_config) as llm:
             task = MMMU(self.MODEL_NAME)
             task.evaluate(llm, sampling_params=self.sampling_params)

From 0d470ec7f72d0717442532e3e5380f747bf15d4e Mon Sep 17 00:00:00 2001
From: John Calderon <jcalderon@nvidia.com>
Date: Fri, 3 Oct 2025 19:46:04 +0000
Subject: [PATCH 18/24] fix test
 L40S-PyTorch-2.test_e2e.test_ptp_quickstart_multimodal[NVILA-8B-FP16-vila/NVILA-8B-video-False]

Signed-off-by: John Calderon <jcalderon@nvidia.com>
---
 tensorrt_llm/_torch/models/modeling_qwen2vl.py | 3 ++-
 tensorrt_llm/_torch/pyexecutor/model_engine.py | 4 +---
 2 files changed, 3 insertions(+), 4 deletions(-)

diff --git a/tensorrt_llm/_torch/models/modeling_qwen2vl.py b/tensorrt_llm/_torch/models/modeling_qwen2vl.py
index ad3a72f3f55b..4371a8759493 100644
--- a/tensorrt_llm/_torch/models/modeling_qwen2vl.py
+++ b/tensorrt_llm/_torch/models/modeling_qwen2vl.py
@@ -96,7 +96,8 @@ def __init__(self,
                  tokenizer: AutoTokenizer,
                  trust_remote_code: bool = True):
         self.model_config = model_config
-        self.tokenizer = tokenizer
+        self.tokenizer = tokenizer if tokenizer is not None else AutoTokenizer.from_pretrained(
+            model_path)
         self.use_fast = True
         self.model_path = model_path
         self.processor = AutoProcessor.from_pretrained(
diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py
index 49eb4e8e5222..8d650469a057 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_engine.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -12,7 +12,6 @@
 
 import torch
 import torch._dynamo.config
-from transformers import AutoTokenizer
 
 import tensorrt_llm.bindings.internal.userbuffers as ub
 from tensorrt_llm._utils import (is_trace_enabled, nvtx_range, release_gc,
@@ -174,8 +173,7 @@ def __init__(
 
         self.attn_runtime_features = attn_runtime_features or AttentionRuntimeFeatures(
         )
-        tokenizer = AutoTokenizer.from_pretrained(model_path)
-        self.input_processor = create_input_processor(model_path, tokenizer)
+        self.input_processor = create_input_processor(model_path, None)
         self.input_processor_with_hash = create_input_processor_with_hash(
             self.input_processor)
         if model is None:

From bcc9004d23c5b25e367986546613357b495a10a1 Mon Sep 17 00:00:00 2001
From: John Calderon <jcalderon@nvidia.com>
Date: Sat, 4 Oct 2025 23:36:47 +0000
Subject: [PATCH 19/24] fix tests
 A10-PyTorch-1.test_e2e.test_openai_chat_multimodal_example and
 A10-PyTorch-1.test_e2e.test_trtllm_serve_multimodal_example

Signed-off-by: John Calderon <jcalderon@nvidia.com>
---
 tests/unittest/llmapi/apps/_test_openai_chat_multimodal.py   | 2 +-
 .../llmapi/apps/_test_trtllm_serve_multimodal_example.py     | 5 ++++-
 2 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/tests/unittest/llmapi/apps/_test_openai_chat_multimodal.py b/tests/unittest/llmapi/apps/_test_openai_chat_multimodal.py
index d92ca0616726..1607e6fd5df3 100644
--- a/tests/unittest/llmapi/apps/_test_openai_chat_multimodal.py
+++ b/tests/unittest/llmapi/apps/_test_openai_chat_multimodal.py
@@ -48,7 +48,7 @@ def server(model_name: str, temp_extra_llm_api_options_file: str):
     model_path = get_model_path(model_name)
     args = [
         "--extra_llm_api_options", temp_extra_llm_api_options_file,
-        "--max_batch_size", "64"
+        "--max_batch_size", "64", "--enable_chunked_prefill"
     ]
     with RemoteOpenAIServer(model_path, args) as remote_server:
         yield remote_server
diff --git a/tests/unittest/llmapi/apps/_test_trtllm_serve_multimodal_example.py b/tests/unittest/llmapi/apps/_test_trtllm_serve_multimodal_example.py
index 5b28e12675cb..e619fa6c21a8 100644
--- a/tests/unittest/llmapi/apps/_test_trtllm_serve_multimodal_example.py
+++ b/tests/unittest/llmapi/apps/_test_trtllm_serve_multimodal_example.py
@@ -46,7 +46,10 @@ def temp_extra_llm_api_options_file(request):
 @pytest.fixture(scope="module")
 def server(model_name: str, temp_extra_llm_api_options_file: str):
     model_path = get_model_path(model_name)
-    args = ["--extra_llm_api_options", temp_extra_llm_api_options_file]
+    args = [
+        "--enable_chunked_prefill", "--extra_llm_api_options",
+        temp_extra_llm_api_options_file
+    ]
     with RemoteOpenAIServer(model_path, port=8000,
                             cli_args=args) as remote_server:
         yield remote_server

From 345954c6dae31e008007898a4b8155c4b1279567 Mon Sep 17 00:00:00 2001
From: John Calderon <jcalderon@nvidia.com>
Date: Sun, 5 Oct 2025 20:52:27 +0000
Subject: [PATCH 20/24] include flag to check chunked prefill flag during
 profiling

Signed-off-by: John Calderon <jcalderon@nvidia.com>
---
 tensorrt_llm/_torch/pyexecutor/_util.py                   | 8 +++++---
 tests/integration/defs/accuracy/test_llm_api_pytorch.py   | 1 -
 .../unittest/llmapi/apps/_test_openai_chat_multimodal.py  | 2 +-
 .../llmapi/apps/_test_trtllm_serve_multimodal_example.py  | 5 +----
 4 files changed, 7 insertions(+), 9 deletions(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py
index e92dfa027b2f..6efd48f6cf11 100644
--- a/tensorrt_llm/_torch/pyexecutor/_util.py
+++ b/tensorrt_llm/_torch/pyexecutor/_util.py
@@ -193,9 +193,11 @@ def _create_dummy_mm_context_request(
     def _create_dummy_context_requests(
             self, input_seq_len: int) -> List[trtllm.Request]:
         requests = []
-        if hasattr(self._model_engine.model,
-                   "original_arch") and MODEL_CLASS_VISION_ENCODER_MAPPING.get(
-                       self._model_engine.model.original_arch, None):
+        if hasattr(
+                self._model_engine.model,
+                "original_arch") and MODEL_CLASS_VISION_ENCODER_MAPPING.get(
+                    self._model_engine.model.original_arch, None
+                ) and self._model_engine.attn_runtime_features.chunked_prefill:
             requests = self._create_dummy_mm_context_request(input_seq_len)
         # if succeed profiling with multimodal requests then return, otherwise profile
         # with default case
diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py
index acbb8a27dca7..213d0597ad8b 100644
--- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py
+++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py
@@ -3550,7 +3550,6 @@ class TestQwen2_VL_7B(LlmapiAccuracyTestHarness):
     def test_auto_dtype(self):
         with LLM(self.MODEL_PATH,
                  max_num_tokens=16384,
-                 enable_chunked_prefill=True,
                  kv_cache_config=self.kv_cache_config) as llm:
             task = MMMU(self.MODEL_NAME)
             task.evaluate(llm, sampling_params=self.sampling_params)
diff --git a/tests/unittest/llmapi/apps/_test_openai_chat_multimodal.py b/tests/unittest/llmapi/apps/_test_openai_chat_multimodal.py
index 1607e6fd5df3..d92ca0616726 100644
--- a/tests/unittest/llmapi/apps/_test_openai_chat_multimodal.py
+++ b/tests/unittest/llmapi/apps/_test_openai_chat_multimodal.py
@@ -48,7 +48,7 @@ def server(model_name: str, temp_extra_llm_api_options_file: str):
     model_path = get_model_path(model_name)
     args = [
         "--extra_llm_api_options", temp_extra_llm_api_options_file,
-        "--max_batch_size", "64", "--enable_chunked_prefill"
+        "--max_batch_size", "64"
     ]
     with RemoteOpenAIServer(model_path, args) as remote_server:
         yield remote_server
diff --git a/tests/unittest/llmapi/apps/_test_trtllm_serve_multimodal_example.py b/tests/unittest/llmapi/apps/_test_trtllm_serve_multimodal_example.py
index e619fa6c21a8..5b28e12675cb 100644
--- a/tests/unittest/llmapi/apps/_test_trtllm_serve_multimodal_example.py
+++ b/tests/unittest/llmapi/apps/_test_trtllm_serve_multimodal_example.py
@@ -46,10 +46,7 @@ def temp_extra_llm_api_options_file(request):
 @pytest.fixture(scope="module")
 def server(model_name: str, temp_extra_llm_api_options_file: str):
     model_path = get_model_path(model_name)
-    args = [
-        "--enable_chunked_prefill", "--extra_llm_api_options",
-        temp_extra_llm_api_options_file
-    ]
+    args = ["--extra_llm_api_options", temp_extra_llm_api_options_file]
     with RemoteOpenAIServer(model_path, port=8000,
                             cli_args=args) as remote_server:
         yield remote_server

From a546ae9084944e5b1b6719ac984e151c3e0ce37b Mon Sep 17 00:00:00 2001
From: John Calderon <jcalderon@nvidia.com>
Date: Thu, 9 Oct 2025 19:06:42 +0000
Subject: [PATCH 21/24] change logic to get initial input_seq_len

Signed-off-by: John Calderon <jcalderon@nvidia.com>
---
 tensorrt_llm/_torch/pyexecutor/_util.py        | 17 ++++++++---------
 tests/unittest/llmapi/test_memory_profiling.py |  4 +---
 2 files changed, 9 insertions(+), 12 deletions(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py
index 6efd48f6cf11..3f55a34fe91a 100644
--- a/tensorrt_llm/_torch/pyexecutor/_util.py
+++ b/tensorrt_llm/_torch/pyexecutor/_util.py
@@ -160,13 +160,10 @@ def _create_dummy_mm_context_request(
 
         max_num_tokens = len(prompt_token_ids)
         assert max_num_tokens > 0, "the length of the prompt of the dummy mm req is less than or equal to 0"
-        remaining_tokens = max(max_num_tokens, input_seq_len)
+        remaining_tokens = min(max_num_tokens, input_seq_len)
         if remaining_tokens > input_seq_len:
             logger.warning(f"Profiling with multimedia prompt which contains more tokens than the allowed input_seq_len. " \
                            f"Multimodal prompt has {remaining_tokens} while the input_seq_len is: {input_seq_len}")
-        ## add + 1 to avoid error: RuntimeError: The max KV cache length of input sequences (X + 1) exceeds the KV cache manager's maximum supported length X.
-        ## at line "/code/tensorrt_llm/tensorrt_llm/_torch/attention_backend/trtllm.py", line 837
-        self._max_seq_len = remaining_tokens + 1
         while remaining_tokens > 0:
             req_mm_input = trtllm.MultimodalInput(
                 multimodal_hashes=multimodal_input.multimodal_hashes,
@@ -181,6 +178,9 @@ def _create_dummy_mm_context_request(
                                      output_config=trtllm.OutputConfig(),
                                      end_id=-1,
                                      multimodal_input=req_mm_input)
+            # TODO:
+            # create_input_processor_with_hash shouldn’t be required during profiling,
+            # but is temporarily needed due to the multimodal input dependency for chunked prefill
             request.py_multimodal_data = multimodal_data
             remaining_tokens -= max_num_tokens
             requests.append(request)
@@ -193,11 +193,10 @@ def _create_dummy_mm_context_request(
     def _create_dummy_context_requests(
             self, input_seq_len: int) -> List[trtllm.Request]:
         requests = []
-        if hasattr(
-                self._model_engine.model,
-                "original_arch") and MODEL_CLASS_VISION_ENCODER_MAPPING.get(
-                    self._model_engine.model.original_arch, None
-                ) and self._model_engine.attn_runtime_features.chunked_prefill:
+        if hasattr(self._model_engine.model,
+                   "original_arch") and MODEL_CLASS_VISION_ENCODER_MAPPING.get(
+                       self._model_engine.model.original_arch, None):
+            input_seq_len = min(self._max_num_tokens, input_seq_len)
             requests = self._create_dummy_mm_context_request(input_seq_len)
         # if succeed profiling with multimodal requests then return, otherwise profile
         # with default case
diff --git a/tests/unittest/llmapi/test_memory_profiling.py b/tests/unittest/llmapi/test_memory_profiling.py
index b62a161a1402..a7b89c483824 100644
--- a/tests/unittest/llmapi/test_memory_profiling.py
+++ b/tests/unittest/llmapi/test_memory_profiling.py
@@ -22,9 +22,7 @@ def test_profile_kvcache():
     VLM_MODEL = "Qwen2.5-VL-7B-Instruct"
     VLM_MODEL_PATH = get_model_path(VLM_MODEL)
 
-    build_config = BuildConfig(max_batch_size=2048,
-                               max_beam_width=1,
-                               max_seq_len=8192)
+    build_config = BuildConfig(max_beam_width=1, max_num_tokens=16384)
     dynamic_batch_config = DynamicBatchConfig(
         enable_batch_size_tuning=True,
         enable_max_num_tokens_tuning=False,

From f011861966c80b61e5f9257b83289feea60e3e32 Mon Sep 17 00:00:00 2001
From: John Calderon <jcalderon@nvidia.com>
Date: Tue, 14 Oct 2025 16:19:52 +0000
Subject: [PATCH 22/24] fix latest rebase from
 7291cdc42287297bf72015e7201fede7985edeae

Signed-off-by: John Calderon <jcalderon@nvidia.com>
---
 tensorrt_llm/_torch/pyexecutor/_util.py        | 3 +--
 tests/unittest/llmapi/test_memory_profiling.py | 6 +++---
 2 files changed, 4 insertions(+), 5 deletions(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py
index 3f55a34fe91a..a1f7cbe39830 100644
--- a/tensorrt_llm/_torch/pyexecutor/_util.py
+++ b/tensorrt_llm/_torch/pyexecutor/_util.py
@@ -416,8 +416,7 @@ def configure_kv_cache_capacity(self, py_executor: PyExecutor) -> None:
         # set max_gpu_total_bytes
         self._kv_cache_config.max_gpu_total_bytes = kv_cache_max_memory
         if isinstance(self._profiling_stage_data, dict):
-            self._profiling_stage_data[
-                "max_gpu_total_bytes"] = kv_cache_max_memory
+            self._profiling_stage_data["activation_bytes"] = activation_bytes
         # ---------------------------handle max_gpu_total_bytes---------------------------------
 
     def _create_kv_cache_manager(
diff --git a/tests/unittest/llmapi/test_memory_profiling.py b/tests/unittest/llmapi/test_memory_profiling.py
index a7b89c483824..f1f036b1fd7d 100644
--- a/tests/unittest/llmapi/test_memory_profiling.py
+++ b/tests/unittest/llmapi/test_memory_profiling.py
@@ -62,7 +62,7 @@ def test_profile_kvcache():
     py_executor = create_py_executor(llm_args=torchllm_args,
                                      checkpoint_dir=VLM_MODEL_PATH,
                                      profiling_stage_data=profiling_data)
-    vlm_max_gpu_total_bytes_with_mm_reqs = profiling_data["max_gpu_total_bytes"]
+    vlm_activation_bytes_with_mm_reqs = profiling_data["activation_bytes"]
     py_executor.shutdown()
     torch.cuda.empty_cache()
 
@@ -71,8 +71,8 @@ def test_profile_kvcache():
     py_executor_2 = create_py_executor(llm_args=torchllm_args,
                                        checkpoint_dir=VLM_MODEL_PATH,
                                        profiling_stage_data=profiling_data)
-    vlm_max_gpu_total_bytes_no_mm_reqs = profiling_data["max_gpu_total_bytes"]
+    vlm_activation_bytes_no_mm_reqs = profiling_data["activation_bytes"]
     py_executor_2.shutdown()
     torch.cuda.empty_cache()
 
-    assert vlm_max_gpu_total_bytes_with_mm_reqs < vlm_max_gpu_total_bytes_no_mm_reqs, f"available KVCache for VLMs is expected to be less when profiling with mm reqs, but got {vlm_max_gpu_total_bytes_with_mm_reqs} for mm reqs and {vlm_max_gpu_total_bytes_no_mm_reqs} without mm reqs"
+    assert vlm_activation_bytes_with_mm_reqs > vlm_activation_bytes_no_mm_reqs, f"Activation bytes should be higher with mm reqs, but got {vlm_activation_bytes_with_mm_reqs} for mm reqs and {vlm_activation_bytes_no_mm_reqs} without mm reqs"

From 5410711e6d48a8a2f32a3987113363126f5dafd8 Mon Sep 17 00:00:00 2001
From: John Calderon <jcalderon@nvidia.com>
Date: Tue, 14 Oct 2025 18:18:22 +0000
Subject: [PATCH 23/24] Fix rebase to 1cdb0b6

Signed-off-by: John Calderon <jcalderon@nvidia.com>
---
 .../_torch/models/modeling_qwen2vl.py         |  7 ++---
 tensorrt_llm/_torch/pyexecutor/_util.py       | 29 +++++++------------
 .../unittest/llmapi/test_memory_profiling.py  |  1 -
 3 files changed, 14 insertions(+), 23 deletions(-)

diff --git a/tensorrt_llm/_torch/models/modeling_qwen2vl.py b/tensorrt_llm/_torch/models/modeling_qwen2vl.py
index 4371a8759493..160ce302b878 100644
--- a/tensorrt_llm/_torch/models/modeling_qwen2vl.py
+++ b/tensorrt_llm/_torch/models/modeling_qwen2vl.py
@@ -322,9 +322,8 @@ def get_dummy_prompt(self, input_seq_len: int):
             # reduce img resolution
             img_max_dim = img_max_dim >> 1
 
-            image = self.get_dummy_image(
-                max_width=img_max_dim, max_height=img_max_dim
-            )  #w, h is sqrt of min_pixels value (3136)
+            image = self.get_dummy_image(max_width=img_max_dim,
+                                         max_height=img_max_dim)
 
             test_mm_prompt = default_multimodal_input_loader(
                 tokenizer=self.tokenizer,
@@ -342,7 +341,7 @@ def get_dummy_prompt(self, input_seq_len: int):
         # on how many tokens we need to complete the input_seq_len, the output of
         # default_multimodal_input_loader may give more tokens then the input_seq_len and this
         # can lead to errors.
-        # That is why we try to clipped the variable text_token_left to a lower threshold
+        # That is why we try to clip the variable text_token_left to a lower threshold
         # but close enough to the actual input_seq_len
         text_generation_perc_threshold = 0.95
         text_token_left = int((input_seq_len - len_prompt_tokens_ids) *
diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py
index a1f7cbe39830..bec3cd119726 100644
--- a/tensorrt_llm/_torch/pyexecutor/_util.py
+++ b/tensorrt_llm/_torch/pyexecutor/_util.py
@@ -61,24 +61,17 @@ def get_kv_cache_manager_cls(model_config: ModelConfig):
 class KvCacheCreator:
     """Groups together logic related to KV cache construction."""
 
-    def __init__(
-        self,
-        *,
-        model_engine: PyTorchModelEngine,
-        draft_model_engine: Optional[PyTorchModelEngine],
-        mapping: Mapping,
-        net_max_seq_len: int,
-        kv_connector_manager: Optional[KvCacheConnectorManager],
-        max_num_tokens: int,
-        max_beam_width: int,
-        tokens_per_block: int,
-        max_seq_len: int,
-        max_batch_size: int,
-        kv_cache_config: KvCacheConfig,
-        pytorch_backend_config: PyTorchConfig,
-        speculative_config: SpeculativeConfig,
-        sparse_attention_config: SparseAttentionConfig,
-    ):
+    def __init__(self, *, model_engine: PyTorchModelEngine,
+                 draft_model_engine: Optional[PyTorchModelEngine],
+                 mapping: Mapping, net_max_seq_len: int,
+                 kv_connector_manager: Optional[KvCacheConnectorManager],
+                 max_num_tokens: int, max_beam_width: int,
+                 tokens_per_block: int, max_seq_len: int, max_batch_size: int,
+                 kv_cache_config: KvCacheConfig,
+                 pytorch_backend_config: PyTorchConfig,
+                 speculative_config: SpeculativeConfig,
+                 sparse_attention_config: SparseAttentionConfig,
+                 profiling_stage_data: Optional[dict]):
         self._model_engine = model_engine
         self._draft_model_engine = draft_model_engine
         self._mapping = mapping
diff --git a/tests/unittest/llmapi/test_memory_profiling.py b/tests/unittest/llmapi/test_memory_profiling.py
index f1f036b1fd7d..57c668e9e153 100644
--- a/tests/unittest/llmapi/test_memory_profiling.py
+++ b/tests/unittest/llmapi/test_memory_profiling.py
@@ -52,7 +52,6 @@ def test_profile_kvcache():
         "postprocess_tokenizer_dir": VLM_MODEL,
         "reasoning_parser": None,
         "fail_fast_on_attention_window_too_large": False,
-        "enable_chunked_prefill": True,
         "cuda_graph_config": cuda_graph_config,
     }
 

From 9919b4c071bd1f134f74d53d9ff10db5ed4b1125 Mon Sep 17 00:00:00 2001
From: John Calderon <jcalderon@nvidia.com>
Date: Wed, 15 Oct 2025 14:17:48 +0000
Subject: [PATCH 24/24] fix format file _util.py

Signed-off-by: John Calderon <jcalderon@nvidia.com>
---
 tensorrt_llm/_torch/pyexecutor/_util.py | 30 ++++++++++++++++---------
 1 file changed, 19 insertions(+), 11 deletions(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py
index bec3cd119726..4099e0e104dd 100644
--- a/tensorrt_llm/_torch/pyexecutor/_util.py
+++ b/tensorrt_llm/_torch/pyexecutor/_util.py
@@ -61,17 +61,25 @@ def get_kv_cache_manager_cls(model_config: ModelConfig):
 class KvCacheCreator:
     """Groups together logic related to KV cache construction."""
 
-    def __init__(self, *, model_engine: PyTorchModelEngine,
-                 draft_model_engine: Optional[PyTorchModelEngine],
-                 mapping: Mapping, net_max_seq_len: int,
-                 kv_connector_manager: Optional[KvCacheConnectorManager],
-                 max_num_tokens: int, max_beam_width: int,
-                 tokens_per_block: int, max_seq_len: int, max_batch_size: int,
-                 kv_cache_config: KvCacheConfig,
-                 pytorch_backend_config: PyTorchConfig,
-                 speculative_config: SpeculativeConfig,
-                 sparse_attention_config: SparseAttentionConfig,
-                 profiling_stage_data: Optional[dict]):
+    def __init__(
+        self,
+        *,
+        model_engine: PyTorchModelEngine,
+        draft_model_engine: Optional[PyTorchModelEngine],
+        mapping: Mapping,
+        net_max_seq_len: int,
+        kv_connector_manager: Optional[KvCacheConnectorManager],
+        max_num_tokens: int,
+        max_beam_width: int,
+        tokens_per_block: int,
+        max_seq_len: int,
+        max_batch_size: int,
+        kv_cache_config: KvCacheConfig,
+        pytorch_backend_config: PyTorchConfig,
+        speculative_config: SpeculativeConfig,
+        sparse_attention_config: SparseAttentionConfig,
+        profiling_stage_data: Optional[dict],
+    ):
         self._model_engine = model_engine
         self._draft_model_engine = draft_model_engine
         self._mapping = mapping