关于大模型推理的常见讨论,往往聚焦于“更高的吞吐量”、“更大的批处理规模”以及“更复杂的服务化系统”。
然而,在机器人控制、实时交互和边缘部署等场景中,我们面临的是一套截然不同的挑战:批量极小、延迟极度敏感、输入形态频繁变化,控制闭环无法容忍数十毫秒的额外调度开销。
TensorRT 擅长将模型编译为冻结的引擎,vLLM 与 SGLang 则专注于服务高并发的 LLM;但当任务变为“相机图像输入,机械臂动作即刻输出”时,推理系统的底层逻辑已然改变。
- FlashRT 是一款专为小批量、低延迟敏感的 AI 工作负载而设计的高性能实时推理引擎。其核心集成应用是为 Pi0、Pi0.5、GROOT N1.6 以及 Pi0-FAST 提供生产级的 VLA(视觉-语言-动作)控制功能,同时也支持通义千问 3.6-270 亿参数等大语言模型。
- 代码仓库:https://github.com/LiangSu8899/FlashRT
- 全文约 9000 字,预计阅读 40 分钟,播客版约 30 分钟。
FlashRT 正是沿着这条差异化路径生长出的工程实践:它不追求将所有工作交给通用编译器,而是将模型的前向推理过程拆解为稳定的 CUDA 内核、指针协议、量化权重和静态 CUDA Graph,并在 PyTorch/JAX 前端之下,构建一条面向实时控制的底层高速通道。
下表对比了 Pi0.5、Pi0、GROOT N1.6 以及优化版 Pi0-FAST 四款 AI 模型,在边缘端 Jetson AGX Thor 与桌面端 RTX 5090 上的推理延迟与吞吐量表现。在 Jetson 平台上,常规模型延迟约为 41-46ms,吞吐量在 22-24Hz 之间;而 Pi0-FAST 经过流程优化后,延迟降至 8.1ms/token,吞吐量达到 123 tok/s。在 RTX 5090 平台上,各模型性能大幅跃升:Pi0 延迟低至 17.58-24.48ms,GROOT 吞吐量最高达 80Hz,Pi0-FAST 更是降至 2.39ms/token,吞吐量高达 418 tok/s。这既直观体现了高端桌面 GPU 的算力优势,也验证了推理优化对端侧部署的关键价值。同时,多视图、多 token 生成场景下的性能差异,也反映出任务复杂度与算力需求之间的强关联,为机器人、嵌入式等场景的模型部署提供了清晰参考。
另一项对比涵盖了 Original openpi、PyTorch naive、torch.compile、Triton-based VLA、NVIDIA VLA-Perf、Isaac GR00T 以及 FlashRT 等多种 AI 推理方案,在 Jetson Thor、RTX 4090/5090 硬件上,针对 Pi0、Pi0.5、GROOT N1.6 模型的推理延迟表现。未优化的 Original openpi 在 Jetson Thor 上延迟高达 714ms,PyTorch naive 在 RTX 4090 上约为 200ms。经过 torch.compile、Triton、TensorRT 等优化后,延迟显著降低。而本次的 FlashRT 方案表现最优:在 RTX 5090 上,Pi0、Pi0.5、GROOT N1.6 的延迟分别低至 21.16ms、17.58ms 和 13.08ms;在 Jetson Thor 上也实现了 46ms、39.78-51.51ms 和 45ms 的低延迟。这既体现了推理优化对边缘与桌面端部署的关键价值,也凸显了 FlashRT 在端侧与桌面 GPU 上的性能优势,为机器人等场景的实时部署提供了高效方案。
unsetunset本文目录unsetunset
- 一、快速上手使用:先跑起来,再理解它为什么快
- 二、FlashRT 到底解决了什么问题
- 2.1 它不是通用 LLM Serving,而是实时控制推理栈
- 2.2 核心性能数据说明了它的目标边界
- 三、总体架构:把模型前向拆成四层契约
- 3.1 公共 API 层:用户只看到 load_model 和 predict
- 3.2 Frontend 层:框架相关的输入、权重和校准
- 3.3 Pipeline 层:框架无关的指针化 forward
- 3.4 Hardware 层:attention 和 GEMM 按硬件分发
- 四、为什么 CUDA Graph 是 FlashRT 的性能中轴
- 4.1 小 batch 场景下,Python 调度本身就是瓶颈
- 4.2 图捕获要求:指针稳定、无动态分配、shape 稳定
- 五、权重加载:把 state_dict 变成声明式规格表
- 5.1 WEIGHT_SPEC 是 FlashRT 的轻量编译描述
- 5.2 Pi0 与 Pi0.5 的差异被编码进 spec,而不是散落在推理逻辑里
- 六、AttentionBackend:让模型不关心 FlashAttention 还是 Thor FMHA
- 6.1 同一 pipeline,背后可以是不同 attention 实现
- 6.2 GROOT 的 attention site 体现了复杂模型的拆分方式
- 七、FP8 与 NVFP4:量化不是压缩,而是控制误差传播
- 7.1 FP8 路径:权重量化与激活校准进入图捕获
- 7.2 NVFP4 路径:Pi0.5 encoder FFN 的更激进尝试
- 八、Qwen3.6 路径:FlashRT 正在把 VLA 经验扩展到 LLM
- 8.1 Qwen3.6 的阶段式迁移计划
- 8.2 KV cache 的指针布局同样服务于图稳定性
- 九、训练反哺:把推理 kernel 直接拿到训练 forward 上
- 9.1 设计原则:推理 kernel = 训练 kernel
- 9.2 数据:训练步速与 VRAM 双降
- 9.3 闭环意义:解决具身智能”频繁微调”的痛点
- 十、与 TensorRT、vLLM、SGLang 的根本差异
- 10.1 FlashRT 拒绝“先导出、再编译、再部署”的重型路径
- 10.2 它也不是 vLLM 式高并发调度器
- 十一、工程取舍:FlashRT 为什么值得研究
- 11.1 优点:把端到端延迟当成第一目标
- 11.2 代价:模型适配需要工程翻译
- 11.3 最大启发:实时 AI 需要另一种推理系统哲学
- 结语:把模型部署还给控制闭环
unsetunset一、快速上手使用:先跑起来,再理解它为什么快unsetunset
根据 FlashRT 的 README 文档,实现最简调用的方式仅需三行 API:加载模型、传入图像与语言指令、获取动作输出。但需注意,该项目依赖已编译好的 .so 内核文件。如果你没有使用预先配置好的构建环境,必须先执行 CMake 构建步骤。
以下是最简的 Python 调用示例:
import flash_rt
model = flash_rt.load_model(
checkpoint="/path/to/pi05_checkpoint",
config="pi05", # 也可以是 "pi0"、"groot"、"pi0fast"
framework="torch", # 也可以是 "jax"
)
actions = model.predict(
images=[base_img, wrist_img],
prompt="pick up the red block",
)
如果你打算使用 Docker,README 中推荐的预构建镜像命令如下:
docker pull ghcr.io/liangsu8899/flashrt:latest
docker run --rm --gpus all -it ghcr.io/liangsu8899/flashrt:latest
不过,仓库中的 docker/README.md 也指出,由于包名从 flash_vla 重构为 flash_rt,公共镜像可能尚未更新推送。因此,更稳妥的做法是在本地自行构建:
git clone https://github.com/LiangSu8899/FlashRT.git
cd FlashRT
docker build -t flashrt:dev -f docker/Dockerfile .
docker run --rm --gpus all -it flashrt:dev
在原生 Linux 环境下,必要的构建命令可概括如下:
python3.12 -m venv .venv
source .venv/bin/activate
pip install torch --index-url https://download.pytorch.org/whl/cu128
pip install pybind11 cmake "numpy>=1.24" safetensors
pip install "transformers<4.56" pandas pillow pyarrow
git clone https://github.com/LiangSu8899/FlashRT.git
cd FlashRT
git clone --depth 1 --branch v4.4.2 https://github.com/NVIDIA/cutlass.git third_party/cutlass
pip install -e ".[torch]"
cmake -B build -S .
cmake --build build -j$(nproc)
验证构建是否成功,可运行以下命令:
python examples/quickstart.py
--checkpoint /path/to/pi05_checkpoint
--benchmark 20
更完整的依赖信息、GPU 架构覆盖范围以及 Docker 使用细节,可参考项目中的 README.md、docs/INSTALL.md 和 docker/README.md。
二、FlashRT 究竟解决了什么问题
2.1 它不是通用的大模型服务引擎,而是专为实时控制设计的推理栈
FlashRT 在 README 开篇就明确了自己的定位:它是一个面向“小批量、低延迟、实时 AI 工作负载”的高性能推理引擎。这个定位至关重要。
传统的推理优化通常默认一个场景:服务端不断接收请求,通过批处理、KV 缓存管理、调度队列和多租户策略来提升整体吞吐量。然而,机器人 VLA 控制更像一个“工业控制回路”:每一次相机观测和文本指令都必须尽快转化为动作序列。这里的瓶颈远不止矩阵乘法,而是涵盖了 Python 调度、框架算子分发、动态形状、内存分配、量化校准、注意力机制实现选择等一整条链路。
FlashRT 的解决方案包括:
- 使用手写 CUDA 内核覆盖归一化、激活函数、旋转位置编码(RoPE)、融合算子以及量化等内存敏感的路径;
- 借助 cuBLASLt、CUTLASS 和 FlashAttention-2 处理计算密集型的 GEMM 和注意力机制;
- 采用 FP8 E4M3 和 NVFP4 来降低权重和激活的内存开销;
- 利用 CUDA Graph 将整个前向传播捕获为静态图,之后每次推理仅需重放该图;
- 通过硬件分发层,将 Jetson Thor、RTX 5090、RTX 4090 等不同目标的差异隐藏在后端。
这不仅仅是“对 PyTorch 模型进行加速”,而是将推理系统重新设计为“前端负责 I/O 和权重,底层负责稳定指针与静态图”的架构。
2.2 核心性能数据揭示了它的目标边界
README 中展示的代表性性能结果如下:
- Pi0.5 在 Jetson AGX Thor 上约 44 毫秒,约 23 Hz;
- Pi0.5 在 RTX 5090 上约 17.58 毫秒,约 57 Hz;
- GROOT N1.6 在 RTX 5090 上约 13.08 毫秒;
- Pi0-FAST 在 RTX 5090 上约 2.39 毫秒/令牌;
- Qwen3.6-27B NVFP4 在 RTX 5090 上典型约 100 令牌/秒,峰值约 129 令牌/秒,并支持 256K 上下文路径。
这些数字的意义不仅在于“速度快”,更在于它们指向了一个工程目标:在小批量场景下,不依赖复杂的批处理来摊薄开销,而是直接消除这些开销本身。
三、总体架构:将模型前向分解为四层契约
3.1 公共 API 层:用户只需接触 load_model 和 predict
FlashRT 的调用入口统一为 flash_rt.load_model(...)。该函数将 checkpoint 路径、框架类型、模型配置、硬件选择、CUDA Graph 策略以及 FP4 开关等众多参数,全部收敛到一个简洁的 API 中。
来源:flash_rt/api.py
def load_model(checkpoint, framework="torch", num_views=2, autotune=3,
recalibrate=False, weight_cache=True, config="pi05", device=None,
decode_cuda_graph=False, decode_graph_steps=80,
max_decode_steps=256,
hardware="auto",
embodiment_tag=None,
action_horizon=None,
use_fp4=False,
fp4_layers=None,
use_awq=None,
awq_alpha=0.5,
use_p1_split_gu=None):
"""Load a FlashRT model.
Args:
checkpoint: path to checkpoint directory.
- torch: safetensors directory
- jax: Orbax checkpoint directory
framework: "torch" or "jax"
num_views: number of camera views (default 2)
autotune: CUDA Graph autotune intensity.
config: model config name: "pi05", "pi0", "groot", "pi0fast"
hardware: GPU backend selection. ``"auto"`` (default) detects the
current CUDA device via compute capability and picks the
best-matching backend.
这段代码背后贯穿着“稳定外壳,内部可替换”的设计哲学。用户无需关心 Pi0.5 在 Thor 上走的是 CUTLASS FMHA,还是在 RTX 上走 FlashAttention-2;也无需知道权重格式是 safetensors 还是 Orbax。用户只需表达“我要加载哪个模型”,系统便会自动选择最优路径。
3.2 Frontend 层:框架相关的输入、权重与校准
FlashRT 全面支持 PyTorch safetensors 与 JAX Orbax 两种格式。框架差异主要集中在前端:包括读取 checkpoint、准备 tokenizer、图像预处理以及触发校准等环节。
在 VLAModel.predict 方法中,可以看到它先将用户传入的图像列表整理成观测字典,随后在必要时触发基于真实数据的校准,最后调用底层 pipeline 的 infer 函数:
来源:flash_rt/api.py
if self._needs_real_data_calibration:
self._pipe.calibrate_with_real_data([obs])
self._needs_real_data_calibration = False
result = self._pipe.infer(obs)
return result['actions']
这里的“懒校准”策略极具工程智慧:API 层面保持三行代码的简洁,但首次真实输入会完成 FP8 激活尺度等准备工作;之后便进入稳定的图重放路径。
3.3 Pipeline 层:框架无关的指针化前向传播
项目模板文档清晰地阐述了 FlashRT 的心智模型:传统模型的
forward()会被翻译成一系列fvk.*内核调用;这串 Python 代码仅在 CUDA Graph 捕获阶段执行一次,真正推理时不再逐行执行,而是直接重放整个图。
来源:flash_rt/frontends/torch/_template/README.md
| Your model code | FlashRT equivalent |
| ----------------------------------------- | ------------------------------------------------------------- |
| `out = model.encoder(images, text)` | `encoder_forward(ctx, fvk, bufs, weights, dims)` |
| `F.rms_norm(x, weight)` | `fvk.rms_norm_fp16(x_ptr, w_ptr, out_ptr, S, D, eps, stream)`|
| `F.scaled_dot_product_attention(q, k, v)` | `attn.run("encoder", layer_idx, q_seq=S, stream=stream)` |
| Subsequent `model(...)` calls | `infer(obs)` → `cudaGraphLaunch(self._enc_ae_graph)` |
**Key insight**: FlashRT replaces your model's `forward()` with a
fully captured CUDA Graph.
这正是 FlashRT 的编译器特质所在:它没有生成传统意义上的 IR,也没有导出 ONNX,而是将模型前向过程“手工降级”为稳定的算子图和指针协议。 从某种意义上说,它就像一个专为少数模型族定制的静态编译器后端。
3.4 Hardware 层:按硬件分发的 attention 与 GEMM
FlashRT 将硬件差异封装在
hardware/目录中。README 中的架构说明显示,Thor、RTX 5090、RTX 4090 会被自动路由到不同的后端。尤其在 attention 方面:Thor 可走 CUTLASS FMHA 或 cuBLAS 分解路径,而 RTX 则走 vendored FlashAttention-2。
这层抽象的核心在于AttentionBackend协议。
来源:flash_rt/hardware/backend.py
class AttentionBackend(Protocol):
“””Hardware provider for attention in a model pipeline.”””
def sites(self) -> tuple[str, …]:
“””Return the tuple of site names this backend was configured for.”””
def get_slot_ptrs(self, site: str, layer_idx: int) -> dict[str, int]:
“””Return raw int device pointers for all slots at (site, layer).
Pointers returned are valid for the lifetime of the backend and
stable across CUDA Graph capture + replay.
“””
…
def run(
self,
site: str,
layer_idx: int,
q_seq: int,
*,
kv_seq: Optional[int] = None,
stream: int = 0,
state_nk: Optional[int] = None,
) -> int:
“””Execute attention for one (site, layer) and return the output ptr.”””
这一协议解决了一个根本性的问题:模型 pipeline 完全不需要关心 attention 的具体实现方式。它所需要的仅仅是一组稳定的 Q/K/V 槽位指针,以及一个统一的 run() 调用接口。后端可以自行管理 buffer,也可以由 pipeline 持有 buffer;只要这些指针在 CUDA Graph 的捕获与重放过程中保持稳定,就可以顺利地被纳入静态图。
四、为何 CUDA Graph 是 FlashRT 的性能核心
4.1 小批量场景下,Python 调度本身就是性能瓶颈
当 batch size 足够大时,CPU 端的调度开销可以被 GPU 计算所掩盖;然而,VLA 控制场景通常为 batch=1 或极小的 batch 大小。在这种情况下,每个 PyTorch 操作的调度、内存分配器的行为、stream 同步以及 Python 循环都会成为可感知的延迟来源。
FlashRT 通过 ctypes 直接调用 CUDA Runtime API,实现了一个与框架无关的 CUDA Graph 封装:
来源:flash_rt/core/cuda_graph.py
class CUDAGraph:
“””Framework-agnostic CUDA Graph using raw CUDA Runtime API.”””
def init(self):
self._graph = ctypes.c_void_p()
self._graph_exec = ctypes.c_void_p()
self._captured = False
def begin_capture(self, stream: ctypes.c_void_p):
# cudaStreamCaptureModeRelaxed=2: only capture ops on THIS stream.
_check(_cudart.cudaStreamBeginCapture(stream, 2), “cudaStreamBeginCapture”)
def end_capture(self, stream: ctypes.c_void_p):
_check(_cudart.cudaStreamEndCapture(stream, ctypes.byref(self._graph)),
“cudaStreamEndCapture”)
_check(_cudart.cudaGraphInstantiate(
ctypes.byref(self._graph_exec), self._graph, 0),
“cudaGraphInstantiate”)
self._captured = True
def replay(self, stream: ctypes.c_void_p):
ifnot self._captured:
raise RuntimeError(“No graph captured”)
_check(_cudart.cudaGraphLaunch(self._graph_exec, stream), “cudaGraphLaunch”)
这里特意使用了 cudaStreamCaptureModeRelaxed=2,注释明确表明只捕获当前 stream 上的操作。这一点对于兼容 JAX/XLA 后端尤为关键,因为全局捕获模式可能会与框架的后台线程产生冲突。
4.2 图捕获的核心要求:指针稳定、无动态分配、shape 稳定
CUDA Graph 并非魔法。它之所以能带来性能提升,代价是“静态性”:捕获时记录的内核序列、指针和内存分配关系,在重放时必须仍然有效。因此,FlashRT 的大量代码都围绕一个核心原则展开:所有中间 buffer 预先分配,前向传播路径只传递裸指针。
CudaBuffer 正是这一设计思想在底层实现中的具体体现。
# 来源:flash_rt/core/cuda_buffer.py
class CudaBuffer:
"""GPU buffer — managed or device memory."""
def __init__(self, nbytes: int, managed: bool = True):
self._ptr = ctypes.c_void_p()
self._managed = managed
if managed:
_check(_cudart.cudaMallocManaged(ctypes.byref(self._ptr), nbytes, 1),
"cudaMallocManaged")
else:
_check(_cudart.cudaMalloc(ctypes.byref(self._ptr), nbytes), "cudaMalloc")
self._nbytes = nbytes
@classmethod
def from_numpy(cls, arr: np.ndarray) -> 'CudaBuffer':
"""Create device buffer, upload via chunked H2D (Thor-safe ≤4MB/chunk)."""
arr = np.ascontiguousarray(arr)
buf = cls(arr.nbytes, managed=False)
chunk = 4 * 1024 * 1024 # 4MB chunks, Thor-safe
for off in range(0, arr.nbytes, chunk):
n = min(chunk, arr.nbytes - off)
_check(_cudart.cudaMemcpy(
ctypes.c_void_p(buf._ptr.value + off),
ctypes.c_void_p(arr.ctypes.data + off), n, 1), "H2D chunk")
return buf
代码中“Thor-safe ≤4MB/chunk”这一注释尤其值得关注。这并非学术论文中的抽象理论,而是来自真实部署场景的硬件约束。FlashRT 的设计哲学在此显露无遗:为了保障实时系统的可靠性,许多优化手段追求的并非“更优雅”,而是“更确定”。
五、权重加载:将 state_dict 转化为声明式规格表
5.1 WEIGHT_SPEC:FlashRT 的轻量化编译描述
不同模型的 checkpoint 命名规则、层数深度、QKV 排列方式以及 Gate/Up 融合策略千差万别。如果为每个前端模型都手动编写大量加载代码,系统将很快变得难以控制。为此,FlashRT 引入了 ModelWeightSpec、LayerBlock 和 Item 这一套声明式机制,将权重加载过程规格化。
# 来源:flash_rt/executors/weight_loader.py
@dataclass
class ModelWeightSpec:
"""Top-level spec for one frontend's weight loading.
A frontend declares exactly one ``WEIGHT_SPEC = ModelWeightSpec(...)``
at class scope and its ``_load_weights`` body becomes::
WeightLoader(source, target=self, spec=self.WEIGHT_SPEC).run()
"""
framework: str
blocks: list[LayerBlock] = field(default_factory=list)
singletons: list[Item] = field(default_factory=list)
buffers: list[BufferSpec] = field(default_factory=list)
dims: dict[str, Any] = field(default_factory=dict)
这套机制类似于模型权重的“链接脚本”:它清晰地告知加载器从何处获取张量、如何进行拼接、是否需要转置或量化,以及最终应该放置到目标对象的哪个属性中。
5.2 Pi0 与 Pi0.5 的差异被编码进 spec,而非散落在推理逻辑中
在 Pi0 模型中,decoder 模块会将 RMSNorm 的权重参数融合到 QKV 以及 Gate/Up 矩阵中;而 Pi0.5 由于采用了不同的 AdaRMSNorm Dense 路径,因此不会执行此类融合操作。这一关键差异在模型规格文件中得到了直接体现。
来源:flash_rt/frontends/torch/_pi0_thor_spec.py
def _decoder_block() -> LayerBlock:
“””Pi0 decoder (18 layers, cuBLASLt path, with norm fuse).”””
dp = “paligemma_with_expert.gemma_expert.model.layers.{i}”
items = [
Item(“qkv_w”,
FusedQKV(q=f”{dp}.self_attn.q_proj.weight”,
k=f”{dp}.self_attn.k_proj.weight”,
v=f”{dp}.self_attn.v_proj.weight”,
norm_fuse=f”{dp}.input_layernorm.weight”,
interleave_q_heads=8,
interleave_k_heads=1),
[tT(), Quant()],
FlatCat(“_dec_qkv_flat”), scale_into=”_ae_w_scales”),
Item(“gu_w”,
FusedGateUp(gate=f”{dp}.mlp.gate_proj.weight”,
up=f”{dp}.mlp.up_proj.weight”,
norm_fuse=f”{dp}.post_attention_layernorm.weight”),
[tT(), Quant()],
FlatCat(“_dec_gu_flat”), scale_into=”_ae_w_scales”),
]
return LayerBlock(prefix_fmt=””, num_layers=18, items=items, name=”decoder”)
来源:flash_rt/frontends/torch/_pi05_thor_spec.py
def _decoder_block() -> LayerBlock:
“””Gemma-expert decoder / action expert (18 layers, cuBLASLt path).
Unlike Pi0, Pi0.5 does not fuse the per-layer norm weight into
QKV/GateUp — modulation happens via AdaRMSNorm Dense layers instead.
“””
dp = “paligemma_with_expert.gemma_expert.model.layers.{i}”
items = [
Item(“qkv_w”,
FusedQKV(q=f”{dp}.self_attn.q_proj.weight”,
k=f”{dp}.self_attn.k_proj.weight”,
v=f”{dp}.self_attn.v_proj.weight”,
interleave_q_heads=8,
interleave_k_heads=1),
[tT(), Quant()],
FlatCat(“_dec_qkv_flat”), scale_into=”_ae_w_scales”),
]
return LayerBlock(prefix_fmt=””, num_layers=18, items=items, name=”decoder”)
这恰恰是工程可维护性的精髓所在:模型之间的差异被局部化到了声明式的权重变换层,而底层的 pipeline 则可以保持一个稳定的结构——只需按槽位读取指针、按层调用内核即可。
unsetunset六、AttentionBackend:让模型不关心 FlashAttention 还是 Thor FMHAunsetunset
6.1 同一 pipeline,背后可以是不同 attention 实现
hardware/backend.py的注释阐述了两种存储模型:RTX 后端可能自行持有 Q/K/V/O 张量;而 Thor 后端则可能由 pipeline 维护 scratch 和 KV cache。该协议允许二者共存,仅要求get_slot_ptrs()返回稳定的指针,run()返回输出指针。
这种设计思路与操作系统中的设备驱动接口非常相似:上层只需发起“在某个 site 和 layer 执行 attention”的请求,底层自行决定如何执行。
6.2 GROOT 的 attention site 体现了复杂模型的拆分方式
GROOT N1.6 模型包含了 SigLIP 视觉模块、Qwen3 编码器、DiT self-attention 以及 DiT cross-attention。RTX 后端将这些模块拆分为不同的 site:
来源:flash_rt/hardware/rtx/attn_backend_groot.py
class RtxFlashAttnBackendGroot:
“””FA2-backed attention for GROOT N1.6 (FP16) on RTX hardware.
Site mapping
————
GROOT has four distinct attention shapes:
* ``"siglip"`` — 27 SigLIP layers, batched per-view self-attn,
16 heads, head_dim=72, 256 tokens/view.
* ``"qwen3"`` — 16 Qwen3 encoder layers, GQA 16Q/8KV, head_dim=128.
* ``"dit_self"`` — 16 DiT self-attention blocks.
* ``"dit_cross"`` — 16 DiT cross-attention blocks.
“””
在 run() 方法中,系统会根据 site 参数将请求分发至不同的注意力计算路径:
# 来源:flash_rt/hardware/rtx/attn_backend_groot.py
def run(
self,
site: str,
layer_idx: int,
q_seq: int,
*,
kv_seq=None,
stream: int = 0,
) -> int:
"""根据给定的 site 分发到对应的 legacy attention 调用。"""
if site == "siglip":
return self.vision_attn(stream=stream)
if site == "qwen3":
return self.qwen3_attn(layer_idx, q_seq, stream=stream)
if site == "dit_self":
return self.dit_self_attn(layer_idx, q_seq, stream=stream)
if site == "dit_cross":
if kv_seq is None:
raise ValueError(
"dit_cross site 表示与 Qwen3 骨干特征的交叉注意力;"
"必须提供 kv_seq 参数"
)
这种设计理念体现了 FlashRT 的系统观:注意力机制并非一个抽象的通用算子,而是一组由 shape、缓存策略、site 标识和硬件实现共同定义的调用协议。
七、FP8 与 NVFP4:量化的本质不是压缩,而是控制误差传播
7.1 FP8 路径:权重量化与激活校准融入图捕获阶段
根据 README 文档,FlashRT 采用 FP8 E4M3 格式,并将校准结果缓存至
~/.flash_rt/calibration/目录。首次调用时可能需要几秒完成校准与 CUDA Graph 捕获,后续调用则直接进入 graph replay 阶段。
这种设计在实时系统中非常合理:将不稳定、成本高昂且数据相关的准备工作前置,在线阶段仅保留确定性执行逻辑。
7.2 NVFP4 路径:Pi0.5 encoder FFN 的激进尝试
README 对 Pi0.5 的 NVFP4 路径描述非常详尽:Gate+Up 和 Down GEMM 跨越 18 层 encoder FFN 以 NVFP4 格式运行;AWQ 用于实现激活感知的预缩放;P1 split-GU 方案将合并的 Gate+Up 拆分为两个独立的 NVFP4 GEMM,并由专用 kernel 进行组合,从而减少 DRAM 往返次数。
这背后的系统意义在于:低比特量化的收益不仅来自“少存储权重”,更关键的是降低内存带宽压力和减少中间激活的搬运开销。然而,其代价是误差可能在多层网络中累积放大。为此,FlashRT 引入了 AWQ 和专用融合 kernel,将精度补偿机制也整合进图结构中。
以下是一个简化的伪代码,用于说明其设计思路:
# 来源:根据 README.md 中 NVFP4 encoder FFN 说明整理
for layer in encoder_layers:
# 1. 残差流保持 fp16 精度,避免整段网络被低比特误差拖垮
x_norm = rms_norm_fp16(x)
# 2. AWQ 预缩放:根据输入通道修正权重量化的敏感性
x_scaled = x_norm * inverse_awq_scale
# 3. Gate 和 Up 拆分为两个 NVFP4 GEMM,直接产出 packed FP4 与 scale
gate_fp4 = nvfp4_gemm(x_scaled, gate_weight_fp4)
up_fp4 = nvfp4_gemm(x_scaled, up_weight_fp4)
# 4. 专用融合 kernel 完成 GEGLU/SiLU 类门控与再量化
hidden_fp4 = geglu_two_mul_fp4_to_fp4(gate_fp4, up_fp4)
# 5. Down 投影后回到残差路径
out = nvfp4_gemm(hidden_fp4, down_weight_fp4)
x = x + out
这里的关键不在于“使用了 FP4”这个事实,而在于理解:低比特区域的边界在哪里、scale 如何融合、残差流是否保持高精度、Gate/Up 是否减少了内存往返。这也是底层推理优化中最容易被忽视的部分——量化本质上是一套误差控制工程,而不是一个简单的 dtype 开关。
八、Qwen3.6 路径:FlashRT 正将 VLA 经验扩展到 LLM 领域
8.1 Qwen3.6 的分阶段迁移计划
Qwen3.6-Next 并非传统的 dense transformer,而是由 48 层 Gated DeltaNet 线性注意力与 16 层 softmax 全注意力组成的混合结构。这意味着 KV cache 仅在 16 层 full-attn 中增长,长上下文吞吐相对稳定——256K 上下文能够正常运行的背后,是架构红利配合 Turboquant 与 FlashRT 静态图叠加的结果;单卡情况下,通常在 30k 左右就会触发 OOM。
因此,仓库中的 flash_rt/models/qwen36/__init__.py 记录了 Qwen3.6 路径的分阶段规划:首先通过 PyTorch eager wrapper 锁定前端契约和 cosine 基线,然后逐步替换 attention、MLP、RMSNorm、RoPE、SwiGLU、Gated DeltaNet、KV cache、采样、CUDA Graph 以及 MTP speculative decoding。
“`python
来源:flash_rt/models/qwen36/init.py
“””FlashRT — Qwen3.6-27B model pipelines.
阶段规划:
Phase 1: 基于 HF AutoModelForCausalLM 的 PyTorch-eager 包装器
Phase 2: 用 fvk kernel 调用替换 full-attn + MLP + RMSNorm + RoPE + SwiGLU
Phase 3: 编写 csrc/kernels/{gated_deltanet,causal_conv1d_qwen36}
Phase 4: KV cache + 采样 + CUDA graph + 干净的 decode 循环
Phase 5: 视觉塔(仅限多模态场景)
Phase 6: MTP speculative 解码(可选)
“””
这说明 FlashRT 并非仅仅局限于机器人领域的 VLA 应用,其核心目标是将“指针化 pipeline + kernel 库 + 静态图”这一整套系统方法论,迁移至大语言模型(LLM)的处理路径中。
8.2 KV cache 的指针布局同样服务于图稳定性
在 Qwen3.6 RTX attention 后端中,KV cache 是按层(layer)连续分块的,并提供了层跨度(layer stride)与行跨度(row stride),使得 pipeline 只需计算当前 token 的写入偏移量即可。
来源:flash_rt/hardware/rtx/attn_backend_qwen36.py
@property
def kv_layer_stride_bytes(self) -> int:
return self._max_seq * self.NUM_KV_HEADS * self.HEAD_DIM * 2
@property
def kv_row_stride_bytes(self) -> int:
return self.NUM_KV_HEADS * self.HEAD_DIM * 2
def get_slot_ptrs(self, site: str, layer_idx: int) -> dict:
layer_off_bytes = layer_idx * self.kv_layer_stride_bytes
return {
“Q”: self.Q_buf.data_ptr(),
“K”: self.K_cache.data_ptr() + layer_off_bytes,
“V”: self.V_cache.data_ptr() + layer_off_bytes,
“kv_layer_stride_bytes”: self.kv_layer_stride_bytes,
“kv_row_stride_bytes”: self.kv_row_stride_bytes,
}
这段代码非常接近硬件底层:它处理的并非张量(tensor)语义,而是直接的字节偏移量。正是这种明确的内存布局设计,使得 KV cache 能够被纳入一个稳定的解码循环中,最终服务于 CUDA Graph 的捕获以及长上下文的推理任务。
九、训练反哺:把推理 kernel 直接拿到训练 forward 上
绝大多数推理引擎与训练框架是两套独立的工程体系:训练阶段使用 PyTorch 或 JAX 的标准算子,推理阶段则再进行一次模型导出(export)、策略搜索(tactic search)和引擎编译(engine compile)。FlashRT 选择了另一条截然不同的道路——让训练的前向传播(forward)直接调用推理时使用的同一组手写 kernel。
9.1 设计原则:推理 kernel = 训练 kernel
FlashRT 的 training/ 目录下提供了一个基于 PyTorch 的 FP8 + LoRA 训练栈,其反向传播(backward)走 BF16 LoRA 路径。关键在于前向传播路径:FP8 cuBLASLt GEMM、融合的 RMSNorm、SwiGLU 等 kernel,与推理路径使用的是同一份动态链接库(.so)。
来源:training/README.md
“the same kernels FlashRT’s inference path uses for FP8 GEMM, fused
RMSNorm, and SwiGLU are reused for the training forward pass, while
backward stays in BF16 through small LoRA adapters.”
工程意义在于:训练完成后,直接执行 merge_lora_into_base 操作,导出为 safetensors 格式,再由同一个推理引擎加载。整个过程无需模型导出,无需引擎重新编译,也无需切换框架。
9.2 数据:训练步速与 VRAM 双降
在 RTX 5090 上,设置 batch size 为 4、lora_rank 为 16,与 lerobot BF16+LoRA 以及 openpi 官方 JAX LoRA 在相同目标集上进行比较:
| 维度 | 基线 | FlashRT FP8+LoRA | 倍数 |
|---|---|---|---|
| 训练步速 vs lerobot BF16 | 1.83 step/s | 4.04 step/s | 2.21× |
| VRAM vs openpi JAX LoRA | > 22.5 GB | 10.05 GB (memory-priority) | 2.24× 缩减 |
| 单 shape FP8 fwd (gate_up) | BF16 baseline | — | 2.00× fwd / 1.67× fwd+bwd |
在 memory-priority 配置下(即 cache_bf16_weight=False),使得拥有 27 亿参数的 Pi0.5 模型的 LoRA 微调,能够在消费级的 24 GB 单卡上轻松完成,不再依赖 A100 或 H100 等专业显卡。
9.3 闭环意义:解决具身智能“频繁微调”的痛点
具身智能场景与通用 LLM 服务最大的区别之一在于,模型并非一次训练完成即可部署:
- 新的机器人本体或新任务 → 需要进行 LoRA 微调
- 真实机器人数据回流 → 需要 RECAP 或 RL 持续迭代
- 部署到不同机型 → 需要重新进行校准
每增加一个环节,传统路径就需要经历一次从“训练框架 → 模型导出 → 推理引擎”的转换循环。FlashRT 的做法是将这条链路压缩成一条直线:
calibrate (FP8 scale) ─┐
├─ 共享同一组 hand-written kernel
train LoRA (FP8 fwd) ─┤
├─ 共享同一份 weight/scale 协议
merge_lora_into_base ─┤
├─ 共享同一个 .so
serve (FP8 inference) ─┘
这不仅仅是“训练变快”那么简单——它的本质是用 kernel 抽象层取代了框架的中间态。对于 RL、频繁的 RECAP 微调以及多机型部署这类具身智能的“高换手”工作流,在这条 pipeline 下,不再需要每次都重复执行模型导出和编译的流程。
十、与 TensorRT、vLLM、SGLang 的根本差异
10.1 FlashRT 拒绝“先导出、再编译、再部署”的重型路径
TensorRT 的优势在于其策略搜索(tactic search)、引擎编译以及强大的图优化能力。但这也意味着它存在模型导出、引擎生成、驱动或 GPU 变更后的兼容性问题,以及对动态输入和迭代开发流程的摩擦。
FlashRT 选择了一条不同的道路:“不导出 ONNX、不生成 .engine 文件、不进行每次编译”。它直接加载 safetensors 或 Orbax 格式的模型,并将模型逻辑直接映射到其自身的 kernel 组合上。这种方式虽然牺牲了一部分通用性,但换来了更短的部署路径、更可控的冷启动时间,以及对特定模型系列的高度优化。
10.2 它也不是 vLLM 式高并发调度器
vLLM 和 SGLang 的核心价值在于高并发的 LLM 服务,尤其是在 KV cache 分页、请求调度和连续批处理方面。然而,FlashRT 的目标场景是小批量(small batch)和实时控制。对它而言,将 10 个请求攒在一起处理并非优化,反而可能成为控制系统的灾难。
因此,FlashRT 的核心并非“调度更多请求”,而是“让单次请求以最快速度完成”。
十一、工程取舍:FlashRT 为何值得深入剖析
11.1 优势:将端到端延迟置于首位
FlashRT 最值得借鉴之处,在于它并未止步于单个 kernel 的基准测试,而是围绕完整的推理链路进行优化:权重加载、量化、注意力机制、通用矩阵乘法、缓冲区管理、CUDA Graph、硬件分发、API 抽象以及基准测试,全部整合进同一套系统。
11.2 代价:模型适配需要工程层面的“翻译”
FlashRT 并非适用于任意 PyTorch 模型的自动加速器。若要引入新模型,开发者需要编写前端、权重规范、注意力站点、流水线,并将源模型的 forward 流程翻译为内核调用序列。这意味着开发者必须深入理解模型结构、权重命名规则、张量形状、量化边界以及硬件路径。
换言之,它更像一个“高性能模型移植框架”,而非“黑盒编译器”。
11.3 最大启示:实时 AI 需要另一种推理系统哲学
过去数年,AI 基础设施主要围绕云端 LLM 服务构建;然而,机器人、AR、边缘 Agent 以及低延迟世界模型将提出截然不同的需求:输入复杂、批量极小、闭环严格、硬件异构、可靠性优先。 FlashRT 展示了一条可能的路径:不追求无限的通用性,而是在明确场景中将系统边界压缩到极致。
结语:让模型部署回归控制闭环
FlashRT 的价值,并不仅仅在于“某个 VLA 模型跑得更快”。其真正引人深思之处在于:它重新定义了实时推理引擎的职责分工。框架负责生态,前端负责加载,规范负责权重变换,流水线负责静态算子序列,硬件后端负责将注意力机制和通用矩阵乘法落实到具体 GPU 上,而 CUDA Graph 则负责将运行时从“逐 op 调度”转变为“一次性重放”。
这套设计并不适用于所有模型,也无意取代 TensorRT、vLLM 或 SGLang。 然而,在小批量、低延迟、机器人控制、边缘 GPU 等场景中,它精准击中了通用推理系统的盲区:真正的瓶颈并非模型能否运行,而是能否稳定、快速、无摩擦地融入控制闭环。
如果说大模型服务的关键词是吞吐量、调度与扩展性,那么 FlashRT 所代表的关键词则是确定性、静态化与端到端延迟。它提醒我们:AI 系统工程的未来不会只有一种形态。面向实时世界的模型,不仅需要更智能的“大脑”,也需要一条足够短、足够硬、足够可预测的“神经通路”。
关注“鲸栖”小程序,掌握最新AI资讯
本文来自网络搜集,不代表鲸林向海立场,如有侵权,联系删除。转载请注明出处:https://www.itsolotime.com/archives/33574

