tensorrtx 手写 TensorRT 实现 Vision Transformer（ViT）：从 .wts 导出、引擎构建到 Pre-LN Transformer 逐层解析

导读

本文以 vit/README.md 为核心，结合 tensorrtx 仓库中 ViT 模块的完整源码，系统讲解如何完全不用任何解析器（parser），仅凭 TensorRT network definition API 手写构建 ViT 推理引擎：涵盖从 HuggingFace 权重导出 .wts、CMake 构建、引擎序列化到动态 batch 推理的完整链路，并逐层剖析 Pre-LayerNorm Transformer 编码器在 TensorRT 中的数学实现与张量形状流转。读完本文，你将掌握 ViT-B/L/H 系列 5 种变体在 TensorRT 10.x 下的部署方法，理解自实现 Swi-GeLU 与 Myelin 融合内核的触发条件，并能读懂 profiler 输出的层级性能报告。

1. 概述：手写 ViT 的 TensorRT 实现

该模块是对 Vision Transformer（论文编号 arxiv.org.2010.11929，即 2020 年 10 月发表的 An Image is Worth 16x16 Words）的手写 TensorRT 实现：不依赖 ONNX 解析器、不使用 nvinfer1 之外的第三方插件，全部网络结构由 vit.cc 中的 createEngine() 与 ViTLayer() 两个函数通过 INetworkDefinition API 逐层搭建。

这种实现方式的价值在于：

• 结构完全可控，每一层的名称、数据类型、形状都可以显式指定，便于做细粒度性能剖析；
• 便于为特定 GPU 架构手工编排融合友好的算子顺序（例如触发 TensorRT 内部的 Myelin 融合通道）；
• 权重来源（HuggingFace 官方 checkpoint）与推理代码解耦，权重以通用 .wts 文本格式保存。

1.1 功能特性清单

• 支持 TensorRT SDK 8.5.1+ ~ 10.15.1+（README 声明范围，实际代码中通过 TRT_VERSION_GE/LT 宏做了版本分支）；
• 支持 Windows 11 操作系统（CMakeLists.txt 中针对 WIN32 设置了 MSVC 运行库）；
• 支持原生或自实现两种 Swi-GeLU（见第 8 节）；
• 支持原生或自实现两种多头自注意力（见第 7 节与第 10 节）；
• 默认启用一个 dummy profiler（VisionTransformerProfiler，见 profiler.h）；
• 默认启用一个 dummy output allocator（CudaOutputAllocator，见 cuda_allocator.h）；
• 默认启用 optimization profile，动态 batch 配置为 min=1 / opt=1 / max=2（源码常量 BUILD_MIN_BATCH / BUILD_OPT_BATCH / BUILD_MAX_BATCH，见 vit.cc 顶部）；
• 支持 ViT-B/16、ViT-B/32、ViT-L/16、ViT-L/32、ViT-H/14 五种官方变体。

1.2 当前已知限制

原文档明确列出如下限制，部署前需注意：

• IAttention 在 TensorRT 10.14 ~ 10.15 之间存在 bug，不能使用；README 建议使用 10.15.1+ 以启用原生融合内核。源码中也有对应编译期告警：当 TRT_VERSION_GE(10,14,0) && TRT_VERSION_LT(10,15,0) 时，构建过程会打印 "IAttention is available in TensorRT 10.14.1 SDK but have bugs, use 10.15.1+ to enable native fused kernel"；
• TensorRT < 8 不受支持：部分算子未在 cuDNN 中实现（utils.h 中 checkTrtEnv 对 TRT<8 场景做了 SM>86 的拦截）；
• SM < 86、TensorRT < 10、CUDA < 12 的组合尚未被完全测试；
• ViT-H/14 使用 HuggingFace 的 ImageNet-21k checkpoint（google/vit-huge-patch14-224-in21k，共 21843 类），公开渠道不存在 1k 微调版 checkpoint——这意味着 ViT-H/14 引擎的分类头是 21843 维，标签映射文件也不适用。

2. 五种变体的配置参数

所有变体共享同一套架构代码，仅尺寸不同。配置表集中在 ViTConfig 结构体与 getVariantConfig() 函数中（vit.cc），也在 gen_wts.py 的 VARIANTS 字典中有一致镜像：

