Chapter 5: CUDA 侧优化：Tensor Cores、Kernel 融合与调度

1. 开篇

在 NVIDIA DRIVE Orin 平台上，GPU（基于 Ampere 架构的 GA10B）不仅仅是 DLA 的替补，它是处理高复杂度 Vision Transformer、不规则算子（如 Deformable Conv）、以及后处理逻辑的绝对核心。

很多算法工程师习惯于在 PyTorch 中以 "Layer" 为单位思考模型（如 nn.Conv2d, nn.Linear），但在 Orin 的 GPU 上，真正的性能决胜点在于 "Kernel" 的粒度。一个在 PyTorch 中看起来很美的模型，如果充斥着碎片化的算子、非对齐的通道数、以及大量的显存读写操作，在 Orin 上运行时可能会遭遇严重的“带宽墙（Memory Wall）”或“调度延迟（Launch Latency）”。

本章的目标是让你学会像 GPU 架构师一样思考：如何通过设计 NHWC 友好、Tensor Core 对齐、高度融合 的网络架构，来榨干 Orin GPU 的每一个 SM（Streaming Multiprocessor）和每一 GB 的显存带宽。

2. 核心论述

2.1 数据的物理形态：Layout 与 Vectorized Access

在深度学习框架中，张量（Tensor）是一个数学概念；但在硬件上，张量是一段连续或非连续的内存。Orin 的 GPU 性能极大程度上取决于数据加载的效率。

2.1.1 为什么 Tensor Cores 唯独偏爱 NHWC？

NVIDIA Ampere 架构的 Tensor Cores 在执行矩阵乘法（GEMM）或卷积时，本质上是在处理 矩阵块（Tiles）。

  • 向量化加载 (Vectorized Loads): GPU 的加载指令（如 LDG.E.128）一次可以加载 128 bits（16字节）甚至更多的数据。
  • NHWC (Channels-Last): 在这种格式下，同一个像素点的 $C$ 个通道在内存中是连续存放的。如果 $C$ 是 32 或 64 的倍数，GPU 的一个线程（Thread）或一个线程束（Warp）可以用一条指令把一个像素的多个通道一次性“捞”进寄存器。
  • NCHW (Channels-First): 这是 PyTorch 的默认格式。同一个像素的 $C$ 个通道被分散在内存的 $C$ 个不同平面（Plane）上。要计算一个像素的卷积，线程需要跨越大段内存地址去收集数据，导致非合并访问（Uncoalesced Access），显存带宽利用率可能低至 20%-30%。

2.1.2 Reformat：看不见的性能杀手

当你在 TensorRT 中部署网络时，如果你强行使用了 NCHW 优化的层（如某些自定义 Plugin）夹杂在 NHWC 优化的层（如 Tensor Core 卷积）中间，TensorRT 会自动插入 Reformat 节点。

  • 表现：在 Nsight Systems profiler 中，你会看到大量的 kPCIE 或 kMEM 类型的 kernel，名字通常包含 reformat、transpose 或 copy。
  • 代价：这些操作不贡献任何 FLOPs，却消耗了最贵的显存带宽和功耗。

ASCII 图解：NHWC 的向量化优势

假设我们使用 FP16 (2 bytes)，通道数 C=8。
我们需要读取 Pixel 0 的所有通道数据。

[ NCHW 内存布局 ]
Addr:  0x00      0x100     0x200     0x300 ... (假设 H*W=128)
Data:  [C0_P0]...[C1_P0]...[C2_P0]...[C3_P0]...
操作:  需要发送 8 次独立的内存请求，每次地址跨度很大。
       Cache 命中率极低。

[ NHWC 内存布局 ]
Addr:  0x00  0x02  0x04  0x06  0x08  0x0A  0x0C  0x0E
Data:  [C0]  [C1]  [C2]  [C3]  [C4]  [C5]  [C6]  [C7]  (均为 Pixel 0)
操作:  GPU 发送 1 次 128-bit 向量加载指令 (LDG.128)。
       数据一次性进入寄存器文件，总线利用率 100%。

