导读

物理智能不仅需要更强的模型，也需要更快的系统。

除了关注基础模型的Scaling外，当模型真正进入机器人、工业设备与真实环境时，推理系统往往成为决定体验与性能的关键环节。一次动作决策背后，可能需要经历视觉理解、语义推理、状态建模与动作生成等多个阶段，而每一次毫秒级延迟都会被放大为真实世界中的响应迟滞。

因此，我们认为，Physical AI 的竞争不仅是模型能力的竞争，也将是 AI Infrastructure 的竞争。

在本工作中，我们以行业较为通用的VLA模型为例进行了系统性的推理优化。从框架重构到 CUDA Graph，从 Triton Kernel 到 Kernel Fusion，我们构建了一套跨 Framework、Runtime 与 Kernel 的协同优化方案，将推理延迟从 130ms 降低至 32ms，实现约 4 倍性能提升。

这不仅是一项针对单一模型的优化实践，更是我们对于下一代 Physical AI Inference Stack 的一次探索。

背景和优化成果

机器人 VLA（Vision-Language-Action）模型将视觉感知、语言理解和动作生成融合为单一网络。推理延迟会直接影响机器人控制的实时性——过长的推理时间会导致动作指令滞后、执行卡顿，在高动态任务中甚至导致任务失败。

原始 StarVLA训练框架的推理流程存在三个层面的效率问题。其一是框架调度开销：HuggingFace Transformers 的通用接口在推理时引入不必要的条件判断和 config 解析。其二是 Kernel launch 次数过多：单次推理约 10,000 次 CUDA kernel launch，调度开销占总时间 40% 以上。其三是通用算子未针对推理形状优化：PyTorch 原生 RMSNorm、MLP 等算子未利用 VLA 推理中 batch=1、序列长度 100-400 等特征。

考虑到实际应用中，模型结构在短期内相对固定，且追求极致的推理延迟是核心目标，我们没有采用 vLLM 或 TensorRT-LLM 等通用推理框架，而是选择自行搭建轻量级的推理框架，配合 Triton 进行算子开发。这一决策让我们能够精细控制推理过程的每一个环节，也为后续的 CUDA Graph 和算子融合等深度优化提供了基础。基于这一框架，我们在 RTX 4090 上完成了从纯语言到视觉语言再到完整 VLA 的三阶段推理适配，覆盖了全部 10 种算子的高效实现。最终在单 Camera 224×224 输入，精度结果差别不超过1%前提下，纯模型平均推理耗时从 130ms 减少到 32ms，实现了 4倍的加速。

性能优化方案

整体方案

StarVLA 模型由三个模块串联组成：Vision Encoder（视觉编码器，24 层 ViT，约 300MB 参数）、LLM Decoder（语言模型，36 层，含投影层 Project Layers，共约 2.75B 参数）和 Action Head（动作生成器，4 步 × 36 层 DiT，约 1.1B 参数）。其中 LLM 采用 GQA 注意力机制（32/8 heads）和 SwiGLU MLP（中间维度 9728），且为纯 prefill 模式（无自回归 decode 阶段）。Action Head 使用 Flow-Matching Euler denoising，从随机噪声开始，经过 4 步去噪逐步生成有意义的动作轨迹。

推理总层数约 240 层，但涉及的算子类型仅 10 种：RMSNorm、LayerNorm、RoPE、GELU、SiLU、MatMul（QKV/MLP）、SDPA、Sin/Cos、Embedding、Element-wise。这一特性使得每个算子都值得手工深度优化。

基于以上特点，我们的优化方案为：自研推理框架，使用cuda graph来减少系统层面的开销；同时基于triton做算子开发和算子融合，在算子层面进一步挖掘硬件潜力，提升推理速度，整体优化后的模型结构和推理过程见上图。

自研推理框架

问题

原始推理基于 HuggingFace Transformers + PyTorch nn.Module 封装。训练框架为了灵活性包含了大量在推理时不需要的逻辑。每层 nn.Module 的 Python 级 forward 调用引入调度开销，在 36-144 层的累积下变得显著。from_pretrained 权重加载接口需要解析模型结构和配置，路径不透明。

方案

