Skip to content

[TRTLLM-13419][refactor] Share encoder input packing across requests #21052

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
yechank-nvidia wants to merge 3 commits into
base: main
Choose a base branch
from
+599 −105
Draft

[TRTLLM-13419][refactor] Share encoder input packing across requests #21052

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
cceab0d
[TRTLLM-13419][refactor] Share encoder input packing across requests
yechank-nvidia Oct 26, 2026
92463e0
[TRTLLM-13419][perf] Skip grouping for uniform encoder inputs
yechank-nvidia Oct 26, 2026
666ee8c
[TRTLLM-13419][perf] Avoid repeated Qwen3-VL encoder input gathering
yechank-nvidia Oct 26, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
170 changes: 81 additions & 89 deletions tensorrt_llm/_torch/models/modeling_gemma4mm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -58,8 +58,10 @@
from .modeling_gemma4_vision import Gemma4VisionModel
from .modeling_multimodal_mixin import (
EncoderCachePartition,
EncoderInputField,
MultimodalModelMixin,
PreparedLlmInputs,
batch_encoder_inputs,
)
from .modeling_multimodal_utils import _MULTIMODAL_ENV_NAME, _is_mm_disagg
from .modeling_utils import ModelConfig, filter_weights, register_auto_model
Expand All @@ -85,6 +87,12 @@
get_aspect_ratio_preserving_size,
)

_GEMMA4_PIXEL_FIELD = {"pixel_values": EncoderInputField(("pixel_values",))}
_GEMMA4_VISION_FIELDS = {
**_GEMMA4_PIXEL_FIELD,
"image_position_ids": EncoderInputField(("image_position_ids",)),
}


class RMSNormNoScale(nn.Module):
"""RMSNorm without learnable scale (for multimodal embedder pre-projection)."""
Expand Down Expand Up @@ -1043,107 +1051,91 @@ def encode_multimodal_inputs(self, multimodal_params: List[MultimodalParams]) ->
"Gemma4 requests containing multiple multimodal input types are not supported."
)

