From 72f828a4d08243e63638fb86aee586de0aa0479c Mon Sep 17 00:00:00 2001 From: Connor Carpenter Date: Sun, 12 Jul 2026 15:34:39 -0700 Subject: [PATCH 1/4] feat(openengine): add sibling protocol server Signed-off-by: Connor Carpenter --- OPENENGINE_COMMIT | 1 + scripts/install_openengine.py | 81 ++ .../_torch/models/modeling_gemma4mm.py | 6 + .../_torch/models/modeling_nemotron_nano.py | 6 + tensorrt_llm/_torch/models/modeling_phi4mm.py | 33 +- .../_torch/models/modeling_qwen3vl.py | 6 + tensorrt_llm/_utils.py | 5 +- tensorrt_llm/commands/serve.py | 150 +- tensorrt_llm/grpc/grpc_request_manager.py | 54 +- tensorrt_llm/inputs/registry.py | 118 +- tensorrt_llm/openengine/README.md | 21 + tensorrt_llm/openengine/__init__.py | 8 + tensorrt_llm/openengine/_schema_pin.py | 8 + tensorrt_llm/openengine/converters.py | 392 ++++++ tensorrt_llm/openengine/lora_registry.py | 120 ++ tensorrt_llm/openengine/server.py | 129 ++ tensorrt_llm/openengine/servicer.py | 1226 +++++++++++++++++ .../kv_cache_manager_v2/_event_manager.py | 9 + tensorrt_llm/serve/kv_event_fanout.py | 134 ++ tensorrt_llm/serve/openai_server.py | 41 +- tensorrt_llm/serve/request_tracker.py | 156 +++ tensorrt_llm/serve/stats_fanout.py | 81 ++ .../_torch/modeling/test_modeling_phi4mm.py | 17 + .../references/trtllm_serve_cli.yaml | 18 + tests/unittest/openengine/test_converters.py | 225 +++ tests/unittest/openengine/test_fanout.py | 105 ++ .../unittest/openengine/test_lora_registry.py | 89 ++ .../openengine/test_native_grpc_compat.py | 49 + .../openengine/test_request_tracker.py | 73 + tests/unittest/openengine/test_server.py | 86 ++ tests/unittest/openengine/test_servicer.py | 775 +++++++++++ 31 files changed, 4117 insertions(+), 105 deletions(-) create mode 100644 OPENENGINE_COMMIT create mode 100644 scripts/install_openengine.py create mode 100644 tensorrt_llm/openengine/README.md create mode 100644 tensorrt_llm/openengine/__init__.py create mode 100644 tensorrt_llm/openengine/_schema_pin.py create mode 100644 tensorrt_llm/openengine/converters.py create mode 100644 tensorrt_llm/openengine/lora_registry.py create mode 100644 tensorrt_llm/openengine/server.py create mode 100644 tensorrt_llm/openengine/servicer.py create mode 100644 tensorrt_llm/serve/kv_event_fanout.py create mode 100644 tensorrt_llm/serve/request_tracker.py create mode 100644 tensorrt_llm/serve/stats_fanout.py create mode 100644 tests/unittest/openengine/test_converters.py create mode 100644 tests/unittest/openengine/test_fanout.py create mode 100644 tests/unittest/openengine/test_lora_registry.py create mode 100644 tests/unittest/openengine/test_native_grpc_compat.py create mode 100644 tests/unittest/openengine/test_request_tracker.py create mode 100644 tests/unittest/openengine/test_server.py create mode 100644 tests/unittest/openengine/test_servicer.py diff --git a/OPENENGINE_COMMIT b/OPENENGINE_COMMIT new file mode 100644 index 000000000000..2fb5de7752ae --- /dev/null +++ b/OPENENGINE_COMMIT @@ -0,0 +1 @@ +cea19cb06acf03c911b84d5c147e519b60dd92a6 diff --git a/scripts/install_openengine.py b/scripts/install_openengine.py new file mode 100644 index 000000000000..89f8f3e3ba93 --- /dev/null +++ b/scripts/install_openengine.py @@ -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() diff --git a/tensorrt_llm/_torch/models/modeling_gemma4mm.py b/tensorrt_llm/_torch/models/modeling_gemma4mm.py index 0f5a340c0a9a..16b2129442eb 100644 --- a/tensorrt_llm/_torch/models/modeling_gemma4mm.py +++ b/tensorrt_llm/_torch/models/modeling_gemma4mm.py @@ -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, diff --git a/tensorrt_llm/_torch/models/modeling_nemotron_nano.py b/tensorrt_llm/_torch/models/modeling_nemotron_nano.py index 386b30716d08..e39376945942 100644 --- a/tensorrt_llm/_torch/models/modeling_nemotron_nano.py +++ b/tensorrt_llm/_torch/models/modeling_nemotron_nano.py @@ -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, diff --git a/tensorrt_llm/_torch/models/modeling_phi4mm.py b/tensorrt_llm/_torch/models/modeling_phi4mm.py index 2fbe96ac03c1..5e45078e47f6 100644 --- a/tensorrt_llm/_torch/models/modeling_phi4mm.py +++ b/tensorrt_llm/_torch/models/modeling_phi4mm.py @@ -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 @@ -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, @@ -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]) diff --git a/tensorrt_llm/_torch/models/modeling_qwen3vl.py b/tensorrt_llm/_torch/models/modeling_qwen3vl.py index 95da74e5d3bd..3943560deef0 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3vl.py +++ b/tensorrt_llm/_torch/models/modeling_qwen3vl.py @@ -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: diff --git a/tensorrt_llm/_utils.py b/tensorrt_llm/_utils.py index cf242cd205c9..c7605056a923 100644 --- a/tensorrt_llm/_utils.py +++ b/tensorrt_llm/_utils.py @@ -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": diff --git a/tensorrt_llm/commands/serve.py b/tensorrt_llm/commands/serve.py index bc949b3db948..8817a918adc6 100644 --- a/tensorrt_llm/commands/serve.py +++ b/tensorrt_llm/commands/serve.py @@ -122,8 +122,7 @@ def _signal_handler_cleanup_child(signum, frame): def is_non_default_or_required(param_name, value, backend, explicit_cli_keys): - """ - Check if a parameter should be explicitly included in llm_args. + """Check if a parameter should be explicitly included in llm_args. Returns True if parameter is either: 1. Always required (core params that must be present), OR @@ -387,7 +386,9 @@ def launch_server( served_model_name: Optional[str] = None, allow_request_chat_template: bool = False, num_input_processor_workers: int = 8, - num_media_load_workers: int = 8): + num_media_load_workers: int = 8, + openengine_host: str = "127.0.0.1", + openengine_port: Optional[int] = None): backend = llm_args["backend"] model = served_model_name or llm_args["model"] @@ -427,6 +428,51 @@ def launch_server( f"{backend} is not a known backend, check help for available options.", param_hint="backend") + request_tracker = None + kv_event_fanout = None + stats_fanout = None + openengine_server_cls = None + resolved_openengine_role = None + if openengine_port is not None: + from fastapi.responses import JSONResponse + + try: + from tensorrt_llm.openengine.server import (OpenEngineServer, + openengine_role) + except ModuleNotFoundError as error: + if error.name and error.name.startswith("openengine"): + llm.shutdown() + raise click.ClickException( + "OpenEngine support requires openengine-proto 0.2.0. " + "Run scripts/install_openengine.py against the pinned " + "local OpenEngine sibling checkout.") from error + raise + from tensorrt_llm.serve.kv_event_fanout import KvEventFanout + from tensorrt_llm.serve.request_tracker import (RequestTracker, + track_http_response) + from tensorrt_llm.serve.stats_fanout import StatsFanout + + request_tracker = RequestTracker(llm) + kv_cache_config = getattr(llm.args, "kv_cache_config", None) + event_buffer_size = getattr(kv_cache_config, + "event_buffer_max_size", 0) + if event_buffer_size > 0: + kv_event_fanout = KvEventFanout(llm, + buffer_size=event_buffer_size) + stats_buffer_size = getattr(llm.args, "iter_stats_max_iterations", + 1000) + if stats_buffer_size is None or stats_buffer_size == 0: + stats_buffer_size = 1000 + elif stats_buffer_size < 0: + stats_buffer_size = None + stats_fanout = StatsFanout(llm, buffer_size=stats_buffer_size) + openengine_server_cls = OpenEngineServer + try: + resolved_openengine_role = openengine_role(server_role) + except ValueError: + llm.shutdown() + raise + server = OpenAIServer( generator=llm, model=model, @@ -438,22 +484,76 @@ def launch_server( chat_template=chat_template, allow_request_chat_template=allow_request_chat_template, input_processor_workers=num_input_processor_workers, - media_load_workers=num_media_load_workers) + media_load_workers=num_media_load_workers, + kv_event_fanout=kv_event_fanout, + stats_fanout=stats_fanout, + shutdown_generator=openengine_server_cls is None) _apply_fastapi_middlewares(server.app, middleware) + if request_tracker is not None: + + @server.app.middleware("http") + async def openengine_process_admission(request, call_next): + if request.method != "POST": + return await call_next(request) + try: + request_tracker.begin_external() + if stats_fanout is not None: + stats_fanout.wake() + except RuntimeError: + return JSONResponse( + status_code=503, + content={"error": "TensorRT-LLM is draining"}) + response = None + try: + response = await call_next(request) + finally: + if response is None: + await request_tracker.finish_external() + return await track_http_response(response, request_tracker) + # Optionally disable GC (default: not disabled) if os.getenv("TRTLLM_SERVER_DISABLE_GC", "0") == "1": gc.disable() - uvloop.run(server(host, port, sockets=[s])) + if openengine_server_cls is None: + uvloop.run(server(host, port, sockets=[s])) + else: + + async def serve_http_and_openengine(): + openengine_server = None + try: + openengine_server = openengine_server_cls( + llm=llm, + model=model, + role=resolved_openengine_role, + host=openengine_host, + port=openengine_port, + tracker=request_tracker, + media_config=multimodal_server_config, + reasoning_parser=getattr(llm.args, "reasoning_parser", + None), + tool_parser=tool_parser, + kv_event_fanout=kv_event_fanout, + stats_fanout=stats_fanout, + ) + await openengine_server.start() + await server(host, port, sockets=[s]) + finally: + try: + if openengine_server is not None: + await openengine_server.stop() + finally: + llm.shutdown() + + uvloop.run(serve_http_and_openengine()) def launch_grpc_server(host: str, port: int, llm_args: dict, served_model_name: Optional[str] = None): - """ - Launch a gRPC server for TensorRT-LLM. + """Launch a gRPC server for TensorRT-LLM. This provides a high-performance gRPC interface designed for external routers (e.g., sgl-router) using pre-tokenized input and raw token ID output. @@ -1075,6 +1175,18 @@ def convert(self, value: Any, param: Optional["click.Parameter"], help="Run gRPC server instead of OpenAI HTTP server. " "gRPC server accepts pre-tokenized requests and returns raw token IDs.", status="prototype") +@stability_option("--openengine-host", + type=str, + default="127.0.0.1", + help="Host for the optional OpenEngine sibling gRPC server.", + status="prototype") +@stability_option( + "--openengine-port", + type=int, + default=None, + help= + "Port for the optional OpenEngine sibling gRPC server. Disabled when absent.", + status="prototype") @stability_option( "--served_model_name", type=str, @@ -1132,13 +1244,18 @@ def serve( telemetry: bool, custom_module_dirs: list[Path], chat_template: Optional[str], allow_request_chat_template: bool, middleware: tuple[str, ...], grpc: bool, enable_visual_gen: bool, - served_model_name: Optional[str], visual_gen_args: Optional[str]): + served_model_name: Optional[str], visual_gen_args: Optional[str], + openengine_host: str, openengine_port: Optional[int]): """Running an OpenAI API compatible server MODEL: model name | HF checkpoint path | TensorRT engine path """ logger.set_level(log_level) + if grpc and openengine_port is not None: + raise click.UsageError( + "--openengine-port cannot be combined with dedicated --grpc mode") + if moe_cluster_parallel_size is not None: logger.warning( "--moe_cluster_parallel_size / --cluster_size is deprecated and " @@ -1277,6 +1394,12 @@ def _serve_llm(): multimodal_server_config = MultimodalServerConfig( media_io_kwargs=parsed_media_io_kwargs) + if openengine_port is not None and server_role is not None: + role_name = getattr(server_role, "name", str(server_role)).upper() + if role_name in ("MM_ENCODER", "VISUAL_GEN"): + raise click.UsageError( + f"OpenEngine does not support server role {role_name!r}") + if grpc: # gRPC mode: launch gRPC server instead of OpenAI HTTP server # Check for unsupported arguments that are silently ignored in gRPC mode @@ -1323,7 +1446,9 @@ def _serve_llm(): served_model_name=served_model_name, allow_request_chat_template=allow_request_chat_template, num_input_processor_workers=num_input_processor_workers, - num_media_load_workers=num_media_load_workers) + num_media_load_workers=num_media_load_workers, + openengine_host=openengine_host, + openengine_port=openengine_port) def _serve_visual_gen(): parsed_visual_gen_args = (VisualGenArgs.from_yaml(visual_gen_args) @@ -1338,6 +1463,9 @@ def _serve_visual_gen(): is_visual_gen = (enable_visual_gen or visual_gen_args is not None or get_is_diffusion_only_model(model)) if is_visual_gen: + if openengine_port is not None: + raise click.UsageError( + "--openengine-port is not supported by the VisualGen server") _serve_visual_gen() else: _serve_llm() @@ -1665,7 +1793,6 @@ def disaggregated( schedule_style: str, ): """Running server in disaggregated mode""" - logger.set_level(log_level) if metrics_log_interval != 0: @@ -1756,7 +1883,6 @@ def set_cuda_device(): status="beta") def disaggregated_mpi_worker(config_file: Optional[str], log_level: str): """Launching disaggregated MPI worker""" - from tensorrt_llm._utils import mpi_rank if os.environ.get(DisaggLauncherEnvs. TLLM_DISAGG_RUN_REMOTE_MPI_SESSION_CLIENT) != "1": @@ -1944,7 +2070,7 @@ def _launch_disaggregated_leader(sub_comm, instance_idx: int, config_file: str, f"Child process {_child_p_global.pid} failed to be killed even after 30s." ) assert _child_p_global.poll( - ) is not None, f"the subprocess should be terminated" + ) is not None, "the subprocess should be terminated" # Check if the process was launched and assert it's terminated if _child_p_global and hasattr(_child_p_global, diff --git a/tensorrt_llm/grpc/grpc_request_manager.py b/tensorrt_llm/grpc/grpc_request_manager.py index ca2e59eb1f36..9cb4868682c6 100644 --- a/tensorrt_llm/grpc/grpc_request_manager.py +++ b/tensorrt_llm/grpc/grpc_request_manager.py @@ -34,6 +34,7 @@ from tensorrt_llm.llmapi.llm_utils import KvCacheRetentionConfig from tensorrt_llm.logger import logger from tensorrt_llm.sampling_params import GuidedDecodingParams, SamplingParams +from tensorrt_llm.serve.request_tracker import RequestTracker from . import trtllm_service_pb2 as pb2 @@ -59,8 +60,9 @@ def __init__(self, llm: Any): llm: The TensorRT-LLM LLM instance (tensorrt_llm.LLM or tensorrt_llm._tensorrt_engine.LLM) """ self.llm = llm - # Track active requests: request_id -> GenerationResult - self._rid_to_result: Dict[str, GenerationResult] = {} + self._request_tracker = RequestTracker(llm) + # Kept as an alias for callers which inspect native gRPC state. + self._rid_to_result = self._request_tracker.active_requests logger.info("GrpcRequestManager initialized") @@ -93,6 +95,7 @@ async def generate( GenerationResult objects containing token IDs (text will be empty because detokenize=False) """ + admitted = False try: # Submit to LLM.generate_async which returns a GenerationResult # that is an async iterator @@ -111,7 +114,15 @@ async def generate( ) # Track the result for potential abort - self._rid_to_result[request_id] = gen_result + try: + self._request_tracker.admit(request_id, gen_result) + except (RuntimeError, ValueError): + try: + gen_result.abort() + except (RuntimeError, AssertionError) as error: + logger.warning("Failed to abort rejected request %s: %s", request_id, error) + raise + admitted = True # Iterate over the async generator # GenerationResult implements __aiter__ and __anext__ @@ -130,7 +141,8 @@ async def generate( raise finally: # Cleanup tracking - self._rid_to_result.pop(request_id, None) + if admitted: + await self._request_tracker.finish(request_id) async def abort(self, request_id: str) -> bool: """Abort a running request. @@ -141,22 +153,13 @@ async def abort(self, request_id: str) -> bool: Returns: True if request was found and aborted, False otherwise """ - gen_result = self._rid_to_result.get(request_id) - - if gen_result is None: + if request_id not in self._rid_to_result: logger.debug(f"Abort: request {request_id} not found (may have already completed)") return False - - try: - # GenerationResult has an abort() method - gen_result.abort() - self._rid_to_result.pop(request_id, None) + aborted = await self._request_tracker.abort(request_id) + if aborted: logger.info(f"Request {request_id} aborted") - return True - except Exception as e: - logger.error(f"Error aborting request {request_id}: {e}") - self._rid_to_result.pop(request_id, None) - return False + return aborted async def health_check(self) -> Tuple[bool, str]: """Check if the engine is healthy. @@ -168,22 +171,7 @@ async def health_check(self) -> Tuple[bool, str]: if self.llm is None: return False, "LLM not initialized" - # Check executor health (includes error queue, MPI worker - # liveness, and fatal error state — not just doing_shutdown). - if hasattr(self.llm, "_executor"): - if self.llm._executor is None: - return False, "Executor is not available" - if not self.llm._executor.check_health(): - error_msg = "Executor is unhealthy" - if self.llm._executor._fatal_error is not None: - exc = self.llm._executor._fatal_error - lines = str(exc).splitlines() - short = (lines[0] if lines else type(exc).__name__)[:200] - error_msg = f"{type(exc).__name__}: {short}" - logger.error(f"Health check fatal error: {repr(exc)}") - return False, error_msg - - return True, "OK" + return await self._request_tracker.health() except Exception as e: logger.error(f"Health check error: {e}") return False, f"Error: {e}" diff --git a/tensorrt_llm/inputs/registry.py b/tensorrt_llm/inputs/registry.py index 58ca102d35be..ec4b6df03478 100644 --- a/tensorrt_llm/inputs/registry.py +++ b/tensorrt_llm/inputs/registry.py @@ -28,9 +28,17 @@ ExtraProcessedInputs = Dict[str, Any] +@dataclass(frozen=True) +class MultimodalLoraSpec: + """Model-required LoRA adapter for a set of input modalities.""" + + name: str + adapter_id: int + path: str + + class InputProcessor(Protocol): - """ - Protocol for InputProcessor classes. + """Protocol for InputProcessor classes. InputProcessor's functions are more relevant to multimodal use cases: - Preprocess: extra steps to manipulate the prompts. - Forward: the main logic to process the inputs. In multimodal cases, this may run a multimodal encoder model. @@ -121,8 +129,7 @@ def __call__( class BaseMultimodalInputProcessor(ABC): - """ - Base class for multimodal input processors with default implementations + """Base class for multimodal input processors with default implementations of get_num_tokens_per_image and get_num_tokens_per_video methods. This class provides default implementations that work with most AutoProcessor-based @@ -150,6 +157,24 @@ class BaseMultimodalInputProcessor(ABC): # inputs to `call_with_token_ids` instead of detokenizing upstream. supports_token_id_mm_expansion: ClassVar[bool] = False + def get_openengine_modalities(self) -> tuple[str, ...]: + """Modalities this processor accepts through aggregated OpenEngine.""" + return ("image", ) + + def get_openengine_prefill_decode_modalities(self) -> tuple[str, ...]: + """Modalities supported in context-first OpenEngine P/D requests.""" + return () + + def get_required_lora_spec( + self, modalities: tuple[str, ...]) -> MultimodalLoraSpec | None: + """Return a model-owned adapter required to process these modalities.""" + del modalities + return None + + def get_model_owned_lora_identities(self) -> dict[str, int]: + """Names and IDs reserved for adapters selected by the model processor.""" + return {} + def __init__(self, model_path, config, @@ -373,16 +398,14 @@ class flag and implement both of the following hooks; otherwise raises @property def use_fast(self) -> bool: - """ - Whether to use fast tokenizer for AutoProcessor. + """Whether to use fast tokenizer for AutoProcessor. Default is True for most multimodal models. """ return self._use_fast @property def multimodal_hashing_supported(self) -> Optional[bool]: - """ - Whether multimodal hashing is supported for this processor. + """Whether multimodal hashing is supported for this processor. Returns None if unknown (will be detected at runtime), True if supported, False if not supported. @@ -474,8 +497,7 @@ def get_mm_special_token_ids(self) -> Optional[Tensor]: @property def get_num_multimodal_tokens(self): - """ - Get the Hugging Face processor's '_get_num_multimodal_tokens' method. + """Get the Hugging Face processor's '_get_num_multimodal_tokens' method. """ if hasattr(self.processor, '_get_num_multimodal_tokens'): return self.processor._get_num_multimodal_tokens @@ -491,8 +513,7 @@ def get_num_tokens_per_image( image: Union[Image.Image, torch.Tensor], **kwargs, ): - """ - Calculate the number of tokens generated for an image. + """Calculate the number of tokens generated for an image. Delegates to the Hugging Face processor's ``_get_num_multimodal_tokens``. @@ -521,8 +542,7 @@ def get_num_tokens_per_video( video_audio: Optional[Any] = None, **kwargs, ): - """ - Calculate the number of tokens generated for a video. + """Calculate the number of tokens generated for a video. Delegates to the Hugging Face processor's ``_get_num_multimodal_tokens``; a fallback treats the video as a stack of frames if the HF processor @@ -643,10 +663,9 @@ def get_dummy_mm_data_for_tokens( class MultimodalPlaceholderPlacement(enum.Enum): - """ - The placement of the multimodal placeholder in the prompt. Valid values are: - - BEFORE_TEXT: the placeholders are placed before the text prompt. - - AFTER_TEXT: the placeholders are placed after the text prompt. + """The placement of the multimodal placeholder in the prompt. Valid values are: + - BEFORE_TEXT: the placeholders are placed before the text prompt. + - AFTER_TEXT: the placeholders are placed after the text prompt. """ INVALID = -1 BEFORE_TEXT = 0 @@ -655,31 +674,30 @@ class MultimodalPlaceholderPlacement(enum.Enum): @dataclass(frozen=True) class MultimodalPlaceholderMetadata: - """ - Metadata for the multimodal placeholder. It has 5 components: - - placeholder_map: - A mapping from modality to placeholder string. - Modality can be "image", "video", "audio", etc. - - placeholder_placement: - The placement of the placeholders, e.g. before or after the text prompt. - Only used when interleave_placeholders is False (the default). - Ignored when interleave_placeholders is True. - - placeholders_separator: - The separator between the placeholders, e.g. some models use "\n" to separate the placeholders. - - content_format: - Optional override for the content format expected by the chat template. - ContentFormat.OPENAI means the template handles multimodal content dicts natively. - ContentFormat.STRING means the template expects plain string content. - ContentFormat.PASSTHROUGH skips chat template rendering entirely. - None means auto-detect at runtime via Jinja AST analysis. - - interleave_placeholders: - When True and content_parts is available, placeholders are inserted - at the exact media positions within the text (interleaved). - In this mode, placeholder_placement is ignored - the position of - each placeholder is determined by where the media appears in the - user's message. - When False (default), placeholders are bulk-prepended or appended - according to placeholder_placement. + """Metadata for the multimodal placeholder. It has 5 components: + - placeholder_map: + A mapping from modality to placeholder string. + Modality can be "image", "video", "audio", etc. + - placeholder_placement: + The placement of the placeholders, e.g. before or after the text prompt. + Only used when interleave_placeholders is False (the default). + Ignored when interleave_placeholders is True. + - placeholders_separator: + The separator between the placeholders, e.g. some models use "\n" to separate the placeholders. + - content_format: + Optional override for the content format expected by the chat template. + ContentFormat.OPENAI means the template handles multimodal content dicts natively. + ContentFormat.STRING means the template expects plain string content. + ContentFormat.PASSTHROUGH skips chat template rendering entirely. + None means auto-detect at runtime via Jinja AST analysis. + - interleave_placeholders: + When True and content_parts is available, placeholders are inserted + at the exact media positions within the text (interleaved). + In this mode, placeholder_placement is ignored - the position of + each placeholder is determined by where the media appears in the + user's message. + When False (default), placeholders are bulk-prepended or appended + according to placeholder_placement. """ placeholder_map: Dict[str, str] = field(default_factory=dict) placeholder_placement: MultimodalPlaceholderPlacement = MultimodalPlaceholderPlacement.AFTER_TEXT @@ -689,8 +707,7 @@ class MultimodalPlaceholderMetadata: class MultimodalPlaceholderRegistry: - """ - Registry for the multimodal models to keep track of the placeholder information. + """Registry for the multimodal models to keep track of the placeholder information. """ def __init__(self) -> None: @@ -806,8 +823,7 @@ def __init__(self) -> None: def support_multimodal_disaggregated(model_cls: Type[nn.Module]): - """ - Model-class decorator to declare support for multimodal disaggregated inputs. + """Model-class decorator to declare support for multimodal disaggregated inputs. Apply this to a model class AFTER its input processor has been registered via @register_input_processor. The decorator will locate the processor class, @@ -840,10 +856,9 @@ def register_input_processor( processor_cls: Type[InputProcessor], model_type: str, placeholder_metadata: MultimodalPlaceholderMetadata = None): - """ - Register an input processor to a model class. + """Register an input processor to a model class. - NOTE: + Note: 1. Since this API is only used for multimodal models, we are checking the model type only for that. 2. If this is used for other models in the future, this logic needs to be @@ -1066,8 +1081,7 @@ def create_input_processor_with_hash( def multimodal_hashing_process( inputs: TextPrompt, sampling_params: SamplingParams ) -> Tuple[List[int], Optional[ExtraProcessedInputs]]: - """ - Process multimodal hashing for media tokens if possible. + """Process multimodal hashing for media tokens if possible. Delegates the raw `(token_ids, extra_processed_inputs)` production to the input processor's `__call__` (which itself dispatches between the diff --git a/tensorrt_llm/openengine/README.md b/tensorrt_llm/openengine/README.md new file mode 100644 index 000000000000..4a7063ea4de6 --- /dev/null +++ b/tensorrt_llm/openengine/README.md @@ -0,0 +1,21 @@ + + + +# TensorRT-LLM OpenEngine sibling server + +`trtllm-serve` can expose OpenEngine alongside its normal HTTP server with +`--openengine-port`. The feature remains disabled when the flag is absent and +keeps the OpenEngine bindings out of TensorRT-LLM's required dependencies. + +The exact sibling source revision is recorded in `OPENENGINE_COMMIT`. Verify +and install its generated Python package with: + +```bash +python scripts/install_openengine.py +python -m pip install -e . +``` + +The installer rejects a different or dirty sibling package/proto checkout. It +prints the required `OPENENGINE_SCHEMA_RELEASE` export; sibling startup fails +closed unless that value exactly matches `OPENENGINE_COMMIT`. No registry +publication is assumed. diff --git a/tensorrt_llm/openengine/__init__.py b/tensorrt_llm/openengine/__init__.py new file mode 100644 index 000000000000..ca230e3fecdd --- /dev/null +++ b/tensorrt_llm/openengine/__init__.py @@ -0,0 +1,8 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""OpenEngine server integration for TensorRT-LLM.""" + +from .server import OpenEngineServer + +__all__ = ["OpenEngineServer"] diff --git a/tensorrt_llm/openengine/_schema_pin.py b/tensorrt_llm/openengine/_schema_pin.py new file mode 100644 index 000000000000..8a5cfa30ece3 --- /dev/null +++ b/tensorrt_llm/openengine/_schema_pin.py @@ -0,0 +1,8 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""OpenEngine source identity embedded in TensorRT-LLM distributions.""" + +OPENENGINE_COMMIT = "cea19cb06acf03c911b84d5c147e519b60dd92a6" + +__all__ = ["OPENENGINE_COMMIT"] diff --git a/tensorrt_llm/openengine/converters.py b/tensorrt_llm/openengine/converters.py new file mode 100644 index 000000000000..e934a5059138 --- /dev/null +++ b/tensorrt_llm/openengine/converters.py @@ -0,0 +1,392 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Conversions between OpenEngine messages and TensorRT-LLM API objects.""" + +import asyncio +import base64 +import hashlib +import json +import re +from dataclasses import asdict, is_dataclass +from typing import Any + +from google.protobuf.json_format import MessageToDict +from openengine.v1 import generation_pb2, input_pb2, kv_pb2 + +from tensorrt_llm.disaggregated_params import DisaggregatedParams, DisaggScheduleStyle +from tensorrt_llm.executor.request import LoRARequest +from tensorrt_llm.executor.result import Logprob +from tensorrt_llm.inputs.media_io import MEDIA_IO_REGISTRY +from tensorrt_llm.inputs.multimodal import MultimodalServerConfig +from tensorrt_llm.sampling_params import GuidedDecodingParams, SamplingParams + +HANDOFF_ATTRIBUTE = "tensorrt_llm.disaggregated_params.v1" +_MODALITY_NAMES = { + input_pb2.MODALITY_UNSPECIFIED: "image", + input_pb2.MODALITY_IMAGE: "image", + input_pb2.MODALITY_VIDEO: "video", + input_pb2.MODALITY_AUDIO: "audio", +} + + +def _optional(message: object, field: str) -> Any | None: + return getattr(message, field) if message.HasField(field) else None + + +def _candidate_count(selection: object, enabled: bool) -> int | None: + if not enabled: + return None + kind = selection.WhichOneof("selection") + if kind in (None, "top_n"): + return selection.top_n if kind == "top_n" else 0 + raise ValueError("TensorRT-LLM OpenEngine supports top_n logprob selection only") + + +def to_sampling_params(request: generation_pb2.GenerateRequest) -> SamplingParams: + """Build TRT-LLM sampling params without inventing wire defaults.""" + sampling = request.sampling + stopping = request.stopping + response = request.response + kwargs: dict[str, Any] = { + "max_tokens": (32 if _optional(stopping, "max_tokens") is None else stopping.max_tokens), + "detokenize": True, + } + field_map = { + "temperature": "temperature", + "top_p": "top_p", + "top_k": "top_k", + "min_p": "min_p", + "frequency_penalty": "frequency_penalty", + "presence_penalty": "presence_penalty", + "repetition_penalty": "repetition_penalty", + "seed": "seed", + } + for proto_name, trt_name in field_map.items(): + value = _optional(sampling, proto_name) + if value is not None: + kwargs[trt_name] = value + num_sequences = _optional(sampling, "num_sequences") + if num_sequences is not None: + kwargs["n"] = num_sequences + kwargs["best_of"] = num_sequences + min_tokens = _optional(stopping, "min_tokens") + if min_tokens is not None: + kwargs["min_tokens"] = min_tokens + ignore_eos = _optional(stopping, "ignore_eos") + if ignore_eos is not None: + kwargs["ignore_eos"] = ignore_eos + include_stop = _optional(stopping, "include_stop_in_output") + if include_stop is not None: + kwargs["include_stop_str_in_output"] = include_stop + stop_text: list[str] = [] + stop_ids: list[int] = [] + for condition in stopping.conditions: + kind = condition.WhichOneof("condition") + if kind == "stop_text": + stop_text.append(condition.stop_text) + elif kind == "stop_token_id": + stop_ids.append(condition.stop_token_id) + if stop_text: + kwargs["stop"] = stop_text + if stop_ids: + kwargs["stop_token_ids"] = stop_ids + + prompt_requested = bool(_optional(response, "return_prompt_logprobs")) + output_requested = bool(_optional(response, "return_output_logprobs")) + prompt_count = _candidate_count(response.prompt_candidates, prompt_requested) + output_count = _candidate_count(response.output_candidates, output_requested) + if prompt_count is not None: + kwargs["prompt_logprobs"] = prompt_count + if output_count is not None: + kwargs["logprobs"] = output_count + prompt_start = _optional(response, "prompt_logprob_start") + if prompt_start not in (None, 0): + raise ValueError("TensorRT-LLM does not support prompt_logprob_start through OpenEngine") + + guide_kind = request.guided.WhichOneof("guide") + if guide_kind: + guide_kwargs: dict[str, Any] + if guide_kind == "json_schema": + guide_kwargs = {"json": request.guided.json_schema} + elif guide_kind == "regex": + guide_kwargs = {"regex": request.guided.regex} + elif guide_kind == "ebnf_grammar": + guide_kwargs = {"grammar": request.guided.ebnf_grammar} + elif guide_kind == "structural_tag": + guide_kwargs = {"structural_tag": request.guided.structural_tag} + elif guide_kind == "choice": + choices = [f"(?:{re.escape(choice)})" for choice in request.guided.choice.choices] + guide_kwargs = {"regex": "^(?:" + "|".join(choices) + ")$"} + elif guide_kind == "json_object": + guide_kwargs = {"json_object": True} + else: + raise ValueError(f"Unsupported guided decoding mode {guide_kind!r}") + kwargs["guided_decoding"] = GuidedDecodingParams(**guide_kwargs) + return SamplingParams(**kwargs) + + +def to_priority(request: generation_pb2.GenerateRequest) -> float: + """Map signed OpenEngine ordering into TRT-LLM's bounded priority domain.""" + if not request.HasField("priority"): + return 0.5 + priority = request.priority + return 0.5 + 0.5 * priority / (1 + abs(priority)) + + +def _message_struct(struct_message: object) -> dict[str, Any]: + if not struct_message.fields: + return {} + return MessageToDict(struct_message, preserving_proto_field_name=True) + + +def validate_media_options(options: dict[str, Any]) -> None: + unknown = set(options).difference(MEDIA_IO_REGISTRY) + if unknown: + raise ValueError(f"Unknown media option modalities: {sorted(unknown)}") + invalid = [name for name, value in options.items() if not isinstance(value, dict)] + if invalid: + raise ValueError(f"Media options must be objects for modalities: {sorted(invalid)}") + + +async def load_media( + media: list[object], + media_options: object, + server_config: MultimodalServerConfig | None, +) -> dict[str, list[Any]] | None: + """Decode ordered OpenEngine media with TRT-LLM's media-I/O merge rules.""" + if not media: + return None + request_options = _message_struct(media_options) + validate_media_options(request_options) + server_options = (server_config.media_io_kwargs if server_config is not None else None) or {} + output: dict[str, list[Any]] = {} + pending: list[tuple[str, asyncio.Future[Any] | asyncio.Task[Any] | Any]] = [] + for item in media: + modality = _MODALITY_NAMES.get(item.modality) + if modality is None: + raise ValueError(f"Unsupported media modality {item.modality}") + media_io = MEDIA_IO_REGISTRY[modality].create( + server_options.get(modality), request_options.get(modality) + ) + source = item.WhichOneof("source") + if source in ("url", "data_uri"): + pending.append((modality, media_io.async_load(getattr(item, source)))) + elif source == "raw_bytes": + pending.append( + (modality, media_io._run_in_executor(media_io.load_bytes, item.raw_bytes)) + ) + else: + raise ValueError("Each media item must carry exactly one source") + decoded = await asyncio.gather(*(awaitable for _, awaitable in pending)) + for (modality, _), value in zip(pending, decoded): + output.setdefault(modality, []).append(value) + return output + + +def stable_request_id(request_id: str) -> int: + """Map arbitrary wire request IDs into TRT-LLM's positive int64 domain.""" + digest = hashlib.sha256(request_id.encode("utf-8")).digest() + return int.from_bytes(digest[:8], "big") & ((1 << 63) - 1) + + +def _json_value(value: Any) -> Any: + if is_dataclass(value): + return _json_value(asdict(value)) + if isinstance(value, dict): + return {str(key): _json_value(item) for key, item in value.items()} + if isinstance(value, (list, tuple)): + return [_json_value(item) for item in value] + if isinstance(value, bytes): + return base64.b64encode(value).decode("ascii") + if hasattr(value, "logprob"): + output = {"logprob": float(value.logprob)} + if getattr(value, "rank", None) is not None: + output["rank"] = int(value.rank) + return output + if hasattr(value, "item"): + return value.item() + return value + + +def encode_handoff( + params: DisaggregatedParams, *, requires_decode_media: bool = False +) -> kv_pb2.KvSessionRef: + """Encode the revisioned context-first TRT handoff profile.""" + unsupported = { + "first-generation logits": params.first_gen_logits, + "multimodal embedding handles": params.multimodal_embedding_handles, + "multimodal hashes": params.multimodal_hashes, + } + if not requires_decode_media: + unsupported.update( + { + "mrope position IDs": params.mrope_position_ids_handle, + "mrope position deltas": params.mrope_position_deltas_handle, + } + ) + present = [name for name, value in unsupported.items() if value is not None] + if present: + raise ValueError("OpenEngine context-first handoff does not support " + ", ".join(present)) + if params.schedule_style not in (None, DisaggScheduleStyle.CONTEXT_FIRST): + raise ValueError("Generation-first handoff is not supported") + payload = { + "first_gen_tokens": _json_value(params.first_gen_tokens), + "first_gen_log_probs": _json_value(params.first_gen_log_probs), + "ctx_request_id": None if params.ctx_request_id is None else str(params.ctx_request_id), + "disagg_request_id": None + if params.disagg_request_id is None + else str(params.disagg_request_id), + "ctx_dp_rank": params.ctx_dp_rank, + "ctx_info_endpoint": params.ctx_info_endpoint, + "draft_tokens": _json_value(params.draft_tokens), + "ctx_usage": _json_value(params.ctx_usage), + "conversation_id": params.conversation_id, + "schedule_style": "context_first", + "requires_decode_media": requires_decode_media, + "opaque_state": ( + None + if params.opaque_state is None + else base64.b64encode(params.opaque_state).decode("ascii") + ), + } + canonical = json.dumps(payload, sort_keys=True, separators=(",", ":"), ensure_ascii=True) + session = kv_pb2.KvSessionRef( + session_id=str(params.disagg_request_id or params.ctx_request_id or ""), + transfer_backend="tensorrt_llm", + dp_rank=params.ctx_dp_rank or 0, + ) + session.attributes_struct[HANDOFF_ATTRIBUTE] = canonical + return session + + +def _decode_handoff_payload(session: kv_pb2.KvSessionRef) -> dict[str, Any]: + if HANDOFF_ATTRIBUTE not in session.attributes_struct: + raise ValueError(f"KV session is missing {HANDOFF_ATTRIBUTE!r}") + encoded = session.attributes_struct[HANDOFF_ATTRIBUTE] + if not isinstance(encoded, str): + raise ValueError("TensorRT-LLM handoff attribute must be a JSON string") + try: + payload = json.loads(encoded) + except json.JSONDecodeError as error: + raise ValueError("TensorRT-LLM handoff is not valid JSON") from error + if not isinstance(payload, dict): + raise ValueError("TensorRT-LLM handoff must contain a JSON object") + return payload + + +def handoff_requires_decode_media(session: kv_pb2.KvSessionRef) -> bool: + """Return whether decode must recompute transient MM state from raw media.""" + payload = _decode_handoff_payload(session) + value = payload.get("requires_decode_media", False) + if not isinstance(value, bool): + raise ValueError("requires_decode_media must be a boolean") + return value + + +def decode_handoff(session: kv_pb2.KvSessionRef) -> DisaggregatedParams: + """Decode and validate the revisioned context-first TRT handoff profile.""" + payload = _decode_handoff_payload(session) + if payload.get("schedule_style", "context_first") != "context_first": + raise ValueError("Generation-first handoff is not supported") + if any( + payload.get(key) is not None + for key in ( + "first_gen_logits", + "multimodal_embedding_handles", + "multimodal_hashes", + "mrope_position_ids_handle", + "mrope_position_deltas_handle", + ) + ): + raise ValueError("Encoder-stage and first-generation logits handles are not supported") + + def decimal_id(name: str) -> int | None: + value = payload.get(name) + if value is None: + return None + if not isinstance(value, str) or not value.isdecimal(): + raise ValueError(f"{name} must be a decimal string") + return int(value) + + opaque = payload.get("opaque_state") + try: + opaque_state = None if opaque is None else base64.b64decode(opaque, validate=True) + except (ValueError, TypeError) as error: + raise ValueError("opaque_state must be canonical base64") from error + + ctx_usage = payload.get("ctx_usage") + if ctx_usage is not None: + if not isinstance(ctx_usage, dict): + raise ValueError("ctx_usage must be an object") + prompt_tokens = ctx_usage.get("prompt_tokens") + if ( + not isinstance(prompt_tokens, int) + or isinstance(prompt_tokens, bool) + or prompt_tokens < 0 + ): + raise ValueError("ctx_usage.prompt_tokens must be a non-negative integer") + prompt_details = ctx_usage.get("prompt_tokens_details") + if prompt_details is not None: + if not isinstance(prompt_details, dict): + raise ValueError("ctx_usage.prompt_tokens_details must be an object") + cached_tokens = prompt_details.get("cached_tokens", 0) + if ( + not isinstance(cached_tokens, int) + or isinstance(cached_tokens, bool) + or cached_tokens < 0 + ): + raise ValueError( + "ctx_usage.prompt_tokens_details.cached_tokens must be a non-negative integer" + ) + + def logprobs() -> list[Any] | None: + encoded_logprobs = payload.get("first_gen_log_probs") + if encoded_logprobs is None: + return None + if not isinstance(encoded_logprobs, list): + raise ValueError("first_gen_log_probs must be a list") + decoded: list[Any] = [] + for position in encoded_logprobs: + if isinstance(position, (int, float)): + decoded.append(float(position)) + continue + if not isinstance(position, dict): + raise ValueError("first_gen_log_probs entries must be candidate objects or numbers") + candidates: dict[int, Logprob] = {} + for token_id, value in position.items(): + if not isinstance(value, dict) or "logprob" not in value: + raise ValueError("first_gen_log_probs contains an invalid candidate") + try: + decoded_token_id = int(token_id) + except (TypeError, ValueError) as error: + raise ValueError( + "first_gen_log_probs contains a non-integer token ID" + ) from error + candidates[decoded_token_id] = Logprob( + logprob=float(value["logprob"]), + rank=(None if value.get("rank") is None else int(value["rank"])), + ) + decoded.append(candidates) + return decoded + + return DisaggregatedParams( + request_type="generation_only", + first_gen_tokens=payload.get("first_gen_tokens"), + first_gen_log_probs=logprobs(), + ctx_request_id=decimal_id("ctx_request_id"), + disagg_request_id=decimal_id("disagg_request_id"), + ctx_dp_rank=payload.get("ctx_dp_rank", session.dp_rank), + ctx_info_endpoint=payload.get("ctx_info_endpoint"), + draft_tokens=payload.get("draft_tokens"), + ctx_usage=ctx_usage, + conversation_id=payload.get("conversation_id"), + schedule_style=DisaggScheduleStyle.CONTEXT_FIRST, + opaque_state=opaque_state, + ) + + +def to_lora_request(adapter: object | None) -> LoRARequest | None: + if adapter is None: + return None + return LoRARequest(adapter.lora_name, adapter.lora_id, adapter.source_path) diff --git a/tensorrt_llm/openengine/lora_registry.py b/tensorrt_llm/openengine/lora_registry.py new file mode 100644 index 000000000000..22ec8af8efdd --- /dev/null +++ b/tensorrt_llm/openengine/lora_registry.py @@ -0,0 +1,120 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Lazy, logical LoRA registry for the OpenEngine control plane.""" + +import asyncio +import json +from pathlib import Path + +from openengine.v1 import lora_pb2 + +from tensorrt_llm.executor.request import LoRARequest + + +class LoraRegistry: + """Validate adapters now and let TRT-LLM load weights on first use.""" + + def __init__(self, reserved_identities: dict[str, int] | None = None) -> None: + reserved_identities = dict(reserved_identities or {}) + if any(not name for name in reserved_identities): + raise ValueError("Reserved LoRA names must not be empty") + if any(adapter_id < 0 for adapter_id in reserved_identities.values()): + raise ValueError("Reserved LoRA IDs must be non-negative") + if len(set(reserved_identities.values())) != len(reserved_identities): + raise ValueError("Reserved LoRA IDs must be unique") + self._by_name: dict[str, lora_pb2.LoraAdapter] = {} + self._identities_by_name: dict[str, lora_pb2.LoraAdapter] = {} + self._names_by_id: dict[int, str] = {} + self._names_by_path: dict[str, str] = {} + self._reserved_by_name = reserved_identities + self._reserved_by_id = { + adapter_id: name for name, adapter_id in reserved_identities.items() + } + self._lock = asyncio.Lock() + + @staticmethod + def _validated(adapter: lora_pb2.LoraAdapter) -> lora_pb2.LoraAdapter: + if not adapter.lora_name: + raise ValueError("lora_name must not be empty") + if adapter.lora_id < 0: + raise ValueError("lora_id must be non-negative") + path = Path(adapter.source_path).expanduser().resolve() + if not path.is_dir(): + raise ValueError(f"LoRA source path is not a directory: {path}") + if not (path / "adapter_config.json").is_file(): + raise ValueError(f"LoRA directory is missing adapter_config.json: {path}") + try: + config = json.loads((path / "adapter_config.json").read_text()) + except (OSError, json.JSONDecodeError) as error: + raise ValueError(f"LoRA adapter_config.json is invalid: {path}") from error + if not isinstance(config, dict): + raise ValueError(f"LoRA adapter_config.json must contain an object: {path}") + if not any( + (path / filename).is_file() + for filename in ("adapter_model.safetensors", "adapter_model.bin") + ): + raise ValueError(f"LoRA directory is missing adapter weights: {path}") + return lora_pb2.LoraAdapter( + lora_id=adapter.lora_id, lora_name=adapter.lora_name, source_path=str(path) + ) + + async def load(self, adapter: lora_pb2.LoraAdapter) -> tuple[lora_pb2.LoraAdapter, bool]: + validated = self._validated(adapter) + async with self._lock: + reserved_id = self._reserved_by_name.get(validated.lora_name) + if reserved_id is not None: + raise ValueError( + f"LoRA name {validated.lora_name!r} is reserved for a model-owned adapter" + ) + reserved_name = self._reserved_by_id.get(validated.lora_id) + if reserved_name is not None: + raise ValueError( + f"LoRA ID {validated.lora_id} is reserved for model-owned adapter " + f"{reserved_name!r}" + ) + identity = self._identities_by_name.get(validated.lora_name) + if identity is not None: + if identity != validated: + raise ValueError( + f"LoRA name {validated.lora_name!r} is permanently bound to different " + "attributes" + ) + if validated.lora_name in self._by_name: + return identity, True + self._by_name[validated.lora_name] = identity + return identity, False + registered_name = self._names_by_id.get(validated.lora_id) + if registered_name is not None: + raise ValueError( + f"LoRA ID {validated.lora_id} is permanently bound to {registered_name!r}" + ) + registered_name = self._names_by_path.get(validated.source_path) + if registered_name is not None: + raise ValueError( + f"LoRA path {validated.source_path!r} is permanently bound to " + f"{registered_name!r}" + ) + self._by_name[validated.lora_name] = validated + self._identities_by_name[validated.lora_name] = validated + self._names_by_id[validated.lora_id] = validated.lora_name + self._names_by_path[validated.source_path] = validated.lora_name + return validated, False + + async def unload(self, name: str) -> lora_pb2.LoraAdapter: + async with self._lock: + adapter = self._by_name.pop(name, None) + if adapter is None: + raise KeyError(name) + return adapter + + async def list(self) -> list[lora_pb2.LoraAdapter]: + async with self._lock: + return [self._by_name[name] for name in sorted(self._by_name)] + + async def request(self, name: str) -> LoRARequest: + async with self._lock: + adapter = self._by_name.get(name) + if adapter is None: + raise KeyError(name) + return LoRARequest(adapter.lora_name, adapter.lora_id, adapter.source_path) diff --git a/tensorrt_llm/openengine/server.py b/tensorrt_llm/openengine/server.py new file mode 100644 index 000000000000..4604047ad1e2 --- /dev/null +++ b/tensorrt_llm/openengine/server.py @@ -0,0 +1,129 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Lifecycle wrapper for the optional OpenEngine sibling gRPC server.""" + +import grpc +from openengine.v1 import engine_pb2, openengine_pb2_grpc + +from tensorrt_llm.inputs.multimodal import MultimodalServerConfig +from tensorrt_llm.logger import logger +from tensorrt_llm.serve.kv_event_fanout import KvEventFanout +from tensorrt_llm.serve.request_tracker import RequestTracker +from tensorrt_llm.serve.stats_fanout import StatsFanout + +from ._schema_pin import OPENENGINE_COMMIT +from .servicer import OpenEngineServicer, schema_release + +_MAX_MESSAGE_LENGTH = 256 * 1024 * 1024 + + +def validate_schema_release(schema_release: str) -> str: + """Require an immutable OpenEngine source identity before binding.""" + if schema_release == OPENENGINE_COMMIT: + return schema_release + raise RuntimeError( + "OPENENGINE_SCHEMA_RELEASE must exactly match the pinned OPENENGINE_COMMIT " + f"({OPENENGINE_COMMIT})" + ) + + +class OpenEngineServer: + """An OpenEngine gRPC server which never owns or shuts down the LLM.""" + + def __init__( + self, + llm: object, + model: str, + role: int, + host: str, + port: int, + tracker: RequestTracker, + media_config: MultimodalServerConfig | None = None, + reasoning_parser: str | None = None, + tool_parser: str | None = None, + kv_event_fanout: KvEventFanout | None = None, + stats_fanout: StatsFanout | None = None, + ) -> None: + validate_schema_release(schema_release()) + self.host = host + self.port = port + self._kv_event_fanout = kv_event_fanout + self._stats_fanout = stats_fanout + self._server = grpc.aio.server( + options=[ + ("grpc.max_send_message_length", _MAX_MESSAGE_LENGTH), + ("grpc.max_receive_message_length", _MAX_MESSAGE_LENGTH), + ("grpc.keepalive_time_ms", 30_000), + ("grpc.keepalive_timeout_ms", 10_000), + ] + ) + self.servicer = OpenEngineServicer( + llm=llm, + model=model, + role=role, + tracker=tracker, + media_config=media_config, + reasoning_parser=reasoning_parser, + tool_parser=tool_parser, + kv_event_fanout=kv_event_fanout, + stats_fanout=stats_fanout, + event_host=host, + event_port=port, + ) + openengine_pb2_grpc.add_OpenEngineServicer_to_server(self.servicer, self._server) + bound = self._server.add_insecure_port(f"{host}:{port}") + if bound == 0: + raise RuntimeError(f"Failed to bind OpenEngine server to {host}:{port}") + self.port = bound if port == 0 else port + advertised_host = { + "0.0.0.0": "127.0.0.1", + "::": "::1", + "[::]": "::1", + }.get(host, host) + self.servicer.event_host = advertised_host + self.servicer.event_port = self.port + + async def start(self) -> None: + if self._kv_event_fanout is not None: + self._kv_event_fanout.start() + if self._stats_fanout is not None: + self._stats_fanout.start() + try: + await self._server.start() + except BaseException: + if self._kv_event_fanout is not None: + await self._kv_event_fanout.stop() + if self._stats_fanout is not None: + await self._stats_fanout.stop() + self.servicer.close() + raise + logger.info("OpenEngine sibling server started on %s:%d", self.host, self.port) + + async def stop(self, grace: float = 5.0) -> None: + try: + await self._server.stop(grace=grace) + finally: + self.servicer.close() + if self._kv_event_fanout is not None: + await self._kv_event_fanout.stop() + if self._stats_fanout is not None: + await self._stats_fanout.stop() + logger.info("OpenEngine sibling server stopped") + + +def openengine_role(server_role: object | None) -> int: + """Map TRT-LLM serve roles onto the supported OpenEngine milestone.""" + if server_role is None: + return engine_pb2.ENGINE_ROLE_AGGREGATED + name = getattr(server_role, "name", str(server_role)).upper() + if name == "CONTEXT": + return engine_pb2.ENGINE_ROLE_PREFILL + if name == "GENERATION": + return engine_pb2.ENGINE_ROLE_DECODE + if name in ("MM_ENCODER", "VISUAL_GEN"): + raise ValueError(f"OpenEngine does not support TRT-LLM server role {name!r}") + raise ValueError(f"Unknown TRT-LLM server role {name!r}") + + +__all__ = ["OpenEngineServer", "openengine_role", "validate_schema_release"] diff --git a/tensorrt_llm/openengine/servicer.py b/tensorrt_llm/openengine/servicer.py new file mode 100644 index 000000000000..0835a0d0ac97 --- /dev/null +++ b/tensorrt_llm/openengine/servicer.py @@ -0,0 +1,1226 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""OpenEngine gRPC servicer backed by one TensorRT-LLM LLM instance.""" + +import asyncio +import os +import time +import uuid +from collections.abc import AsyncGenerator +from dataclasses import replace +from pathlib import Path +from typing import Any + +import grpc +from openengine import MINIMUM_CLIENT_REVISION, SCHEMA_RELEASE, SCHEMA_REVISION +from openengine.v1 import ( + engine_pb2, + error_pb2, + generation_pb2, + input_pb2, + kv_pb2, + lifecycle_pb2, + lora_pb2, + model_pb2, + observability_pb2, + openengine_pb2_grpc, +) + +import tensorrt_llm +from tensorrt_llm.disaggregated_params import DisaggregatedParams, DisaggScheduleStyle +from tensorrt_llm.executor.request import LoRARequest +from tensorrt_llm.inputs.multimodal import MultimodalServerConfig +from tensorrt_llm.inputs.registry import BaseMultimodalInputProcessor +from tensorrt_llm.logger import logger +from tensorrt_llm.runtime.kv_cache_hash import ( + KV_CACHE_HASH_ALGO_V1, + KV_CACHE_HASH_ALGO_V2_SHA256_64, + get_effective_kv_cache_event_hash_algo, +) +from tensorrt_llm.scheduling_params import SchedulingParams +from tensorrt_llm.serve.kv_event_fanout import KvEventFanout +from tensorrt_llm.serve.request_tracker import RequestTracker +from tensorrt_llm.serve.stats_fanout import StatsFanout + +from .converters import ( + decode_handoff, + encode_handoff, + handoff_requires_decode_media, + load_media, + stable_request_id, + to_priority, + to_sampling_params, +) +from .lora_registry import LoraRegistry + + +def schema_release() -> str: + """Return the immutable OpenEngine source identity configured at startup.""" + return os.getenv("OPENENGINE_SCHEMA_RELEASE", SCHEMA_RELEASE) + + +def _arg(llm: object, name: str, default: Any = None) -> Any: + args = getattr(llm, "args", None) + value = getattr(args, name, None) + if value is not None: + return value + parallel = getattr(args, "parallel_config", None) + return getattr(parallel, name, default) + + +def _token_text(llm: object, token_id: int) -> str: + tokenizer = getattr(llm, "tokenizer", None) + if tokenizer is None: + return "" + try: + return tokenizer.decode([token_id], skip_special_tokens=False) + except (ValueError, TypeError, AttributeError): + return "" + + +def _signed_i64(value: int) -> int: + if value >= 1 << 63: + return value - (1 << 64) + if value < -(1 << 63): + return ((value + (1 << 63)) % (1 << 64)) - (1 << 63) + return value + + +def _block_hash(value: object, hash_algo: str | None) -> kv_pb2.KvBlockHash: + if isinstance(value, bytes): + return kv_pb2.KvBlockHash(value=value, encoding=hash_algo or "bytes") + if isinstance(value, int): + return kv_pb2.KvBlockHash( + value=str(_signed_i64(value)).encode("ascii"), encoding="decimal_int64" + ) + text = str(value) + try: + encoded = bytes.fromhex(text) + except ValueError: + return kv_pb2.KvBlockHash(value=text.encode("utf-8"), encoding="utf8") + return kv_pb2.KvBlockHash(value=encoded, encoding=hash_algo or "hex") + + +class OpenEngineServicer(openengine_pb2_grpc.OpenEngineServicer): + """Engine-neutral service surface over TensorRT-LLM's Python LLM API.""" + + def __init__( + self, + llm: object, + model: str, + role: int, + tracker: RequestTracker, + media_config: MultimodalServerConfig | None = None, + reasoning_parser: str | None = None, + tool_parser: str | None = None, + kv_event_fanout: KvEventFanout | None = None, + stats_fanout: StatsFanout | None = None, + instance_id: str | None = None, + event_host: str = "127.0.0.1", + event_port: int = 0, + kv_session_ttl_seconds: float = 300.0, + post_abort_cleanup_timeout_seconds: float = 1.0, + ) -> None: + self.llm = llm + self.model = model + self.role = role + self.tracker = tracker + self.media_config = media_config + self.reasoning_parser = reasoning_parser or "" + self.tool_parser = tool_parser or "" + processor = getattr(llm, "input_processor", None) + identity_getter = getattr(processor, "get_model_owned_lora_identities", None) + model_owned_loras = identity_getter() if callable(identity_getter) else {} + self.loras = LoraRegistry(model_owned_loras) + self._lora_names_by_id: dict[int, str] = { + adapter_id: name for name, adapter_id in model_owned_loras.items() + } + self.kv_events = kv_event_fanout + self.stats = stats_fanout + self._kv_session_requests: dict[str, str] = {} + self._kv_session_timers: dict[str, asyncio.TimerHandle] = {} + if kv_session_ttl_seconds <= 0: + raise ValueError("kv_session_ttl_seconds must be positive") + self._kv_session_ttl_seconds = kv_session_ttl_seconds + if post_abort_cleanup_timeout_seconds <= 0: + raise ValueError("post_abort_cleanup_timeout_seconds must be positive") + self._post_abort_cleanup_timeout_seconds = post_abort_cleanup_timeout_seconds + self._partial_block_hashes: dict[int, set[object]] = {} + self.instance_id = instance_id or str(uuid.uuid4()) + self.event_host = event_host + self.event_port = event_port + + async def Generate( + self, request: generation_pb2.GenerateRequest, context: grpc.aio.ServicerContext + ) -> AsyncGenerator[generation_pb2.GenerateResponse, None]: + result = None + consumed_session_id = None + try: + self._validate_generate(request) + params = to_sampling_params(request) + media = await load_media(list(request.media), request.media_options, self.media_config) + input_kind = request.WhichOneof("input") + inputs: dict[str, Any] + if input_kind == "prompt": + inputs = {"prompt": request.prompt} + else: + inputs = {"prompt_token_ids": list(request.token_ids.ids)} + if media: + inputs["multi_modal_data"] = media + lora_request = self._required_multimodal_lora(request) + if lora_request is not None and request.lora_name: + raise ValueError( + f"Modalities in this request require the model-owned " + f"{lora_request.lora_name!r} adapter and cannot be combined with " + "a user-selected LoRA" + ) + if lora_request is None and request.lora_name: + lora_request = await self.loras.request(request.lora_name) + disaggregated = self._disaggregated_params(request) + context_usage = ( + disaggregated.ctx_usage + if disaggregated is not None and disaggregated.request_type == "generation_only" + else None + ) + scheduling = self._scheduling_params(request) + priority = to_priority(request) + trace_headers = { + key.lower(): value + for key, value in request.metadata.items() + if key.lower() in ("traceparent", "tracestate", "baggage") + } + result = self.llm.generate_async( + inputs=inputs, + sampling_params=params, + lora_request=lora_request, + streaming=True, + disaggregated_params=disaggregated, + scheduling_params=scheduling, + cache_salt=(request.kv.cache_salt if request.kv.HasField("cache_salt") else None), + trace_headers=trace_headers or None, + priority=priority, + ) + self.tracker.admit(request.request_id, result) + if self.stats is not None: + self.stats.wake() + if request.kv.HasField("session") and request.kv.session.session_id: + consumed_session_id = request.kv.session.session_id + self._track_kv_session(consumed_session_id, request.request_id) + except KeyError as error: + await context.abort( + grpc.StatusCode.FAILED_PRECONDITION, f"LoRA is not loaded: {error.args[0]}" + ) + return + except (ValueError, TypeError, RuntimeError) as error: + if ( + result is not None + and self.tracker.active_requests.get(request.request_id) is not result + ): + try: + result.abort() + except (RuntimeError, AssertionError): + logger.warning( + "Failed to abort rejected duplicate request %s", + request.request_id, + ) + code = ( + grpc.StatusCode.UNAVAILABLE + if self.tracker.draining + else grpc.StatusCode.INVALID_ARGUMENT + ) + await context.abort(code, str(error)) + return + + sent_tokens: dict[int, int] = {} + sent_text: dict[int, int] = {} + prompt_sent = False + completed = False + try: + async for current in result: + if context.cancelled(): + await self.tracker.abort(request.request_id) + return + if not prompt_sent: + prompt = self._prompt_output(current) + if prompt is not None: + prompt_sent = True + yield generation_pb2.GenerateResponse( + request_id=request.request_id, prompt=prompt + ) + if self.role == engine_pb2.ENGINE_ROLE_PREFILL: + if current.finished: + handoff = current.disaggregated_params + if handoff is None and current.outputs: + handoff = current.outputs[0].disaggregated_params + if handoff is None: + raise RuntimeError( + "Context-only result did not return disaggregated parameters" + ) + handoff = replace(handoff, ctx_usage=self._usage_payload(current)) + session = encode_handoff(handoff, requires_decode_media=bool(request.media)) + self._track_kv_session(session.session_id, request.request_id) + yield generation_pb2.GenerateResponse( + request_id=request.request_id, + prefill_ready=generation_pb2.PrefillReady(kv_session=session), + usage=self._usage(current), + ) + completed = True + continue + for output in current.outputs: + token_start = sent_tokens.get(output.index, 0) + text_start = sent_text.get(output.index, 0) + token_ids = list(output.token_ids or []) + delta_ids = token_ids[token_start:] + delta_text = (output.text or "")[text_start:] + if delta_ids or delta_text: + logprobs = list(output.logprobs or [])[token_start:] + yield generation_pb2.GenerateResponse( + request_id=request.request_id, + token=generation_pb2.TokenOutput( + output_index=output.index, + tokens=self._token_infos(delta_ids, logprobs), + text=delta_text, + ), + ) + sent_tokens[output.index] = len(token_ids) + sent_text[output.index] = len(output.text or "") + if current.finished: + outputs = current.outputs + for index, output in enumerate(outputs): + response = generation_pb2.GenerateResponse( + request_id=request.request_id, + finished=self._finished(output), + ) + if index == len(outputs) - 1: + response.usage.CopyFrom(self._usage(current, context_usage)) + yield response + completed = True + except asyncio.CancelledError: + await self.tracker.abort(request.request_id) + raise + except (RuntimeError, ValueError, TypeError) as error: + logger.error("OpenEngine request %s failed: %s", request.request_id, error) + yield generation_pb2.GenerateResponse( + request_id=request.request_id, + error=error_pb2.EngineError( + code=error_pb2.ERROR_CODE_INTERNAL, + message=str(error), + retryable=False, + ), + ) + finally: + if consumed_session_id is not None: + self._release_kv_session(consumed_session_id) + if not completed and request.request_id in self.tracker.active_requests: + await self.tracker.abort(request.request_id) + else: + await self.tracker.finish(request.request_id) + if self.stats is not None: + self.stats.wake() + + def _track_kv_session(self, session_id: str, request_id: str) -> None: + self._release_kv_session(session_id) + self._kv_session_requests[session_id] = request_id + loop = asyncio.get_running_loop() + self._kv_session_timers[session_id] = loop.call_later( + self._kv_session_ttl_seconds, self._release_kv_session, session_id + ) + + def _release_kv_session(self, session_id: str) -> bool: + released = self._kv_session_requests.pop(session_id, None) is not None + timer = self._kv_session_timers.pop(session_id, None) + if timer is not None: + timer.cancel() + return released + + def _release_all_kv_sessions(self) -> None: + for timer in self._kv_session_timers.values(): + timer.cancel() + self._kv_session_timers.clear() + self._kv_session_requests.clear() + + def close(self) -> None: + """Release protocol-owned timers without shutting down the shared LLM.""" + self._release_all_kv_sessions() + + def _validate_generate(self, request: generation_pb2.GenerateRequest) -> None: + if not request.request_id: + raise ValueError("request_id must not be empty") + if request.model and request.model != self.model: + raise ValueError(f"Unknown model {request.model!r}") + if request.WhichOneof("input") is None: + raise ValueError("Generate requires prompt or token_ids input") + if request.request_id in self.tracker.active_requests: + raise ValueError(f"Request {request.request_id!r} is already active") + if self.tracker.draining: + raise RuntimeError("TensorRT-LLM is draining") + configured_backend = getattr( + getattr(self.llm, "args", None), "guided_decoding_backend", None + ) + if request.guided.backend and request.guided.backend != configured_backend: + raise ValueError("Per-request guided decoding backend selection is not supported") + if request.kv.HasField("bypass_prefix_cache") and request.kv.bypass_prefix_cache: + raise ValueError("Prefix-cache bypass is not supported by TensorRT-LLM") + if request.kv.HasField("data_parallel_rank"): + self._scheduling_params(request) + if request.media: + processor = getattr(self.llm, "input_processor", None) + aggregate = self._available_modalities(processor) + prefill_decode = self._available_modalities(processor, prefill_decode=True) + names = { + input_pb2.MODALITY_UNSPECIFIED: "image", + input_pb2.MODALITY_IMAGE: "image", + input_pb2.MODALITY_VIDEO: "video", + input_pb2.MODALITY_AUDIO: "audio", + } + allowed = ( + aggregate if self.role == engine_pb2.ENGINE_ROLE_AGGREGATED else prefill_decode + ) + unsupported = { + names.get(item.modality, "unspecified") + for item in request.media + if names.get(item.modality) not in allowed + } + if unsupported: + raise ValueError( + f"Media modalities are not supported for this role: {sorted(unsupported)}" + ) + if isinstance(processor, BaseMultimodalInputProcessor): + requested_modalities = tuple( + dict.fromkeys(names[item.modality] for item in request.media) + ) + processor.get_required_lora_spec(requested_modalities) + if self.role == engine_pb2.ENGINE_ROLE_DECODE: + if not request.kv.HasField("session"): + raise ValueError("Decode requests require a prefill KV session") + requires_media = handoff_requires_decode_media(request.kv.session) + if requires_media and not request.media: + raise ValueError("Decode request must resend the ordered context-phase media") + if request.media and not requires_media: + raise ValueError("Decode KV session does not require raw media") + elif self.role == engine_pb2.ENGINE_ROLE_PREFILL: + if request.kv.HasField("session"): + raise ValueError("Prefill requests cannot consume a KV session") + elif request.kv.HasField("session"): + raise ValueError("Aggregated requests cannot consume a KV session") + + def _disaggregated_params( + self, request: generation_pb2.GenerateRequest + ) -> DisaggregatedParams | None: + if self.role == engine_pb2.ENGINE_ROLE_AGGREGATED: + return None + if self.role == engine_pb2.ENGINE_ROLE_PREFILL: + return DisaggregatedParams( + request_type="context_only", + disagg_request_id=stable_request_id(request.request_id), + ctx_dp_rank=( + request.kv.data_parallel_rank + if request.kv.HasField("data_parallel_rank") + else None + ), + schedule_style=DisaggScheduleStyle.CONTEXT_FIRST, + conversation_id=request.metadata.get("conversation_id"), + ) + return decode_handoff(request.kv.session) + + def _scheduling_params( + self, request: generation_pb2.GenerateRequest + ) -> SchedulingParams | None: + if not request.kv.HasField("data_parallel_rank"): + return None + rank = request.kv.data_parallel_rank + data_parallel_size = int(_arg(self.llm, "data_parallel_size", 1)) + if rank >= data_parallel_size: + raise ValueError( + f"data_parallel_rank {rank} is outside the configured DP size {data_parallel_size}" + ) + attention_dp = bool(_arg(self.llm, "enable_attention_dp", False)) + if not attention_dp or getattr(self.llm, "_on_trt_backend", False): + if rank == 0: + return None + raise ValueError( + "A nonzero data_parallel_rank requires the PyTorch backend with attention DP " + "enabled" + ) + return SchedulingParams(attention_dp_rank=rank, attention_dp_relax=False) + + def _required_multimodal_lora( + self, request: generation_pb2.GenerateRequest + ) -> LoRARequest | None: + if not request.media: + return None + processor = getattr(self.llm, "input_processor", None) + if not isinstance(processor, BaseMultimodalInputProcessor): + return None + names = { + input_pb2.MODALITY_UNSPECIFIED: "image", + input_pb2.MODALITY_IMAGE: "image", + input_pb2.MODALITY_VIDEO: "video", + input_pb2.MODALITY_AUDIO: "audio", + } + modalities = tuple( + dict.fromkeys(names[item.modality] for item in request.media if item.modality in names) + ) + spec = processor.get_required_lora_spec(modalities) + if spec is None: + return None + if not self._supports_lora(): + raise ValueError( + f"Modalities in this request require the model-owned {spec.name!r} adapter, " + "but TensorRT-LLM was not configured with LoRA support" + ) + self._lora_names_by_id[spec.adapter_id] = spec.name + return LoRARequest(spec.name, spec.adapter_id, spec.path) + + def _available_modalities( + self, processor: object, *, prefill_decode: bool = False + ) -> tuple[str, ...]: + if not isinstance(processor, BaseMultimodalInputProcessor): + return () + declared = ( + processor.get_openengine_prefill_decode_modalities() + if prefill_decode + else processor.get_openengine_modalities() + ) + available = [] + for modality in declared: + spec = processor.get_required_lora_spec((modality,)) + if spec is not None and (not self._supports_lora() or not Path(spec.path).is_dir()): + continue + available.append(modality) + return tuple(available) + + def _token_infos( + self, token_ids: list[int], logprobs: list[Any] + ) -> list[generation_pb2.TokenInfo]: + infos: list[generation_pb2.TokenInfo] = [] + for index, token_id in enumerate(token_ids): + info = generation_pb2.TokenInfo( + token_id=token_id, + token=_token_text(self.llm, token_id), + ) + if index < len(logprobs): + value = logprobs[index] + if isinstance(value, dict): + sampled = value.get(token_id) + if sampled is not None: + info.logprob = sampled.logprob + if sampled.rank is not None: + info.rank = sampled.rank + for candidate_id, candidate in value.items(): + candidate_proto = info.candidates.add( + token_id=candidate_id, + logprob=candidate.logprob, + token=_token_text(self.llm, candidate_id), + ) + if candidate.rank is not None: + candidate_proto.rank = candidate.rank + elif isinstance(value, (int, float)): + info.logprob = float(value) + infos.append(info) + return infos + + def _prompt_output(self, result: object) -> generation_pb2.PromptOutput | None: + if not result.outputs: + return None + prompt_logprobs = result.outputs[0].prompt_logprobs + if not prompt_logprobs: + return None + token_ids = list(result.prompt_token_ids) + return generation_pb2.PromptOutput( + tokens=self._token_infos(token_ids, list(prompt_logprobs)) + ) + + @staticmethod + def _usage(result: object, context_usage: dict[str, Any] | None = None) -> generation_pb2.Usage: + prompt = len(result.prompt_token_ids) + completion = sum(len(output.token_ids or []) for output in result.outputs) + cached = getattr(result, "cached_tokens", None) + if context_usage is not None: + prompt = int(context_usage["prompt_tokens"]) + details = context_usage.get("prompt_tokens_details") or {} + cached = int(details.get("cached_tokens", 0)) + usage = generation_pb2.Usage( + prompt_tokens=prompt, completion_tokens=completion, total_tokens=prompt + completion + ) + if cached is not None: + usage.cached_prompt_tokens = cached + return usage + + @classmethod + def _usage_payload(cls, result: object) -> dict[str, Any]: + usage = cls._usage(result) + return { + "prompt_tokens": usage.prompt_tokens, + "completion_tokens": usage.completion_tokens, + "total_tokens": usage.total_tokens, + "prompt_tokens_details": {"cached_tokens": usage.cached_prompt_tokens}, + } + + @staticmethod + def _finished(output: object) -> generation_pb2.GenerationFinished: + reason_map = { + "stop": generation_pb2.FINISH_REASON_STOP, + "length": generation_pb2.FINISH_REASON_LENGTH, + "cancelled": generation_pb2.FINISH_REASON_CANCELLED, + "timeout": generation_pb2.FINISH_REASON_CANCELLED, + } + finished = generation_pb2.GenerationFinished( + output_index=output.index, + reason=reason_map.get(output.finish_reason, generation_pb2.FINISH_REASON_STOP), + ) + if isinstance(output.stop_reason, int): + finished.stop_match.stop_token_id = output.stop_reason + elif isinstance(output.stop_reason, str): + finished.stop_match.stop_text = output.stop_reason + return finished + + async def GetEngineInfo( + self, request: engine_pb2.GetEngineInfoRequest, context: grpc.aio.ServicerContext + ) -> engine_pb2.EngineInfo: + del request, context + tp = _arg(self.llm, "tensor_parallel_size", 1) + pp = _arg(self.llm, "pipeline_parallel_size", 1) + dp = _arg(self.llm, "data_parallel_size", 1) + return engine_pb2.EngineInfo( + engine_name="tensorrt_llm", + engine_version=getattr(tensorrt_llm, "__version__", "unknown"), + role=self.role, + instance_id=self.instance_id, + supported_models=[self.model], + parallelism=engine_pb2.ParallelismInfo( + tensor_parallel_size=tp, + pipeline_parallel_size=pp, + data_parallel_size=dp, + data_parallel_rank=_arg(self.llm, "data_parallel_rank", 0), + ), + kv_connector=self._kv_connector_info(), + schema_revision=SCHEMA_REVISION, + minimum_client_revision=MINIMUM_CLIENT_REVISION, + schema_release=schema_release(), + ) + + async def GetModelInfo( + self, request: model_pb2.GetModelInfoRequest, context: grpc.aio.ServicerContext + ) -> model_pb2.ModelInfo: + if request.model and request.model != self.model: + await context.abort(grpc.StatusCode.NOT_FOUND, f"Unknown model {request.model!r}") + raise ValueError(f"Unknown model {request.model!r}") + input_processor = getattr(self.llm, "input_processor", None) + aggregate = self._available_modalities(input_processor) + prefill_decode = self._available_modalities(input_processor, prefill_decode=True) + modality_enum = { + "image": input_pb2.MODALITY_IMAGE, + "video": input_pb2.MODALITY_VIDEO, + "audio": input_pb2.MODALITY_AUDIO, + } + args = getattr(self.llm, "args", None) + guided_backend = getattr(args, "guided_decoding_backend", None) + guided_modes = [] + if guided_backend is not None: + guided_modes = [ + model_pb2.GUIDED_DECODING_MODE_JSON_SCHEMA, + model_pb2.GUIDED_DECODING_MODE_REGEX, + model_pb2.GUIDED_DECODING_MODE_EBNF_GRAMMAR, + model_pb2.GUIDED_DECODING_MODE_CHOICE, + model_pb2.GUIDED_DECODING_MODE_JSON_OBJECT, + ] + if guided_backend == "xgrammar": + guided_modes.append(model_pb2.GUIDED_DECODING_MODE_STRUCTURAL_TAG) + max_context = getattr(args, "max_seq_len", None) or getattr(args, "max_input_len", None) + max_requests = getattr(args, "max_batch_size", None) + max_tokens = getattr(args, "max_num_tokens", None) + kv_capacity = self._kv_capacity() + info = model_pb2.ModelInfo( + model_id=self.model, + served_model_name=self.model, + served_model_aliases=[self.model], + tokenizer_modes=["auto"], + supports_text_input=True, + supports_token_ids_input=True, + supports_lora=self._supports_lora(), + supports_multimodal=bool(aggregate), + reasoning_parser=self.reasoning_parser, + tool_call_parser=self.tool_parser, + generation=model_pb2.GenerationCapabilities( + prompt_logprobs=model_pb2.LogprobCapabilities( + supported=True, + candidate_selection_modes=[model_pb2.CANDIDATE_TOKEN_SELECTION_MODE_TOP_N], + max_top_n=20, + ), + output_logprobs=model_pb2.LogprobCapabilities( + supported=True, + candidate_selection_modes=[model_pb2.CANDIDATE_TOKEN_SELECTION_MODE_TOP_N], + max_top_n=20, + ), + guided_decoding=model_pb2.GuidedDecodingCapabilities( + supported=guided_backend is not None, + modes=guided_modes, + ), + max_num_sequences=getattr(args, "max_beam_width", 1), + supports_priority=True, + supports_stop_in_output=True, + supports_cache_salt=True, + supports_prefix_cache_bypass=False, + ), + multimodal_capabilities=model_pb2.MultimodalCapabilities( + aggregate_modalities=[modality_enum[name] for name in aggregate], + prefill_decode_modalities=[modality_enum[name] for name in prefill_decode], + source_types=[ + input_pb2.MEDIA_SOURCE_TYPE_URL, + input_pb2.MEDIA_SOURCE_TYPE_DATA_URI, + input_pb2.MEDIA_SOURCE_TYPE_RAW_BYTES, + ], + supports_per_request_media_options=True, + ), + ) + if max_context is not None: + info.max_context_length = max_context + if max_requests is not None: + info.max_running_requests = max_requests + if max_tokens is not None: + info.max_batched_tokens = max_tokens + if kv_capacity.get("tokensPerBlock") is not None: + info.kv_block_size = kv_capacity["tokensPerBlock"] + if kv_capacity.get("maxNumBlocks") is not None: + info.total_kv_blocks = kv_capacity["maxNumBlocks"] + return info + + def _kv_capacity(self) -> dict[str, int]: + capacity: dict[str, int] = {} + config = getattr(getattr(self.llm, "args", None), "kv_cache_config", None) + tokens_per_block = getattr(config, "tokens_per_block", None) + if tokens_per_block is not None: + capacity["tokensPerBlock"] = int(tokens_per_block) + getter = getattr(self.llm, "get_kv_cache_capacity", None) + if callable(getter): + try: + discovered = getter() + except (RuntimeError, AttributeError, TypeError) as error: + logger.debug("KV capacity discovery unavailable: %s", error) + else: + if isinstance(discovered, dict): + for key in ("maxNumBlocks", "tokensPerBlock", "maxNumTokens"): + if discovered.get(key) is not None: + capacity[key] = int(discovered[key]) + return capacity + + async def GetLoad( + self, request: observability_pb2.GetLoadRequest, context: grpc.aio.ServicerContext + ) -> observability_pb2.LoadInfo: + del context + latest = self.stats.latest_by_rank() if self.stats is not None else {} + rank_infos = [] + running_requests = 0 + queued_requests = 0 + used_kv_blocks = 0 + total_kv_blocks = 0 + prefill_batch_size = 0 + decode_batch_size = 0 + attributes = {"source": "shared_stats_fanout" if latest else "shared_request_tracker"} + tokens_per_block: set[int] = set() + for rank, stat in sorted(latest.items()): + kv_stats = stat.get("kvCacheStats") or {} + ifb_stats = stat.get("inflightBatchingStats") or {} + rank_running = int(stat.get("numActiveRequests", 0)) + rank_queued = int(stat.get("numQueuedRequests", 0)) + rank_total = int(kv_stats.get("maxNumBlocks", 0)) + rank_used = int( + kv_stats.get( + "usedNumBlocks", + max(0, rank_total - int(kv_stats.get("freeNumBlocks", rank_total))), + ) + ) + rank_prefill = int(ifb_stats.get("numContextRequests", 0)) + rank_decode = int(ifb_stats.get("numGenRequests", 0)) + rank_tokens_per_block = kv_stats.get("tokensPerBlock") + if rank_tokens_per_block is not None: + rank_tokens_per_block = int(rank_tokens_per_block) + tokens_per_block.add(rank_tokens_per_block) + attributes[f"rank.{rank}.kv_tokens_per_block"] = str(rank_tokens_per_block) + running_requests += rank_running + queued_requests += rank_queued + used_kv_blocks += rank_used + total_kv_blocks += rank_total + prefill_batch_size += rank_prefill + decode_batch_size += rank_decode + if request.include_per_rank: + rank_infos.append( + observability_pb2.RankLoadInfo( + data_parallel_rank=rank, + running_requests=rank_running, + queued_requests=rank_queued, + used_kv_blocks=rank_used, + total_kv_blocks=rank_total, + prefill_batch_size=rank_prefill, + decode_batch_size=rank_decode, + ) + ) + + capacity = self._kv_capacity() + running_requests = max(running_requests, self.tracker.active_count) + if not latest: + total_kv_blocks = capacity.get("maxNumBlocks", 0) + if capacity.get("tokensPerBlock") is not None: + tokens_per_block.add(capacity["tokensPerBlock"]) + if len(tokens_per_block) == 1: + attributes["kv_tokens_per_block"] = str(next(iter(tokens_per_block))) + response = observability_pb2.LoadInfo( + instance_id=self.instance_id, + timestamp_unix_nanos=time.time_ns(), + running_requests=running_requests, + queued_requests=queued_requests, + active_kv_sessions=len(self._kv_session_requests), + ranks=rank_infos, + attributes=attributes, + ) + if used_kv_blocks or latest: + response.used_kv_blocks = used_kv_blocks + if total_kv_blocks: + response.total_kv_blocks = total_kv_blocks + if prefill_batch_size or latest: + response.prefill_batch_size = prefill_batch_size + if decode_batch_size or latest: + response.decode_batch_size = decode_batch_size + return response + + async def Health( + self, request: lifecycle_pb2.HealthRequest, context: grpc.aio.ServicerContext + ) -> lifecycle_pb2.HealthResponse: + if request.model and request.model != self.model: + await context.abort(grpc.StatusCode.NOT_FOUND, f"Unknown model {request.model!r}") + raise ValueError(f"Unknown model {request.model!r}") + if request.role not in (engine_pb2.ENGINE_ROLE_UNSPECIFIED, self.role): + await context.abort( + grpc.StatusCode.INVALID_ARGUMENT, + "Requested health role does not match this engine", + ) + raise ValueError("Requested health role does not match this engine") + if request.include_inference_probe: + await context.abort( + grpc.StatusCode.UNIMPLEMENTED, + "Role-safe inference probes are not implemented by this server", + ) + raise NotImplementedError("Role-safe inference probes are not implemented") + healthy, message = await self.tracker.health() + state = ( + lifecycle_pb2.HEALTH_STATE_DRAINING + if self.tracker.draining + else lifecycle_pb2.HEALTH_STATE_READY + if healthy + else lifecycle_pb2.HEALTH_STATE_NOT_READY + ) + checks = [lifecycle_pb2.HealthCheck(name="scheduler", state=state, message=message)] + return lifecycle_pb2.HealthResponse(state=state, checks=checks) + + async def Abort( + self, request: lifecycle_pb2.AbortRequest, context: grpc.aio.ServicerContext + ) -> lifecycle_pb2.AbortResponse: + target = request.WhichOneof("target") + if target is None: + await context.abort(grpc.StatusCode.INVALID_ARGUMENT, "Abort target must be specified") + raise ValueError("Abort target must be specified") + if target == "all_requests": + count = await self.tracker.abort_all() + self._release_all_kv_sessions() + return lifecycle_pb2.AbortResponse( + status=lifecycle_pb2.ABORT_STATUS_ABORTED, message=f"Aborted {count} requests" + ) + request_id = request.request_id + released_session = False + if target == "kv_session": + session_id = request.kv_session.session_id + request_id = self._kv_session_requests.get(session_id, session_id) + released_session = self._release_kv_session(session_id) + aborted = await self.tracker.abort(request_id) + return lifecycle_pb2.AbortResponse( + status=( + lifecycle_pb2.ABORT_STATUS_ABORTED + if aborted or released_session + else lifecycle_pb2.ABORT_STATUS_ALREADY_FINISHED + ), + message=( + f"Aborted {request_id}" + if aborted + else f"Released KV session {request.kv_session.session_id}" + if released_session + else f"Request {request_id} is not active" + ), + ) + + async def Drain( + self, request: lifecycle_pb2.DrainRequest, context: grpc.aio.ServicerContext + ) -> AsyncGenerator[lifecycle_pb2.DrainResponse, None]: + del context + if request.stop_accepting_new_requests: + await self.tracker.start_drain() + yield lifecycle_pb2.DrainResponse( + state=lifecycle_pb2.DRAIN_STATE_STARTED, + in_flight_requests=self.tracker.active_count, + open_kv_sessions=len(self._kv_session_requests), + message="Process-wide drain started", + ) + timeout = request.deadline_ms / 1000.0 if request.HasField("deadline_ms") else None + empty = await self.tracker.wait_empty(timeout) + if not empty and request.abort_after_deadline: + await self.tracker.abort_all() + empty = await self.tracker.wait_empty(self._post_abort_cleanup_timeout_seconds) + if not empty: + yield lifecycle_pb2.DrainResponse( + error=error_pb2.EngineError( + code=error_pb2.ERROR_CODE_INTERNAL, + message="Drain deadline expired with active requests", + retryable=False, + ), + in_flight_requests=self.tracker.active_count, + open_kv_sessions=len(self._kv_session_requests), + ) + return + self._release_all_kv_sessions() + yield lifecycle_pb2.DrainResponse( + state=lifecycle_pb2.DRAIN_STATE_COMPLETE, + in_flight_requests=0, + open_kv_sessions=0, + message="Drain complete", + ) + + async def LoadLora( + self, request: lora_pb2.LoadLoraRequest, context: grpc.aio.ServicerContext + ) -> lora_pb2.LoadLoraResponse: + if not self._supports_lora(): + await context.abort( + grpc.StatusCode.FAILED_PRECONDITION, + "TensorRT-LLM was not configured with LoRA support", + ) + raise RuntimeError("TensorRT-LLM was not configured with LoRA support") + try: + adapter, already_loaded = await self.loras.load(request.adapter) + except ValueError as error: + await context.abort(grpc.StatusCode.INVALID_ARGUMENT, str(error)) + raise + self._lora_names_by_id[adapter.lora_id] = adapter.lora_name + return lora_pb2.LoadLoraResponse(adapter=adapter, already_loaded=already_loaded) + + async def UnloadLora( + self, request: lora_pb2.UnloadLoraRequest, context: grpc.aio.ServicerContext + ) -> lora_pb2.UnloadLoraResponse: + if not self._supports_lora(): + await context.abort( + grpc.StatusCode.FAILED_PRECONDITION, + "TensorRT-LLM was not configured with LoRA support", + ) + raise RuntimeError("TensorRT-LLM was not configured with LoRA support") + try: + adapter = await self.loras.unload(request.lora_name) + except KeyError: + await context.abort( + grpc.StatusCode.NOT_FOUND, f"LoRA {request.lora_name!r} is not loaded" + ) + raise + return lora_pb2.UnloadLoraResponse(adapter=adapter) + + async def ListLoras( + self, request: lora_pb2.ListLorasRequest, context: grpc.aio.ServicerContext + ) -> lora_pb2.ListLorasResponse: + del request + if not self._supports_lora(): + await context.abort( + grpc.StatusCode.FAILED_PRECONDITION, + "TensorRT-LLM was not configured with LoRA support", + ) + raise RuntimeError("TensorRT-LLM was not configured with LoRA support") + return lora_pb2.ListLorasResponse(adapters=await self.loras.list()) + + def _supports_lora(self) -> bool: + args = getattr(self.llm, "args", None) + if getattr(args, "lora_config", None) is not None or bool( + getattr(args, "enable_lora", False) + ): + return True + build_config = getattr(args, "build_config", None) + plugin_config = getattr(build_config, "plugin_config", None) + return bool(getattr(plugin_config, "lora_plugin", False)) + + def _kv_connector_info(self) -> kv_pb2.KvConnectorInfo: + enabled = self.role in (engine_pb2.ENGINE_ROLE_PREFILL, engine_pb2.ENGINE_ROLE_DECODE) + return kv_pb2.KvConnectorInfo( + enabled=enabled, + transfer_backend="tensorrt_llm" if enabled else "", + supported_protocols=["nixl"] if enabled else [], + supports_remote_prefill=enabled, + supports_decode_pull=enabled, + supports_abort_cleanup=enabled, + supports_drain=enabled, + schema_version=1, + ) + + def _kv_events_enabled(self) -> bool: + args = getattr(self.llm, "args", None) + config = getattr(args, "kv_cache_config", None) + if self.kv_events is None or getattr(config, "event_buffer_max_size", 0) <= 0: + return False + effective_hash_algo = get_effective_kv_cache_event_hash_algo( + getattr(config, "kv_cache_event_hash_algo", "auto"), + bool(getattr(config, "use_kv_cache_manager_v2", False)), + ) + return effective_hash_algo in (KV_CACHE_HASH_ALGO_V1, KV_CACHE_HASH_ALGO_V2_SHA256_64) + + async def GetKvConnectorInfo( + self, request: kv_pb2.GetKvConnectorInfoRequest, context: grpc.aio.ServicerContext + ) -> kv_pb2.KvConnectorInfo: + del request, context + return self._kv_connector_info() + + async def GetKvEventSources( + self, request: kv_pb2.GetKvEventSourcesRequest, context: grpc.aio.ServicerContext + ) -> kv_pb2.GetKvEventSourcesResponse: + if not self._kv_events_enabled(): + return kv_pb2.GetKvEventSourcesResponse() + data_parallel_size = int(_arg(self.llm, "data_parallel_size", 1)) + available_ranks = set(range(data_parallel_size)) + requested_ranks = set(request.data_parallel_ranks) + invalid_ranks = requested_ranks - available_ranks + if invalid_ranks: + await context.abort( + grpc.StatusCode.INVALID_ARGUMENT, + f"Unknown data-parallel ranks: {sorted(invalid_ranks)}", + ) + return kv_pb2.GetKvEventSourcesResponse() + selected_ranks = sorted(requested_ranks or available_ranks) + return kv_pb2.GetKvEventSourcesResponse( + sources=[ + kv_pb2.KvEventSource( + transport="grpc", + endpoint_addr=kv_pb2.KvEndpoint( + host=self.event_host, port=self.event_port, protocol="grpc" + ), + data_parallel_rank=rank, + encoding="protobuf", + schema_version=1, + max_queue_size=self.kv_events.buffer_size, + ) + for rank in selected_ranks + ] + ) + + async def SubscribeKvEvents( + self, request: kv_pb2.SubscribeKvEventsRequest, context: grpc.aio.ServicerContext + ) -> AsyncGenerator[kv_pb2.SubscribeKvEventsResponse, None]: + if not self._kv_events_enabled(): + await context.abort( + grpc.StatusCode.FAILED_PRECONDITION, "KV event collection is disabled" + ) + return + if request.include_snapshot or request.start_sequence_number: + await context.abort( + grpc.StatusCode.UNIMPLEMENTED, + "KV event snapshot and replay subscriptions are not implemented", + ) + return + selected_ranks = set(request.data_parallel_ranks) + available_ranks = set(range(int(_arg(self.llm, "data_parallel_size", 1)))) + invalid_ranks = selected_ranks - available_ranks + if invalid_ranks: + await context.abort( + grpc.StatusCode.INVALID_ARGUMENT, + f"Unknown data-parallel ranks: {sorted(invalid_ranks)}", + ) + return + async for sequence, raw in self.kv_events.subscribe(selected_ranks): + if context.cancelled(): + return + rank = raw.get("attention_dp_rank", 0) + if selected_ranks and rank not in selected_ranks: + continue + yield kv_pb2.SubscribeKvEventsResponse(batch=self._kv_batch(raw, sequence)) + + def _kv_batch(self, raw: dict[str, Any], sequence: int) -> kv_pb2.KvEventBatch: + rank = raw.get("attention_dp_rank", 0) + group_idx = raw.get("layer_group_id", 0) + block_size = self._kv_capacity().get("tokensPerBlock") + hash_algo = raw.get("hash_algo") + data = raw.get("data", raw) + event_type = data.get("type") + events: list[kv_pb2.KvEvent] = [] + if event_type == "stored": + if block_size is None or block_size <= 0: + raise ValueError("Cannot bridge KV events without a configured tokens_per_block") + parent = data.get("parent_hash") + blocks = data.get("blocks", []) + complete_blocks = [] + for block in blocks: + block_hash = block.get("block_hash") + tokens = block.get("tokens", []) + if len(tokens) > block_size: + raise ValueError("TRT-LLM KV event block exceeds tokens_per_block") + if len(tokens) < block_size: + self._partial_block_hashes.setdefault(rank, set()).add(block_hash) + break + if block.get("cache_salt") not in (None, ""): + return self._kv_fail_closed_batch( + rank, + sequence, + "Salted TRT-LLM KV event cannot be represented by OpenEngine", + ) + if any(token.get("token_extra_id", 0) for token in tokens): + return self._kv_fail_closed_batch( + rank, + sequence, + "Eagle TRT-LLM KV event cannot be represented by OpenEngine", + ) + token_ids = [token.get("token_id") for token in tokens] + if any( + not isinstance(token_id, int) or not 0 <= token_id < (1 << 32) + for token_id in token_ids + ): + return self._kv_fail_closed_batch( + rank, + sequence, + "Non-text TRT-LLM KV event contains non-integer tokens", + ) + complete_blocks.append(block) + if complete_blocks: + lora_ids = { + int(block["lora_id"]) + for block in complete_blocks + if block.get("lora_id") is not None + } + lora_names = { + str(block.get("lora_name") or data.get("lora_name")) + for block in complete_blocks + if block.get("lora_name") or data.get("lora_name") + } + if len(lora_ids) > 1 or len(lora_names) > 1: + return self._kv_fail_closed_batch( + rank, + sequence, + "TRT-LLM KV event mixes LoRA identities", + ) + lora_id = next(iter(lora_ids), None) + lora_name = next( + iter(lora_names), + self._lora_names_by_id.get(lora_id, "") if lora_id is not None else "", + ) + mm_keys = [] + for block_index, block in enumerate(complete_blocks): + for key in block.get("mm_keys", []): + mm_hash = key.get("hash") + if ( + key.get("type") != "mm_key" + or not isinstance(mm_hash, str) + or not mm_hash + ): + return self._kv_fail_closed_batch( + rank, + sequence, + "TRT-LLM KV event contains malformed multimodal metadata", + ) + try: + mm_hash_prefix = str(int(mm_hash[:16], 16)) + except ValueError: + return self._kv_fail_closed_batch( + rank, + sequence, + "TRT-LLM KV event contains a non-hex multimodal hash", + ) + mm_keys.append( + kv_pb2.OpaqueKeyTuple( + values=[ + "trt_mm_v1", + str(block_index), + mm_hash_prefix, + str(key.get("start_offset", 0)), + ] + ) + ) + stored = kv_pb2.BlockStored( + block_hashes=[ + _block_hash(block.get("block_hash"), hash_algo) for block in complete_blocks + ], + token_ids=[ + token["token_id"] + for block in complete_blocks + for token in block.get("tokens", []) + ], + block_size=block_size, + lora_id=0 if lora_id is None else int(lora_id), + lora_name=lora_name, + medium=kv_pb2.STORAGE_MEDIUM_GPU, + extra_keys=mm_keys, + group_idx=group_idx, + ) + if parent is not None: + stored.parent_block_hash.CopyFrom(_block_hash(parent, hash_algo)) + events.append(kv_pb2.KvEvent(block_stored=stored)) + else: + stored = kv_pb2.BlockStored( + block_size=block_size, + medium=kv_pb2.STORAGE_MEDIUM_GPU, + group_idx=group_idx, + ) + if parent is not None: + stored.parent_block_hash.CopyFrom(_block_hash(parent, hash_algo)) + events.append(kv_pb2.KvEvent(block_stored=stored)) + elif event_type == "removed": + partial_hashes = self._partial_block_hashes.setdefault(rank, set()) + removed_hashes = [] + for value in data.get("block_hashes", []): + if value in partial_hashes: + partial_hashes.discard(value) + else: + removed_hashes.append(value) + events.append( + kv_pb2.KvEvent( + block_removed=kv_pb2.BlockRemoved( + block_hashes=[_block_hash(value, hash_algo) for value in removed_hashes], + medium=kv_pb2.STORAGE_MEDIUM_GPU, + group_idx=group_idx, + ) + ) + ) + elif event_type == "all_cleared": + self._partial_block_hashes.pop(rank, None) + events.append(kv_pb2.KvEvent(all_blocks_cleared=kv_pb2.AllBlocksCleared())) + return kv_pb2.KvEventBatch( + sequence_number=sequence, + timestamp_unix_nanos=time.time_ns(), + data_parallel_rank=rank, + events=events, + ) + + @staticmethod + def _kv_fail_closed_batch(rank: int, sequence: int, message: str) -> kv_pb2.KvEventBatch: + logger.error("%s; clearing the advertised KV index", message) + return kv_pb2.KvEventBatch( + sequence_number=sequence, + timestamp_unix_nanos=time.time_ns(), + data_parallel_rank=rank, + events=[kv_pb2.KvEvent(all_blocks_cleared=kv_pb2.AllBlocksCleared())], + ) + + async def Embed(self, request: object, context: grpc.aio.ServicerContext) -> None: + del request + await context.abort( + grpc.StatusCode.UNIMPLEMENTED, "Embedding is not implemented by this server" + ) + + async def Classify(self, request: object, context: grpc.aio.ServicerContext) -> None: + del request + await context.abort( + grpc.StatusCode.UNIMPLEMENTED, "Classification is not implemented by this server" + ) + + async def Score(self, request: object, context: grpc.aio.ServicerContext) -> None: + del request + await context.abort( + grpc.StatusCode.UNIMPLEMENTED, "Scoring is not implemented by this server" + ) + + async def SubscribeRuntimeEvents( + self, request: object, context: grpc.aio.ServicerContext + ) -> None: + del request + await context.abort( + grpc.StatusCode.UNIMPLEMENTED, "Runtime event streaming is not implemented" + ) diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_event_manager.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_event_manager.py index c308790f9c29..114e414ed808 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_event_manager.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_event_manager.py @@ -91,6 +91,8 @@ class KVCacheStoredBlockData: priority: int mm_keys: list[MmKey] = field(default_factory=list) cache_salt: str | None = None + lora_id: int | None = None + lora_name: str | None = None @dataclass(slots=True, frozen=True) @@ -552,12 +554,19 @@ def _stored_block_from_radix_block( if life_cycle_ids is not None and not found_page: return None + root = block + while self._is_radix_block(root): + root = root.prev + lora_id, cache_salt = self._root_attrs_from_root_block(root) + return KVCacheStoredBlockData( block_hash=self._hash_from_radix_block(block), tokens=[self._normalize_token(token) for token in block.tokens], cache_level=int(cache_level), priority=int(priority), mm_keys=[], + cache_salt=None if cache_salt is None else str(cache_salt), + lora_id=lora_id, ) @staticmethod diff --git a/tensorrt_llm/serve/kv_event_fanout.py b/tensorrt_llm/serve/kv_event_fanout.py new file mode 100644 index 000000000000..60fce03c8a31 --- /dev/null +++ b/tensorrt_llm/serve/kv_event_fanout.py @@ -0,0 +1,134 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Single-consumer TRT-LLM KV-event pump with in-process fanout.""" + +import asyncio +from collections import deque +from collections.abc import AsyncGenerator +from typing import Any + +from tensorrt_llm.logger import logger + + +class KvEventFanout: + """Own TRT-LLM's event queue and copy events to every protocol.""" + + def __init__(self, llm: object, buffer_size: int = 1024) -> None: + if buffer_size <= 0: + raise ValueError("buffer_size must be positive") + self._llm = llm + self._buffer: deque[dict[str, Any]] = deque(maxlen=buffer_size) + self._subscribers: dict[asyncio.Queue[tuple[int, dict[str, Any]]], frozenset[int]] = {} + self._task: asyncio.Task[None] | None = None + self._sequence_numbers: dict[int, int] = {} + self._last_engine_event_ids: dict[int, int] = {} + self._max_window_size: int | None = None + self._processing_initial_created_events = True + + @property + def buffer_size(self) -> int: + return self._buffer.maxlen or 0 + + def start(self) -> None: + if self._task is None: + self._task = asyncio.create_task(self._pump()) + + async def stop(self) -> None: + if self._task is None: + return + self._task.cancel() + try: + await self._task + except asyncio.CancelledError: + pass + self._task = None + + def drain_http_buffer(self) -> list[dict[str, Any]]: + events = list(self._buffer) + self._buffer.clear() + return events + + async def subscribe( + self, data_parallel_ranks: set[int] | None = None + ) -> AsyncGenerator[tuple[int, dict[str, Any]], None]: + self.start() + queue: asyncio.Queue[tuple[int, dict[str, Any]]] = asyncio.Queue(maxsize=self.buffer_size) + self._subscribers[queue] = frozenset(data_parallel_ranks or ()) + try: + while True: + yield await queue.get() + finally: + self._subscribers.pop(queue, None) + + def _is_routable_event(self, event: dict[str, Any]) -> bool: + event_type = event.get("data", event).get("type") + if event_type == "created" and self._processing_initial_created_events: + window_size = event.get("window_size") + if window_size is not None: + self._max_window_size = max(self._max_window_size or 0, int(window_size)) + return False + self._processing_initial_created_events = False + if event_type not in ("stored", "removed"): + return False + window_size = event.get("window_size") + return ( + window_size is None + or self._max_window_size is None + or int(window_size) == self._max_window_size + ) + + def _publish(self, event: dict[str, Any]) -> None: + rank = int(event.get("attention_dp_rank", 0)) + event_id = event.get("event_id") + if isinstance(event_id, int) and not isinstance(event_id, bool): + previous_event_id = self._last_engine_event_ids.get(rank) + if (previous_event_id is None and event_id != 0) or ( + previous_event_id is not None and event_id != previous_event_id + 1 + ): + logger.warning( + "TRT-LLM KV event gap on DP rank %d: expected %d, received %d; " + "clearing the advertised KV index", + rank, + 0 if previous_event_id is None else previous_event_id + 1, + event_id, + ) + self._publish_routable( + { + "attention_dp_rank": rank, + "data": {"type": "all_cleared"}, + }, + include_http_raw_event=False, + ) + self._last_engine_event_ids[rank] = event_id + if not self._is_routable_event(event): + return + self._publish_routable(event) + + def _publish_routable( + self, event: dict[str, Any], *, include_http_raw_event: bool = True + ) -> None: + rank = int(event.get("attention_dp_rank", 0)) + sequence = self._sequence_numbers.get(rank, 0) + 1 + self._sequence_numbers[rank] = sequence + sequenced_event = (sequence, event) + if include_http_raw_event: + self._buffer.append(event) + for queue, selected_ranks in tuple(self._subscribers.items()): + if selected_ranks and rank not in selected_ranks: + continue + try: + queue.put_nowait(sequenced_event) + except asyncio.QueueFull: + logger.warning("Dropping KV event for a slow OpenEngine subscriber") + + async def _pump(self) -> None: + while True: + try: + async for event in self._llm.get_kv_cache_events_async(1): + self._publish(event) + except (IndexError, asyncio.QueueEmpty): + await asyncio.sleep(0) + except (RuntimeError, AttributeError) as error: + logger.debug("KV event pump unavailable: %s", error) + await asyncio.sleep(1) diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index fd5fb55c25c0..030f80dfb552 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -65,6 +65,7 @@ from tensorrt_llm.serve.disagg_auto_scaling import DisaggClusterWorker from tensorrt_llm.serve.encode_batcher import (EncodeBatcher, InputTooLongError, QueueFullError) +from tensorrt_llm.serve.kv_event_fanout import KvEventFanout from tensorrt_llm.serve.metadata_server import create_metadata_server from tensorrt_llm.serve.openai_protocol import ( ChatCompletionRequest, ChatCompletionResponse, ChatCompletionResponseChoice, @@ -92,6 +93,7 @@ from tensorrt_llm.serve.responses_utils import get_steady_clock_now_in_seconds from tensorrt_llm.serve.responses_utils import \ request_preprocess as responses_api_request_preprocess +from tensorrt_llm.serve.stats_fanout import StatsFanout from tensorrt_llm.serve.tool_parser.tool_parser_factory import ToolParserFactory from tensorrt_llm.serve.visual_gen_metrics import \ build_visual_gen_timing_headers @@ -265,7 +267,10 @@ def __init__( embedding_max_queue_delay: float = 0.005, embedding_max_queue_size: int = 2048, input_processor_workers: int = 8, - media_load_workers: int = 8): + media_load_workers: int = 8, + kv_event_fanout: Optional[KvEventFanout] = None, + stats_fanout: Optional[StatsFanout] = None, + shutdown_generator: bool = True): self.generator = generator self._is_visual_gen = isinstance(generator, VisualGen) self._embedding_max_queue_delay = embedding_max_queue_delay @@ -290,6 +295,9 @@ def __init__( self.multimodal_server_config = cfg self.allow_request_chat_template = allow_request_chat_template self.server_role = server_role + self.kv_event_fanout = kv_event_fanout + self.stats_fanout = stats_fanout + self.shutdown_generator = shutdown_generator # Will be set in __call__ self.binding_addr = None self.host = None @@ -346,6 +354,8 @@ def __init__( @asynccontextmanager async def lifespan(app: FastAPI): + if self.kv_event_fanout is not None: + self.kv_event_fanout.start() if self.metadata_server is not None: metadata = { "model": self.model, @@ -384,7 +394,12 @@ async def lifespan(app: FastAPI): # Start background iteration stats collector if metrics are enabled # The args for pytorch and autodeploy backend has attribute `enable_iter_perf_stats` while # tensorrt backend does not have this attribute but it always has iter stats enabled. - if self.metrics_collector and getattr( + if self.stats_fanout is not None and getattr( + self.generator.args, "enable_iter_perf_stats", True): + consumer = (self.metrics_collector.log_iteration_stats + if self.metrics_collector is not None else None) + self.stats_fanout.start(consumer) + elif self.metrics_collector and getattr( self.generator.args, "enable_iter_perf_stats", True): # The background loop becomes the sole consumer of the # engine stats queue; /metrics reads from a tee buffer @@ -432,7 +447,12 @@ async def lifespan(app: FastAPI): await self.disagg_cluster_worker.deregister_worker() if self.resource_governor is not None: self.resource_governor.close() - self.generator.shutdown() + if self.kv_event_fanout is not None: + await self.kv_event_fanout.stop() + if self.stats_fanout is not None: + await self.stats_fanout.stop() + if self.shutdown_generator: + self.generator.shutdown() self.app = FastAPI(lifespan=lifespan) if _MSGSPEC_ENABLED: @@ -1193,6 +1213,9 @@ async def get_iteration_stats(self) -> JSONResponse: stats.append(stat) return JSONResponse(content=stats) + if self.stats_fanout is not None: + return JSONResponse(content=self.stats_fanout.drain_http_buffer()) + # When the background collector loop is active it is the sole # consumer of the engine stats queue; serve /metrics from the tee # buffer it populates so we do not race it for queue items. Racing @@ -1316,6 +1339,9 @@ async def get_perf_metrics(self) -> JSONResponse: return JSONResponse(content=list(perf_metrics)) async def get_kv_cache_events(self) -> JSONResponse: + if self.kv_event_fanout is not None: + return JSONResponse( + content=self.kv_event_fanout.drain_http_buffer()) events = [] try: async for event in self.generator.get_kv_cache_events_async(0): @@ -1337,10 +1363,13 @@ async def _extract_metrics(self, res: RequestOutput, raw_request: Request): # (e.g. PostprocWorker). self.metrics_collector.log_request_metrics_dict( res.metrics_dict) - # Note: Iteration stats are collected by the background _iteration_stats_collector_loop task - # Wake up the stats collector to drain iteration stats - if getattr(self.generator.args, "enable_iter_perf_stats", True): + # Note: Iteration stats are collected by the background collector. + if (self.stats_fanout is None and getattr( + self.generator.args, "enable_iter_perf_stats", True)): self._iteration_stats_wakeup_event.set() + if self.stats_fanout is not None and getattr( + self.generator.args, "enable_iter_perf_stats", True): + self.stats_fanout.wake() if self.generator.args.return_perf_metrics: output = res.outputs[0] item = { diff --git a/tensorrt_llm/serve/request_tracker.py b/tensorrt_llm/serve/request_tracker.py new file mode 100644 index 000000000000..8c6ebe3c8674 --- /dev/null +++ b/tensorrt_llm/serve/request_tracker.py @@ -0,0 +1,156 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Protocol-neutral tracking for requests submitted to an LLM instance.""" + +import asyncio +from collections.abc import Iterable +from typing import Protocol + +from tensorrt_llm.logger import logger + + +class AbortableResult(Protocol): + def abort(self) -> None: + """Abort the associated engine request.""" + + +class RequestTracker: + """Track active engine results for admission, cancellation, and drain.""" + + def __init__(self, llm: object) -> None: + self.llm = llm + self._requests: dict[str, AbortableResult] = {} + self._external_requests = 0 + self._draining = False + self._changed = asyncio.Condition() + + @property + def draining(self) -> bool: + return self._draining + + @property + def active_count(self) -> int: + return len(self._requests) + self._external_requests + + @property + def active_requests(self) -> dict[str, AbortableResult]: + """Compatibility view for protocol request managers.""" + return self._requests + + def admit(self, request_id: str, result: AbortableResult) -> None: + if self._draining: + raise RuntimeError("TensorRT-LLM is draining") + if request_id in self._requests: + raise ValueError(f"Request {request_id!r} is already active") + self._requests[request_id] = result + + async def finish(self, request_id: str) -> None: + async with self._changed: + self._requests.pop(request_id, None) + self._changed.notify_all() + + def begin_external(self) -> None: + if self._draining: + raise RuntimeError("TensorRT-LLM is draining") + self._external_requests += 1 + + async def finish_external(self) -> None: + async with self._changed: + self._external_requests = max(0, self._external_requests - 1) + self._changed.notify_all() + + async def abort(self, request_id: str) -> bool: + result = self._requests.get(request_id) + if result is None: + return False + try: + result.abort() + except (RuntimeError, AssertionError) as error: + logger.warning("Failed to abort request %s: %s", request_id, error) + return False + finally: + await self.finish(request_id) + return True + + async def abort_all(self) -> int: + tracked_by_identity = {id(result): result for result in self._requests.values()} + executor = getattr(self.llm, "_executor", None) + executor_results = getattr(executor, "_results", None) + if isinstance(executor_results, dict): + tracked_by_identity.update( + (id(result), result) for result in tuple(executor_results.values()) + ) + + aborted = 0 + for result in tracked_by_identity.values(): + try: + result.abort() + except (RuntimeError, AssertionError) as error: + logger.warning("Failed to abort engine request: %s", error) + else: + aborted += 1 + + async with self._changed: + self._requests.clear() + self._changed.notify_all() + return aborted + + async def start_drain(self) -> int: + async with self._changed: + self._draining = True + self._changed.notify_all() + return self.active_count + + async def wait_empty(self, timeout: float | None = None) -> bool: + async def _wait() -> None: + async with self._changed: + await self._changed.wait_for(lambda: self.active_count == 0) + + if self.active_count == 0: + return True + try: + await asyncio.wait_for(_wait(), timeout=timeout) + except asyncio.TimeoutError: + return False + return True + + async def health(self) -> tuple[bool, str]: + if not hasattr(self.llm, "_executor"): + return True, "OK" + executor = getattr(self.llm, "_executor", None) + if executor is None: + return False, "Executor is not available" + try: + healthy = executor.check_health() + except (RuntimeError, AttributeError) as error: + return False, f"Executor health check failed: {error}" + if healthy: + return True, "OK" + fatal_error = getattr(executor, "_fatal_error", None) + if fatal_error is None: + return False, "Executor is unhealthy" + lines = str(fatal_error).splitlines() + short = (lines[0] if lines else type(fatal_error).__name__)[:200] + return False, f"{type(fatal_error).__name__}: {short}" + + def iter_results(self) -> Iterable[AbortableResult]: + return tuple(self._requests.values()) + + +async def track_http_response(response: object, tracker: RequestTracker) -> object: + """Keep HTTP admission active until a streaming response body closes.""" + body_iterator = getattr(response, "body_iterator", None) + if body_iterator is None: + await tracker.finish_external() + return response + + async def tracked_body(): + try: + async for chunk in body_iterator: + yield chunk + finally: + await tracker.finish_external() + + response.body_iterator = tracked_body() + return response diff --git a/tensorrt_llm/serve/stats_fanout.py b/tensorrt_llm/serve/stats_fanout.py new file mode 100644 index 000000000000..7b2093efe3eb --- /dev/null +++ b/tensorrt_llm/serve/stats_fanout.py @@ -0,0 +1,81 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Single-consumer TRT-LLM iteration-stat pump with shared snapshots.""" + +import asyncio +from collections import deque +from collections.abc import Callable +from typing import Any + +from tensorrt_llm.logger import logger + + +class StatsFanout: + """Own the engine stats queue and tee it to HTTP, metrics, and OpenEngine.""" + + def __init__(self, llm: object, buffer_size: int | None = 1000) -> None: + self._llm = llm + self._buffer: deque[dict[str, Any]] = deque(maxlen=buffer_size) + self._latest_by_rank: dict[int, dict[str, Any]] = {} + self._consumer: Callable[[dict[str, Any]], None] | None = None + self._wake_event: asyncio.Event | None = None + self._task: asyncio.Task[None] | None = None + + def start(self, consumer: Callable[[dict[str, Any]], None] | None = None) -> None: + self._consumer = consumer + if self._task is None: + self._wake_event = asyncio.Event() + self._task = asyncio.create_task(self._pump()) + + def wake(self) -> None: + if self._wake_event is not None: + self._wake_event.set() + + async def stop(self) -> None: + if self._task is None: + return + self._task.cancel() + try: + await self._task + except asyncio.CancelledError: + pass + self._task = None + self._wake_event = None + + def drain_http_buffer(self) -> list[dict[str, Any]]: + stats = list(self._buffer) + self._buffer.clear() + return stats + + def latest_by_rank(self) -> dict[int, dict[str, Any]]: + return {rank: dict(stat) for rank, stat in self._latest_by_rank.items()} + + def _publish(self, stat: dict[str, Any]) -> None: + rank = int(stat.get("attentionDpRank", stat.get("attention_dp_rank", 0))) + self._latest_by_rank[rank] = stat + self._buffer.append(stat) + if self._consumer is not None: + try: + self._consumer(stat) + except (RuntimeError, ValueError, TypeError) as error: + logger.error("Iteration-stat consumer failed: %s", error) + + async def _pump(self) -> None: + assert self._wake_event is not None + while True: + await self._wake_event.wait() + self._wake_event.clear() + try: + async for stat in self._llm.get_stats_async(timeout=0.5): + self._publish(stat) + except ( + RuntimeError, + AttributeError, + IndexError, + TypeError, + ValueError, + asyncio.QueueEmpty, + ) as error: + logger.error("Error collecting iteration stats: %s", error) + await asyncio.sleep(0.1) diff --git a/tests/unittest/_torch/modeling/test_modeling_phi4mm.py b/tests/unittest/_torch/modeling/test_modeling_phi4mm.py index 58c7de77f938..42013d661070 100644 --- a/tests/unittest/_torch/modeling/test_modeling_phi4mm.py +++ b/tests/unittest/_torch/modeling/test_modeling_phi4mm.py @@ -2,6 +2,7 @@ from dataclasses import dataclass from typing import List +import pytest import torch import transformers from test_modeling_multimodal import MultimodalScenario, TestModelingMultimodal, llm_models_root @@ -10,9 +11,25 @@ _AUDIO_SPECIAL_TOKEN_ID, _IMAGE_SPECIAL_TOKEN_ID, Phi4MMForCausalLM, + Phi4MMInputProcessor, ) from tensorrt_llm.inputs import default_multimodal_input_loader, prompt_inputs + +def test_phi4mm_required_multimodal_lora_selection(tmp_path) -> None: + processor = object.__new__(Phi4MMInputProcessor) + processor._model_path = str(tmp_path) + + assert processor.get_model_owned_lora_identities() == { + "vision-lora": 0, + "speech-lora": 1, + } + assert processor.get_required_lora_spec(("image",)).name == "vision-lora" + assert processor.get_required_lora_spec(("audio",)).name == "speech-lora" + with pytest.raises(ValueError, match="cannot combine"): + processor.get_required_lora_spec(("image", "audio")) + + PHI4MM_CONFIG = { "_name_or_path": str(os.path.join(llm_models_root(), "multimodals/Phi-4-multimodal-instruct")), "architectures": ["Phi4MMForCausalLM"], diff --git a/tests/unittest/api_stability/references/trtllm_serve_cli.yaml b/tests/unittest/api_stability/references/trtllm_serve_cli.yaml index 4a80a5095240..67c595b6f8fe 100644 --- a/tests/unittest/api_stability/references/trtllm_serve_cli.yaml +++ b/tests/unittest/api_stability/references/trtllm_serve_cli.yaml @@ -321,6 +321,24 @@ commands: is_flag: false flags: - "--num_media_load_workers" + openengine_host: + type: str + default: 127.0.0.1 + status: prototype + required: false + multiple: false + is_flag: false + flags: + - "--openengine-host" + openengine_port: + type: int + default: null + status: prototype + required: false + multiple: false + is_flag: false + flags: + - "--openengine-port" otlp_traces_endpoint: type: str default: null diff --git a/tests/unittest/openengine/test_converters.py b/tests/unittest/openengine/test_converters.py new file mode 100644 index 000000000000..b1dc4236eb4f --- /dev/null +++ b/tests/unittest/openengine/test_converters.py @@ -0,0 +1,225 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import json + +import pytest +from openengine.v1 import generation_pb2, input_pb2, kv_pb2 + +from tensorrt_llm.disaggregated_params import DisaggregatedParams, DisaggScheduleStyle +from tensorrt_llm.executor.result import Logprob +from tensorrt_llm.openengine.converters import ( + HANDOFF_ATTRIBUTE, + decode_handoff, + encode_handoff, + handoff_requires_decode_media, + load_media, + to_priority, + to_sampling_params, +) + + +def test_context_first_handoff_round_trip_large_ids_and_binary() -> None: + params = DisaggregatedParams( + request_type="context_only", + first_gen_tokens=[11, 12], + first_gen_log_probs=[ + { + 11: Logprob(logprob=-0.25, rank=1), + 91: Logprob(logprob=-2.0, rank=2), + }, + {12: Logprob(logprob=-0.5, rank=1)}, + ], + ctx_request_id=2**63 - 2, + disagg_request_id=2**63 - 1, + ctx_dp_rank=7, + ctx_info_endpoint="tcp://context:1234", + draft_tokens=[21, 22], + ctx_usage={"prompt_tokens": 9}, + conversation_id="conversation", + schedule_style=DisaggScheduleStyle.CONTEXT_FIRST, + opaque_state=b"\x00\xffbinary", + ) + + session = encode_handoff(params) + payload = json.loads(session.attributes_struct[HANDOFF_ATTRIBUTE]) + assert payload["ctx_request_id"] == str(2**63 - 2) + assert payload["disagg_request_id"] == str(2**63 - 1) + + restored = decode_handoff(session) + assert restored.request_type == "generation_only" + assert restored.ctx_request_id == 2**63 - 2 + assert restored.disagg_request_id == 2**63 - 1 + assert restored.opaque_state == b"\x00\xffbinary" + assert restored.first_gen_log_probs[0][11] == Logprob(-0.25, 1) + assert restored.first_gen_log_probs[0][91] == Logprob(-2.0, 2) + assert restored.first_gen_log_probs[1][12] == Logprob(-0.5, 1) + + +def test_raw_media_handoff_omits_transient_mrope_and_requires_decode_media() -> None: + params = DisaggregatedParams( + request_type="context_only", + schedule_style=DisaggScheduleStyle.CONTEXT_FIRST, + mrope_position_ids_handle={"tensor": "ids"}, + mrope_position_deltas_handle={"tensor": "deltas"}, + ) + + with pytest.raises(ValueError, match="mrope"): + encode_handoff(params) + + session = encode_handoff(params, requires_decode_media=True) + assert handoff_requires_decode_media(session) + restored = decode_handoff(session) + assert restored.mrope_position_ids_handle is None + assert restored.mrope_position_deltas_handle is None + + +@pytest.mark.parametrize( + "params, message", + [ + ( + DisaggregatedParams( + request_type="context_only", + schedule_style=DisaggScheduleStyle.GENERATION_FIRST, + ), + "Generation-first", + ), + (DisaggregatedParams(request_type="context_only", first_gen_logits=[object()]), "logits"), + ( + DisaggregatedParams( + request_type="context_only", + multimodal_embedding_handles=[{"handle": "encoder"}], + ), + "embedding handles", + ), + ], +) +def test_handoff_rejects_out_of_scope_topologies(params: DisaggregatedParams, message: str) -> None: + with pytest.raises(ValueError, match=message): + encode_handoff(params) + + +def test_decode_rejects_non_decimal_identifier() -> None: + session = kv_pb2.KvSessionRef() + session.attributes_struct[HANDOFF_ATTRIBUTE] = json.dumps( + { + "schedule_style": "context_first", + "ctx_request_id": 9007199254740993, + } + ) + with pytest.raises(ValueError, match="decimal string"): + decode_handoff(session) + + +def test_decode_rejects_multimodal_hashes_even_without_embedding_handles() -> None: + session = kv_pb2.KvSessionRef() + session.attributes_struct[HANDOFF_ATTRIBUTE] = json.dumps( + { + "schedule_style": "context_first", + "multimodal_hashes": [[1, 2, 3, 4, 5, 6, 7, 8]], + } + ) + + with pytest.raises(ValueError, match="handles"): + decode_handoff(session) + + +def test_decode_rejects_malformed_context_usage() -> None: + session = kv_pb2.KvSessionRef() + session.attributes_struct[HANDOFF_ATTRIBUTE] = json.dumps( + { + "schedule_style": "context_first", + "ctx_usage": {"prompt_tokens": "2"}, + } + ) + + with pytest.raises(ValueError, match="prompt_tokens"): + decode_handoff(session) + + +def test_sampling_conversion_preserves_requested_controls() -> None: + request = generation_pb2.GenerateRequest( + sampling=generation_pb2.SamplingParams( + temperature=0.25, + top_p=0.8, + top_k=16, + seed=2**63, + num_sequences=2, + ), + stopping=generation_pb2.StoppingOptions( + max_tokens=8, + min_tokens=2, + conditions=[ + generation_pb2.StopCondition(stop_text="done"), + generation_pb2.StopCondition(stop_token_id=42), + ], + ignore_eos=True, + include_stop_in_output=True, + ), + response=generation_pb2.ResponseOptions( + return_output_logprobs=True, + output_candidates=generation_pb2.CandidateTokenSelection(top_n=4), + ), + guided=generation_pb2.GuidedDecoding(regex="[a-z]+"), + ) + params = to_sampling_params(request) + assert params.max_tokens == 8 + assert params.n == 2 + assert params.best_of == 2 + assert params.stop == ["done"] + assert params.stop_token_ids == [42] + assert params.logprobs == 4 + assert params.guided_decoding.regex == "[a-z]+" + + +def test_choice_guidance_escapes_literal_regex_characters() -> None: + request = generation_pb2.GenerateRequest() + request.guided.choice.choices.extend(["a.b", "c+(d)"]) + params = to_sampling_params(request) + assert params.guided_decoding.regex == r"^(?:(?:a\.b)|(?:c\+\(d\)))$" + + +def test_priority_mapping_is_centered_bounded_and_strictly_monotonic() -> None: + absent = generation_pb2.GenerateRequest() + negative = generation_pb2.GenerateRequest(priority=-1000) + zero = generation_pb2.GenerateRequest(priority=0) + positive = generation_pb2.GenerateRequest(priority=1000) + + assert 0 < to_priority(negative) < to_priority(zero) + assert to_priority(zero) == to_priority(absent) == 0.5 + assert to_priority(zero) < to_priority(positive) < 1 + + +@pytest.mark.asyncio +async def test_media_preserves_per_modality_order_and_merge_inputs(monkeypatch) -> None: + calls = [] + + class _MediaIO: + @classmethod + def create(cls, defaults, request): + calls.append((defaults, request)) + return cls() + + async def _run_in_executor(self, function, data): + return function(data) + + def load_bytes(self, data): + return data.decode() + + monkeypatch.setattr( + "tensorrt_llm.openengine.converters.MEDIA_IO_REGISTRY", + {"image": _MediaIO, "video": _MediaIO, "audio": _MediaIO}, + ) + options = generation_pb2.GenerateRequest().media_options + options.update({"image": {"format": "pil"}}) + config = type("Config", (), {"media_io_kwargs": {"image": {"device": "cpu"}}})() + media = [ + input_pb2.MediaItem(modality=input_pb2.MODALITY_IMAGE, raw_bytes=b"one"), + input_pb2.MediaItem(modality=input_pb2.MODALITY_VIDEO, raw_bytes=b"video"), + input_pb2.MediaItem(modality=input_pb2.MODALITY_IMAGE, raw_bytes=b"two"), + ] + + decoded = await load_media(media, options, config) + + assert decoded == {"image": ["one", "two"], "video": ["video"]} + assert calls[0] == ({"device": "cpu"}, {"format": "pil"}) diff --git a/tests/unittest/openengine/test_fanout.py b/tests/unittest/openengine/test_fanout.py new file mode 100644 index 000000000000..623910b7c0fa --- /dev/null +++ b/tests/unittest/openengine/test_fanout.py @@ -0,0 +1,105 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import asyncio + +import pytest + +from tensorrt_llm.serve.kv_event_fanout import KvEventFanout +from tensorrt_llm.serve.stats_fanout import StatsFanout + + +def test_kv_fanout_sequence_exposes_subscriber_queue_drops() -> None: + fanout = KvEventFanout(object(), buffer_size=1) + queue = asyncio.Queue(maxsize=1) + fanout._subscribers[queue] = frozenset({0}) + + first = {"data": {"type": "stored"}, "event": 1, "attention_dp_rank": 0} + second = {"data": {"type": "stored"}, "event": 2, "attention_dp_rank": 0} + third = {"data": {"type": "stored"}, "event": 3, "attention_dp_rank": 0} + other_rank = {"data": {"type": "stored"}, "event": 20, "attention_dp_rank": 1} + fanout._publish(first) + fanout._publish(second) + assert queue.get_nowait() == (1, first) + + fanout._publish(other_rank) + fanout._publish(third) + assert queue.get_nowait() == (3, third) + assert fanout.drain_http_buffer() == [third] + assert fanout._sequence_numbers == {0: 3, 1: 1} + + +def test_kv_fanout_filters_non_global_attention_without_sequence_gaps() -> None: + fanout = KvEventFanout(object()) + fanout._publish({"window_size": 128, "data": {"type": "created"}}) + fanout._publish({"window_size": 4096, "data": {"type": "created"}}) + fanout._publish({"window_size": 128, "data": {"type": "stored"}}) + global_event = {"window_size": 4096, "data": {"type": "stored"}} + fanout._publish(global_event) + + assert fanout._sequence_numbers == {0: 1} + assert fanout.drain_http_buffer() == [global_event] + + +def test_kv_fanout_detects_raw_event_id_gap_before_attention_filtering() -> None: + fanout = KvEventFanout(object()) + queue = asyncio.Queue() + fanout._subscribers[queue] = frozenset({0}) + fanout._publish({"event_id": 0, "window_size": 128, "data": {"type": "created"}}) + fanout._publish({"event_id": 1, "window_size": 4096, "data": {"type": "created"}}) + fanout._publish({"event_id": 2, "window_size": 128, "data": {"type": "stored"}}) + global_event = {"event_id": 3, "window_size": 4096, "data": {"type": "stored"}} + fanout._publish(global_event) + fanout._publish({"event_id": 5, "window_size": 128, "data": {"type": "stored"}}) + next_global = {"event_id": 6, "window_size": 4096, "data": {"type": "removed"}} + fanout._publish(next_global) + + reset = {"attention_dp_rank": 0, "data": {"type": "all_cleared"}} + assert fanout._last_engine_event_ids == {0: 6} + assert fanout._sequence_numbers == {0: 3} + assert [queue.get_nowait() for _ in range(3)] == [ + (1, global_event), + (2, reset), + (3, next_global), + ] + assert fanout.drain_http_buffer() == [global_event, next_global] + + +def test_synthetic_gap_reset_does_not_change_http_raw_event_contract() -> None: + fanout = KvEventFanout(object()) + queue = asyncio.Queue() + fanout._subscribers[queue] = frozenset() + fanout._publish({"event_id": 0, "data": {"type": "created"}}) + stored = {"event_id": 2, "data": {"type": "stored"}} + fanout._publish(stored) + + assert queue.get_nowait()[1]["data"]["type"] == "all_cleared" + assert queue.get_nowait() == (2, stored) + assert fanout.drain_http_buffer() == [stored] + + +@pytest.mark.asyncio +async def test_stats_fanout_is_single_consumer_for_http_metrics_and_load() -> None: + class _Llm: + calls = 0 + + async def get_stats_async(self, timeout: float): + assert timeout == 0.5 + self.calls += 1 + yield {"attentionDpRank": 0, "numActiveRequests": 1} + yield {"attentionDpRank": 1, "numActiveRequests": 2} + + llm = _Llm() + consumed = [] + fanout = StatsFanout(llm, buffer_size=4) + fanout.start(consumed.append) + fanout.wake() + for _ in range(20): + if len(consumed) == 2: + break + await asyncio.sleep(0) + await fanout.stop() + + assert llm.calls == 1 + assert consumed == fanout.drain_http_buffer() + assert set(fanout.latest_by_rank()) == {0, 1} diff --git a/tests/unittest/openengine/test_lora_registry.py b/tests/unittest/openengine/test_lora_registry.py new file mode 100644 index 000000000000..81477be16f47 --- /dev/null +++ b/tests/unittest/openengine/test_lora_registry.py @@ -0,0 +1,89 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import asyncio +from pathlib import Path + +import pytest +from openengine.v1 import lora_pb2 + +from tensorrt_llm.openengine.lora_registry import LoraRegistry + + +def _adapter(path: Path, name: str = "adapter", adapter_id: int = 7) -> lora_pb2.LoraAdapter: + path.mkdir() + (path / "adapter_config.json").write_text("{}") + (path / "adapter_model.safetensors").write_bytes(b"weights") + return lora_pb2.LoraAdapter(lora_name=name, lora_id=adapter_id, source_path=str(path)) + + +@pytest.mark.asyncio +async def test_lora_registry_is_lazy_idempotent_and_logically_unloads( + tmp_path: Path, +) -> None: + registry = LoraRegistry() + adapter = _adapter(tmp_path / "adapter") + + registered, already_loaded = await registry.load(adapter) + assert not already_loaded + repeated, already_loaded = await registry.load(adapter) + assert already_loaded + assert repeated == registered + + request = await registry.request("adapter") + assert request.lora_int_id == 7 + assert request.lora_path == str((tmp_path / "adapter").resolve()) + + await registry.unload("adapter") + with pytest.raises(KeyError): + await registry.request("adapter") + + reloaded, already_loaded = await registry.load(adapter) + assert not already_loaded + assert reloaded == registered + + +@pytest.mark.asyncio +async def test_lora_registry_rejects_conflicting_identity(tmp_path: Path) -> None: + registry = LoraRegistry() + await registry.load(_adapter(tmp_path / "one", "one", 1)) + with pytest.raises(ValueError, match="ID 1"): + await registry.load(_adapter(tmp_path / "two", "two", 1)) + + +@pytest.mark.asyncio +async def test_lora_registry_tombstones_name_id_and_path_after_unload(tmp_path: Path) -> None: + registry = LoraRegistry() + original = _adapter(tmp_path / "one", "one", 1) + await registry.load(original) + await registry.unload("one") + + with pytest.raises(ValueError, match="name 'one'.*permanently"): + await registry.load(_adapter(tmp_path / "renamed-path", "one", 2)) + with pytest.raises(ValueError, match="ID 1.*permanently"): + await registry.load(_adapter(tmp_path / "reused-id", "two", 1)) + with pytest.raises(ValueError, match="path.*permanently"): + await registry.load( + lora_pb2.LoraAdapter(lora_name="two", lora_id=2, source_path=original.source_path) + ) + + +@pytest.mark.asyncio +async def test_lora_registry_reserves_model_owned_names_and_ids(tmp_path: Path) -> None: + registry = LoraRegistry({"vision-lora": 0, "speech-lora": 1}) + + with pytest.raises(ValueError, match="name 'vision-lora'.*model-owned"): + await registry.load(_adapter(tmp_path / "reserved-name", "vision-lora", 7)) + with pytest.raises(ValueError, match="ID 0.*model-owned"): + await registry.load(_adapter(tmp_path / "reserved-id", "user-adapter", 0)) + + +@pytest.mark.asyncio +async def test_lora_registry_serializes_concurrent_idempotent_loads(tmp_path: Path) -> None: + registry = LoraRegistry() + adapter = _adapter(tmp_path / "adapter") + + results = await asyncio.gather(registry.load(adapter), registry.load(adapter)) + + assert sorted(already_loaded for _, already_loaded in results) == [False, True] + assert len(await registry.list()) == 1 diff --git a/tests/unittest/openengine/test_native_grpc_compat.py b/tests/unittest/openengine/test_native_grpc_compat.py new file mode 100644 index 000000000000..5c9a0ebc1151 --- /dev/null +++ b/tests/unittest/openengine/test_native_grpc_compat.py @@ -0,0 +1,49 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest + +from tensorrt_llm.grpc.grpc_request_manager import GrpcRequestManager + + +class _Result: + def __init__(self) -> None: + self.aborted = False + + def abort(self) -> None: + self.aborted = True + + def __aiter__(self): + return self + + async def __anext__(self): + raise StopAsyncIteration + + +class _Llm: + def __init__(self) -> None: + self.last_result = None + + def generate_async(self, *args, **kwargs): + del args, kwargs + self.last_result = _Result() + return self.last_result + + +@pytest.mark.asyncio +async def test_native_grpc_duplicate_aborts_rejected_engine_result() -> None: + llm = _Llm() + manager = GrpcRequestManager(llm) + manager._request_tracker.admit("duplicate", _Result()) + + with pytest.raises(ValueError, match="already active"): + _ = [ + result + async for result in manager.generate( + request_id="duplicate", + prompt_token_ids=[1], + sampling_params=object(), + ) + ] + + assert llm.last_result.aborted diff --git a/tests/unittest/openengine/test_request_tracker.py b/tests/unittest/openengine/test_request_tracker.py new file mode 100644 index 000000000000..5ab44e5a3e5d --- /dev/null +++ b/tests/unittest/openengine/test_request_tracker.py @@ -0,0 +1,73 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import asyncio + +import pytest + +from tensorrt_llm.serve.request_tracker import RequestTracker, track_http_response + + +class _Result: + def __init__(self) -> None: + self.aborted = False + + def abort(self) -> None: + self.aborted = True + + +@pytest.mark.asyncio +async def test_tracker_abort_drain_and_external_admission() -> None: + tracker = RequestTracker(object()) + result = _Result() + tracker.admit("request", result) + tracker.begin_external() + assert tracker.active_count == 2 + + await tracker.start_drain() + with pytest.raises(RuntimeError, match="draining"): + tracker.begin_external() + assert await tracker.abort("request") + assert result.aborted + assert not await tracker.wait_empty(0) + await tracker.finish_external() + assert await tracker.wait_empty(0) + + +@pytest.mark.asyncio +async def test_streaming_http_admission_lives_until_body_closes() -> None: + tracker = RequestTracker(object()) + tracker.begin_external() + + async def body(): + yield b"first" + yield b"second" + + response = type("Response", (), {})() + response.body_iterator = body() + await track_http_response(response, tracker) + + assert tracker.active_count == 1 + chunks = [chunk async for chunk in response.body_iterator] + assert chunks == [b"first", b"second"] + assert tracker.active_count == 0 + + +@pytest.mark.asyncio +async def test_abort_all_reaches_untracked_http_engine_results() -> None: + result = _Result() + llm = type("Llm", (), {})() + llm._executor = type("Executor", (), {"_results": {1: result}})() + tracker = RequestTracker(llm) + tracker.begin_external() + + async def finish_http() -> None: + while not result.aborted: + await asyncio.sleep(0) + await tracker.finish_external() + + finish_task = asyncio.create_task(finish_http()) + assert await tracker.abort_all() == 1 + assert await tracker.wait_empty() + await finish_task + assert result.aborted diff --git a/tests/unittest/openengine/test_server.py b/tests/unittest/openengine/test_server.py new file mode 100644 index 000000000000..213046595d45 --- /dev/null +++ b/tests/unittest/openengine/test_server.py @@ -0,0 +1,86 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from pathlib import Path + +import pytest + +from tensorrt_llm.openengine._schema_pin import OPENENGINE_COMMIT +from tensorrt_llm.openengine.server import OpenEngineServer, validate_schema_release +from tensorrt_llm.openengine.servicer import schema_release + + +def test_schema_release_accepts_exact_pinned_identity() -> None: + assert validate_schema_release(OPENENGINE_COMMIT) == OPENENGINE_COMMIT + + +def test_packaged_schema_pin_matches_repository_pin() -> None: + repository_pin = ( + (Path(__file__).resolve().parents[3] / "OPENENGINE_COMMIT") + .read_text(encoding="utf-8") + .strip() + ) + assert OPENENGINE_COMMIT == repository_pin + + +@pytest.mark.parametrize( + "release", + [ + "", + "unreleased", + "main", + "cea19cb", + "latest", + "v0.2.0", + "signed-tag:1.2.3", + "a" * 40, + "a" * 64, + ], +) +def test_schema_release_rejects_mutable_or_unknown_identity(release: str) -> None: + with pytest.raises(RuntimeError, match="exactly match"): + validate_schema_release(release) + + +def test_schema_release_reads_launch_environment(monkeypatch) -> None: + commit = "cea19cb06acf03c911b84d5c147e519b60dd92a6" + monkeypatch.setenv("OPENENGINE_SCHEMA_RELEASE", commit) + assert schema_release() == commit + + +@pytest.mark.asyncio +async def test_server_stop_cleans_protocol_state_and_stats_on_grpc_failure() -> None: + calls = [] + + class _GrpcServer: + async def stop(self, grace: float) -> None: + calls.append(("grpc", grace)) + raise RuntimeError("stop failed") + + class _Servicer: + def close(self) -> None: + calls.append(("servicer", None)) + + class _Stats: + async def stop(self) -> None: + calls.append(("stats", None)) + + class _KvEvents: + async def stop(self) -> None: + calls.append(("kv_events", None)) + + server = object.__new__(OpenEngineServer) + server._server = _GrpcServer() + server.servicer = _Servicer() + server._kv_event_fanout = _KvEvents() + server._stats_fanout = _Stats() + + with pytest.raises(RuntimeError, match="stop failed"): + await server.stop(grace=1.5) + + assert calls == [ + ("grpc", 1.5), + ("servicer", None), + ("kv_events", None), + ("stats", None), + ] diff --git a/tests/unittest/openengine/test_servicer.py b/tests/unittest/openengine/test_servicer.py new file mode 100644 index 000000000000..2c7dac52bbe0 --- /dev/null +++ b/tests/unittest/openengine/test_servicer.py @@ -0,0 +1,775 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import asyncio +from types import SimpleNamespace + +import pytest +from openengine.v1 import ( + engine_pb2, + generation_pb2, + input_pb2, + kv_pb2, + lifecycle_pb2, + lora_pb2, + model_pb2, + observability_pb2, +) + +from tensorrt_llm.disaggregated_params import DisaggregatedParams, DisaggScheduleStyle +from tensorrt_llm.inputs.registry import MultimodalLoraSpec +from tensorrt_llm.openengine.converters import encode_handoff +from tensorrt_llm.openengine.servicer import OpenEngineServicer +from tensorrt_llm.serve.kv_event_fanout import KvEventFanout +from tensorrt_llm.serve.request_tracker import RequestTracker +from tensorrt_llm.serve.stats_fanout import StatsFanout + + +class _Tokenizer: + def decode(self, token_ids: list[int], **_: object) -> str: + return "".join(chr(96 + token_id) for token_id in token_ids) + + +class _Result: + def __init__(self, handoff: DisaggregatedParams | None = None) -> None: + self.prompt_token_ids = [1, 2] + self.cached_tokens = 1 + self.disaggregated_params = handoff + self.outputs = [ + SimpleNamespace( + index=0, + token_ids=[3, 4], + text="cd", + logprobs=[], + prompt_logprobs=[], + finish_reason="length", + stop_reason=None, + disaggregated_params=handoff, + ) + ] + self.finished = True + self._yielded = False + self.aborted = False + + def __aiter__(self) -> "_Result": + return self + + async def __anext__(self) -> "_Result": + if self._yielded: + raise StopAsyncIteration + self._yielded = True + return self + + def abort(self) -> None: + self.aborted = True + + +class _Llm: + def __init__(self, handoff: DisaggregatedParams | None = None) -> None: + self.tokenizer = _Tokenizer() + self.args = SimpleNamespace() + self.handoff = handoff + self.kwargs = None + self.result = None + + def generate_async(self, **kwargs: object) -> _Result: + self.kwargs = kwargs + self.result = _Result(self.handoff) + return self.result + + +class _Context: + def __init__(self, cancelled: bool = False) -> None: + self._cancelled = cancelled + + def cancelled(self) -> bool: + return self._cancelled + + async def abort(self, code: object, message: str) -> None: + raise AssertionError(f"unexpected abort {code}: {message}") + + +async def _collect_async(iterator): + return [item async for item in iterator] + + +@pytest.mark.asyncio +async def test_aggregate_streams_deltas_finish_and_terminal_usage() -> None: + llm = _Llm() + tracker = RequestTracker(llm) + servicer = OpenEngineServicer(llm, "model", engine_pb2.ENGINE_ROLE_AGGREGATED, tracker) + request = generation_pb2.GenerateRequest( + request_id="request", + model="model", + prompt="hello", + stopping=generation_pb2.StoppingOptions(max_tokens=2), + metadata={ + "traceparent": "00-0123456789abcdef0123456789abcdef-0123456789abcdef-01", + "tracestate": "vendor=value", + "ignored": "not-a-trace-header", + }, + ) + + responses = [response async for response in servicer.Generate(request, _Context())] + + assert [response.WhichOneof("event") for response in responses] == [ + "token", + "finished", + ] + assert list(responses[0].token.tokens)[0].token_id == 3 + assert responses[0].token.text == "cd" + assert responses[-1].usage.prompt_tokens == 2 + assert responses[-1].usage.completion_tokens == 2 + assert llm.kwargs["trace_headers"] == { + "traceparent": "00-0123456789abcdef0123456789abcdef-0123456789abcdef-01", + "tracestate": "vendor=value", + } + assert tracker.active_count == 0 + + +@pytest.mark.asyncio +async def test_generate_cancellation_aborts_and_cleans_tracking() -> None: + llm = _Llm() + tracker = RequestTracker(llm) + service = OpenEngineServicer(llm, "model", engine_pb2.ENGINE_ROLE_AGGREGATED, tracker) + request = generation_pb2.GenerateRequest(request_id="cancel", model="model", prompt="hello") + + responses = [response async for response in service.Generate(request, _Context(True))] + + assert responses == [] + assert llm.result.aborted + assert tracker.active_count == 0 + + +def test_generate_rejects_unsupported_cache_bypass_and_nonzero_dp_placement() -> None: + llm = _Llm() + service = OpenEngineServicer( + llm, "model", engine_pb2.ENGINE_ROLE_AGGREGATED, RequestTracker(llm) + ) + request = generation_pb2.GenerateRequest(request_id="request", model="model", prompt="hello") + request.kv.bypass_prefix_cache = True + with pytest.raises(ValueError, match="Prefix-cache bypass"): + service._validate_generate(request) + + request.kv.bypass_prefix_cache = False + request.kv.data_parallel_rank = 0 + service._validate_generate(request) + assert service._scheduling_params(request) is None + + llm.args = SimpleNamespace(data_parallel_size=2) + request.kv.data_parallel_rank = 1 + with pytest.raises(ValueError, match="attention DP"): + service._validate_generate(request) + + +@pytest.mark.asyncio +async def test_generate_applies_strict_attention_dp_placement() -> None: + llm = _Llm() + llm.args = SimpleNamespace(enable_attention_dp=True, data_parallel_size=4) + llm._on_trt_backend = False + service = OpenEngineServicer( + llm, "model", engine_pb2.ENGINE_ROLE_AGGREGATED, RequestTracker(llm) + ) + request = generation_pb2.GenerateRequest(request_id="request", model="model", prompt="hello") + request.kv.data_parallel_rank = 2 + + _ = [response async for response in service.Generate(request, _Context())] + scheduling = llm.kwargs["scheduling_params"] + + assert scheduling.attention_dp_rank == 2 + assert not scheduling.attention_dp_relax + + +def test_prefill_preserves_routing_rank_in_context_handoff() -> None: + llm = _Llm() + llm.args = SimpleNamespace(enable_attention_dp=True, data_parallel_size=4) + llm._on_trt_backend = False + service = OpenEngineServicer(llm, "model", engine_pb2.ENGINE_ROLE_PREFILL, RequestTracker(llm)) + request = generation_pb2.GenerateRequest(request_id="prefill", model="model", prompt="hello") + request.kv.data_parallel_rank = 3 + + params = service._disaggregated_params(request) + + assert params.ctx_dp_rank == 3 + + +@pytest.mark.asyncio +async def test_generate_selects_model_owned_multimodal_lora(monkeypatch, tmp_path) -> None: + class _Processor: + def get_openengine_modalities(self) -> tuple[str, ...]: + return ("audio",) + + def get_openengine_prefill_decode_modalities(self) -> tuple[str, ...]: + return () + + def get_required_lora_spec(self, modalities: tuple[str, ...]) -> MultimodalLoraSpec | None: + assert modalities == ("audio",) + return MultimodalLoraSpec("speech-lora", 1, str(tmp_path)) + + async def _load_media(*_args, **_kwargs): + return {"audio": [object()]} + + monkeypatch.setattr("tensorrt_llm.openengine.servicer.BaseMultimodalInputProcessor", _Processor) + monkeypatch.setattr("tensorrt_llm.openengine.servicer.load_media", _load_media) + llm = _Llm() + llm.args = SimpleNamespace(lora_config=object()) + llm.input_processor = _Processor() + service = OpenEngineServicer( + llm, "model", engine_pb2.ENGINE_ROLE_AGGREGATED, RequestTracker(llm) + ) + request = generation_pb2.GenerateRequest(request_id="audio", model="model", prompt="transcribe") + request.media.add(modality=input_pb2.MODALITY_AUDIO, raw_bytes=b"audio") + + _ = [response async for response in service.Generate(request, _Context())] + + assert llm.kwargs["lora_request"].lora_name == "speech-lora" + assert llm.kwargs["lora_request"].lora_int_id == 1 + assert await service.loras.list() == [] + + +@pytest.mark.asyncio +async def test_model_info_suppresses_modalities_when_required_lora_is_unavailable( + monkeypatch, tmp_path +) -> None: + class _Processor: + def get_openengine_modalities(self) -> tuple[str, ...]: + return ("audio",) + + def get_openengine_prefill_decode_modalities(self) -> tuple[str, ...]: + return () + + def get_required_lora_spec(self, modalities: tuple[str, ...]) -> MultimodalLoraSpec | None: + assert modalities == ("audio",) + return MultimodalLoraSpec("speech-lora", 1, str(tmp_path / "missing")) + + monkeypatch.setattr("tensorrt_llm.openengine.servicer.BaseMultimodalInputProcessor", _Processor) + llm = _Llm() + llm.input_processor = _Processor() + service = OpenEngineServicer( + llm, "model", engine_pb2.ENGINE_ROLE_AGGREGATED, RequestTracker(llm) + ) + + info = await service.GetModelInfo(model_pb2.GetModelInfoRequest(), _Context()) + + assert not info.supports_multimodal + assert list(info.multimodal_capabilities.aggregate_modalities) == [] + + +@pytest.mark.asyncio +async def test_context_then_generation_round_trips_handoff() -> None: + context_handoff = DisaggregatedParams( + request_type="context_only", + first_gen_tokens=[7], + ctx_request_id=100, + disagg_request_id=101, + ctx_info_endpoint="tcp://context:1", + schedule_style=DisaggScheduleStyle.CONTEXT_FIRST, + opaque_state=b"state", + ) + context_llm = _Llm(context_handoff) + context_service = OpenEngineServicer( + context_llm, + "model", + engine_pb2.ENGINE_ROLE_PREFILL, + RequestTracker(context_llm), + ) + context_request = generation_pb2.GenerateRequest( + request_id="request", + model="model", + token_ids=input_pb2.TokenIds(ids=[1, 2]), + ) + context_responses = [ + response async for response in context_service.Generate(context_request, _Context()) + ] + assert len(context_responses) == 1 + assert context_responses[0].WhichOneof("event") == "prefill_ready" + assert len(context_service._kv_session_requests) == 1 + + decode_llm = _Llm() + decode_service = OpenEngineServicer( + decode_llm, + "model", + engine_pb2.ENGINE_ROLE_DECODE, + RequestTracker(decode_llm), + ) + decode_request = generation_pb2.GenerateRequest( + request_id="decode", + model="model", + token_ids=input_pb2.TokenIds(ids=[1, 2]), + ) + decode_request.kv.session.CopyFrom(context_responses[0].prefill_ready.kv_session) + decode_responses = [ + response async for response in decode_service.Generate(decode_request, _Context()) + ] + decoded = decode_llm.kwargs["disaggregated_params"] + assert decoded.request_type == "generation_only" + assert decoded.disagg_request_id == 101 + assert decoded.opaque_state == b"state" + assert decoded.ctx_usage == { + "prompt_tokens": 2, + "completion_tokens": 2, + "total_tokens": 4, + "prompt_tokens_details": {"cached_tokens": 1}, + } + terminal_usage = decode_responses[-1].usage + assert terminal_usage.prompt_tokens == 2 + assert terminal_usage.cached_prompt_tokens == 1 + assert terminal_usage.completion_tokens == 2 + assert terminal_usage.total_tokens == 4 + + abort_response = await context_service.Abort( + lifecycle_pb2.AbortRequest(kv_session=context_responses[0].prefill_ready.kv_session), + _Context(), + ) + assert abort_response.status == lifecycle_pb2.ABORT_STATUS_ABORTED + assert not context_service._kv_session_requests + + +@pytest.mark.asyncio +async def test_prefill_session_expires_and_close_cancels_timers() -> None: + handoff = DisaggregatedParams( + request_type="context_only", + disagg_request_id=101, + schedule_style=DisaggScheduleStyle.CONTEXT_FIRST, + ) + llm = _Llm(handoff) + service = OpenEngineServicer( + llm, + "model", + engine_pb2.ENGINE_ROLE_PREFILL, + RequestTracker(llm), + kv_session_ttl_seconds=0.01, + ) + request = generation_pb2.GenerateRequest(request_id="prefill", model="model", prompt="hello") + + _ = [response async for response in service.Generate(request, _Context())] + assert service._kv_session_requests + assert service._kv_session_timers + await asyncio.sleep(0.02) + assert not service._kv_session_requests + assert not service._kv_session_timers + + service._track_kv_session("second", "request") + timer = service._kv_session_timers["second"] + service.close() + assert timer.cancelled() + assert not service._kv_session_requests + + +def test_decode_media_is_required_only_for_marked_handoff(monkeypatch) -> None: + class _Processor: + def get_openengine_modalities(self) -> tuple[str, ...]: + return ("image", "video") + + def get_openengine_prefill_decode_modalities(self) -> tuple[str, ...]: + return ("image", "video") + + def get_required_lora_spec(self, modalities: tuple[str, ...]) -> MultimodalLoraSpec | None: + del modalities + return None + + monkeypatch.setattr("tensorrt_llm.openengine.servicer.BaseMultimodalInputProcessor", _Processor) + llm = _Llm() + llm.input_processor = _Processor() + service = OpenEngineServicer(llm, "model", engine_pb2.ENGINE_ROLE_DECODE, RequestTracker(llm)) + session = encode_handoff( + DisaggregatedParams( + request_type="context_only", + schedule_style=DisaggScheduleStyle.CONTEXT_FIRST, + mrope_position_ids_handle={"tensor": "ids"}, + mrope_position_deltas_handle={"tensor": "deltas"}, + ), + requires_decode_media=True, + ) + request = generation_pb2.GenerateRequest( + request_id="decode", model="model", token_ids=input_pb2.TokenIds(ids=[1, 2]) + ) + request.kv.session.CopyFrom(session) + + with pytest.raises(ValueError, match="must resend"): + service._validate_generate(request) + + request.media.add(modality=input_pb2.MODALITY_IMAGE, raw_bytes=b"image") + service._validate_generate(request) + + +def test_kv_batch_preserves_chain_geometry_and_source_sequence() -> None: + llm = _Llm() + llm.args = SimpleNamespace(kv_cache_config=SimpleNamespace(tokens_per_block=2)) + service = OpenEngineServicer( + llm, "model", engine_pb2.ENGINE_ROLE_AGGREGATED, RequestTracker(llm) + ) + service._lora_names_by_id[7] = "adapter" + batch = service._kv_batch( + { + "attention_dp_rank": 2, + "layer_group_id": 3, + "window_size": 4096, + "hash_algo": "v2_sha256", + "data": { + "type": "stored", + "parent_hash": "00" * 32, + "blocks": [ + { + "block_hash": "11" * 32, + "tokens": [{"token_id": 1}, {"token_id": 2}], + "mm_keys": [ + { + "type": "mm_key", + "hash": "0123456789abcdef0011", + "start_offset": 4, + } + ], + "lora_id": 7, + }, + { + "block_hash": "22" * 32, + "tokens": [{"token_id": 3}, {"token_id": 4}], + "lora_id": 7, + }, + ], + }, + }, + sequence=9, + ) + + assert batch.sequence_number == 9 + assert batch.data_parallel_rank == 2 + assert len(batch.events) == 1 + stored = batch.events[0].block_stored + assert len(stored.block_hashes) == 2 + assert stored.parent_block_hash.value == bytes.fromhex("00" * 32) + assert list(stored.token_ids) == [1, 2, 3, 4] + assert stored.block_size == 2 + assert stored.lora_id == 7 + assert stored.lora_name == "adapter" + assert stored.group_idx == 3 + assert list(stored.extra_keys[0].values) == [ + "trt_mm_v1", + "0", + "81985529216486895", + "4", + ] + + +def test_kv_batch_stops_at_partial_block_and_uses_decimal_int_hashes() -> None: + llm = _Llm() + llm.args = SimpleNamespace(kv_cache_config=SimpleNamespace(tokens_per_block=2)) + service = OpenEngineServicer( + llm, "model", engine_pb2.ENGINE_ROLE_AGGREGATED, RequestTracker(llm) + ) + batch = service._kv_batch( + { + "attention_dp_rank": 0, + "data": { + "type": "stored", + "parent_hash": 2**64 - 1, + "blocks": [ + { + "block_hash": 2**63, + "tokens": [{"token_id": 1}, {"token_id": 2}], + }, + {"block_hash": 9, "tokens": [{"token_id": 3}]}, + { + "block_hash": 10, + "tokens": [{"token_id": 4}, {"token_id": 5}], + }, + ], + }, + }, + sequence=3, + ) + + assert len(batch.events) == 1 + stored = batch.events[0].block_stored + assert stored.parent_block_hash.encoding == "decimal_int64" + assert stored.parent_block_hash.value == b"-1" + assert stored.block_hashes[0].value == str(-(2**63)).encode() + assert 9 in service._partial_block_hashes[0] + + +def test_kv_batch_preserves_model_owned_lora_id_zero_and_gap_reset() -> None: + class _Processor: + @staticmethod + def get_model_owned_lora_identities() -> dict[str, int]: + return {"vision-lora": 0} + + llm = _Llm() + llm.input_processor = _Processor() + llm.args = SimpleNamespace(kv_cache_config=SimpleNamespace(tokens_per_block=2)) + service = OpenEngineServicer( + llm, "model", engine_pb2.ENGINE_ROLE_AGGREGATED, RequestTracker(llm) + ) + batch = service._kv_batch( + { + "data": { + "type": "stored", + "blocks": [ + { + "block_hash": 2, + "tokens": [{"token_id": 1}, {"token_id": 2}], + "lora_id": 0, + } + ], + } + }, + sequence=1, + ) + service._partial_block_hashes[0] = {3} + reset = service._kv_batch({"attention_dp_rank": 0, "data": {"type": "all_cleared"}}, sequence=2) + + assert batch.events[0].block_stored.lora_id == 0 + assert batch.events[0].block_stored.lora_name == "vision-lora" + assert reset.events[0].WhichOneof("event") == "all_blocks_cleared" + assert 0 not in service._partial_block_hashes + + +@pytest.mark.parametrize( + "unsupported", + [ + {"cache_salt": "private"}, + {"tokens": [{"token_id": 1, "token_extra_id": 1}, {"token_id": 2}]}, + {"mm_keys": [{"type": "mm_key", "hash": "not-hex"}]}, + ], +) +def test_kv_batch_fails_closed_for_unrepresentable_cache_namespace(unsupported) -> None: + llm = _Llm() + llm.args = SimpleNamespace(kv_cache_config=SimpleNamespace(tokens_per_block=2)) + service = OpenEngineServicer( + llm, "model", engine_pb2.ENGINE_ROLE_AGGREGATED, RequestTracker(llm) + ) + block = { + "block_hash": 2, + "tokens": [{"token_id": 1}, {"token_id": 2}], + **unsupported, + } + + batch = service._kv_batch( + {"data": {"type": "stored", "parent_hash": 1, "blocks": [block]}}, + sequence=1, + ) + + assert len(batch.events) == 1 + assert batch.events[0].WhichOneof("event") == "all_blocks_cleared" + + +@pytest.mark.asyncio +async def test_discovery_and_load_use_config_and_shared_stats() -> None: + llm = _Llm() + llm.args = SimpleNamespace( + guided_decoding_backend=None, + lora_config=None, + enable_lora=False, + kv_cache_config=SimpleNamespace(tokens_per_block=64), + ) + llm.get_kv_cache_capacity = lambda: { + "maxNumBlocks": 100, + "tokensPerBlock": 64, + "maxNumTokens": 6400, + } + stats = StatsFanout(llm) + stats._publish( + { + "attentionDpRank": 1, + "numActiveRequests": 3, + "numQueuedRequests": 2, + "kvCacheStats": { + "maxNumBlocks": 100, + "freeNumBlocks": 25, + "tokensPerBlock": 64, + }, + "inflightBatchingStats": {"numContextRequests": 1, "numGenRequests": 2}, + } + ) + service = OpenEngineServicer( + llm, + "model", + engine_pb2.ENGINE_ROLE_AGGREGATED, + RequestTracker(llm), + stats_fanout=stats, + ) + + model_info = await service.GetModelInfo(model_pb2.GetModelInfoRequest(), _Context()) + assert model_info.kv_block_size == 64 + assert model_info.total_kv_blocks == 100 + assert not model_info.generation.guided_decoding.supported + assert not model_info.supports_lora + + load = await service.GetLoad( + observability_pb2.GetLoadRequest(include_per_rank=True), _Context() + ) + assert load.running_requests == 3 + assert load.queued_requests == 2 + assert load.used_kv_blocks == 75 + assert load.total_kv_blocks == 100 + assert load.prefill_batch_size == 1 + assert load.decode_batch_size == 2 + assert load.ranks[0].data_parallel_rank == 1 + assert load.attributes["source"] == "shared_stats_fanout" + assert load.attributes["kv_tokens_per_block"] == "64" + assert load.attributes["rank.1.kv_tokens_per_block"] == "64" + + llm.args.guided_decoding_backend = "xgrammar" + guided_info = await service.GetModelInfo(model_pb2.GetModelInfoRequest(), _Context()) + assert guided_info.generation.guided_decoding.supported + assert ( + model_pb2.GUIDED_DECODING_MODE_STRUCTURAL_TAG + in guided_info.generation.guided_decoding.modes + ) + + +@pytest.mark.asyncio +async def test_load_never_undercounts_live_tracker_from_stale_stats() -> None: + llm = _Llm() + stats = StatsFanout(llm) + stats._publish({"attentionDpRank": 0, "numActiveRequests": 0}) + tracker = RequestTracker(llm) + tracker.admit("active", _Result()) + service = OpenEngineServicer( + llm, + "model", + engine_pb2.ENGINE_ROLE_AGGREGATED, + tracker, + stats_fanout=stats, + ) + + load = await service.GetLoad(observability_pb2.GetLoadRequest(), _Context()) + + assert load.running_requests == 1 + + +@pytest.mark.asyncio +async def test_kv_event_sources_are_rank_scoped() -> None: + llm = _Llm() + llm.args = SimpleNamespace( + data_parallel_size=2, + kv_cache_config=SimpleNamespace(event_buffer_max_size=8), + ) + fanout = KvEventFanout(llm, buffer_size=8) + service = OpenEngineServicer( + llm, + "model", + engine_pb2.ENGINE_ROLE_AGGREGATED, + RequestTracker(llm), + kv_event_fanout=fanout, + ) + + response = await service.GetKvEventSources(kv_pb2.GetKvEventSourcesRequest(), _Context()) + + assert [source.data_parallel_rank for source in response.sources] == [0, 1] + + +@pytest.mark.asyncio +async def test_kv_event_discovery_requires_decimal_compatible_hashes() -> None: + llm = _Llm() + llm.args = SimpleNamespace( + kv_cache_config=SimpleNamespace( + event_buffer_max_size=8, + kv_cache_event_hash_algo="v2_sha256", + use_kv_cache_manager_v2=True, + ) + ) + service = OpenEngineServicer( + llm, + "model", + engine_pb2.ENGINE_ROLE_AGGREGATED, + RequestTracker(llm), + kv_event_fanout=KvEventFanout(llm), + ) + + response = await service.GetKvEventSources(kv_pb2.GetKvEventSourcesRequest(), _Context()) + assert list(response.sources) == [] + + llm.args.kv_cache_config.kv_cache_event_hash_algo = "v2_sha256_64" + response = await service.GetKvEventSources(kv_pb2.GetKvEventSourcesRequest(), _Context()) + assert len(response.sources) == 1 + + +@pytest.mark.asyncio +async def test_direct_kv_subscription_rejects_unknown_rank() -> None: + llm = _Llm() + llm.args = SimpleNamespace( + data_parallel_size=2, + kv_cache_config=SimpleNamespace(event_buffer_max_size=8), + ) + service = OpenEngineServicer( + llm, + "model", + engine_pb2.ENGINE_ROLE_AGGREGATED, + RequestTracker(llm), + kv_event_fanout=KvEventFanout(llm), + ) + request = kv_pb2.SubscribeKvEventsRequest(data_parallel_ranks=[2]) + + with pytest.raises(AssertionError, match="INVALID_ARGUMENT"): + await _collect_async(service.SubscribeKvEvents(request, _Context())) + + +@pytest.mark.asyncio +async def test_drain_deadline_does_not_wait_forever_for_external_http() -> None: + llm = _Llm() + tracker = RequestTracker(llm) + tracker.begin_external() + service = OpenEngineServicer( + llm, + "model", + engine_pb2.ENGINE_ROLE_AGGREGATED, + tracker, + post_abort_cleanup_timeout_seconds=0.01, + ) + request = lifecycle_pb2.DrainRequest( + stop_accepting_new_requests=True, + deadline_ms=1, + abort_after_deadline=True, + ) + + responses = await asyncio.wait_for( + _collect_async(service.Drain(request, _Context())), timeout=0.1 + ) + + assert responses[-1].WhichOneof("event") == "error" + await tracker.finish_external() + + +@pytest.mark.asyncio +async def test_model_info_omits_unknown_kv_geometry() -> None: + llm = _Llm() + service = OpenEngineServicer( + llm, "model", engine_pb2.ENGINE_ROLE_AGGREGATED, RequestTracker(llm) + ) + info = await service.GetModelInfo(model_pb2.GetModelInfoRequest(), _Context()) + assert not info.HasField("kv_block_size") + assert not info.HasField("total_kv_blocks") + + +@pytest.mark.asyncio +async def test_model_info_unknown_model_maps_to_not_found() -> None: + llm = _Llm() + service = OpenEngineServicer( + llm, "model", engine_pb2.ENGINE_ROLE_AGGREGATED, RequestTracker(llm) + ) + + with pytest.raises(AssertionError, match="NOT_FOUND"): + await service.GetModelInfo(model_pb2.GetModelInfoRequest(model="other"), _Context()) + + +@pytest.mark.asyncio +async def test_health_probe_and_unconfigured_lora_fail_explicitly(tmp_path) -> None: + llm = _Llm() + service = OpenEngineServicer( + llm, "model", engine_pb2.ENGINE_ROLE_AGGREGATED, RequestTracker(llm) + ) + with pytest.raises(AssertionError, match="UNIMPLEMENTED"): + await service.Health(lifecycle_pb2.HealthRequest(include_inference_probe=True), _Context()) + + adapter_dir = tmp_path / "adapter" + adapter_dir.mkdir() + with pytest.raises(AssertionError, match="FAILED_PRECONDITION"): + await service.LoadLora( + lora_pb2.LoadLoraRequest( + adapter=lora_pb2.LoraAdapter( + lora_id=1, lora_name="adapter", source_path=str(adapter_dir) + ) + ), + _Context(), + ) From 461b120a4b2144ad684113b9409c2f3d60e0f3a6 Mon Sep 17 00:00:00 2001 From: Connor Carpenter Date: Sun, 12 Jul 2026 22:58:43 -0700 Subject: [PATCH 2/4] build(ucx): avoid conflicting target link features Signed-off-by: Connor Carpenter --- .../executor/cache_transmission/ucx_utils/CMakeLists.txt | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/cpp/tensorrt_llm/executor/cache_transmission/ucx_utils/CMakeLists.txt b/cpp/tensorrt_llm/executor/cache_transmission/ucx_utils/CMakeLists.txt index a59d85d77923..dae43e989694 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/ucx_utils/CMakeLists.txt +++ b/cpp/tensorrt_llm/executor/cache_transmission/ucx_utils/CMakeLists.txt @@ -31,7 +31,9 @@ if(ENABLE_UCX) target_link_libraries(${UCX_WRAPPER_TARGET} PRIVATE $) - 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}) From 90386d69b6d6c65e9ee3dd5a0d6963fed962e652 Mon Sep 17 00:00:00 2001 From: Connor Carpenter Date: Sun, 12 Jul 2026 22:58:54 -0700 Subject: [PATCH 3/4] fix(disaggregation): prepare lazy LoRA on decode Signed-off-by: Connor Carpenter --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 3 ++- .../_torch/executor/test_benchmark_disagg.py | 27 +++++++++++++++++++ 2 files changed, 29 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index f69041fa7ab8..33a8675b9efb 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -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): diff --git a/tests/unittest/_torch/executor/test_benchmark_disagg.py b/tests/unittest/_torch/executor/test_benchmark_disagg.py index dabb2b910d55..b1d2c829f86d 100644 --- a/tests/unittest/_torch/executor/test_benchmark_disagg.py +++ b/tests/unittest/_torch/executor/test_benchmark_disagg.py @@ -30,6 +30,7 @@ import pytest from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState +from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManagerType from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests # --------------------------------------------------------------------------- @@ -649,6 +650,32 @@ def test_fetch_called_once_even_in_benchmark_disagg(self): mock_fetch.assert_called_once() +class TestPrepareDisaggGenInitResources: + def test_prepares_peft_cache_for_lazy_lora(self): + from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor + + ex = object.__new__(PyExecutor) + managers = { + kind: Mock() + for kind in ( + ResourceManagerType.KV_CACHE_MANAGER, + ResourceManagerType.SPEC_RESOURCE_MANAGER, + ResourceManagerType.DRAFT_KV_CACHE_MANAGER, + ResourceManagerType.PEFT_CACHE_MANAGER, + ) + } + ex.resource_manager = Mock(resource_managers=managers) + ex._recv_disagg_gen_cache = Mock() + request = Mock() + + ex._prepare_disagg_gen_init([request]) + + for manager in managers.values(): + prepared = manager.prepare_resources.call_args.args[0] + assert prepared.context_requests == [request] + ex._recv_disagg_gen_cache.assert_called_once_with([request]) + + # --------------------------------------------------------------------------- # Benchmark fill admission flow control # --------------------------------------------------------------------------- From cee72f7da8e6f112237e2d8a59ccc37f42e0b63b Mon Sep 17 00:00:00 2001 From: Connor Carpenter Date: Sun, 12 Jul 2026 22:59:55 -0700 Subject: [PATCH 4/4] fix(openengine): complete context handoff accounting Signed-off-by: Connor Carpenter --- tensorrt_llm/openengine/servicer.py | 37 ++++++++++++++++++++-- tests/unittest/openengine/test_servicer.py | 34 +++++++++++++++++++- 2 files changed, 68 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/openengine/servicer.py b/tensorrt_llm/openengine/servicer.py index 0835a0d0ac97..effd562b4495 100644 --- a/tensorrt_llm/openengine/servicer.py +++ b/tensorrt_llm/openengine/servicer.py @@ -257,7 +257,13 @@ async def Generate( raise RuntimeError( "Context-only result did not return disaggregated parameters" ) - handoff = replace(handoff, ctx_usage=self._usage_payload(current)) + handoff = replace( + handoff, + ctx_usage=self._usage_payload(current), + ctx_info_endpoint=( + handoff.ctx_info_endpoint or self._context_info_endpoint() + ), + ) session = encode_handoff(handoff, requires_decode_media=bool(request.media)) self._track_kv_session(session.session_id, request.request_id) yield generation_pb2.GenerateResponse( @@ -334,12 +340,39 @@ def _release_kv_session(self, session_id: str) -> bool: timer.cancel() return released + def _context_info_endpoint(self) -> str | None: + """Return the context worker endpoint needed for direct NIXL discovery.""" + params = getattr(self.llm, "disaggregated_params", None) + if not isinstance(params, dict): + return None + endpoint = params.get("ctx_info_endpoint") + if isinstance(endpoint, str): + return endpoint or None + if isinstance(endpoint, (list, tuple)): + return next( + (item for item in endpoint if isinstance(item, str) and item), + None, + ) + return None + def _release_all_kv_sessions(self) -> None: for timer in self._kv_session_timers.values(): timer.cancel() self._kv_session_timers.clear() self._kv_session_requests.clear() + def _active_kv_session_count(self) -> int: + """Count sessions whose owning request is still engine-active. + + Prefill handoff handles remain abortable until their TTL expires, but + their context request has already finished and must not inflate load. + Drain continues to report and release all open handles separately. + """ + return sum( + request_id in self.tracker.active_requests + for request_id in self._kv_session_requests.values() + ) + def close(self) -> None: """Release protocol-owned timers without shutting down the shared LLM.""" self._release_all_kv_sessions() @@ -771,7 +804,7 @@ async def GetLoad( timestamp_unix_nanos=time.time_ns(), running_requests=running_requests, queued_requests=queued_requests, - active_kv_sessions=len(self._kv_session_requests), + active_kv_sessions=self._active_kv_session_count(), ranks=rank_infos, attributes=attributes, ) diff --git a/tests/unittest/openengine/test_servicer.py b/tests/unittest/openengine/test_servicer.py index 2c7dac52bbe0..6c8e34bd3502 100644 --- a/tests/unittest/openengine/test_servicer.py +++ b/tests/unittest/openengine/test_servicer.py @@ -18,7 +18,7 @@ from tensorrt_llm.disaggregated_params import DisaggregatedParams, DisaggScheduleStyle from tensorrt_llm.inputs.registry import MultimodalLoraSpec -from tensorrt_llm.openengine.converters import encode_handoff +from tensorrt_llm.openengine.converters import decode_handoff, encode_handoff from tensorrt_llm.openengine.servicer import OpenEngineServicer from tensorrt_llm.serve.kv_event_fanout import KvEventFanout from tensorrt_llm.serve.request_tracker import RequestTracker @@ -284,6 +284,10 @@ async def test_context_then_generation_round_trips_handoff() -> None: assert len(context_responses) == 1 assert context_responses[0].WhichOneof("event") == "prefill_ready" assert len(context_service._kv_session_requests) == 1 + context_load = await context_service.GetLoad( + observability_pb2.GetLoadRequest(), _Context() + ) + assert context_load.active_kv_sessions == 0 decode_llm = _Llm() decode_service = OpenEngineServicer( @@ -325,6 +329,34 @@ async def test_context_then_generation_round_trips_handoff() -> None: assert not context_service._kv_session_requests +@pytest.mark.asyncio +async def test_prefill_adds_context_endpoint_from_llm_discovery() -> None: + handoff = DisaggregatedParams( + request_type="context_only", + first_gen_tokens=[7], + disagg_request_id=101, + schedule_style=DisaggScheduleStyle.CONTEXT_FIRST, + ) + llm = _Llm(handoff) + llm.disaggregated_params = {"ctx_info_endpoint": ["tcp://context:1234"]} + service = OpenEngineServicer( + llm, + "model", + engine_pb2.ENGINE_ROLE_PREFILL, + RequestTracker(llm), + ) + request = generation_pb2.GenerateRequest( + request_id="prefill", + model="model", + token_ids=input_pb2.TokenIds(ids=[1, 2]), + ) + + responses = [response async for response in service.Generate(request, _Context())] + decoded = decode_handoff(responses[0].prefill_ready.kv_session) + + assert decoded.ctx_info_endpoint == "tcp://context:1234" + + @pytest.mark.asyncio async def test_prefill_session_expires_and_close_cancels_timers() -> None: handoff = DisaggregatedParams(