ARTICLE · INTELLIGENCE

战地情报 · 详情页

来自尧图项目组的一线实战观察与深度解析

FlashRT 架构深度剖析:从内核库到 Serving 层的 8 大核心组件全解

FlashRT 架构深度剖析:从内核库到 Serving 层的 8 大核心组件全解

【免费下载链接】FlashRT

FlashRT is a high-performance realtime inference engine for small-batch, latency-sensitive AI workloads. The flagship integration is production VLA control for Pi0, Pi0.5, GROOT N1.6, and Pi0-FAST. Also support llm e.g, qwen3.6-27B

项目地址:https://gitcode.com/gh_mirrors/fl/FlashRT
点击查看免费下载

FlashRT 是一款面向小批量、延迟敏感场景的高性能实时推理引擎,也是 Pi0、Pi0.5、GROOT N1.6 等机器人 VLA 模型的生产级控制引擎。这篇文章带你一次讲透FlashRT 的完整架构:从最底层的 CUDA 内核库,到中间的 Python 前端、执行契约,再到顶部的 Serving 场景宿主层,共 8 大核心组件 + 2 个扩展层,全部用大白话讲清楚。

无论你是想评估它能否加速自己的模型,还是想读懂源码再下手,这份组件地图都能帮你快速定位每一块"积木"的位置与职责。


先看懂全局:四层运行时架构 🧱

FlashRT不是编译器或图重写器,它的定位很朴素:一个内核库 + 8 个基础设施组件,由薄薄的每模型流水线拼装起来。整体分层如下:

层级目录职责一句话理解
场景宿主层serving/会话、调度器、协议、机器人循环、OpenAI 兼容 HTTP/SSE"怎么接客"
Python 前端层flash_rt/权重、校准、CUDA Graph 捕获、模型 API"怎么搭车"
执行契约层exec/C ABI:Buffer / Graph / Plan,replay 期机制"轨道标准"
内核库层csrc/CUDA/C++ 内核与注意力/GEMM 基础组件"发动机本体"

这套设计的精髓在于:Serving 层刻意放在模型运行时之上。它只管会话策略、流式协议、回合生命周期、中断/重置行为;核心运行时只负责图与内核机制。两者各司其职,互不纠缠。

完整分层说明见 docs/architecture.md 与 docs/cpp_runtime_design.md。


组件 1:公共 API —— 三行代码的入口 🚪

一切从两个函数开始:

model = flash_rt.load_model(checkpoint, config="pi05", framework="torch") actions = model.predict(images=..., prompt=...)

load_model内部只做三件事:

  1. 探测 GPU 算力(compute capability),解析出arch(如thor、rtx_sm120、rtx_sm89);
  2. 拿着(config, framework, arch)三元组去硬件分派表里查对应的前端;
  3. 导入模块、实例化前端,包在一个轻量的VLAModel适配器里返回。

API 被刻意设计得小而稳定——不暴露注意力后端、校准内部状态、图捕获状态这些复杂概念,新手只面对两行代码。

  • 源码入口:flash_rt/api.py
  • 稳定 API 面文档:docs/stable_api.md

组件 2:硬件分派表 —— 全仓库唯一的"路由字典" 🗺️

这是整个代码库中唯一知道"哪个前端处理哪种硬件"的地方,本质就是一个 Python 字典:

