使用CUDA内核重写模型推理：瓶颈不仅仅是GEMM [P]
摘要
作者描述了构建FlashRT的过程，这是一个以CUDA为核心的推理运行时，通过使用C++/CUDA内核重写模型推理路径，来解决小批量/实时工作负载中超出GEMM的瓶颈，在Jetson Thor和RTX 5090上实现了显著的延迟改进。文章讨论了关于精度的经验（FP8有帮助，FP4好坏参半）以及绕过通用运行时进行实时推理的必要性。
我一直在开发一个面向小批量/实时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)
相似文章
从 RTX 到 Spark：NVIDIA 加速 Gemma 4 赋能本地智能体 AI
NVIDIA 与谷歌合作优化 Gemma 4 模型，以实现在 RTX GPU、DGX Spark 和 Jetson 设备上的本地部署，从而支持高效的端侧智能体 AI，具备推理、编程、多模态能力以及 35 多种语言的支持。
Gemma 4 26B 在单块 RTX 5090 上达到 600 Tok/s
一项基准测试显示，使用 vLLM 搭配 DFlash 投机解码，在单块 RTX 5090 上将 Gemma 4 26B 的推理速度提升至约 578 tokens/s，相比基线实现了 2.56 倍的加速。
一个可定制的编译器，用于为AI模型生成高效的融合GPU内核 [P]
作者介绍了一款用 Python 编写、高度可定制且易于修改的 ML 编译器。该编译器通过多级 IR 流水线将 LLMs 转换为优化的 CUDA 内核，在特定操作上实现了与 PyTorch 相当甚至更优的性能。文章详细阐述了该编译器的优化过程、降级规则以及用于生成高效融合 GPU 内核的 CLI 用法。
Luce Megakernel: 为什么没有人谈论这个？
Lucebox Hub 为本地 LLM 推理提供优化的 CUDA 内核（Megakernel、DFlash、PFlash），在各种模型和 GPU 上相比 llama.cpp 实现了显著的加速（2-10 倍）。
Apple M3 Ultra上实时扩散模型推理的系统优化
本文对Apple M3 Ultra上的实时扩散模型推理进行了系统优化研究，通过CoreML转换和蒸馏模型在512x512分辨率下达到了22.7 FPS，揭示了针对CUDA优化的技术无法直接迁移到Apple统一内存架构。