使用CUDA内核重写模型推理：瓶颈不仅仅是GEMM [P]

我一直在开发一个面向小批量/实时ML工作负载的CUDA优先推理运行时。核心思想很简单：不再将PyTorch/TensorRT/通用图运行时作为主要执行路径，而是直接用C++/CUDA内核重写模型推理路径。这最初源于机器人/VLA相关的工作负载，但问题更为普遍。在小批量推理中，瓶颈往往不仅仅是单个缓慢的GEMM。大量延迟来自于数学计算周围的运行时胶水：* 碎片化的小内核 * 归一化/残差/激活边界 * 量化和反量化开销 * 布局转换 * Python/运行时调度 * 图编译器融合失败 * FP8/FP4区域的精度转换。对于云端LLM服务，批处理可以隐藏很多这些问题。但对于机器人、VLA、世界模型以及其他实时工作负载，批量大小通常为1，无处可藏。每一次启动、同步和格式边界都会直接体现在延迟中。 以下是我实现的一些当前结果： |模型/工作负载|硬件|FlashRT延迟| |:-|:-|:-| |Pi0.5|Jetson Thor|\~44 ms| |Pi0|Jetson Thor|\~46 ms| |GROOT N1.6|Jetson Thor|\~41–45 ms| |Pi0.5|RTX 5090|\~17.6 ms| |GROOT N1.6|RTX 5090|\~12.5–13.1 ms| |Pi0-FAST|RTX 5090|\~2.39 ms/token| |Qwen3.6 27B|RTX 5090|\~129 tok/s with NVFP4| |Motus / Wan风格世界模型|RTX 5090|\~1.3s基线 → 目标\~100ms端到端| Motus/世界模型的情况尤其有趣。基线路径大约需要1.3秒端到端。目标是\~100毫秒端到端，但难点并非仅仅是“使用更快的GEMM”。瓶颈在于VAE、联合注意力、启动碎片化以及实际计算周围的大量胶水代码。 这项工作得出的一条经验：低精度并不自动就是胜利。FP8一直很有用。FP4/NVFP4则好坏参半。它有助于减少内存占用和改善某些大型GEMM区域，但如果FP4区域很小、不连续或周围有转换/缩放开销，那么端到端加速可能微乎其微。例如，在某些VLA/世界模型路径中，除非区域很大且深度融合，否则FP4相对于FP8仅能带来几个百分点的延迟改善。 这改变了我对推理优化的思考方式。对于大批量云端服务，通用运行时和批处理通常已经足够。但对于实时小批量推理，运行时开销本身就成为了工作负载。好奇其他人是否在使用torch.compile、TensorRT、XLA、Triton或自定义CUDA内核时观察到类似行为。在什么情况下你会放弃让通用编译器优化模型，而直接重写推理路径？ 实现代码：[https://github.com/LiangSu8899/FlashRT](https://github.com/LiangSu8899/FlashRT)