[TRTLLM-12467][feat] EPD improvements

.claude/skills/trtllm-model-onboard-multimodal/SKILL.md

@@ -83,7 +83,7 @@ metadata:
 When `@support_multimodal_disaggregated` is set and the deployment uses `TLLM_MULTIMODAL_DISAGGREGATED=1`:
 
 - **Encoder worker:** runs as a standalone `MultimodalEncoder` (`mm_encoder_only=True`). It executes only the multimodal encoder and ships `mm_embeddings` (+ mRoPE position ids/deltas) to prefill+decode workers as shared-tensor handles.
-- **Prefill+decode worker:** the model's `__init__` skips constructing `self.mm_encoder` when `_is_disagg()` is true; the input processor's `_attach_multimodal_embeddings_impl()` override binds the encoder handles into the request (the base `attach_multimodal_embeddings` wrapper detokenizes tokenized inputs for non-fast-path VLMs, then delegates to your impl). For context-only requests, the engine re-clones mrope tensors so IPC handles outlive the encoder worker's freed memory — replicate that pattern for any new GPU-resident mm tensors.
+- **Prefill+decode worker:** the model's `__init__` skips constructing `self.mm_encoder` when `_is_mm_disagg()` is true; the input processor's `attach_multimodal_embeddings()` override binds the encoder handles into the request. For context-only requests, the engine re-clones mrope tensors so IPC handles outlive the encoder worker's freed memory — replicate that pattern for any new GPU-resident mm tensors.
 
 ### Templates to study
 
@@ -229,7 +229,7 @@ class {Name}Model(PreTrainedModel):
         if hasattr(self, "llm"):
             return  # idempotency guard — re-entry from `post_config` etc.
 
-        if not _is_disagg():
+        if not _is_mm_disagg():
             self.mm_encoder = {Name}VisionModel(model_config)
         else:
             self.mm_encoder = None
@@ -269,7 +269,7 @@ class {Name}Model(PreTrainedModel):
 
         multimodal_params = kwargs.get("multimodal_params", [])
         mm_embeds = []