2.2 32 字节对齐：Orin 的“强迫症”

仅仅使用 NHWC 是不够的，通道数 $C$ 的具体数值决定了是否需要 Padding。 Orin 的 Tensor Cores 在 INT8 推理时，最佳吞吐模式要求输入 Tensor 在 Channel 维度上对齐到 32 字节（即 32 个 INT8 通，或 16 个 FP16 通道）。

  • 黄金数字: $32, 64, 96, 128, 256, \dots$
  • 毒药数字: $3, 17, 30, 80$ (80 虽然是偶数，但对于某些 INT8 pattern 可能不是最优，不过通常 16 的倍数在 FP16 下已足够)。
  • Padding 的代价: 如果你的网络定义是 C=30，TensorRT 往往会隐式地将其 Pad 到 32。这不仅浪费了 2/32 的显存，还可能引入额外的 Masking 计算逻辑。在 Backbone 设计中，永远向上取整到 32 的倍数。

2.3 Attention 机制的硬件化改造

Transformer 的核心是 Attention，在 Orin 上优化 Attention 的关键是打破 Memory Bound。

2.3.1 标准 Attention 的痛点

标准公式 $Softmax(QK^T)V$ 涉及三个巨大的矩阵操作。

  1. HBM 读写风暴: 计算 $S = QK^T$ 后，需要将 $S$ ($N \times N$) 写回显存。
  2. Softmax 瓶颈: 读取 $S$，计算 exp 和 sum，再写回。Softmax 是 Elementwise 操作，计算密度极低，纯粹消耗带宽。
  3. Dropout: 训练时的 Dropout 在推理时虽然没有，但如果网络结构未剪枝干净，残留的 Mask 操作也会拖慢速度。

2.3.2 Flash Attention / FMHA

Orin 的 TensorRT 强力支持 FMHA (Fused Multi-Head Attention) 插件。

  • SRAM Tiling: 它不将中间结果 $N \times N$ 写回 HBM，而是在 SM 的 L1/Shared Memory 中分块完成 Softmax。
  • 架构约束 (Rule of Thumb):
    • Head Size: 必须设计为 32, 64, 128。最推荐 64。如果设为 40 或 80，FMHA kernel 可能无法启动，回退到慢速实现。
    • Seq Length: 虽然支持动态长度，但最好设有明确的上限（如 2048 或 4096）以优化显存预分配。

2.4 Kernel 融合：消灭“微操作”

GPU 的启动开销（Launch Overhead）约为 5-10 微秒。如果一个算子（如 Add）的执行时间只有 1 微秒，那么 GPU 实际上 90% 的时间在空转等待 CPU 的命令。

2.4.1 垂直融合 (Vertical Fusion)

这是 TensorRT 的拿手戏。它会将 Conv + Bias + BN + ReLU 合并成一个单一的 CUDA Kernel（CBR）。

  • 设计陷阱: 在 Conv 和 ReLU 之间插入 Print、Save、或者分支结构（Branching），会打断融合。
  • Add 融合: 残差连接的 Add 操作通常可以融合进前一个 Conv 的 kernel 中（作为 output 阶段的操作）。

2.4.2 水平融合 (Horizontal Fusion)

如果你有多个结构相同的头（例如 BiFPN 的不同层级，或者检测头的 Class/Box 分支），且它们的输入源相同或可推导：

  • 手动优化: 在网络定义阶段，尝试用 Group Conv 或者单一的大 Conv 替代多个小 Conv，然后通过 Slice 分开输出。
  • 理由: 一个大 Kernel 的并行度（Occupancy）远高于多个小 Kernel，且只读取一次输入数据。

ASCII 图解：水平融合示例

[原始设计：低效]
Input (C=256) ---> [Conv 1x1, Out=64] ---> Class Head
              |--> [Conv 1x1, Out=64] ---> Box Head
              |--> [Conv 1x1, Out=32] ---> Dir Head