我们从权重加载到算子调度全链路自建，不依赖 HuggingFace Transformers。具体包括三方面工作：第一，直接权重加载——从 PyTorch checkpoint 或 safetensors 直接读取权重张量，跳过 config 解析环节。第二，模块级代码组织——每个模型组件独立为单一文件（vision_encoder.py / language_model.py / dit_model.py / action_head.py），通过统一的 dispatch 层控制算子调度。第三，三阶段递进——从纯语言模型到视觉语言模型，再到完整 VLA 模型，下层被上层复用，逐步验证。

收益

自研推理框架消除了框架调度开销，使控制流完全透明。贡献了约 15% 的加速（130ms → 115ms）。这一优化的通用性最高——任何使用 HuggingFace Transformers 进行生产部署的模型，都可以通过裁剪框架层获得相近比例的收益，且改造成本与模型复杂度无关。

CUDA Graph

问题

GPU 推理中，CPU 需通过 driver 为每个 GPU kernel 发起一次 launch 调用。单次 launch 约 4μs，而单次推理的 10,000 多次 launch 中约 5,900 次为 element-wise、reshape、copy 等小 kernel，其 launch 开销可与计算本身相当。这意味着大量时间花在了 CPU 向 GPU 下达指令上，而非真正的计算。

方案

CUDA Graph 将一组 kernel launch 录制为一张图（graph），后续推理通过一次 replay 完成所有 kernel 的调度，从而消除单次 launch 开销。实现采用标准的 warmup → capture → replay 模式：先做一次完整推理预热，然后在 torch.cuda.graph() 上下文中重新执行以录制 kernel 序列，后续推理直接 replay 已录制的序列。我们对三个核心组件分别应用了 CUDA Graph：Vision Graph 覆盖 PatchEmbed 到 24 个 ViT blocks 再到 Merger 的全链路，LLM Graph 覆盖 36 层 Transformer 到 final RMSNorm，Action Graph 覆盖 4 步 denoising × 36 层 DiT。实现中有两个关键细节需要处理。其一是显存预分配：graph capture 前预分配输入/输出 buffer，replay 时通过 copy_() 更新输入数据，避免 graph 内动态显存分配。其二是序列长度自适应：LLM Graph 在输入序列长度变化时自动回退 eager 模式并重新 capture，通过缓存序列长度做匹配检测。

收益

CUDA Graph 将 CPU 侧的 kernel launch 调用量从 10,605 次降至 3 次（每个 graph 一次 replay()），彻底消除了 CPU launch 瓶颈。nsys 观测到的 GPU 侧 kernel 执行次数从 10,605 降至 5,139（减少 52%），其中小 kernel 因为图优化消除或融合上的优化个数从 5,913 降至 499（减少 92%）。剩余的大计算 kernel（matmul、attention 等）被录制为一张图，由 GPU 自行调度执行。整体贡献了约 60% 的加速（115ms → 57ms），是四项优化中收益最大的单项。CUDA Graph 对任何推理时计算图形状固定的模型都有效。形状固定程度越高，收益越大——本模型恰好是纯 prefill 场景（无自回归 decode，序列长度在推理过程中不变），因此收益尤为突出。

算子开发

问题

PyTorch 的通用算子实现针对广泛输入形状做权衡，未利用 VLA 推理的特有模式（batch=1、序列长度 100-400、hidden 大小固定为 1024 或 2560）。部分算子调用次数极多——例如 RMSNorm 在 LLM pre-norm 和 QK norm 中被调用 144 次，微小优化的累积收益非常显著。

方案

我们使用 Triton 语言针对 VLA 推理形状开发自定义 kernel。Triton 相比 CUDA 开发效率更高，同时保持接近手工优化的性能。部分 kernel 使用了 @triton.autotune 机制，在编译时自动搜索最优的 block size 和 warp 数配置。以下以 RMSNorm 和 RoPE 为例说明算子开发的核心思路。

RMSNorm（均方根归一化）在 LLM 中被调用 72 次（pre-norm）加上 72 次（QK norm），共 144 次。PyTorch 的实现是一条分解的算子链：x.float() → x*x → .mean(-1) → torch.rsqrt() → * x → * w.float()，每个步骤独立 kernel，中间结果经 HBM 读写约 3-4 轮。我们实现了 single-pass 版本，一次加载整行到 SRAM，在寄存器中完成 sum_sq、rstd 和 final scaling 的全部计算，再写回 HBM——整个 kernel 只读一次、写一次 HBM。在 LLM prefill（M=182, N=2560）场景下加速比为 1.81×。