变体	HuggingFace 模型	hidden (D)	层数	头数 (H)	FFN 维度 (4D)	patch	输入尺寸	类别数
ViT-B/16	google/vit-base-patch16-224	768	12	12	3072	16	224	1000
ViT-B/32	google/vit-base-patch32-384	768	12	12	3072	32	384	1000
ViT-L/16	google/vit-large-patch16-224	1024	24	16	4096	16	224	1000
ViT-L/32	google/vit-large-patch32-384	1024	24	16	4096	32	384	1000
ViT-H/14	google/vit-huge-patch14-224-in21k	1280	32	16	5120	14	224	21843

模型类型命名支持别名：b16、b32、l16、l32、h14 等（源码 normalizeModelType() 会去掉 -///空格并转大写后归一化，例如 vit-b/16、B/16 均映射到 ViT-B/16）。patch 数、序列长度的换算见源码：

num_patches = (img_size / patch)²，seq_len = num_patches + 1（含 cls token）

例如 ViT-B/16：(224/16)² = 196 个 patch，序列长度 196 + 1 = 197。

3. 完整使用流程（四步走）

原文档给出完整的四步部署链路，每一步都对应仓库中的真实脚本/命令：

第 1 步：用 gen_wts.py 生成 .wts 权重文件

# 从 ViT-B/16、ViT-B/32、ViT-L/16、ViT-L/32、ViT-H/14 中任选其一
# 别名 b16、b32、l16、l32、h14 同样被接受
python gen_wts.py ViT-B/16
# -> 写入 models/ViT-B-16.wts

脚本（vit/gen_wts.py）内部逻辑：