-        if len(multimodal_params) > 0 and not _is_disagg():
+        if len(multimodal_params) > 0 and not _is_mm_disagg():
             mm_embeds = get_multimodal_embeddings(
                 encoder_forward_fn=self.mm_encoder.forward,
                 multimodal_params=multimodal_params[:num_context_requests],
@@ -341,7 +341,7 @@ class {Name}Model(PreTrainedModel): ...
 
 ```python
 def load_weights(self, weights, weight_mapper):
-    if not _is_disagg():
+    if not _is_mm_disagg():
         self.mm_encoder.load_weights(weights)
         # Release mmap pages backing the encoder weights as soon as we're done.
         if hasattr(weights, "mark_consumed"):

docs/source/models/supported-models.md

@@ -81,6 +81,7 @@ Note: Support for other models may vary. Features marked "N/A" are not applicabl
 [^7]: Text-only support via the [AutoDeploy](../features/auto_deploy/auto-deploy.md) backend.
 [^8]: Supports text and image inputs. The vision tower runs in BF16 even when the text decoder is quantized (FP8 block-scale or NVFP4). The text decoder is also usable standalone (text-only) via the `Step3p5ForCausalLM` architecture.
 [^9]: Audio modality only supported on E2B/E4B variants.
+[^10]: Audio requires a checkpoint with a `sound_config` and is supported only on the full (non-disaggregated) model path, not the EPD disaggregated path.
 
 # Multimodal Feature Support Matrix (PyTorch Backend)
 
@@ -94,7 +95,7 @@ Note: Support for other models may vary. Features marked "N/A" are not applicabl
 | `LlavaNextForConditionalGeneration`  | Yes               | Yes        | Yes             | Yes           | Yes              | Yes            | Yes                   | Yes                       | L + I     |
 | `Llama4ForConditionalGeneration`     | Yes               | Yes        | No              | Yes           | Yes              | No             | Yes                   | No                        | L + I     |
 | `Mistral3ForConditionalGeneration`   | Yes               | Yes        | Yes             | Yes           | Yes              | Yes            | Yes                   | No                        | L + I     |
-| `NemotronH_Nano_VL_V2`               | Yes               | Yes        | Yes             | Yes           | Yes              | N/A            | Yes                   | No                        | L + I + V |
+| `NemotronH_Nano_VL_V2`               | Yes               | Yes        | Yes             | Yes           | Yes              | N/A            | Yes                   | Yes                       | L + I + V + A [^10] |
 | `Phi4MMForCausalLM`                  | Yes               | Yes        | Yes             | Yes           | Yes              | Yes            | Yes                   | No                        | L + I + A |
 | `Qwen2VLForConditionalGeneration`    | Yes               | Yes        | Yes             | Yes           | Yes              | Yes            | Yes                   | No                        | L + I + V |
 | `Qwen2_5_VLForConditionalGeneration` | Yes               | Yes        | Yes             | Yes           | Yes              | Yes            | Yes                   | Yes                       | L + I + V |

tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py

@@ -66,6 +66,8 @@
     {
         "layout_metadata",
         "mm_bidirectional_blocks",
+        "multimodal_embedding",
+        "multimodal_embedding_lengths",
         "special_token_offsets",
         "multimodal_embed_mask_cumsum",
     }

tensorrt_llm/_torch/models/modeling_exaone4_5.py

@@ -9,7 +9,7 @@
 from transformers.models.auto import CONFIG_MAPPING
 
 from tensorrt_llm._torch.models.checkpoints.base_weight_mapper import BaseWeightMapper
-from tensorrt_llm._torch.models.modeling_multimodal_utils import _is_disagg
+from tensorrt_llm._torch.models.modeling_multimodal_utils import _is_mm_disagg
 
 from ...inputs import (
     ContentFormat,
@@ -198,7 +198,7 @@ def __init__(
         llm_model_config.pretrained_config = llm_model_config.pretrained_config.text_config
         self.llm = AutoModelForCausalLM.from_config(llm_model_config)
 
-        if not _is_disagg():
+        if not _is_mm_disagg():
             mm_encoder_config = copy.deepcopy(model_config)
             self.mm_encoder = Exaone4_5_VisionModel(mm_encoder_config, Qwen2_5_VisionModel)
         else:
@@ -231,7 +231,7 @@ def forward(
         mm_multimodal_params = self._get_requests_with_mm_data(multimodal_params)
 
         if len(mm_multimodal_params) > 0:
-            if not _is_disagg():
+            if not _is_mm_disagg():
                 mm_embeds = get_multimodal_embeddings(
                     encoder_forward_fn=self.mm_encoder.forward,
                     multimodal_params=mm_multimodal_params,
@@ -262,6 +262,6 @@ def forward(
     def load_weights(self, weights, weight_mapper: BaseWeightMapper):
         assert isinstance(weight_mapper, Exaone4_5HfWeightMapper)
         weights = weight_mapper.preprocess_weights(weights)
-        if not _is_disagg():
+        if not _is_mm_disagg():
             self.mm_encoder.load_weights(weights)
         self.llm.load_weights(weights, weight_mapper)

tensorrt_llm/_torch/models/modeling_gemma3vl.py

@@ -1,6 +1,5 @@
 import copy
 import dataclasses
-import os
 from typing import List, Optional, Tuple
 
 import torch
@@ -23,17 +22,11 @@
 from ..modules.linear import Linear
 from ..modules.rms_norm import RMSNorm
 from .modeling_gemma3 import Gemma3ForCausalLM
-from .modeling_multimodal_utils import fuse_input_embeds
+from .modeling_multimodal_utils import (_MULTIMODAL_ENV_NAME, _is_mm_disagg,
+                                        fuse_input_embeds)
 from .modeling_siglip import SiglipVisionModel
 from .modeling_utils import ModelConfig, filter_weights, register_auto_model
 
-_MULTIMODAL_ENV_NAME = "TLLM_MULTIMODAL_DISAGGREGATED"
-
-
-# Make this a runtime lookup rather than a module-wide constant for easier unit testing.
-def _is_disagg() -> bool:
-    return os.getenv(_MULTIMODAL_ENV_NAME, "0") == "1"
-
 
 class Gemma3InputProcessor(BaseMultimodalInputProcessor,
                            BaseMultimodalDummyInputsBuilder):
@@ -185,7 +178,7 @@ def forward(self, vision_outputs: torch.Tensor):
 class Gemma3VLM(PreTrainedModel):
 
     def __init__(self, model_config: ModelConfig[Gemma3Config]):
-        if _is_disagg():
+        if _is_mm_disagg():
             raise NotImplementedError(
                 "Gemma3VLM does not support disaggregated inference yet. Please unset "
                 f"the {_MULTIMODAL_ENV_NAME} environment variable, or set it to '0'."

tensorrt_llm/_torch/models/modeling_gemma4mm.py

@@ -23,7 +23,6 @@
 import copy
 import dataclasses
 import math
-import os
 from typing import Dict, List, Optional, Tuple
 
 import torch
@@ -51,7 +50,12 @@
 from .modeling_gemma4 import Gemma4ForCausalLM
 from .modeling_gemma4_audio import Gemma4AudioModel
 from .modeling_gemma4_vision import Gemma4VisionModel
-from .modeling_multimodal_utils import find_input_mm_embeds, fuse_input_embeds
+from .modeling_multimodal_utils import (
+    _MULTIMODAL_ENV_NAME,
+    _is_mm_disagg,
+    find_input_mm_embeds,
+    fuse_input_embeds,
+)
 from .modeling_utils import ModelConfig, filter_weights, register_auto_model
 
 _MIN_TRANSFORMERS_FOR_GEMMA4 = "5.5.0"
@@ -69,12 +73,6 @@
     PreTrainedModel,
 )
 
-_MULTIMODAL_ENV_NAME = "TLLM_MULTIMODAL_DISAGGREGATED"
-
-
-def _is_disagg() -> bool:
-    return os.getenv(_MULTIMODAL_ENV_NAME, "0") == "1"
-
 
 class RMSNormNoScale(nn.Module):
     """RMSNorm without learnable scale (for multimodal embedder pre-projection)."""
@@ -602,7 +600,7 @@ def _check_and_adjust_experts_implementation(self, *args, **kwargs):
         return None
 
     def __init__(self, model_config: ModelConfig[Gemma4Config]):
-        if _is_disagg():
+        if _is_mm_disagg():
             raise NotImplementedError(
                 "Gemma4ForConditionalGeneration does not support "
                 "disaggregated inference yet. Please unset the "

tensorrt_llm/_torch/models/modeling_kimi_k25.py

@@ -49,7 +49,7 @@
     PreTrainedTokenizerBase,
 )
 
-from tensorrt_llm.inputs.multimodal import MultimodalParams
+from tensorrt_llm.inputs.multimodal import DisaggPrefillMultimodalInputs, MultimodalParams
 from tensorrt_llm.mapping import Mapping
 
 from ..._utils import prefer_pinned
@@ -1492,20 +1492,22 @@ def call_with_text_prompt(
             "multimodal_data": multimodal_data,
         }
 
-    def get_prompt_token_ids(
+    def build_disagg_prefill_multimodal_inputs(
         self,
         inputs: TextPrompt,
         mm_handles: List[Dict[str, Any]],
-    ) -> Tuple[List[int], List[int], List[int]]:
-        """Build token IDs with multimodal placeholders expanded for disaggregated serving.
+    ) -> DisaggPrefillMultimodalInputs:
+        """Build disaggregated prefill inputs from multimodal embedding handles.
 
         Args:
             inputs: Text prompt input container.
             mm_handles: List of multimodal embedding handles from the
                 context phase, each containing ``tensor_size``.
 
         Returns:
-            Tuple of (expanded_ids, mm_token_lengths, mm_token_offsets).
+            DisaggPrefillMultimodalInputs containing expanded token IDs,
+            prompt-side MM positions/lengths, exact runs, and encoder-output
+            embedding lengths.
         """
         text_prompt = inputs.get("prompt")
         if not text_prompt:
@@ -1556,7 +1558,15 @@ def get_prompt_token_ids(
                 expanded_ids[write_pos] = input_ids[read_pos]
                 write_pos += 1
 
-        return (expanded_ids.to(torch.int32).tolist(), mm_token_length, mm_token_offsets)
+        return DisaggPrefillMultimodalInputs(
+            prompt_token_ids=expanded_ids.to(torch.int32).tolist(),
+            multimodal_lengths=mm_token_length,
+            multimodal_positions=mm_token_offsets,
+            multimodal_embedding_lengths=[mm_handle["tensor_size"][0] for mm_handle in mm_handles],
+            multimodal_item_run_cu_offsets=list(range(len(mm_token_length) + 1)),
+            multimodal_run_positions=mm_token_offsets,
+            multimodal_run_lengths=mm_token_length,
+        )
 
 
 # ---------------------------------------------------------------------------

tensorrt_llm/_torch/models/modeling_llava_next.py

@@ -1,5 +1,4 @@
 import copy
-import os
 from typing import Any, ClassVar, Dict, List, Optional, Tuple, Union
 
 import numpy as np
@@ -15,7 +14,9 @@
     BaseWeightMapper
 from tensorrt_llm._torch.models.checkpoints.hf.llava_next_weight_mapper import \
     LlavaNextHfWeightMapper
-from tensorrt_llm.inputs.multimodal import MultimodalParams
+from tensorrt_llm._torch.models.modeling_multimodal_utils import _is_mm_disagg
+from tensorrt_llm.inputs.multimodal import (DisaggPrefillMultimodalInputs,
+                                            MultimodalParams)
 
 from ...inputs import (BaseMultimodalDummyInputsBuilder,
                        BaseMultimodalInputProcessor, ContentFormat,
@@ -29,11 +30,10 @@
 from .modeling_auto import AutoModelForCausalLM
 from .modeling_clip import CLIPVisionModel
 from .modeling_multimodal_utils import (find_input_mm_embeds, fuse_input_embeds,
+                                        get_attached_multimodal_embeddings,
                                         get_multimodal_embeddings)
 from .modeling_utils import register_auto_model, register_vision_encoder
 
-DISAGG = os.getenv('TLLM_MULTIMODAL_DISAGGREGATED', '0') == '1'
-
 
 class LlavaNextInputProcessor(BaseMultimodalInputProcessor,
                               BaseMultimodalDummyInputsBuilder):
@@ -111,7 +111,8 @@ def _expand_image_placeholders_in_token_ids(
         num_mm_tokens_per_placeholder: List[int],
     ) -> Tuple[List[int], List[int], List[int]]:
         """
-        Shared logic (called by expand_prompt_token_ids_for_mm and get_prompt_token_ids):
+        Shared logic (called by expand_prompt_token_ids_for_mm and
+        build_disagg_prefill_multimodal_inputs):
         replace each image placeholder token in prompt_token_ids
         with placeholder_id repeated num_mm_tokens_per_placeholder[i] times.
 
@@ -268,12 +269,11 @@ def _postprocess(
         mm_features = mm_features.view(-1, mm_features.shape[-1])
         return fused_input_ids, mm_features
 
-    def get_prompt_token_ids(
-        self, inputs: Union[TextPrompt, TokensPrompt],
-        mm_handles: List[Dict[str,
-                              Any]]) -> Tuple[List[int], List[int], List[int]]:
+    def build_disagg_prefill_multimodal_inputs(
+            self, inputs: Union[TextPrompt, TokensPrompt],
+            mm_handles: List[Dict[str, Any]]) -> DisaggPrefillMultimodalInputs:
         """
-        Build input token ids with multimodal placeholders expanded to the number of MM tokens.
+        Build disaggregated prefill inputs from multimodal embedding handles.
 
         Uses an already tokenized prompt or tokenizes the txt prompt first.
 
@@ -282,10 +282,9 @@ def get_prompt_token_ids(
             mm_handles: List of multimodal embedding handles.
 
         Returns:
-            Tuple[List[int], List[int], List[int]]:
-                - expanded_ids: token ids with each image token expanded to a placeholder repeated per MM token
-                - mm_token_length: per-image MM token lengths
-                - mm_token_offsets: start offsets (positions) for each image's MM tokens within expanded_ids
+            DisaggPrefillMultimodalInputs containing expanded token IDs,
+            prompt-side MM positions/lengths, exact runs, and encoder-output
+            embedding lengths.
         """
         # TODO: Move this function to the base input processor class when extending for more models
         text_prompt = inputs.get("prompt")
@@ -327,7 +326,18 @@ def get_prompt_token_ids(
                 f"({mm_token_length[-1] + mm_token_offsets[-1]}) should be less "
                 f"than or equal to final_length ({final_length})")
 
-        return expanded_ids, mm_token_length, mm_token_offsets
+        return DisaggPrefillMultimodalInputs(
+            prompt_token_ids=expanded_ids,
+            multimodal_lengths=mm_token_length,
+            multimodal_positions=mm_token_offsets,
+            multimodal_embedding_lengths=[
+                mm_handle["tensor_size"][0] for mm_handle in mm_handles
+            ],
+            multimodal_item_run_cu_offsets=list(range(len(mm_token_length) +
+                                                      1)),
+            multimodal_run_positions=mm_token_offsets,
+            multimodal_run_lengths=mm_token_length,
+        )
 
     def _attach_multimodal_embeddings_impl(
         self, inputs: TextPrompt,
@@ -619,7 +629,7 @@ def __init__(self, model_config: ModelConfig[PretrainedConfig], *args,
         super().__init__(config)
         if hasattr(self, "llm"):
             return
-        if not DISAGG:
+        if not _is_mm_disagg():
             self.mm_encoder = LlavaNextVisionModel(model_config)
         else:
             self.mm_encoder = None
@@ -694,15 +704,14 @@ def forward(
         multimodal_params = kwargs.get("multimodal_params", [])
         mm_embeds = []
         if len(multimodal_params) > 0:
-            if not DISAGG:
+            if self.mm_encoder is not None:
                 mm_embeds = get_multimodal_embeddings(
                     encoder_forward_fn=self.mm_encoder.forward,
                     multimodal_params=multimodal_params[:num_context_requests])
             else:
-                raise NotImplementedError(
-                    "LlavaNextModel does not support disaggregated inference yet. Please unset "
-                    f"the TLLM_MULTIMODAL_DISAGGREGATED environment variable, or set it to '0'."
-                )
+                # E/P prefill: encoder already ran; use attached embeddings.
+                mm_embeds = get_attached_multimodal_embeddings(
+                    multimodal_params[:num_context_requests])
             mm_embeds = find_input_mm_embeds(
                 mm_embeds, multimodal_params[:num_context_requests])
         input_ids, inputs_embeds = fuse_input_embeds(

tensorrt_llm/_torch/models/modeling_mistral.py

@@ -25,7 +25,7 @@
 from tensorrt_llm._torch.models.modeling_multimodal_mixin import (
     MultimodalEncoderOutput, MultimodalModelMixin, PreparedLlmInputs)
 from tensorrt_llm._torch.models.modeling_multimodal_utils import (
-    _MULTIMODAL_ENV_NAME, _is_disagg)
+    _MULTIMODAL_ENV_NAME, _is_mm_disagg)
 from tensorrt_llm._torch.models.modeling_utils import (DecoderModel,
                                                        DecoderModelForCausalLM,
                                                        _load_weights_impl,
@@ -567,7 +567,8 @@ def __init__(
         self,
         model_config: ModelConfig[Mistral3Config],
     ):
-        if _is_disagg():
+        # No MM E/P handoff here yet. Fail before partial model setup.
+        if _is_mm_disagg():
             raise NotImplementedError(
                 "Mistral3VLM does not support disaggregated inference yet. Please unset "
                 f"the {_MULTIMODAL_ENV_NAME} environment variable, or set it to '0'."