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 的读取速度。

Prefill 阶段(计算密集) 输入 Prompt 并行计算 生成 KV 矩阵乘法密集 · GPU 算力利用率高 首个 token Decode 阶段(带宽密集) 读取 KV 单步推理 输出 token 循环直到生成 EOS 或达到最大长度 嵌入式设备瓶颈分析 ① 算力不足 Prefill 阶段耗时显著增加 ② 带宽受限 Decode 阶段 token/s 大幅下降 ③ 内存容量小 KV Cache 受限 → 上下文长度受限

2.4 嵌入式移植的核心权衡

基于上述硬件特性,嵌入式移植需要做出一系列与服务器端不同的决策。以下是需要重点关注的权衡点:

  1. 模型规模 vs 上下文长度:7B 模型 FP16 权重约 14GB,在 16GB 内存设备上几乎无法运行。通常需要 INT4 量化(约 4GB)并限制最大上下文长度(如 2048 token),才能同时容纳权重和 KV Cache。
  2. 精度 vs 推理质量:INT4 量化会带来一定精度损失。对于嵌入式场景,需通过困惑度(perplexity)或下游任务指标评估量化影响,必要时采用混合精度(部分层保持 FP16)。
  3. 批处理大小 vs 延迟:嵌入式设备通常面向单用户或低并发场景,batch size 设为 1 即可。增大 batch size 会显著增加显存占用,但吞吐提升有限。
  4. 引擎加载时间 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 量级模型开始评估,优先保证可用性。