* 3次 Kernel Launch
* Input 被重复读取 3 次

[优化设计：高效]
Input (C=256) ---> [Conv 1x1, Out=160] ---> Output (C=160)

* 1次 Kernel Launch
* Input 只读 1 次
* 后续通过 Pointer 偏移 (View) 即可获得各 Head 数据，无需物理 Slice。

2.5 算力与带宽的平衡：Arithmetic Intensity

算术强度 (Arithmetic Intensity) = $\frac{\text{FLOPs}}{\text{Bytes Access}}$。 Orin 的 Tensor Cores 算力极强，这使得大多数 Vision Transformer 网络变成了 Bandwidth Bound（带宽受限）。

  • Depthwise Conv 的尴尬: 深度可分离卷积（DW-Conv）在 CPU/DLA 上很棒，但在 GPU 上，由于其算术强度极低（计算量小，但读写量并未同比例减少），其加速比往往不如标准卷积明显。
  • Linear/MLP 的优势: 在 Transformer 的 FFN 块中，Linear 层（Dense Matrix Multiply）具有极高的算术强度，是 GPU 最喜欢的算子。
  • 建议: 在 GPU 侧，要过分迷信 MobileNet 风格的 DW-Conv。适当使用 Group Conv 或标准 Conv，利用 Tensor Cores 的暴力计算能力，反而可能获得更低的 Latency。

3. 本章小结

  1. Layout 为王: 坚决拥抱 NHWC。所有自定义算子或前后处理都应尽量在 NHWC 格式下进行，避免 TensorRT 插入 Reformat 节点。
  2. 对齐是硬指标: Channel 数和 Head Size 必须是 32 或 64 的倍数。Head Size = 64 是 Flash Attention 的甜点区。
  3. 融合提升效率: 通过水平融合（合并分支）和垂直融合（Conv-BN-Act）减少 HBM 访问次数。避免在主干网络中使用导致断图的非标准算子。
  4. 警惕“小算子”: 避免网络中充斥着大量的 Reshape, Transpose, Slice, Concat。这些是纯粹的带宽消耗者。
  5. 认识硬件特性: Depthwise Conv 在 GPU 上未必快，标准 Dense 计算（Linear/Conv）更能发挥 Tensor Cores 的效能。

4. 练习题

基础题 (50%)

Q1: 维度对齐的量化影响 在设计一个用于 Orin 的 CNN Backbone 时，你将某层的输出通道数设为 C=150。

  1. 在 FP16 模式下，这对 Tensor Cores 的影响是什么？
  2. 在 INT8 模式下（假设使用 TensorRT 的 qdq 模式），这会带来什么额外的内存开销或计算开销？
点击查看答案提示
  1. FP16: Tensor Cores 通常以 16字节或32字节为 chunk 工作。150 不是 8 或 16 的倍数（150 / 8 = 18.75）。这意味着 GPU 在处理每行数据的末尾时，需要处理非对齐的内存访问，或者 TensorRT 会隐式 pad 到 152/160，浪费计算力。
  2. INT8: INT8 对齐要求更严（通常 32 字节 / 32 channels）。TensorRT 极大概率会将 150 pad 到 160 (32 * 5)。这意味着显存占用增加了 10/150 $\approx$ 6.7%，且计算量也相应增加。更重要的是，这可能导致后续所有层都需要为了对齐而进行 padding。

Q2: Kernel Launch 瓶颈 你编写了一个包含 50 层连续 1x1 Conv + ReLU 的网络，每层的 feature map 很小（如 $8 \times 8$）。在 Orin 上 profile 发现，GPU 利用率（SM Occupancy）只有 10%，且有大量空隙。 请解释原因，并给出两种架构层面的解决方案。

