Skip to content

[#17016][refactor] Isolate SMG gRPC protocol integration #17221

New issue

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

Sign up for GitHub

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

Already on GitHub? Sign in to your account

Closed
connorcarpenter15 wants to merge 1 commit into from
Closed

[#17016][refactor] Isolate SMG gRPC protocol integration #17221

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
153 changes: 21 additions & 132 deletions tensorrt_llm/commands/serve.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -57,9 +57,6 @@
# Global variable to store the Popen object of the child process
_child_p_global: Optional[subprocess.Popen] = None

# Bound gRPC messages while leaving room for multimodal image payloads.
_GRPC_MAX_MESSAGE_LENGTH_BYTES = 32 * 1024 * 1024


def _pop_bool_config_option(config: dict[str, Any], key: str) -> bool:
return validate_config_bool(config.pop(key, False), key)
Expand Down Expand Up @@ -619,129 +616,6 @@ def launch_server(
_terminate_attached_frontends(frontend_children)


def launch_grpc_server(host: str,
port: int,
llm_args: dict,
served_model_name: Optional[str] = None):
"""
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.

Args:
host: Host to bind to
port: Port to bind to
llm_args: Arguments for LLM initialization (from get_llm_args)
served_model_name: Custom model name for API responses (defaults to model path)
"""
import grpc

try:
from grpc_reflection.v1alpha import reflection
REFLECTION_AVAILABLE = True
except ImportError:
REFLECTION_AVAILABLE = False

from tensorrt_llm.grpc import trtllm_service_pb2, trtllm_service_pb2_grpc
from tensorrt_llm.grpc.grpc_request_manager import GrpcRequestManager
from tensorrt_llm.grpc.grpc_servicer import TrtllmServiceServicer

async def serve_grpc_async():
logger.info("Initializing TensorRT-LLM gRPC server...")

backend = llm_args.get("backend")
model_path = served_model_name or llm_args.get("model", "")

if backend == "pytorch":
llm_args.pop("build_config", None)
llm = PyTorchLLM(**llm_args)
elif backend == "_autodeploy":
from tensorrt_llm._torch.auto_deploy import LLM as AutoDeployLLM
llm_args.pop("build_config", None)
llm = AutoDeployLLM(**llm_args)
else:
raise click.BadParameter(
f"{backend} is not a known backend, check help for available options.",
param_hint="backend")

logger.info("Model loaded successfully")

# Create request manager
request_manager = GrpcRequestManager(llm)

# Create servicer
servicer = TrtllmServiceServicer(request_manager, model_path=model_path)

# Create gRPC server
server = grpc.aio.server(
options=[
("grpc.max_send_message_length",
_GRPC_MAX_MESSAGE_LENGTH_BYTES),
("grpc.max_receive_message_length",
_GRPC_MAX_MESSAGE_LENGTH_BYTES),
("grpc.keepalive_time_ms", 30000), # 30s keepalive
("grpc.keepalive_timeout_ms", 10000), # 10s timeout
("grpc.keepalive_permit_without_calls", True),
("grpc.http2.min_recv_ping_interval_without_data_ms", 10000),
], )

# Add servicer to server
trtllm_service_pb2_grpc.add_TrtllmServiceServicer_to_server(
servicer, server)

# Enable reflection for grpcurl and other tools
if REFLECTION_AVAILABLE:
service_names = (
trtllm_service_pb2.DESCRIPTOR.services_by_name["TrtllmService"].
full_name,
reflection.SERVICE_NAME,
)
reflection.enable_server_reflection(service_names, server)
logger.info("gRPC reflection enabled")

# Bind to address
address = f"{host}:{port}"
server.add_insecure_port(address)

# Start server
await server.start()
logger.info(f"TensorRT-LLM gRPC server started on {address}")
logger.info("Server is ready to accept requests")

# Handle shutdown signals
loop = asyncio.get_running_loop()
stop_event = asyncio.Event()

def signal_handler():
logger.info("Received shutdown signal")
stop_event.set()

for sig in (signal.SIGTERM, signal.SIGINT):
loop.add_signal_handler(sig, signal_handler)

# Serve until shutdown signal
try:
await stop_event.wait()
except KeyboardInterrupt:
logger.info("Interrupted by user")
finally:
logger.info("Shutting down TensorRT-LLM gRPC server...")

# Stop gRPC server
await server.stop(grace=5.0)
logger.info("gRPC server stopped")

# Shutdown LLM
if hasattr(llm, "shutdown"):
llm.shutdown()
logger.info("LLM engine stopped")

logger.info("Shutdown complete")

uvloop.run(serve_grpc_async())


def launch_mm_encoder_server(
host: str,
port: int,
Expand Down Expand Up @@ -1228,6 +1102,11 @@ def launch_visual_gen_server(
help="Run gRPC server instead of OpenAI HTTP server. "
"gRPC server accepts pre-tokenized requests and returns raw token IDs.",
status="prototype")
@stability_option("--grpc-protocol",
type=click.Choice(["smg", "openengine"]),
default="smg",
help="Protocol used when --grpc is enabled.",
status="prototype")
@stability_option(
"--served_model_name",
type=str,
Expand Down Expand Up @@ -1285,14 +1164,21 @@ def serve(model: str, tokenizer: Optional[str], custom_tokenizer: Optional[str],
agent_types: Optional[str], video_pruning_rate: Optional[float],
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]):
middleware: tuple[str, ...], grpc: bool, grpc_protocol: str,
enable_visual_gen: bool, served_model_name: Optional[str],
visual_gen_args: Optional[str]):
"""Running an OpenAI API compatible server

MODEL: model name | HF checkpoint path | TensorRT engine path
"""
logger.set_level(log_level)

if not grpc and grpc_protocol != "smg":
raise click.UsageError("--grpc-protocol requires --grpc")
if grpc and grpc_protocol == "openengine":
raise click.UsageError(
"OpenEngine gRPC support is not available in this build")

if moe_cluster_parallel_size is not None:
logger.warning(
"--moe_cluster_parallel_size / --cluster_size is deprecated and "
Expand Down Expand Up @@ -1462,10 +1348,13 @@ def _serve_llm():
f"Argument '{name}' is not supported when running in gRPC mode. "
f"The gRPC server is designed for use with external routers that handle "
f"these features (e.g., tool parsing, chat templates).")
launch_grpc_server(host,
port,
llm_args,
served_model_name=served_model_name)
from tensorrt_llm.grpc.smg.server import \
launch_server as launch_smg_server

launch_smg_server(host,
port,
llm_args,
served_model_name=served_model_name)
else:
# Default: launch OpenAI HTTP server
launch_server(
Expand Down
87 changes: 3 additions & 84 deletions tensorrt_llm/grpc/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,87 +1,6 @@
# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

r"""TensorRT-LLM gRPC module for high-performance communication with external routers.
"""Protocol-specific gRPC integrations for TensorRT-LLM."""

This module provides a gRPC server interface that accepts pre-tokenized requests
and returns raw token IDs, enabling efficient binary communication with Rust-based
routers like sgl-router.

Key Features:
- Pre-tokenized input (no Python tokenization overhead)
- Raw token ID output (no Python detokenization overhead)
- Streaming support with delta tokens
- Full sampling parameter support
- Guided decoding (JSON schema, regex, grammar)
- LoRA and prompt tuning support
- Disaggregated inference support

Proto definitions are provided by the smg-grpc-proto package (pip install smg-grpc-proto).

Usage:
python -m tensorrt_llm.commands.serve /path/to/model \
--grpc \
--host 0.0.0.0 \
--port 50051
"""

# Try to import generated protobuf modules from smg-grpc-proto package
try:
from smg_grpc_proto.generated import trtllm_service_pb2, trtllm_service_pb2_grpc

PROTOS_AVAILABLE = True
except ImportError:
PROTOS_AVAILABLE = False
trtllm_service_pb2 = None
trtllm_service_pb2_grpc = None

# Try to import request manager
try:
from .grpc_request_manager import (
GrpcRequestManager,
create_disaggregated_params_from_proto,
create_lora_request_from_proto,
create_sampling_params_from_proto,
)

REQUEST_MANAGER_AVAILABLE = True
except ImportError:
REQUEST_MANAGER_AVAILABLE = False
GrpcRequestManager = None
create_sampling_params_from_proto = None
create_lora_request_from_proto = None
create_disaggregated_params_from_proto = None

# Try to import servicer
try:
from .grpc_servicer import TrtllmServiceServicer

SERVICER_AVAILABLE = True
except ImportError:
SERVICER_AVAILABLE = False
TrtllmServiceServicer = None

__all__ = [
"PROTOS_AVAILABLE",
"REQUEST_MANAGER_AVAILABLE",
"SERVICER_AVAILABLE",
"trtllm_service_pb2",
"trtllm_service_pb2_grpc",
"GrpcRequestManager",
"TrtllmServiceServicer",
"create_sampling_params_from_proto",
"create_lora_request_from_proto",
"create_disaggregated_params_from_proto",
]
__all__ = []
87 changes: 87 additions & 0 deletions tensorrt_llm/grpc/smg/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,87 @@
# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

r"""SMG integration for high-performance communication with external routers.

This module provides a gRPC server interface that accepts pre-tokenized requests
and returns raw token IDs, enabling efficient binary communication with Rust-based
routers like sgl-router.

Key Features:
- Pre-tokenized input (no Python tokenization overhead)
- Raw token ID output (no Python detokenization overhead)
- Streaming support with delta tokens
- Full sampling parameter support
- Guided decoding (JSON schema, regex, grammar)
- LoRA and prompt tuning support
- Disaggregated inference support

Proto definitions are provided by the smg-grpc-proto package (pip install smg-grpc-proto).

Usage:
python -m tensorrt_llm.commands.serve /path/to/model \
--grpc \
--host 0.0.0.0 \
--port 50051
"""

# Try to import generated protobuf modules from smg-grpc-proto package
try:
from smg_grpc_proto.generated import trtllm_service_pb2, trtllm_service_pb2_grpc

PROTOS_AVAILABLE = True
except ImportError:
PROTOS_AVAILABLE = False
trtllm_service_pb2 = None
trtllm_service_pb2_grpc = None

# Try to import request manager
try:
from .request_manager import (
GrpcRequestManager,
create_disaggregated_params_from_proto,
create_lora_request_from_proto,
create_sampling_params_from_proto,
)

REQUEST_MANAGER_AVAILABLE = True
except ImportError:
REQUEST_MANAGER_AVAILABLE = False
GrpcRequestManager = None
create_sampling_params_from_proto = None
create_lora_request_from_proto = None
create_disaggregated_params_from_proto = None

# Try to import servicer
try:
from .servicer import TrtllmServiceServicer

SERVICER_AVAILABLE = True
except ImportError:
SERVICER_AVAILABLE = False
TrtllmServiceServicer = None
Comment on lines +40 to +74

@coderabbitai coderabbitai Bot Aug 3, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🩺 Stability & Availability | 🟠 Major

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
# Inspect how request_manager.py and servicer.py reference the optional protobuf module,
# and confirm whether launch_server guards against missing smg-grpc-proto.
echo "=== smg/__init__.py ==="
cat -n tensorrt_llm/grpc/smg/__init__.py

echo "=== request_manager.py head (imports and type hints) ==="
sed -n '1,60p' tensorrt_llm/grpc/smg/request_manager.py

echo "=== servicer.py head (imports) ==="
sed -n '1,45p' tensorrt_llm/grpc/smg/servicer.py

echo "=== server.py (full) ==="
cat -n tensorrt_llm/grpc/smg/server.py

Repository: NVIDIA/TensorRT-LLM

Length of output: 12433

🏁 Script executed:

#!/bin/bash
set -eu

echo "=== protobuf references in request_manager.py and servicer.py ==="
rg -n -C 2 'pb2|trtllm_service_pb2' \
  tensorrt_llm/grpc/smg/request_manager.py \
  tensorrt_llm/grpc/smg/servicer.py

echo "=== definitions with protobuf annotations ==="
rg -n -C 3 'def .*pb2|-> .*pb2|: .*pb2' \
  tensorrt_llm/grpc/smg/request_manager.py \
  tensorrt_llm/grpc/smg/servicer.py

echo "=== package import dependency edges ==="
rg -n -C 2 'from \\. import|from \\.request_manager|from \\.servicer|class .*trtllm_service' \
  tensorrt_llm/grpc/smg/__init__.py \
  tensorrt_llm/grpc/smg/request_manager.py \
  tensorrt_llm/grpc/smg/servicer.py \
  tensorrt_llm/grpc/smg/server.py

Repository: NVIDIA/TensorRT-LLM

Length of output: 24118

🏁 Script executed:

#!/bin/bash
set -eu

python3 - <<'PY'
from pathlib import Path
import ast

for name in (
    "tensorrt_llm/grpc/smg/request_manager.py",
    "tensorrt_llm/grpc/smg/servicer.py",
):
    tree = ast.parse(Path(name).read_text())
    print(f"=== {name} ===")
    for node in ast.walk(tree):
        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
            annotations = []
            for arg in [*node.args.posonlyargs, *node.args.args, *node.args.kwonlyargs]:
                if arg.annotation is not None:
                    annotations.append(ast.unparse(arg.annotation))
            if node.args.vararg and node.args.vararg.annotation:
                annotations.append(ast.unparse(node.args.vararg.annotation))
            if node.args.kwarg and node.args.kwarg.annotation:
                annotations.append(ast.unparse(node.args.kwarg.annotation))
            if node.returns is not None:
                annotations.append(ast.unparse(node.returns))
            if any("pb2" in annotation for annotation in annotations):
                print(node.lineno, node.name, annotations)
PY

Repository: NVIDIA/TensorRT-LLM

Length of output: 1904

🏁 Script executed:

#!/bin/bash
set -eu

echo "=== protobuf references in request_manager.py and servicer.py ==="
rg -n -C 2 'pb2|trtllm_service_pb2' \
  tensorrt_llm/grpc/smg/request_manager.py \
  tensorrt_llm/grpc/smg/servicer.py

echo "=== definitions with protobuf annotations ==="
rg -n -C 3 'def .*pb2|-> .*pb2|: .*pb2' \
  tensorrt_llm/grpc/smg/request_manager.py \
  tensorrt_llm/grpc/smg/servicer.py

echo "=== package import dependency edges ==="
rg -n -C 2 'from \\. import|from \\.request_manager|from \\.servicer|class .*trtllm_service' \
  tensorrt_llm/grpc/smg/__init__.py \
  tensorrt_llm/grpc/smg/request_manager.py \
  tensorrt_llm/grpc/smg/servicer.py

Repository: NVIDIA/TensorRT-LLM

Length of output: 24118

🏁 Script executed:

#!/bin/bash
set -eu

python3 - <<'PY'
from pathlib import Path
import ast

for name in (
    "tensorrt_llm/grpc/smg/request_manager.py",
    "tensorrt_llm/grpc/smg/servicer.py",
):
    tree = ast.parse(Path(name).read_text())
    print(f"=== {name} ===")
    for node in ast.walk(tree):
        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
            annotations = []
            for arg in [
                *node.args.posonlyargs,
                *node.args.args,
                *node.args.kwonlyargs,
            ]:
                if arg.annotation is not None:
                    annotations.append(ast.unparse(arg.annotation))
            if node.args.vararg and node.args.vararg.annotation:
                annotations.append(ast.unparse(node.args.vararg.annotation))
            if node.args.kwarg and node.args.kwarg.annotation:
                annotations.append(ast.unparse(node.args.kwarg.annotation))
            if node.returns is not None:
                annotations.append(ast.unparse(node.returns))
            if any("pb2" in annotation for annotation in annotations):
                print(node.lineno, node.name, annotations)
PY

Repository: NVIDIA/TensorRT-LLM

Length of output: 1904

Fail clearly when smg-grpc-proto is missing. request_manager.py evaluates pb2.* annotations during import, so the None fallback raises AttributeError before launch_server() runs. Make these imports lazy or raise a clear error before importing them, such as Install smg-grpc-proto to use --grpc.
[medium_effort_and_high_reward]

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tensorrt_llm/grpc/smg/__init__.py` around lines 40 - 74, The optional
protobuf fallback in the module-level imports causes request_manager and
servicer imports to fail with an unclear AttributeError when smg-grpc-proto is
unavailable. Update the import flow around PROTOS_AVAILABLE,
REQUEST_MANAGER_AVAILABLE, and SERVICER_AVAILABLE so protobuf-dependent modules
are imported only after validating the generated protobuf package, or raise a
clear installation error before those imports; preserve normal imports when
smg-grpc-proto is installed.


__all__ = [
"PROTOS_AVAILABLE",
"REQUEST_MANAGER_AVAILABLE",
"SERVICER_AVAILABLE",
"trtllm_service_pb2",
"trtllm_service_pb2_grpc",
"GrpcRequestManager",
"TrtllmServiceServicer",
"create_sampling_params_from_proto",
"create_lora_request_from_proto",
"create_disaggregated_params_from_proto",
]
4 changes: 2 additions & 2 deletions tensorrt_llm/grpc/grpc_request_manager.py → tensorrt_llm/grpc/smg/request_manager.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
Expand All @@ -13,7 +13,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.

"""gRPC Request Manager for TensorRT-LLM.
"""SMG request manager for TensorRT-LLM.

Manages request lifecycle for gRPC requests, converting between protobuf
and TensorRT-LLM types. Designed for high-performance communication with
Expand Down
Loading