Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
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
1 change: 1 addition & 0 deletions OPENENGINE_COMMIT
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
cea19cb06acf03c911b84d5c147e519b60dd92a6
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,9 @@ if(ENABLE_UCX)

target_link_libraries(${UCX_WRAPPER_TARGET}
PRIVATE $<LINK_LIBRARY:WHOLE_ARCHIVE,ucxx::ucxx>)
target_link_libraries(${UCX_WRAPPER_TARGET} PUBLIC ucxx::ucxx ucx::ucs)
# ucxx is already linked with WHOLE_ARCHIVE above. Repeating the same target
# without a feature is rejected by newer CMake versions.
target_link_libraries(${UCX_WRAPPER_TARGET} PUBLIC ucx::ucs)
target_link_libraries(${UCX_WRAPPER_TARGET} PUBLIC ${CUDA_RT_LIB})
target_link_libraries(${UCX_WRAPPER_TARGET} PUBLIC ${TORCH_LIBRARIES})
target_link_libraries(${UCX_WRAPPER_TARGET} PRIVATE ${ZMQ_LIBRARIES})
Expand Down
81 changes: 81 additions & 0 deletions scripts/install_openengine.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,81 @@
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

"""Verify and install the exact local OpenEngine Python sibling package."""

import argparse
import re
import runpy
import subprocess
import sys
from pathlib import Path


def _run(*args: str, cwd: Path) -> str:
return subprocess.run(
args,
cwd=cwd,
check=True,
capture_output=True,
text=True,
).stdout.strip()


def main() -> None:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument(
"--sibling",
type=Path,
help="OpenEngine checkout (default: ../openengine-trtllm)",
)
parser.add_argument(
"--verify-only",
action="store_true",
help="Verify the sibling without invoking pip",
)
args = parser.parse_args()

root = Path(__file__).resolve().parents[1]
sibling = (args.sibling or root.parent / "openengine-trtllm").resolve()
expected = (root / "OPENENGINE_COMMIT").read_text(encoding="utf-8").strip()
if re.fullmatch(r"[0-9a-f]{40}", expected) is None:
raise RuntimeError("OPENENGINE_COMMIT must contain one full lowercase Git SHA")
packaged = runpy.run_path(str(root / "tensorrt_llm" / "openengine" / "_schema_pin.py"))[
"OPENENGINE_COMMIT"
]
if packaged != expected:
raise RuntimeError(
f"Packaged OpenEngine pin is {packaged}, but OPENENGINE_COMMIT contains {expected}"
)

actual = _run("git", "rev-parse", "HEAD", cwd=sibling)
if actual != expected:
raise RuntimeError(f"OpenEngine sibling is at {actual}, but TensorRT-LLM pins {expected}")
dirty = _run(
"git",
"status",
"--porcelain",
"--",
"packages/python",
"proto",
cwd=sibling,
)
if dirty:
raise RuntimeError("OpenEngine Python/proto sources have uncommitted changes")

package = sibling / "packages" / "python"
if not (package / "pyproject.toml").is_file():
raise RuntimeError(f"OpenEngine Python package is missing: {package}")
if not args.verify_only:
subprocess.run(
[sys.executable, "-m", "pip", "install", "-e", str(package)],
cwd=root,
check=True,
)

print(f"Verified OpenEngine {expected}")
print(f"export OPENENGINE_SCHEMA_RELEASE={expected}")


