add trtllm disagg scripts

examples/python_rs/llm/trtllm/README.md

@@ -0,0 +1,74 @@
+<!--
+SPDX-FileCopyrightText: Copyright (c) 2025 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.
+-->
+
+# TensorRT-LLM Integration with Triton Distributed
+
+This example demonstrates how to use Triton Distributed to serve large language models with the tensorrt_llm engine, enabling efficient model serving with monolithic option.
+
+## Prerequisites
+
+Start required services (etcd and NATS):
+
+   Option A: Using [Docker Compose](/runtime/rust/docker-compose.yml) (Recommended)
+   ```bash
+   docker-compose up -d
+   ```
+
+   Option B: Manual Setup
+
+    - [NATS.io](https://docs.nats.io/running-a-nats-service/introduction/installation) server with [Jetstream](https://docs.nats.io/nats-concepts/jetstream)
+        - example: `nats-server -js --trace`
+    - [etcd](https://etcd.io) server
+        - follow instructions in [etcd installation](https://etcd.io/docs/v3.5/install/) to start an `etcd-server` locally
+
+*Note*: This example is work in progress.
+
+## Building the Environment [WIP]
+
+Use the following image with tritonserver, triton distributed and tensorrt_llm.
+
+```bash
+nvcr.io/pfteb4cqjzrs/playground/tritondistributed:trtllm_disagg-shreyas
+```
+
+## Launching the Environment
+```
+# Run image interactively
+docker run -it --network host --shm-size=64G --ulimit memlock=-1 --ulimit stack=67108864 -e HF_HOME=/home/scratch.shreyasm_gpu/hf_cache --gpus=all -v .:/workspace nvcr.io/pfteb4cqjzrs/playground/tritondistributed:trtllm_disagg-shreyas
+```
+
+## Deployment Options
+
+### 1. Disaggregated Deployment
+
+NOTE: This will start both context and generation servers.
+```bash
+mpirun --allow-run-as-root -n 2 python3 -m disagg.worker --engine_args model.json &
+```
+
+Run router. This is still a WIP. It needs information about which endpoints are ctx/gen.
+```bash
+python3 -m disagg.router
+```
+
+Run client
+```bash
+python3 -m common.client \
+    --prompt "Describe the capital of France" \
+    --max-tokens 10 \
+    --temperature 0.5
+```

examples/python_rs/llm/trtllm/common/client.py

@@ -0,0 +1,81 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 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.
+
+
+import argparse
+import asyncio
+
+import uvloop
+from triton_distributed_rs import DistributedRuntime, triton_worker
+
+from .protocol import Request
+
+
+@triton_worker()
+async def worker(
+    runtime: DistributedRuntime,
+    prompt: str,
+    max_tokens: int,
+    temperature: float,
+    streaming: bool,
+):
+    """
+    Instantiate a `backend` client and call the `generate` endpoint
+    """
+    # get endpoint
+    endpoint = (
+        runtime.namespace("triton-init").component("router").endpoint("generate")
+    )
+
+    # create client
+    client = await endpoint.client()
+
+    # list the endpoints
+    print(client.endpoint_ids())
+
+    # issue request
+    tasks = []
+    for _ in range(1):
+        tasks.append(
+            client.generate(
+                Request(
+                    prompt=prompt,
+                    sampling_params={
+                        "temperature": temperature,
+                        "max_tokens": max_tokens,
+                    },
+                    streaming=streaming,
+                ).model_dump_json()
+            )
+        )
+    streams = await asyncio.gather(*tasks)
+
+    # process response
+    for stream in streams:
+        async for resp in stream:
+            print(resp)
+
+
+if __name__ == "__main__":
+    uvloop.install()
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--prompt", type=str, default="what is the capital of france?")
+    parser.add_argument("--max-tokens", type=int, default=10)
+    parser.add_argument("--temperature", type=float, default=0.5)
+    parser.add_argument("--streaming", type=bool, default=True)
+    args = parser.parse_args()
+
+    asyncio.run(worker(args.prompt, args.max_tokens, args.temperature, args.streaming))

examples/python_rs/llm/trtllm/common/parser.py

@@ -0,0 +1,86 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 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.
+
+import argparse
+import json
+import os
+from typing import Any, Dict, Tuple
+
+# Define the expected keys for each config
+# TODO: Add more keys as needed
+PYTORCH_CONFIG_KEYS = {
+    "use_cuda_graph",
+    "cuda_graph_batch_sizes",
+    "cuda_graph_max_batch_size",
+    "cuda_graph_padding_enabled",
+    "enable_overlap_scheduler",
+    "kv_cache_dtype",
+    "torch_compile_enabled",
+    "torch_compile_fullgraph",
+    "torch_compile_inductor_enabled",
+}
+
+LLM_ENGINE_KEYS = {
+    "model",
+    "tokenizer",
+    "tokenizer_model",
+    "skip_tokenizer_init",
+    "trust_remote_code",
+    "tensor_parallel_size",
+    "dtype",
+    "revision",
+    "tokenizer_revision",
+    "speculative_model",
+    "enable_chunked_prefill",
+    "disagg_config",
+}
+
+
+def _get_llm_args(args_dict):
+    pytorch_config_args = {
+        k: v for k, v in args_dict.items() if k in PYTORCH_CONFIG_KEYS and v is not None
+    }
+    llm_engine_args = {
+        k: v for k, v in args_dict.items() if k in LLM_ENGINE_KEYS and v is not None
+    }
+    if "model" not in llm_engine_args:
+        raise ValueError("Model name is required in the TRT-LLM engine config.")
+
+    return (pytorch_config_args, llm_engine_args)
+
+
+def _init_engine_args(engine_args_filepath):
+    """Initialize engine arguments from config file."""
+    if not os.path.isfile(engine_args_filepath):
+        raise ValueError(
+            f"'{engine_args_filepath}' containing TRT-LLM engine args must be provided in when launching the worker"
+        )
+
+    try:
+        with open(engine_args_filepath) as file:
+            trtllm_engine_config = json.load(file)
+    except json.JSONDecodeError as e:
+        raise RuntimeError(f"Failed to parse engine config: {e}")
+
+    return _get_llm_args(trtllm_engine_config)
+
+
+def parse_tensorrt_llm_args() -> Tuple[Dict[str, Any], Dict[str, Any]]:
+    parser = argparse.ArgumentParser(description="A TensorRT-LLM Worker parser")
+    parser.add_argument(
+        "--engine_args", type=str, required=True, help="Path to the engine args file"
+    )
+    args = parser.parse_args()
+    return _init_engine_args(args.engine_args)

examples/python_rs/llm/trtllm/common/protocol.py

@@ -0,0 +1,29 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 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.
+
+from tensorrt_llm.llmapi import DisaggregatedParams
+from pydantic import BaseModel
+
+
+class Request(BaseModel):
+    prompt: str
+    sampling_params: dict
+    streaming: bool = True
+    disaggregated_params: DisaggregatedParams = {}
+
+
+class Response(BaseModel):
+    text: str
+    disaggregated_params: DisaggregatedParams = {}

examples/python_rs/llm/trtllm/disagg/disagg_config.yaml

@@ -0,0 +1,18 @@
+model: TinyLlama/TinyLlama-1.1B-Chat-v1.0
+hostname: localhost
+port: 8000
+backend: "pytorch"
+context_servers:
+  num_instances: 1
+  gpu_fraction: 0.25
+  tp_size: 1
+  pp_size: 1
+  urls:
+      - "localhost:8001"
+generation_servers:
+  num_instances: 1
+  gpu_fraction: 0.25
+  tp_size: 1
+  pp_size: 1
+  urls:
+      - "localhost:8002"