2. 基础知识:TensorRT-LLM核心概念、嵌入式设备硬件特性与LLM推理基础
2.1 TensorRT-LLM 核心概念
TensorRT-LLM 是 NVIDIA 推出的 LLM 推理加速框架,它将大语言模型的计算图划分为多个阶段,并通过预编译的 TensorRT 引擎在 GPU 上高效执行。理解其核心概念是移植工作的前提。
| 核心概念 | 说明 | 对嵌入式移植的意义 |
|---|---|---|
| Engine(引擎) | 经过 TensorRT 编译后的序列化推理模型,包含优化后的 kernel 与权重 | 引擎文件需与目标 GPU 架构严格匹配,跨设备不可复用 |
| Builder(构建器) | 负责将 ONNX 或权重文件编译为 Engine 的组件 | 嵌入式设备上通常不运行 Builder,需在宿主机交叉编译 |
| Plugin(插件) | 自定义算子,用于支持 TensorRT 原生不支持的层(如 FlashAttention 变体) | 嵌入式 GPU 需确认插件是否有对应架构的 kernel 实现 |
| KV Cache | 推理过程中缓存历史 token 的 Key/Value 张量,避免重复计算 | 显存容量直接决定可支持的最大上下文长度 |
| In-flight Batching | 动态批处理机制,允许不同请求在不同解码阶段并行执行 | 嵌入式场景并发低,可考虑关闭以节省显存 |
关键认知:TensorRT-LLM 的编译产物(Engine)是硬件绑定的。在 Jetson Orin 上编译的引擎无法在 Xavier 上运行,甚至同系列不同显存版本也需要重新校准。因此,移植工作的第一步是明确目标设备的 GPU 架构代号(如 Ampere、Ada Lovelace、Hopper)。
2.2 嵌入式设备硬件特性
嵌入式设备与服务器 GPU 的差异是移植中所有技术决策的根源。以 NVIDIA Jetson 系列为例,其硬件特性直接影响模型选择、精度设置和内存管理策略。
2.2.1 GPU/NPU 计算单元
- GPU:Jetson Orin 搭载 Ampere 架构 GPU,支持 FP16/BF16/INT8 计算,但 INT8 需要通过 TensorRT 的 PTQ(训练后量化)校准获得缩放因子。
- NPU(如有):部分嵌入式 SoC 集成 NPU(如瑞芯微 RK3588 的 6 TOPS NPU),但 TensorRT-LLM 目前不支持 NPU 后端,仅支持 NVIDIA GPU。若目标平台非 NVIDIA,需考虑其他推理框架。
- 张量核心:TensorRT-LLM 的 kernel 依赖张量核心加速矩阵乘法,嵌入式 GPU 的张量核心数量远少于服务器 GPU,导致计算吞吐受限。
2.2.2 内存与存储
嵌入式设备采用统一内存架构(CPU 与 GPU 共享物理内存),这与服务器 GPU 的独立显存有本质区别。这意味着:
# 查看 Jetson 设备内存信息
sudo jetson_clocks --show
cat /proc/meminfo | grep MemTotal
# 典型 Orin NX 16GB 模块:CPU/GPU 共享 16GB LPDDR5
- 显存上限 = 系统内存上限:模型权重 + KV Cache + 激活值必须全部容纳在共享内存中,无法像服务器那样使用 CPU 内存作为显存溢出后备。
- 存储带宽:Jetson Orin 的内存带宽约 100-200 GB/s(视型号),而 A100 达 2 TB/s。带宽不足会严重制约 decode 阶段的速度。
- 存储介质:eMMC 或 NVMe SSD 的读写速度影响引擎加载时间。一个 2GB 的 Engine 文件从 eMMC 加载可能需要数秒,需考虑启动时预加载策略。
常见误区:很多开发者将服务器上的“模型显存占用”直接套用到嵌入式设备。实际上,由于统一内存架构,CPU 进程的内存占用也会挤压 GPU 可用空间。例如,运行 ROS 或视觉处理管线时,可用显存会显著减少。移植时必须预留 20%-30% 的内存余量给系统和其他进程。
2.3 LLM 推理基础
LLM 推理分为两个阶段:预填充(Prefill)和解码(Decode)。理解这两个阶段的计算特征,才能针对嵌入式设备做正确的优化取舍。
2.3.1 Prefill 阶段
当用户输入一段 prompt 时,模型并行处理所有输入 token,生成第一个输出 token 的隐藏状态。此阶段是计算密集型(compute-bound),需要大量矩阵乘法,对 GPU 算力要求高。
2.3.2 Decode 阶段
模型逐个生成后续 token,每步只处理一个 token,但需要读取全部 KV Cache。此阶段是内存带宽密集型(memory-bound),GPU 算力利用率低,瓶颈在于权重和 KV Cache 的读取速度。
2.4 嵌入式移植的核心权衡
基于上述硬件特性,嵌入式移植需要做出一系列与服务器端不同的决策。以下是需要重点关注的权衡点:
- 模型规模 vs 上下文长度:7B 模型 FP16 权重约 14GB,在 16GB 内存设备上几乎无法运行。通常需要 INT4 量化(约 4GB)并限制最大上下文长度(如 2048 token),才能同时容纳权重和 KV Cache。
- 精度 vs 推理质量:INT4 量化会带来一定精度损失。对于嵌入式场景,需通过困惑度(perplexity)或下游任务指标评估量化影响,必要时采用混合精度(部分层保持 FP16)。
- 批处理大小 vs 延迟:嵌入式设备通常面向单用户或低并发场景,batch size 设为 1 即可。增大 batch size 会显著增加显存占用,但吞吐提升有限。
- 引擎加载时间 vs 首次响应:大型 Engine 文件从闪存加载并反序列化需要时间。若对首 token 延迟敏感,可考虑将 Engine 常驻内存或使用更快的存储介质。
实操建议:在开始移植前,先使用 tegrastats 工具监控设备在空载和简单负载下的内存占用与 GPU 频率,建立性能基线。同时用 nvidia-smi(Jetson 上可用)确认 GPU 型号与驱动版本,确保与 TensorRT-LLM 的版本兼容性。
2.5 常见误区与复盘
误区一:忽略共享内存特性。在服务器上,CPU 内存与 GPU 显存独立,模型加载后 CPU 内存可释放。但在 Jetson 上,即使模型已加载到 GPU,CPU 进程的内存分配仍会减少可用显存。移植时需用 nvpmodel -q 检查当前电源模式,不同模式下的 GPU 频率和内存带宽差异巨大。
误区二:盲目追求大模型。嵌入式设备的功耗和散热限制决定了其无法长时间满负荷运行。一个 13B 模型即使量化后勉强装入内存,推理速度也可能低至 1-2 token/s,实际不可用。建议从 3B-8B 量级模型开始评估,优先保证可用性。