trtllm_async_server.py — TensorRT-LLM 异步服务器

文件概述

TensorRT-LLM 推理 HTTP 服务器和副本管理的实现（约 379 行）。

核心类

TRTLLMHttpServer (Ray Actor)

@ray.remote
class TRTLLMHttpServer:
    """TRT-LLM HTTP 推理服务器（单节点）"""

    async def launch_server(self):
        """启动 TRT-LLM 推理服务

        使用 tensorrt_llm 的 AsyncLLM 引擎:
        - backend="pytorch"
        - orchestrator_type="ray"
        - 支持 KV Cache 配置
        - 支持 CUDA Graph 优化
        """
        self.llm = await AsyncLLM(
            model=self.model_config.local_path,
            backend="pytorch",
            orchestrator_type="ray",
            kv_cache_config=kv_cache_config,
            ...
        )

    async def generate(self, prompt_ids, sampling_params, request_id):
        """使用 TRT-LLM 生成序列"""
        outputs = await self.llm.generate_async(inputs=prompt_ids, ...)
        return TokenOutput(token_ids=outputs.outputs[0].token_ids, ...)

    async def wake_up(self):
        """恢复推理引擎"""
        await self.llm.resume(tags=ServerAdapter.get_full_tags())

    async def sleep(self):
        """释放推理引擎资源"""
        await self.llm.release(tags=ServerAdapter.get_full_tags())

TRTLLMReplica(RolloutReplica)

class TRTLLMReplica(RolloutReplica):
    """管理 TRT-LLM 服务器副本"""

    def get_pgs_and_bundle_indices(self):
        """计算 placement group 和 bundle 索引

        确保 TRT-LLM worker 被正确分配到指定的 GPU 上
        """

    async def launch_servers(self):
        """启动 TRT-LLM 服务器

        使用 NodeAffinitySchedulingStrategy 确保
        服务器和 worker 在同一节点上
        """

与其他模块的关系

• TRTLLMReplica 继承 rollout/replica.py 的 RolloutReplica
• 使用 TensorRT-LLM 的 AsyncLLM 和 OpenAIServer
• 被推理管理器调度

小结

TRT-LLM 服务器利用 TensorRT-LLM 的异步引擎提供高性能推理，支持 CUDA Graph 和 KV Cache 优化。