RoPE（旋转位置编码）在 LLM 的 36 层中对 Q/K 张量做分半旋转变换，每个 token 都需要计算。PyTorch 的实现需要创建 q_rot / k_rot 新 tensor，经过 split、逐元素运算、cat 和 dtype 转换等多个 kernel 步骤。我们实现了 in-place 的 Triton kernel，直接从原始 Q/K 张量加载前后两半段到寄存器，完成旋转变换后写回同一位置。在 B=1, S=182, H=32, D=128 场景下加速比为 1.52×。

以上两个算子是收益为正的典型代表。总体而言，大多数算子（RMSNorm、RoPE、GELU MLP、SwiGLU MLP）在 Triton 下都有 1.3-7× 不等的加速，但并非所有算子都适合 Triton——例如 SDPA（attention 计算）在 cuDNN 下有深度优化，Triton 版本反而更慢，因此attention算子在实际推理中仍使用 cuDNN 后端。

收益

算子开发贡献了约 10% 的加速（57ms → 47ms）。注意此处的加速是在 CUDA Graph 已将 kernel launch 开销消除后的 kernel 计算层面收益，因此 10% 是算子计算效率上的真实提升。值得强调的是，这些算子是所有 Transformer 模型的通用算子——LLaMA、Qwen、Gemma、GPT 等主流模型都使用相同的归一化和位置编码逻辑。因此，本文开发的Triton kernel 可以直接迁移到其他模型的推理优化中，无需重新开发。

算子融合

问题

RTX 4090 的 HBM 带宽为 1 TB/s，而片上 SRAM 带宽约 20 TB/s，二者有数量级差距。PyTorch 中连续算子链的每个中间结果都写回 HBM，带宽瓶颈突出。以 GELU MLP 为例：先计算 up = x @ W_up 并将结果写入 HBM，再计算 gelu = GELU(up) 读写 HBM，最后计算 out = gelu @ W_down 从 HBM 读取 gelu。共 4 次 HBM 访问（读 x、写 up、读 up、写 out），其中写 up 和读 up 可通过 SRAM 避免。

方案

在 Triton kernel 的分块循环中保持中间结果在 SRAM。具体流程为：从 HBM 加载 x 的一个分块到 SRAM，计算 up = x @ W_up（保持在 SRAM 寄存器），计算 gelu = GELU(up)（保持在 SRAM），累加 out += gelu @ W_down（分块部分和），最终将 out 写回 HBM。通过这种方式，HBM 访问从 4 次降至 2 次。

收益

基于算子独立基准测试结果（PyTorch eager 模式 vs Triton fused，不包含 CUDA Graph 环境），Fused GELU MLP 的加速效果与矩阵规模正相关：Merger 场景（M=49, C=4096）实现 7.33× 加速，VisionMLP（M=512, C=1024）实现 3.82× 加速。类似原理应用于 Fused SwiGLU MLP（LLM 场景，1.27× 加速）和 Fused QKV（将三个独立 matmul 合并为一次）。整体上，算子融合贡献了约 15% 的加速（47ms → 32ms）。

后续工作

在后续各种优化工作中，预计量化是收益最明显的工作。FP8 和 INT4 两种方案在不同模块上有不同的适用性：LLM Decoder 以 matmul 为主，量化收益最大且行业实践最成熟；Action Head 的去噪迭代过程对精度敏感，需要谨慎评估量化误差的累积效应；Vision Encoder 占比小，量化优先级较低。综合来看，量化技术有望将整体推理延迟进一步降低，从 32ms 降至 20ms 以内。

引用

如果在研究或工程中使用本工作，请引用如下：

@article{rimbot2026aivlaqwen3,
  author = {日冕研究团队 (RimBot Research Team)},
  title = {AI Infra如何优化VLA - 模型推理性能提升技术拆解},
  year = {2026},
  howpublished = {Technical Report},
  note = {An engineering report on optimizing Qwen3-VLA inference with NVIDIA RTX 4090 through custom inference framework, CUDA Graph, Triton kernels, and operator fusion, achieving 4x acceleration.}
}