点击查看答案提示
  • 原因: Kernel Launch Bound。Feature map 太小，计算量极低，GPU 瞬间就做完了，大部分时间花在 CPU 下发指令和 GPU 启动 Kernel 的 overhead 上。GPU 处于“饿死”状态。
  • 解决方案:
    1. 增大 Batch Size: 比如通过多摄合并推理，让 $8 \times 8$ 变成 $8 \times 8 \times 6$，增加计算密度。
    2. 算子融合: 这种结构是否可以用一个更深的 MLP 代替？或者是否可以通过 CUDA Graph 捕获整个图来消除 CPU Launch 开销。
    3. 架构调整: 避免在网络深层（低分辨率）使用过深、过窄的卷积链，改为更宽（Wide）的结构。

Q3: Reformat 节点定位 在 TensorRT 生成的 engine graph 中，你发现在 Input -> Conv1 之间有一个名为 Reformat_0 的节点。Input 是 NCHW 格式，Conv1 是 Tensor Core 优化的。 请问这个 Reformat 节点在做什么？如何修改 Input 端的代码来消除它？

点击查看答案提示
  • 作用: 该节点正在执行 NCHW -> NHWC 的转置操作，这是因为 Tensor Core 卷积核需要 NHWC 布局。
  • 消除方法: 在预处理阶段（CUDA Kernel 或 CPU 代码中），直接将图像数据写入为 NHWC 格式（Interleaved）。然后在 TensorRT 构建期，显式告诉 Parser 输入的格式是 NHWC。这样 Input -> Conv1 就是零拷贝直连。

挑战题 (50%)

Q4: Attention 的 Head Size 设计 你的 Transformer 模型当前配置是：Embedding=768, Num_Heads=12。这意味每个 Head 的维度是 $768/12 = 64$。 为了提升精度，团队决定将 Num_Heads 增加 16，导致 Head_Dim = 48。 请预测在 Orin 上这会对推理时延产生什么影响？为什么？（假设使用 TensorRT FMHA）。

点击查看答案提示
  • 影响: 时延可能会不降反升，或者提升幅度远低于预期，甚至导致 FMHA 失效。
  • 原因:
    1. FMHA 限制: TensorRT 的优化 FMHA kernel 通常针对特定的 Head Size 编译（如 32, 64, 128）。48 是一个非典型尺寸。
    2. Kernel Fallback: 如果不支持 Size=48，TRT 可能会回退到通用的、基于显存读写的 Attention 实现（即 $QK^T$ 写回 HBM），这将导致巨大的带宽开销。
    3. Padding: 即使支持，内核内部也可能将其 Pad 到 64 进行计算，导致计算资源的浪费。
  • 建议: 保持 Head_Dim=64，调整 Embedding Size 到 $16 \times 64 = 1024$，或者保持 Embedding=768 但调整 Head 数。

Q5: 深度思考 - Depthwise Separable Conv 在 Ampere 上的陷阱 MobileNetV2/V3 在手机端 CPU 上很快，因为它们大量使用了 Depthwise Separable Conv (DW-Conv + Pointwise-Conv)。 在 Orin GPU 上，直接把 ResNet50 替换成 MobileNetV3 往往不能获得预期的加速比（例如 FLOPs 降低了 10 倍，但 Latency 只降低了 2 倍）。请从 算术强度 (Arithmetic Intensity) 和 显存带宽 的角度解释原因。

点击查看答案提示
  • 算术强度低: DW-Conv 的计算量（FLOPs）非常小，但它仍然需要读取完整的 Input Feature Map 并写入 Output。这意味着它的算术强度（FLOPs/Bytes）很低。
  • Bandwidth Bound: Orin GPU 的 Tensor Cores 算力极大，对于 DW-Conv 这种低计算密度的算子，瓶颈完全在显存带宽（HBM Bandwidth）。GPU 的 SM 大部分时间在等待数据，而不是在计算。
  • Launch Overhead: MobileNet 往往层数很深且单层计算量小，导致 Launch 开销占比高。
  • 结论: 在高算力 GPU ，有时使用标准卷积（Standard Conv）虽然 FLOPs 高，但能更好地利用 Tensor Cores 的密集计算能力，且减少了算子总数，综合效率可能更高。

