面向低延迟大语言模型推理的混合 JIT-CUDA 图优化
机器学习
2026-04-28 v1 人工智能
硬件体系结构
摘要
大语言模型 (LLM) 在自然语言和多模态任务上取得了显著性能，但其实际部署仍受到推理延迟和内核启动开销的限制，尤其是在交互式、短序列场景中。本文提出一种混合运行时框架，将 Just-In-Time (JIT) 编译与 CUDA Graph 执行相结合，以减少启动开销，同时保持自回归解码期间的运行时灵活性。该框架将 Transformer 推理划分为通过 CUDA Graph 重放执行的静态组件和通过 JIT 编译内核处理的动态组件，实现了解码步骤间的异步图捕获与复用。我们在 LLaMA-2 7B 上进行评估，使用单 GPU、批量大小为 1 的推理，测试提示长度从 10 到 500 个 token。实验结果表明，该混合运行时在该场景下将时间到第一个 token (TTFT) 降低最高可达 66.0%，并在 P99 延迟方面低于 TensorRT-LLM。这些结果表明，混合 JIT-CUDA Graph 执行能够有效降低短序列 LLM 工作负载的推理延迟和方差，使其成为面向延迟敏感型 AI 应用的实用优化策略。
引用
@article{arxiv.2604.23467,
title = {Hybrid JIT-CUDA Graph Optimization for Low-Latency Large Language Model Inference},
author = {Divakar Kumar Yadav and Tian Zhao},
journal= {arXiv preprint arXiv:2604.23467},
year = {2026}
}