# The cache helper splits this tensor by request. Preserve that order while still batching
# consecutive requests with the same modality.
if (
all(len(modalities) == 1 for modalities in param_modalities)
and len({modalities[0] for modalities in param_modalities}) > 1
):
embeddings = []
for _, param_group in groupby(
zip(param_modalities, multimodal_params), key=lambda item: item[0][0]
):
embeddings.append(
self.encode_multimodal_inputs(
[multimodal_param for _, multimodal_param in param_group]
)
# The cache helper splits this tensor by request. Preserve consecutive
# modality runs and their encoder call composition without rechecking
# each run through a recursive encoder entry.
single_modality_inputs = all(len(modalities) == 1 for modalities in param_modalities)
mixed_modality_runs = (
single_modality_inputs and len({modalities[0] for modalities in param_modalities}) > 1
)
if mixed_modality_runs:
input_groups = [
(modality, [param.multimodal_data[modality] for _, param in param_group])
for modality, param_group in groupby(
zip(param_modalities, multimodal_params), key=lambda item: item[0][0]
)
]
elif single_modality_inputs and multimodal_params:
modality = param_modalities[0][0]
input_groups = [
(modality, [param.multimodal_data[modality] for param in multimodal_params])
]
else:
input_groups = [
(
modality,
[
data
for param in multimodal_params
if (data := param.multimodal_data.get(modality, {})).get(input_field)
is not None
],
)
return torch.cat(embeddings, dim=0)

pixel_values_list = []
image_position_ids_list = []
image_seq_lens_extended: List[int] = []
audio_features_list = []
audio_features_mask_list: List[Optional[torch.Tensor]] = []
video_pixel_values_list = []
video_position_ids_list = []
video_seq_lens_extended: List[int] = []

for multimodal_param in multimodal_params:
multimodal_data = multimodal_param.multimodal_data
image_data = multimodal_data.get("image", {})
pixel_values = image_data.get("pixel_values")
if pixel_values is not None:
pixel_values_list.append(pixel_values)
image_position_ids = image_data.get("image_position_ids")
if image_position_ids is not None:
image_position_ids_list.append(image_position_ids)
image_seq_lens = image_data.get("image_seq_lens")
if image_seq_lens is not None:
image_seq_lens_extended.extend(image_seq_lens)

audio_data = multimodal_data.get("audio", {})
audio_features = audio_data.get("audio_features")
if audio_features is not None:
audio_features_list.append(audio_features)
audio_features_mask_list.append(audio_data.get("audio_features_mask"))

video_data = multimodal_data.get("video", {})
video_pixel_values = video_data.get("pixel_values")
if video_pixel_values is not None:
video_pixel_values_list.append(video_pixel_values)
video_position_ids = video_data.get("image_position_ids")
if video_position_ids is not None:
video_position_ids_list.append(video_position_ids)
video_seq_lens = video_data.get("image_seq_lens")
if video_seq_lens is not None:
video_seq_lens_extended.extend(video_seq_lens)
for modality, input_field in modality_inputs
]

multimodal_embeddings = []
if pixel_values_list:
pixel_values = torch.cat(pixel_values_list)
image_position_ids = (
torch.cat(image_position_ids_list)
if len(image_position_ids_list) == len(pixel_values_list)
else None
for modality, inputs in input_groups:
if not inputs:
continue
if modality == "audio":
if self.embed_audio is not None:
multimodal_embeddings.append(
torch.cat(
[
self._get_audio_features(
data["audio_features"], data.get("audio_features_mask")
)
for data in inputs
],
dim=0,
)
)
if mixed_modality_runs:
multimodal_embeddings[-1] = multimodal_embeddings[-1].clone(
memory_format=torch.contiguous_format
)
elif mixed_modality_runs:
raise ValueError(
"Gemma4 received active multimodal parameters without encoder inputs"
)
continue
fields = (
_GEMMA4_VISION_FIELDS
if all(data.get("image_position_ids") is not None for data in inputs)
else _GEMMA4_PIXEL_FIELD
)
batched = batch_encoder_inputs(inputs, fields)
image_seq_lens = []
for data in inputs:
if data.get("image_seq_lens") is not None:
image_seq_lens.extend(data["image_seq_lens"])
multimodal_embeddings.append(
self._get_image_features(
pixel_values=pixel_values,
image_position_ids=image_position_ids,
pixel_values=batched["pixel_values"],
image_position_ids=batched.get("image_position_ids"),
image_seq_lens=(
image_seq_lens_extended
if len(image_seq_lens_extended) == pixel_values.shape[0]
image_seq_lens
if len(image_seq_lens) == batched["pixel_values"].shape[0]
else None
),
)
)

if video_pixel_values_list:
video_pixel_values = torch.cat(video_pixel_values_list)
video_position_ids = (
torch.cat(video_position_ids_list)
if len(video_position_ids_list) == len(video_pixel_values_list)
else None
)
multimodal_embeddings.append(
self._get_image_features(
pixel_values=video_pixel_values,
image_position_ids=video_position_ids,
image_seq_lens=(
video_seq_lens_extended
if len(video_seq_lens_extended) == video_pixel_values.shape[0]
else None
),
if mixed_modality_runs:
# Own this run's storage before another call can replay a graph.
multimodal_embeddings[-1] = multimodal_embeddings[-1].clone(
memory_format=torch.contiguous_format
)
)

if audio_features_list and self.embed_audio is not None:
per_audio_embeddings = [
self._get_audio_features(audio_features, audio_features_mask_list[index])
for index, audio_features in enumerate(audio_features_list)
]
multimodal_embeddings.append(torch.cat(per_audio_embeddings, dim=0))

if not multimodal_embeddings:
raise ValueError("Gemma4 received active multimodal parameters without encoder inputs")
Expand Down
116 changes: 116 additions & 0 deletions tensorrt_llm/_torch/models/modeling_multimodal_mixin.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,10 @@
import contextlib
import copy
import itertools
from collections.abc import Mapping
from dataclasses import dataclass
from dataclasses import field as dataclass_field
from functools import partial
from typing import (
TYPE_CHECKING,
Any,
Expand Down Expand Up @@ -62,6 +65,119 @@ class MultimodalEncoderContractError(ValueError):
"""A request-local MM encoder input or output contract violation."""


def _batch_encoder_input(
values: Sequence[torch.Tensor], *, dim: int, pad_dims: tuple[int, ...], pad_value: float
) -> torch.Tensor:
if pad_dims and len(values) > 1:
maxima = {axis: max(value.shape[axis] for value in values) for axis in pad_dims}
padded = []
for value in values:
padding = [0] * (2 * value.ndim)
for axis, maximum in maxima.items():
padding[2 * (value.ndim - axis) - 1] = maximum - value.shape[axis]
padded.append(
torch.nn.functional.pad(value, padding, value=pad_value) if any(padding) else value
)
values = padded
return torch.cat(values, dim=dim)


@dataclass(frozen=True)
class EncoderInputField:
"""How one encoder argument combines already ordered input payloads.

``source_names`` lists aliases in lookup order; the first non-None value
is used. Only ``pad_dims`` may vary besides the concatenation dimension.
Padding extends the right side with ``pad_value``. Callers supply compatible
tensor ranks, dtypes, devices and remaining dimensions, and retain ownership
of modality order, optional-field policy and attention-mask semantics.
``batch`` applies the same padding and owned concatenation to a gathered column.
"""

source_names: tuple[str, ...]
dim: int = 0
pad_dims: tuple[int, ...] = ()
pad_value: float = 0.0
batch: Callable[[Sequence[torch.Tensor]], torch.Tensor] = dataclass_field(
init=False, repr=False, compare=False
)

def __post_init__(self) -> None:
if not self.source_names:
raise ValueError("Encoder input fields require at least one source name")
if self.dim < 0 or any(axis < 0 or axis == self.dim for axis in self.pad_dims):
raise ValueError("Encoder padding axes must be nonnegative and exclude the concat axis")
object.__setattr__(
self,
"batch",
partial(
_batch_encoder_input,
dim=self.dim,
pad_dims=self.pad_dims,
pad_value=self.pad_value,
)
if self.pad_dims or self.dim
else torch.cat,
)


def batch_encoder_inputs(
inputs: Sequence[Mapping[str, Any]],
fields: Mapping[str, EncoderInputField],
) -> dict[str, torch.Tensor]:
"""Gather, explicitly pad and concatenate encoder kwargs without reordering.

Inputs are read-only and the result always owns concatenated storage, even
for one input. Models may use this inside ``EncoderGroup.build_batched_input``
while existing grouped output restoration continues to use item metadata.
Lengths, grids and masks are never inferred from tensor shapes.
"""
if not inputs:
raise ValueError("Cannot batch empty encoder inputs")
result: dict[str, torch.Tensor] = {}
only = inputs[0] if len(inputs) == 1 else None
for name, field in fields.items():
source_name = field.source_names[0]
if len(field.source_names) == 1:
try:
values = (
[only[source_name]]
if only is not None
else [payload[source_name] for payload in inputs]
)
except KeyError as exc:
raise ValueError(
f"Missing encoder input {name!r}: expected {field.source_names}"
) from exc
elif len(field.source_names) == 2:
alternate = field.source_names[1]
try:
values = [
value if (value := payload.get(source_name)) is not None else payload[alternate]
for payload in inputs
]
except KeyError as exc:
raise ValueError(
f"Missing encoder input {name!r}: expected {field.source_names}"
) from exc
else:
values = []
for payload in inputs:
value = payload.get(source_name)
if value is None:
for alias in field.source_names[1:]:
value = payload.get(alias)
if value is not None:
break
if value is None:
raise ValueError(
f"Missing encoder input {name!r}: expected {field.source_names}"
)
values.append(value)
result[name] = field.batch(values)
return result


def _assemble_multimodal_encoder_embeddings(
item_tensors: Dict[int, torch.Tensor], total_items: int
) -> torch.Tensor:
Expand Down
Loading
Loading