Q6: 多任务头的水平融合实践 设计一个检测网络，有 3 个输出头分别预测 Class (80 ch), Box (4 ch), Centerness (1 ch)。 目前的实现是 3 个独立的 nn.Conv2d。 请设计一个具体的融合方案，并说明在 TensorRT 中如何通过 ISliceLayer (Slice) 实现 Zero-Copy 的输出拆分。

点击查看答案提示
  • 方案: 定义一个单一的 nn.Conv2d，输出通道数 $K = 80 + 4 + 1 = 85$（实际上为了对齐建议 Pad 到 96）。
  • 实现:
    1. 执行一次大卷积，得到 $B \times 96 \times H \times W$ 的张量。
    2. 在 TensorRT (Network Definition) 中，添加 3 个 Slice Layer：
      • Slice 1: Start=0, Size=80 (Class)
      • Slice 2: Start=80, Size=4 (Box)
      • Slice 3: Start=84, Size=1 (Center)
  • Zero-Copy 原理: 如果后续操作允许（且维度对齐），Slice 在 TensorRT 内部往往是指针偏移（Pointer Arithmetic），并不会真正发生 D2D 内存拷贝。但要注意，如果这三个分支后续在不同 Stream 消费，可能还是会发生拷贝。最重要的是，我们只做了一次 Input Feature Map 的读取和一次 Kernel Launch。

5. 常见陷阱与错误 (Gotchas)

5.1 "FP16 里的 FP32 卧底"

  • 现象: 你开启了 FP16 模式，但速度没有翻倍。查看 Profiler，发现大量 Cast（类型转换）算子。
  • 原因: 某些算子（如 LayerNorm, GELU 的特定实现, 或者是 Softmax 在某些 axis 上）在 TensorRT 中如果没有对应的 FP16 实现，或者被策略层认为精度风险过大，会强制回退到 FP32。
  • 对策:
    • 检查 LayerNorm 是否使用了 TensorRT 的 Plugin（如 LayerNormPlugin 通常支持 FP16）。
    • 在 PyTorch 导出 ONNX ，尽量使用标准算子组合，避免自定义复杂的数学运算。
    • 强制指定 Layer 类型（但这可能带来精度溢出风险，需配合 QAT 验证）。

5.2 动态形状 (Dynamic Shapes) 的代价

  • 现象: 为了通用性，将输入维度设为 -1, 3, -1, -1。推理时发现性能不稳定，且比静态形状慢。
  • 原因:
    • TensorRT 在构建引擎时，会针对具体的维度选择最优的 Kernel。如果是动态形状，TRT 必须选择一套“通吃”但非最优的方案，或者在运行时动态选择（造成抖动）。
    • 显存分配无法在初始化时固定，可能导致运行时的 malloc 开销。
  • 建议: 在自动驾驶场景中，相机分辨率通常是固定的。尽量使用 Static Shape。如果必须支持多尺度，考虑编译多个 Profile（如 Profile A: $640 \times 640$, Profile B: $1280 \times 1280$）。

5.3 忽视了 Padding 的副作用

  • 现象: 为了对齐，手动给 Feature Map 做了 Padding，结果后面跟着一堆 Masked_Fill 或者复杂的索引逻辑。
  • 陷阱: 额外的 Masking 逻辑（Elementwise操作）消耗的带宽可能比 Padding 带来的对齐收益还大。
  • Rule of Thumb: 尽量让 Padding 发生在“无感”的地方（如 Conv 的 stride/padding 参数中），而不是显式的数据拷贝 Padding。

5.4 滥用 Group Conv

  • 现象: 为了减少参数量，将 Group 数设得非常大（如 Group = Channel，即 Depthwise）。
  • 陷阱: 当 Group 数很大时，Tensor Core 无法高效工作（因为无法在 Channel 方向累积足够的乘加运算）。在 GPU 上，Group Size 最好不要小于 4 或 8。Group=1 (标准卷积) 通常是吞吐量最高的选择。