• 强制要求设置缓存环境变量：TORCH_HOME 与 HF_HOME，未设置会直接 RuntimeError。README 提到脚本会在导入 torch、transformers 前设置下载缓存（例如 TORCH_HOME=/mnt/data/storage/torch、HF_HOME=/mnt/data/storage/huggingface），实际脚本通过 require_cache_env() 校验环境变量存在；
• 通过 AutoModelForImageClassification.from_pretrained(hub_id) 加载官方 checkpoint，并强制 _attn_implementation = "eager"，保证导出的是与代码实现一致的显式注意力算子；
• 对 in21k checkpoint（ViT-H/14）强制设置 num_labels = 21843：因为这类 checkpoint 未携带分类头，HF 默认会退化成 num_labels=2 的随机小头，与引擎配置表不匹配；
• .wts 格式：首行是张量个数，随后每行 <a href="https://link.gitcode.com/i/a5404f53e0abaed02edb4ec3d5440960" target="_blank">name] [size] <size 个 float32 大端序十六进制>，与 [utils.h 中 loadWeights() 的解析逻辑完全对应；
• 默认输出路径为 models/<safe_name>.wts，其中 / 替换为 -（如 ViT-B/16 -> models/ViT-B-16.wts）；也可通过第二参数指定输出路径，如 python gen_wts.py ViT-L/16 models/vit_l16.wts。

第 2 步：构建 C++ 代码

pushd tensorrtx/vit
cmake -S . -B build -G Ninja --fresh
cmake --build build

构建系统要点（vit/CMakeLists.txt）：

• 要求 CMake 3.17+、C++20 / CUDA C++20 标准；
• 可执行文件 vit 由 vit.cc + cuda_allocator.cc + profiler.cc 构成，链接 CUDA runtime、CUDA driver 与 TensorRT、OpenCV（用于图像读取与预处理）；
• 默认 CMAKE_CUDA_ARCHITECTURES 覆盖 80/86/89/90/100/120，覆盖 Ampere 到 Blackwell；
• 通过仓库自带的 FindTensorRT.cmake 定位 TensorRT（若 TensorRT::TensorRT 目标未被外部 find 到）；
• 从 IDE（如 CLion）运行时若 PATH 中没有 nvcc，会自动在 CUDA_PATH / CUDA_HOME、/usr/local/cuda*/bin 等路径探测 CUDA 编译器。

第 3 步：将 .wts 序列化为 engine 文件

./build/vit -s <wts_path> <engine_path> <model_type>
# 示例：
./build/vit -s models/ViT-B-16.wts models/ViT-B-16.engine ViT-B/16

-s 模式（源码 main() 中 argc == 5 分支）会先打印构建配置摘要（[cfg] ... hidden=768 layers=12 heads=12 ff=3072 patch=16 img=224 classes=1000 batch[min/opt/max]=1/1/2 precision=fp16），然后调用 APIToModel() -> createEngine() 完成图构建、序列化并写出 engine 文件。注意 precision=fp16：BUILD_PRECISION 编译期常量默认为 DataType::kHALF，如需 FP32 需修改 vit.cc 顶部的该常量后重新编译。

构建阶段关键配置（见 createEngine()）：

• TensorRT 10+ 使用 kSTRONGLY_TYPED 强类型网络（createNetworkV2），输入 tensor 数据类型与权重类型决定精度；TensorRT < 10 使用 kEXPLICIT_BATCH 弱类型网络并显式开启 BuilderFlag::kFP16；
• workspace 内存池上限 16 MiB（WORKSPACE_SIZE），builder_optimization_level = 5；
• 为输入 input（形状 [-1, 3, img_size, img_size]）创建 optimization profile：min=1 / opt=1 / max=2。

第 4 步：运行推理

./build/vit -d <engine_path> <image_or_image_dir>
# 示例：
./build/vit -d models/ViT-B-16.engine assets/cats.jpg

-d 模式（argc == 4 或 5）细节：

• 输入既可以是单张图片，也可以是目录；目录会被递归收集（严格说是遍历目录）所有 .jpg/.jpeg/.png/.bmp 文件并排序（collectImages()）；
• batch 自动取 min(图片数, 引擎 max_batch=2)，README 指出"传入含两张图片的目录会以单 batch 跑两个样本"；
• 每次推理前会以默认流（stream）执行，输出通过 CudaOutputAllocator 管理；
• 先跑 5 次 warmup，再跑 1 次输出 Top-3，最后跑 200 次基准（kBenchmarkRuns，见 utils.h）并输出 AVG / P50 / P90 / P95 / P99 百分位延迟；
• 可选追加 --profile 参数启用 profiler 输出，例如：

  ./build/vit -d models/ViT-B-16.engine assets/cats.jpg --profile
  

• 标签映射读取 assets/imagenet1000_clsidx_to_labels.txt（1000 类模型才加载标签；ViT-H/14 输出 21843 类无标签）。

4. 推理输出示例与 Myelin 融合验证

原文档给出了在 RTX 4080、TensorRT 10.15.1 SDK 上的真实输出形态：

====
1880us
-1.125, 0.4623, -0.1215, ...
====
prediction result:
Top: 0 idx: 285, logits: 8.262, label: Egyptian cat
Top: 1 idx: 281, logits: 7.872, label: tabby, tabby cat
Top: 2 idx: 282, logits: 6.477, label: tiger cat
========== VisionTransformerProfiler ==========
                                       TensorRT layer name    Runtime, %  Invocations Runtime, ms
Reformatting CopyNode for Input Tensor 0 to patch embedding     3.2%          20        0.95
                                            patch embedding     1.5%          20        0.45
...
              _gemm_mha_v2_myl3_6                             0.5%          20        0.14
...
========== VisionTransformerProfiler total runtime = 29.67 ms ==========

该输出验证了一个关键事实：手写网络成功触发了 TensorRT 内部的 MHA 融合通道（"Myelin"，简称 "myl"），尤其是融合内核 _gemm_mha_v2_myl3_6。层名中的 myl3 即 Myelin pass 第 3 轮的产物，_gemm_mha_v2 是融合后的 MHA 矩阵乘法内核。这意味着手写图与官方融合器兼容，无需退化为逐算子执行。

5. 整体网络图：patch embedding 与分类头

除 ViTLayer() 编码器层外，createEngine() 还构建了网络的"头尾"：

1. Patch embedding（vit.cc 的 createEngine()）：以卷积实现 patch 投影——addConvolutionNd，hidden 个输出通道、卷积核 patch x patch、步长 patch，对应 ViT 的 vit.embeddings.patch_embeddings.projection.weight/bias；输出经 Shuffle reshape 为 (N, hidden, num_patches) 再转置为 (N, num_patches, hidden)；
2. cls token 与位置编码：cls_token、position_embeddings 按 max_batch=2 复制为常量（makeRepeatedBatchWeights），再通过 addBatchedConstantSlice 按实际 batch 动态裁剪；在序列轴（axis=1）上拼接 cls token 后与位置编码逐元素相加（position_embed）；
3. 末尾 LayerNorm：vit.layernorm.weight/bias；
4. 分类头：用 Gather 取序列第 0 维（CLS token），reshape 为 (-1, hidden)，与 classifier.weight/bias 做线性层，输出命名为 logits。

输入输出 tensor 名称分别为 input 与 logits（NAMES 数组），预处理与 TensorRT 保持一致：OpenCV 读取 BGR 图后 resize 到 img_size，除以 255，减去 mean (0.5,0.5,0.5)、除以 std (0.5,0.5,0.5)，再打包为 FP16 CHW 布局（见 utils.h 的 preprocess_img()）。

6. Transformer 编码器层（Pre-LayerNorm）总体结构

ViTLayer() 用 TensorRT 原语构建一个 ViT 编码器块，与 ViT 论文一致采用 Pre-LayerNorm（先归一化、后做子层）结构，完整包含：

• 注意力前的 LayerNorm（layernorm_before）
• 多头自注意力（MHSA）：QKV 投影 → 缩放点积注意力 → 输出投影
• 残差连接
• 注意力后的 LayerNorm（layernorm_after）
• 前馈网络（FFN/MLP）：dense → GeLU → dense
• 残差连接

函数返回最终残差输出张量。其数学骨架为：

X' = LN1(X)
A  = MHSA(X')
Y  = X + A
Y' = LN2(Y)
F  = FFN(Y')
Z  = Y + F

6.1 符号与张量形状约定

设输入 X ∈ R^{N×L×D}：

• N：batch 大小（代码中的动态维 -1）
• L：序列长度（patch 数 + 1，代码中为动态维）
• D：隐藏维度，由所选变体决定

注意力头配置：H = param.head_num，每个头的维度 d = D/H。

6.2 权重形状（概念层）

• Q/K/V 投影权重：W_Q, W_K, W_V ∈ R^{D×D}
• Q/K/V 偏置：b_Q, b_K, b_V ∈ R^D（注意：原生 nvidia 接口（IAttention 路径）不使用偏置）
• 输出投影：W_O ∈ R^{D×D}, b_O ∈ R^D
• FFN（扩张比 4）：W_1 ∈ R^{D×4D}, b_1 ∈ R^{4D}，W_2 ∈ R^{4D×D}, b_2 ∈ R^{D}（FFN 维度由变体决定）

6.3 LayerNorm 定义

LayerNorm 沿**最后一维（隐藏维度 D）**独立作用于每个 (n, ℓ) 位置：

LN(x) = γ ⊙ (x - μ)/√(σ² + ε) + β
μ = (1/D) Σ x_i
σ² = (1/D) Σ (x_i - μ)²

其中 γ 对应 .weight、β 对应 .bias、ε = param.lnorm_eps（默认 1e-12）。源码实现 addLinearNorm() 会根据 TRT 版本选择 addNormalizationV2（10.15+，新 API）或 addNormalization（旧 API），并统一 setEpsilon(1e-12f)；LayerNorm 的 axes 掩码取 1U << (nbDims-1)，即仅归一化最后一维。

7. 多头自注意力（MHSA）逐算子拆解

7.1 线性投影（代码 2.1 节）

对 X' = LN1(X)，分别做三个线性变换并加偏置：

Q = X'W_Qᵀ + b_Q
K = X'W_Kᵀ + b_K
V = X'W_Vᵀ + b_V
Q,K,V ∈ R^{N×L×D}

源码中 QKV 权重以 Dims3{1,H,H} 常量加入，addMatrixMultiply(..., M::kNONE, ..., M::kTRANSPOSE) 完成投影，随后 addElementWise(..., E::kSUM) 加偏置，并分别命名为 query/key/value。随后用 Shuffle 将每个张量 reshape 成 (N, L, H, d)，再二次转置 {0,2,1,3} 得到 (N, H, L, d) 的 head-first 布局：

q_s->setReshapeDimensions(Dims4{0, 0, param.head_num, attn_head_size});
q_s->setSecondTranspose({0, 2, 1, 3}); // (N,H,L,d)

7.2 缩放点积注意力（SDPA）

• Attention logits：S = Q Kᵀ ∈ R^{N×H×L×L}，对应 addMatrixMultiply(q_s, kNONE, k_s, kTRANSPOSE)；
• 缩放：α = 1/√d，S̃ = α·S。源码中 scale_f = 1.0f/sqrt(attn_head_size) 会按输入 dtype 生成 FP16/FP32 常量（attn_name + ".scale"），再以 E::kPROD 逐元素乘到 logits 上；
• Softmax：沿**最后一维（key 索引）**归一化，P ∈ R^{N×H×L×L}：

qk_softmax = SoftMax(attn_qk);
qk_softmax->setAxes(1U << (nbDims-1)); // last axis

• 加权求和：O = P V ∈ R^{N×H×L×d}，对应 addMatrixMultiply(qk_softmax, kNONE, v_s, kNONE)。

7.3 合并多头与输出投影

• 先 setFirstTranspose({0,2,1,3}) 转回 (N,L,H,d)，再 reshape 回 (N,L,D)（如 ViT-B 的 768）；
• 输出投影：A = O_merged W_Oᵀ + b_O（源码中权重为 Dims3{1,H,H} 常量，out_proj_w 参与 kTRANSPOSE 乘法，再加 out_proj_b，层名 attn_fcb）。

7.4 注意力残差

Y = X + A，源码为 addElementWise(input, attn_fcb, kSUM)，命名 attn_residual。原文档特别指出：这条恒等通路对梯度流动与数值稳定至关重要，推理时即便注意力变得尖锐或含噪，它仍保留了"直达"信号路径。

7.5 注意力后 LayerNorm

Y' = LN2(Y)，源码调用 addLinearNorm(net, *_res, post_ln_scale, post_ln_bias, axes)，权重来自 layernorm_after.weight/bias。

8. FFN（前馈网络）与 GeLU 的两种实现路径

ViT 的 FFN 是扩张比为 4 的两层 MLP，中间夹 GeLU：

1. 第一层（扩张到 4D=3072，对 ViT-B）：H = Y'W_1ᵀ + b_1 ∈ R^{N×L×4D}；
2. GeLU 激活；
3. 第二层（投影回 D=768）：F = GeLU(H)·W_2ᵀ + b_2 ∈ R^{N×L×D}。

GeLU 定义与近似

严格定义 GeLU(x) = x·Φ(x)（Φ 为标准正态 CDF）。工程实现普遍使用 tanh 近似：

GeLU(x) ≈ x·(1 + tanh(√(2/π)·(x + 0.044715·x³))) / 2

源码 addGeLU() 按 TRT 版本选择路径：

• TensorRT < 10.0：用 ElementWise + TANH 逐算子拼出上述 tanh 近似式（kPROD 算 x²、x³、乘以 0.044715、加 x、乘以 √(2/π)、TANH、+1、×0.5x 等一长串层），这正是 README 所说"我们采用 TensorRT 同样的近似方式"；
• TensorRT >= 10.0：直接调用 addActivation(input, ActivationType::kGELU_ERF) 使用原生 erf 版本 GeLU，即 README 提到的"Swi-GeLU 自 TensorRT 10.0+ 起原生支持"。

最终残差

Z = Y + F，源码 addElementWise(out1, attn_residual, kSUM)，命名为 .output_residual 并返回。createEngine() 对全部 num_layers 层循环调用 ViTLayer() 串接。

9. 紧凑形状追踪（以 ViT-B/16 为例，动态 L）

原文档给出了逐算子形状流转表，可直接用于排错：

阶段	形状
输入 X	(N, L, 768)
Pre-LN X'	(N, L, 768)
Q/K/V 投影	(N, L, 768)
reshape+transpose 多头 Q_h/K_h/V_h	(N, H, L, d)
注意力 logits S	(N, H, L, L)
Softmax 权重 P	(N, H, L, L)
头输出 O	(N, H, L, d)
合并多头 O_merged	(N, L, 768)
输出投影 A	(N, L, 768)
残差 Y	(N, L, 768)
Post-LN Y'	(N, L, 768)
FFN 扩张 H	(N, L, 3072)
FFN 投影 F	(N, L, 768)
最终残差 Z	(N, L, 768)

对 ViT-L（D=1024, H=16, d=64, FFN=4096）与 ViT-H（D=1280, H=16, d=80, FFN=5120）只需替换相应维度数值。

10. 原生 IAttention 与 Myelin 融合通道

在 TRT_VERSION_GE(10,15,0) 分支，源码使用 TensorRT 原生 addAttention() 接口替代手写 SDPA：

auto* q_scaled = addElementWise(q_s, qk_scale_w, kPROD);
auto* attn = net->addAttention(q_scaled, k_s, v_s, ANO::kSOFTMAX, false);
attn->setDecomposable(false);

注意两点设计：

• 缩放前置：将 1/√d 的缩放先乘到 Q 上，而不是乘在 QKᵀ 结果上——这与 Myelin 融合内核（_gemm_mha_v2）的输入约定兼容；
• setDecomposable(false) 要求 TensorRT 必须保留融合的 MHA 算子，而不是把它拆解回独立的 matmul/softmax 层，这是确保 profile 输出中出现 myl 层名的关键。

当 TRT 版本在 [10.14, 10.15) 区间时走手写 SDPA 分支（同时打印 bug 警告），< 10.14 同样走手写分支。推理输出中 Reformatting CopyNode、patch embedding 与 __myl_* 命名层（如 _gemm_mha_v2_myl3_6）直接印证了 Myelin 融合 pass 被成功触发。

11. Profiler 与输出分配器

11.1 Profiler（profiler.h / profiler.cc）

Profiler 继承 nvinfer1::IProfiler，通过 reportLayerTime() 累积每个层的耗时与调用次数，operator<< 输出按层名对齐的表格：层名、Runtime %、Invocations、累计毫秒，末尾给出总 runtime。-d 推理时用 context->setProfiler(&profiler) 挂载，便于定位热点层（例如在示例输出中可见 patch embedding 占 1.5%）。

11.2 CudaOutputAllocator（cuda_allocator.h / cuda_allocator.cc）

实现 nvinfer1::IOutputAllocator，按 CUDA 运行时/驱动版本自动选择三种输出缓冲分配策略（Create() 中的自动降级逻辑）：

• CUDA runtime & driver >= 12.0：优先用 cuMemCreate/cuMemMap（kCuMem） 虚拟内存管理；
• runtime >= 11.2：退化为 cudaMallocAsync（kCudaMallocAsync）；
• 兜底：cudaMallocManaged（kCudaMallocManaged） 托管内存。

TRT 10+ 使用异步重分配接口 reallocateOutputAsync，TRT < 10 使用同步接口 reallocateOutput，内存按 tensor 名缓存复用，销毁时统一释放。这是 doInference() 中 CudaOutputAllocator::Create(stream) + context.setOutputAllocator(...) 的底层支撑。

12. 进阶：多变体对比实验脚本

仓库附带 exp/run_experiment.py，可对同一变体同时对比 4 条推理后端：

1. PyTorch FP32（HuggingFace 参考 logits）
2. ONNX Runtime CUDA EP（FP32 ONNX）
3. TensorRT 从 ONNX 构建（FP16，动态 batch profile 1/2/4）
4. TensorRT 从本仓库 .wts 构建（FP16，动态 batch profile 1/2/4）

运行方式：

python exp/run_experiment.py                                   # 全部 5 个变体，batch 1 2 4
python exp/run_experiment.py --variants ViT-B/16 ViT-L/16
python exp/run_experiment.py --batches 1 4

脚本会：先用样本图 assets/cats.jpg 做与 vit.cc 一致的预处理（bilinear resize、/255、mean=std=0.5、NCHW）；计算每个后端的 logits 与 torch 参考的余弦相似度；用 CUDA Event 统计 enqueue 延迟、尝试 CUDA Graph 捕获并统计 graph 延迟；输出 markdown 风格汇总表（含 speedup eq (wts/onnx) 加速比列），并把每 (变体, batch) 的 logits 存为 exp/<safe>__bs<B>.npz。该脚本依赖 models/<safe>.wts 与 models/<safe>.engine 已按第 3 节流程先生成，是验证手写实现与 PyTorch/ONNX 一致性的现成工具。

13. 小结

本文围绕 vit/README.md 完整还原了 tensorrtx 中 ViT 手写 TensorRT 实现的部署全流程与逐层原理：从 gen_wts.py 权重导出、CMake 构建、-s/-d 两段式命令，到 ViTLayer() 中 Pre-LN Transformer 的每个算子（LayerNorm、QKV 投影、SDPA、softmax、残差、GeLU、FFN）的 TensorRT API 表达与形状流转。结合 vit/vit.cc 源码可以看到，工程上通过版本宏在"手写 SDPA"与"原生 IAttention + Myelin 融合"两条路径间切换，并借助强类型网络、FP16 权重转换、优化 profile 与自定义 profiler/输出分配器获得可观测、可对比的部署体验。该实现覆盖 ViT-B/L/H 共 5 种官方变体（含 21843 类的 ViT-H/14 in21k 特例），对希望深入理解 Transformer 在 TensorRT 中的低层映射、或需要在不使用 parser 的前提下手写部署 ViT 类模型的开发者，是一份可直接照搬的完整参考。