_PIPELINE_MAP = { ("pi05", "torch", "thor"): ("flash_rt.frontends.torch.pi05_thor", "Pi05TorchFrontendThor"), ("pi05", "torch", "rtx_sm120"): ("flash_rt.frontends.torch.pi05", "Pi05TorchFrontendRtx"), # … 每个 (模型, 框架, 架构) 三元组一行 }

设计哲学是"极简分派":没有插件清单、没有 entry-points 扫描、没有 manifest YAML,就是一张显式的 Python 字典。外部插件可以在 import 时直接追加新行,无需 fork 仓库——加一个新模型,往往就是往这张表里多写一行。

  • 源码:flash_rt/hardware/init.py
  • 插件扩展模板:docs/plugin_model_template.md

组件 3:Frontend 前端 —— 每个模型唯一要写的文件 📄

Frontend 是"每模型一个文件"的组件,位于flash_rt/frontends/{torch,jax}/<model>_<arch>.py。它是新模型接入时唯一需要新写的部分,其余全部复用共享基础设施。

一个前端负责五件事,每件事都"外包"给共享组件:

职责交给谁
读 checkpoint → 填充模块属性权重加载器(组件 4)
声明注意力形状 → 拿到后端Attention Backend(组件 5)
计算/加载 FP8 激活缩放校准框架(组件 6)
把 forward 捕获成静态图CUDA Graph 捕获(组件 7)
在forward()里调用内核内核库(组件 8)

前端文件通常 800–1500 行,全部是线性代码,没有复杂的继承层级——只有鸭子类型的set_prompt、infer等方法供公共 API 调用。同一模型的 Torch/JAX、Thor/RTX 前端之间不共享 forward 路径代码,但通过组件 4–8 共享一切基础设施。

  • 目录:flash_rt/frontends/
  • 添加新模型指南:docs/adding_new_model.md

组件 4:声明式权重加载 —— 用规格表描述每一块权重 📦

权重加载靠一份声明式规格表WEIGHT_SPEC完成:它描述 checkpoint 中每个权重张量如何映射到模块属性,包含转置、融合 Norm、FP8 量化等变换与量化缩放。

关键点在于框架无关:具体WeightSource有 safetensors(Torch)和 Orbax dict(JAX)两种实现,同一份WEIGHT_SPEC换一下 source 就能同时工作于两个框架——这也是"PyTorch 和 JAX 前端共享同一份内核二进制、结果等价(cosine ≥ 0.999)"的底层保障。

  • 核心代码:flash_rt/executors/weight_loader.py
  • 规格表详解:docs/extension/weight_spec.md

组件 5:Attention Backend —— 同一份流水线,两种硬件都跑得动 🎯

注意力是模型里最难跨硬件的部分。FlashRT 用协议模式解耦:前端只声明一个AttentionSpec(一组 site,如 SigLIP 视觉、PaliGemma 编码器、Pi0.5 解码器各自的 head 数、head 维度、最大序列长度),框架据此返回对应的后端实例。

于是同一份流水线源码,在Thor上走 CUTLASS FMHA / 分解 cuBLAS 路径,在RTX上走内置 FlashAttention-2 路径——切换的只是背后接线的后端实现,前端的 pipeline 代码一字不改。

  • 协议定义:flash_rt/hardware/backend.py
  • 完整文档:docs/extension/attention_backend.md

组件 6:校准框架 —— FP8 量化不是脚本,是框架 ⚖️

FP8 校准在 FlashRT 里是一个框架而非每模型各写一遍的脚本。它统一管理:

  • 缓存键推导:{ckpt_hash}_Se{N}.json,换 checkpoint 或改序列长度时缓存不会误命中;
  • 缓存位置:固定~/.flash_rt/calibration/,跨模型、跨配置路径一致;
  • 两阶段流程:(a) 测量 pass,记录每个 FP8 GEMM 输入的 amax;(b) refit pass,复合alpha = act_scale × weight_scale并把它们作为宿主标量烤进捕获的图里。

前端只需按协议写一个_calibrate(self, sample_obs)方法,缓存查询、保存、文件命名等周边逻辑全部共享。

  • 框架代码:flash_rt/core/quant/calibrator.py
  • 机制详解:docs/calibration.md · docs/extension/calibration.md

组件 7:CUDA Graph 捕获 —— 只传指针的静态图 📸

组件 7 是"快"的来源。它对torch.cuda.CUDAGraph做了两个 FlashRT 专属增强:

  1. 仅指针 forward:被捕获的 forward 只接受int设备指针,不接受 tensor/array。捕获时强制校验,杜绝 replay 期间的意外内存分配;
  2. tactic 稳定 autotune:可选地重捕获最多 N 次并保留最快调度,规避 Thor 上 Myelin / cuBLASLt tactic 的非确定性(每次捕获约 2ms 方差)。

捕获后的图不落盘:没有.engine文件、没有序列化,图在进程内存里活到进程结束;重启即重捕获(热缓存下约 50–500 ms)。这也是它能"扛住 CUDA 驱动升级、换 GPU、换 prompt 都不需重建"的原因。

  • 源码:flash_rt/core/cuda_graph.py

组件 8:内核库 —— 手写的发动机 🏁

最底层是三个动态库,全部由import flash_rt加载,源码在csrc/下:

模块构建条件内容
flash_rt_kernels.so始终构建约 98 个 pybind 入口:norm / 激活 / 融合 / 量化 / GEMM / 注意力 / RoPE / 工具
flash_rt_fp4.so仅 SM100+约 23 个入口:NVFP4 权重预处理 + SM120 block-scaled GEMM
flash_rt_fa2.so仅 RTX2 个入口:FA2 fp16 / bf16 前向

分工清晰:内存受限的算子(norm、激活、残差+norm+量化融合、qkv 拆分+RoPE、patch embed 等)全部手写;算力受限的算子(GEMM、注意力)委托给 cuBLASLt FP8 GEMM、CUTLASS、内置 FA2 与 Thor 的 CUTLASS FMHA。内核库 API稳定——已发布内核的签名跨版本保持不变,新内核可自由新增。

  • 全量内核清单:docs/kernel_catalog.md
  • 融合模式与命名规范:docs/kernel_fusion.md
  • 源码目录:csrc/

串起来看:一次 predict 调用发生了什么 🔗

predict(images, prompt) │ 仅首次调用: │ set_prompt(prompt) │ ├─ 前端._calibrate() ← 校准框架(6):命中缓存或跑 1 次测量 │ ├─ 构建 AttentionSpec → 后端 ← 注意力后端(5) │ ├─ WeightLoader.run() ← 权重加载(4) │ └─ cuda_graph.capture() ← 图捕获(7) │ │ 之后每次调用: └─ graph.replay() + 三次指针拷贝 ← 执行内核库(8)

首次调用慢(校准 + 捕获),之后每次调用只是一次graph.replay()。不编译、不导出、prompt 变化不重捕获——因为所有会变化的东西(prompt token、观测像素、动作噪声种子)都是图的输入,而非图拓扑。

组合流程详见 docs/architecture.md §4。


8 大组件之上:exec 执行契约与 Serving 层 🎬

八个组件之上还有两个扩展层,把"一次性的快"变成"可运营的服务"。

exec/:C ABI 执行契约

exec/把内核+图 replay 这套机制固化为一个极简 C ABI:只有 3 个对象 + 1 把钥匙——

  • Buffer:命名的设备内存,是唯一的"状态"(KV 缓存、缩放因子、噪声、子目标 embedding);图之间靠共享同一个 Buffer实现零拷贝交接;
  • Graph:ShapeKey → 图变体表(精确键 + LRU),支持外部捕获后 adopt,replay 就是cudaGraphLaunch本身,零额外开销;
  • Plan:跨流的多图 DAG,只有数据依赖,不含优先级/抢占——那是上层策略的事。

契约只定机制,不定策略,且对csrc/内核层零依赖。Python 前端负责冷启动(加载、校准、autotune、捕获),宿主进程只驱动热路径 replay。

  • 契约规范:docs/exec_contract.md
  • 源码:exec/

serving/:场景宿主层

libflashrt_exec之上是一组场景宿主,每个宿主为具体场景补齐"策略":

宿主场景
robot_recap/RL rollout 宿主:回合状态机、键盘启停、episode 录制、跨 chunk 中断重置(无需重新捕获)
robot_handoff/分层 planner→actor 的 buffer 交接、多速率、运行中子任务打断
robot_host/C++ 实时 VLA 部署模式(Plan、并发流、buffer 覆写中断),面向上机 / ROS2
qwen36_agent/Qwen3.6 长会话 Agent 服务:精确 token 前缀复用、OpenAI 兼容工具调用、真 SSE 流式
llm_agent/Rust 参考宿主:per-token replay over C ABI 的高并发控制面实验

其标志性设计是graph-replay 原生的执行状态胶囊(快照/恢复/fork/时间旅行),它填补了 vLLM/SGLang 的 block/radix KV 缓存不擅长的场景——混合循环状态、全图捕获、LLM+VLA+机器人同机制运行。

  • 宿主总览:serving/README.md
  • 设计文档:docs/serving_design.md · docs/serving_production.md

它和 TensorRT / vLLM / SGLang 有什么不同?🆚

一句话概括:FlashRT 是按"工作负载形状"(小批量实时)设计的,而不是按"模型类别"设计的。

  • 运行时直接执行前端写好的内核序列,没有图编译器 / tactic 搜索——前端作者本人就是优化器;
  • KV 缓存是静态布局的每层设备 buffer,支持图内小 batch(CFG、多策略、多帧),但不做跨请求连续 batching;
  • 它可以选择部署成"一个捕获图一个进程",再由 Serving 层扇出到多个 FlashRT worker——这是显式的扩展路径。

详细对比见 docs/inference_engine_differences.md。


总结:一张表记住 FlashRT 架构 ✅

#组件位置一句话
1公共 APIflash_rt/api.py两行入口,稳定且极小
2硬件分派表flash_rt/hardware/init.py全仓库唯一的路由字典
3Frontendflash_rt/frontends/每模型唯一要写的文件
4权重加载flash_rt/executors/weight_loader.py声明式 WEIGHT_SPEC,框架无关
5Attention Backendflash_rt/hardware/backend.py协议化,Thor/RTX 一键换线
6校准框架flash_rt/core/quant/calibrator.pyFP8 校准是框架不是脚本
7CUDA Graph 捕获flash_rt/core/cuda_graph.py仅指针 forward + 稳定 autotune
8内核库csrc/ →flash_rt/*.so手写内存受限算子 + 委托算力受限算子
+执行契约exec/C ABI:Buffer/Graph/Plan
+Serving 宿主serving/机器人 rollout、LLM Agent 场景策略

想动手上手?建议路线:先读 docs/INSTALL.md 装好环境,跑通 examples/quickstart.py 的第一次predict,再顺着本文的 8 个组件逐个对照源码。FlashRT 的基础设施侧(不含内核与前端)总共只有约 3000 行——读懂它,比想象中轻松。🚀

【免费下载链接】FlashRT

FlashRT is a high-performance realtime inference engine for small-batch, latency-sensitive AI workloads. The flagship integration is production VLA control for Pi0, Pi0.5, GROOT N1.6, and Pi0-FAST. Also support llm e.g, qwen3.6-27B

项目地址:https://gitcode.com/gh_mirrors/fl/FlashRT
点击查看免费下载

相关推荐

创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考

RELATED READING

延伸阅读

更多一线实战笔记与深度复盘,助您持续精进