if __name__ == "__main__":
main()
6 changes: 6 additions & 0 deletions tensorrt_llm/_torch/models/modeling_gemma4mm.py
Original file line number Diff line number Diff line change
Expand Up @@ -209,6 +209,12 @@ class Gemma4InputProcessor(BaseMultimodalInputProcessor, BaseMultimodalDummyInpu
# image, when present) can be split across chunks safely.
mm_bidirectional_blocks = False

def get_openengine_modalities(self) -> tuple[str, ...]:
return ("image", "audio")

def get_openengine_prefill_decode_modalities(self) -> tuple[str, ...]:
return ()

def __init__(
self,
model_path: str,
Expand Down
6 changes: 6 additions & 0 deletions tensorrt_llm/_torch/models/modeling_nemotron_nano.py
Original file line number Diff line number Diff line change
Expand Up @@ -954,6 +954,12 @@ def forward(self, multimodal_params: List[MultimodalParams]) -> List[torch.Tenso
class NanoV2VLInputProcessor(BaseMultimodalInputProcessor, BaseMultimodalDummyInputsBuilder):
supports_token_id_mm_expansion: ClassVar[bool] = True

def get_openengine_modalities(self) -> tuple[str, ...]:
return ("image", "video", "audio")

def get_openengine_prefill_decode_modalities(self) -> tuple[str, ...]:
return ()

def __init__(
self,
model_path: str,
Expand Down
33 changes: 32 additions & 1 deletion tensorrt_llm/_torch/models/modeling_phi4mm.py
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@
ExtraProcessedInputs, MultimodalPlaceholderMetadata,
MultimodalPlaceholderPlacement, TextPrompt,
register_input_processor)
from ...inputs.registry import MultimodalLoraSpec
from ...logger import logger
from ...lora_helper import LoraConfig
from ...sampling_params import SamplingParams
Expand Down Expand Up @@ -760,6 +761,36 @@ def forward(self, multimodal_params: List[MultimodalParams],
class Phi4MMInputProcessor(BaseMultimodalInputProcessor,
BaseMultimodalDummyInputsBuilder):

def get_model_owned_lora_identities(self) -> dict[str, int]:
return {"vision-lora": 0, "speech-lora": 1}

def get_openengine_modalities(self) -> tuple[str, ...]:
return ("image", "audio")

def get_openengine_prefill_decode_modalities(self) -> tuple[str, ...]:
return ()

def get_required_lora_spec(
self, modalities: tuple[str, ...]) -> MultimodalLoraSpec | None:
requested = set(modalities)
if {"image", "audio"}.issubset(requested):
raise ValueError(
"Phi-4 multimodal requests cannot combine image and audio because they "
"require different built-in LoRA adapters")
if "image" in requested:
return MultimodalLoraSpec(
name="vision-lora",
adapter_id=0,
path=os.path.join(self._model_path, "vision-lora"),
)
if "audio" in requested:
return MultimodalLoraSpec(
name="speech-lora",
adapter_id=1,
path=os.path.join(self._model_path, "speech-lora"),
)
return None

def __init__(self,
model_path: str,
config: transformers.PretrainedConfig,
Expand Down Expand Up @@ -1092,7 +1123,7 @@ def forward(
else:
raise NotImplementedError(
"Phi-4-multimodal does not support disaggregated inference yet. Please unset "
f"the TLLM_MULTIMODAL_DISAGGREGATED environment variable, or set it to '0'."
"the TLLM_MULTIMODAL_DISAGGREGATED environment variable, or set it to '0'."
)
mm_embedding = find_input_mm_embeds(
mm_embedding, multimodal_params[:num_context_requests])
Expand Down
6 changes: 6 additions & 0 deletions tensorrt_llm/_torch/models/modeling_qwen3vl.py
Original file line number Diff line number Diff line change
Expand Up @@ -402,6 +402,12 @@ def get_preferred_media_io_kwargs(self) -> Dict[str, Dict[str, Any]]:
# the per-frame CHW-float conversion in the IO loader.
return {"video": {"format": "np"}}

def get_openengine_modalities(self) -> tuple[str, ...]:
return ("image", "video")

def get_openengine_prefill_decode_modalities(self) -> tuple[str, ...]:
return ("image", "video")

def build_disagg_prefill_multimodal_inputs(
self, inputs: TextPrompt, mm_handles: List[Dict[str, Any]]
) -> DisaggPrefillMultimodalInputs:
Expand Down
3 changes: 2 additions & 1 deletion tensorrt_llm/_torch/pyexecutor/py_executor.py
Original file line number Diff line number Diff line change
Expand Up @@ -5405,7 +5405,8 @@ def _prepare_disagg_gen_init(self, fitting_disagg_gen_init_requests):
for resource_mgr_type in (
ResourceManagerType.KV_CACHE_MANAGER,
ResourceManagerType.SPEC_RESOURCE_MANAGER,
ResourceManagerType.DRAFT_KV_CACHE_MANAGER):
ResourceManagerType.DRAFT_KV_CACHE_MANAGER,
ResourceManagerType.PEFT_CACHE_MANAGER):
if (resource_mgr_type in self.resource_manager.resource_managers
and self.resource_manager.
resource_managers[resource_mgr_type] is not None):
Expand Down
5 changes: 4 additions & 1 deletion tensorrt_llm/_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -1223,7 +1223,10 @@ def _stored_block_to_json(data):
KVCacheEventSerializer._unique_tokens_to_json(token)
for token in data.tokens
],
# "lora_id": data.lora_id, # TODO (shreyasm): enable serialization of lora_id
"lora_id":
getattr(data, "lora_id", None),
"lora_name":
getattr(data, "lora_name", None),
"cache_salt":
data.cache_salt,
"cache_level":
Expand Down
Loading