Vulkan: how much of each LeRobot policy the Vulkan backend takes, and what blocks the rest

Abierto
  • ¿Algo no está bien?
#24,688 0 comentarios 0 reacciones 1 asignado Ver en GitHub

Los mantenedores suelen responder en 1 día

@SS-JIA ya está trabajando en esto.

Desde el 5/11/2026.

Evaluación

Este issue todavía no se ha evaluado.

Descripción

module: vulkan

This issue tracks LeRobot policies exported to ExecuTorch with its Vulkan GPU backend (VulkanPartitioner): how much of each policy Vulkan takes, whether the program exports, and what each project still needs to look at. Other backends have their own issues: TensorRT in #24353, OpenVINO in #24325, XNNPACK in #24667, Qualcomm in #24668, Core ML and MLX in #24669, the CUDA backend in #24684, and Arm in #24697.

Why this matters for LeRobot: Vulkan is a GPU path for boards with no CUDA and no NPU that ExecuTorch supports, for example the Arm Mali GPU on RK3588 boards. ExecuTorch's Vulkan backend is built mainly for Android GPUs, so a board running Linux is on the less tested path. It can also use the Adreno GPU in Qualcomm chips, which also have a Hexagon NPU (see #24668). On NVIDIA GPUs, TensorRT or the CUDA backend is the main path, so the H100 column here only checks correctness. No Mali or Adreno device has been tested yet.

Coverage

Policy Ops on Vulkan Heavy math on Vulkan Vulkan programs Exports to .pte H100 (correctness)
act 1556 / 1561 (100%) 100% 3 $\color{#2da44e}{\checkmark}$ $\color{#2da44e}{\checkmark}$ 88.7 (1.0x)
diffusion 534 / 619 (86%) 94% 50 $\color{#2da44e}{\checkmark}$ ⏳
dm05 ⏳ ⏳ ⏳ ⏳ ⏳
eo1 ⏳ ⏳ ⏳ ⏳ ⏳
evo1 13904 / 16601 (84%) 26% 61 $\color{#d4a017}{\textsf{slow}}$ ⏳
fastwam ⏳ ⏳ ⏳ ⏳ ⏳
fineart_vla 21670 / 21704 (100%) 100% 6 $\color{#d4a017}{\textsf{slow}}$ ⏳
flux3 ⏳ ⏳ ⏳ ⏳ ⏳
gaussian_actor 114 / 127 (90%) 100% 5 $\color{#2da44e}{\checkmark}$ $\color{#2da44e}{\checkmark}$ 1.14 (1.9x)
groot ⏳ ⏳ ⏳ ⏳ ⏳
lawam 15409 / 18124 (85%) 44% at least 198 $\color{#d4a017}{\textsf{slow}}$ ⏳
lingbot_va ⏳ ⏳ ⏳ ⏳ ⏳
lingbot_vla_v2 ⏳ ⏳ ⏳ ⏳ ⏳
molmoact2 ⏳ ⏳ ⏳ ⏳ ⏳
multi_task_dit 1460 / 1463 (100%) 100% 5 $\color{#2da44e}{\checkmark}$ $\color{#2da44e}{\checkmark}$ 774 (3.9x)
pi0 5624 / 5654 (99%) 100% 17 $\color{#2da44e}{\checkmark}$ ⏳
pi05 ❌ ❌ ❌ ❌ -
pi0_fast 9059 / 9082 (100%) 100% 11 $\color{#2da44e}{\checkmark}$ ⏳
smolvla 2933 / 7163 (41%) 16% 209 $\color{#2da44e}{\checkmark}$ ⏳
tdmpc 5477 / 5895 (93%) 100% 140 $\color{#2da44e}{\checkmark}$ $\color{#2da44e}{\checkmark}$ 494 (0.9x)
vla_jepa 7049 / 10817 (65%) 6% 191 $\color{#d4a017}{\textsf{slow}}$ ⏳
vqbet 567 / 594 (95%) 23% 22 $\color{#2da44e}{\checkmark}$ $\color{#2da44e}{\checkmark}$ 1241 (0.02x)
wall_oss_05 ⏳ ⏳ ⏳ ⏳ ⏳
wall_x ⏳ ⏳ ⏳ ⏳ ⏳
xvla 16565 / 16643 (100%) 99% 61 $\color{#d4a017}{\textsf{slow}}$ ⏳

Legend:

  • "Ops on Vulkan": the graph ops that VulkanPartitioner() takes, out of all ops, over all programs of the policy. Ops are counted after to_edge, so a linear layer counts as a matmul and a transpose.
  • "Heavy math on Vulkan": the share of multiply-adds in matmuls, linear layers, convolutions and attention that Vulkan takes. A program counts as many times as one action chunk runs it (for example, once per denoising step).
  • "Vulkan programs": how many Vulkan partitions the policy is split into. Between 2 partitions, the data goes back to the CPU. "At least N" means partitioning itself did not finish within 15 minutes; N is the number of connected groups of ops that Vulkan takes, a lower bound.
  • "Exports to .pte": $\color{#2da44e}{\checkmark}$ means to_edge_transform_and_lower(...).to_executorch() finished with the default options. $\color{#d4a017}{\textsf{slow}}$ means the export did not finish within 15 minutes; a longer run is queued. ❌ means it fails; the cause is in the tables below.
  • "H100 (correctness)": the bundle run with lerobot-rollout --bundle.path on an H100 through its Vulkan driver, with PyTorch never loaded. $\color{#2da44e}{\checkmark}$ means it passed the export's own accuracy check and, for policies without a frame history, a 16-frame comparison against eager PyTorch. The number is milliseconds per action chunk (batch size 1, median). In brackets: how many times faster than eager PyTorch (lerobot-rollout --policy.path --device=cpu) on 16 CPU cores of the same machine, because a Vulkan bundle takes its inputs and returns its outputs on the CPU. Ops that Vulkan does not take run on ExecuTorch's portable CPU kernels, which are slow, so a bundle with heavy ops left on the CPU can be much slower than eager (vqbet).
  • ⏳: this cell is being measured (running or queued).

Notes:

  • The ExecuTorch 1.6.0 wheel has the Vulkan partitioner but not the Vulkan runtime. For the H100 column, ExecuTorch 1.6.0 was built from source with the Vulkan backend on (EXECUTORCH_BUILD_VULKAN=1, shaders compiled with glslc from conda-forge's shaderc, as ExecuTorch's CI does) and with the event tracer off (ExecuTorch table).
  • LeRobot has no Vulkan backend for lerobot-export yet (LeRobot table). So each policy is traced with the trace step of lerobot-export (same checkpoint, task and export flags as the XNNPACK rows of #24667 and the OpenVINO rows of #24325), and each traced program is lowered the way lerobot-export lowers for its ExecuTorch backends, with VulkanPartitioner() and its default options.
  • diffusion, dm05, eo1, fineart_vla, gaussian_actor, lingbot_vla_v2, molmoact2, multi_task_dit, tdmpc, vla_jepa, vqbet and wall_x use random weights, made by the same script and seed as in #24353. The other policies use their real weights.
  • With force_fp16 (run inside the backend in fp16), act, diffusion, gaussian_actor, multi_task_dit, pi0_fast, smolvla and vqbet also export. It does not change what the partitioner takes.
  • evo1 keeps its bf16 weights, so its vision-language part stays on the CPU (first ExecuTorch row).
  • pi0 and pi05 checkpoints use the SigLIP weight names of transformers before 5.22, which the branch cannot load yet (LeRobot table). The pi0 row uses the same pi0 weights re-saved under the new names, which does not change the graph. pi05 is not measured yet.

Built on

  • ExecuTorch: the official 1.6.0 CPU wheel (the final release). Its backends/vulkan Python code is what was measured. No source build and no patches.
  • torch 2.15.0, transformers 5.22.1.
  • LeRobot: the lerobot-export-stack branch of shoumikhin/lerobot, as it was on Nov 3 at 10:53 AM Pacific (see #24353).
  • GPU (for the correctness column): H100, NVIDIA driver 580, Vulkan 1.4.

Known issues

"Fix" is an open pull request unless it says merged. "Needs a decision" is a design question for the maintainers. "Needs an owner" means nobody is fixing it yet.

ExecuTorch
  • Pick up the rows marked "needs an owner". bf16 (first row) and batched convolutions (second row) cost the most heavy math.
  • Every row has a short repro below the table. Repros 1 to 4 need only the ExecuTorch 1.6.0 wheel and no GPU. Repro 5 needs ExecuTorch built with Vulkan and a GPU.
Problem Policies Fix Status
No bf16. torch.bfloat16 is not in the backend's dtype table (_VULKAN_DTYPES in backends/vulkan/utils.py), so every bf16 op stays on the CPU with "dtype not supported". evo1 (74% of its heavy math stays on the CPU), smolvla (84%), vla_jepa (94%) needs a decision: support bf16, or run bf16 programs in fp16 or fp32 inside the backend Repro ready (repro 1)
A 4D convolution is taken only with batch size 1 (check_conv_node in backends/vulkan/op_registry.py). Policies that encode 2 or more frames or cameras as one batch keep their image encoder on the CPU. vqbet (77% of its heavy math stays on the CPU), diffusion, xvla, multi_task_dit needs an owner Repro ready (repro 2)
No Vulkan kernel for elu, log1p (used by Mish through softplus), var, index_put, or transposed 1D convolution. Each one cuts the graph into more Vulkan programs. tdmpc (elu 181 times, log1p 150, var 36), diffusion (log1p 38, transposed conv1d 2), xvla (index_put 21) needs an owner Repro ready (repro 3)
Partitioning big graphs is very slow. VulkanPartitioner uses torch.fx's CapabilityBasedPartitioner, and its merge step grows much faster than the graph. tdmpc (5,895 ops) takes 603 s to partition and ends with 140 Vulkan programs, while its supported ops form only 41 connected groups. tdmpc, evo1, xvla, and the larger policies needs an owner Repro ready (repro 4): 200 blocks take 13 s, 400 blocks take 95 s
With the event tracer on, which the Python wheel's build preset turns on (tools/cmake/preset/pybind.cmake), the Vulkan runtime writes 2 GPU timestamps per dispatch into a query pool of 4,096 (VULKAN_QUERY_POOL_SIZE). A program with more than about 2,000 dispatches stops with "Vulkan QueryPool: Exceeded the maximum number of queries". act (from Python) needs an owner: use the query pool only when profiling is asked for, or let it grow Repro ready (repro 5): 1,000 steps run, 1,500 steps stop; with -DEXECUTORCH_ENABLE_EVENT_TRACER=OFF, 4,000 steps run
Repros (export only, ExecuTorch 1.6.0, torch 2.15.0) 
import time

import torch
from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner
from executorch.exir import to_edge, to_edge_transform_and_lower


def lowered(model, *inputs):
    """How many Vulkan programs the model lowers to, and which ops stay on the CPU."""
    edge = to_edge_transform_and_lower(torch.export.export(model.eval(), inputs), partitioner=[VulkanPartitioner()])
    nodes = [n for n in edge.exported_program().graph.nodes if n.op == "call_function"]
    names = [getattr(n.target, "__name__", str(n.target)) for n in nodes]
    vulkan = names.count("executorch_call_delegate")
    cpu = sorted({name for name in names if name not in ("executorch_call_delegate", "getitem")})
    return vulkan, cpu


class Fn(torch.nn.Module):
    def __init__(self, fn):
        super().__init__()
        self.fn = fn

    def forward(self, x):
        return self.fn(x)


def case(label, model, *inputs):
    print(f"{label:44s} -> Vulkan programs, CPU ops: {lowered(model, *inputs)}")


# 1. bf16: no op is delegated.
case("Linear float32", torch.nn.Linear(64, 64), torch.randn(4, 64))
case("Linear bfloat16", torch.nn.Linear(64, 64).to(torch.bfloat16), torch.randn(4, 64, dtype=torch.bfloat16))

# 2. A 4D convolution is taken only with batch size 1.
case("Conv2d batch 1", torch.nn.Conv2d(3, 8, 3), torch.randn(1, 3, 32, 32))
case("Conv2d batch 2", torch.nn.Conv2d(3, 8, 3), torch.randn(2, 3, 32, 32))

# 3. Ops with no Vulkan kernel.
x = torch.randn(1, 16, 32)
case("ELU", torch.nn.ELU(), x)
case("Mish (softplus -> log1p)", torch.nn.Mish(), x)
case("var", Fn(lambda t: t.var(dim=-1, keepdim=True)), x)
case("ConvTranspose1d", torch.nn.ConvTranspose1d(16, 8, 4, stride=2, padding=1), x)


class Put(torch.nn.Module):
    def forward(self, t, idx):
        y = t.clone()
        y[idx] = 0.0
        return y


case("index_put (y[idx] = 0)", Put(), torch.randn(8, 32), torch.tensor([1, 3]))


# 4. Partitioning time grows fast with graph size when CPU ops sit on side branches.
class Blocks(torch.nn.Module):
    def __init__(self, n):
        super().__init__()
        self.layers = torch.nn.ModuleList(torch.nn.Linear(32, 32) for _ in range(n))

    def forward(self, x):
        for layer in self.layers:
            a = layer(x)
            x = a * torch.sigmoid(a) + torch.nn.functional.elu(a)  # elu stays on the CPU
        return x


for n in (50, 100, 200, 400):
    ep = to_edge(torch.export.export(Blocks(n).eval(), (torch.randn(1, 32),))).exported_program()
    start = time.time()
    tags = VulkanPartitioner().partition(ep).partition_tags
    print(f"{n:4d} blocks: {len(tags):4d} Vulkan partitions, partition() took {time.time() - start:6.1f} s")

Output:

Linear float32                               -> Vulkan programs, CPU ops: (1, [])
Linear bfloat16                              -> Vulkan programs, CPU ops: (0, ['aten.addmm.default', 'aten.permute_copy.default'])
Conv2d batch 1                               -> Vulkan programs, CPU ops: (1, [])
Conv2d batch 2                               -> Vulkan programs, CPU ops: (0, ['aten.convolution.default'])
ELU                                          -> Vulkan programs, CPU ops: (0, ['aten.elu.default'])
Mish (softplus -> log1p)                     -> Vulkan programs, CPU ops: (2, ['aten.log1p.default'])
var                                          -> Vulkan programs, CPU ops: (0, ['aten.var.correction'])
ConvTranspose1d                              -> Vulkan programs, CPU ops: (0, ['aten.convolution.default'])
index_put (y[idx] = 0)                       -> Vulkan programs, CPU ops: (1, ['aten.index_put.default'])
  50 blocks:   51 Vulkan partitions, partition() took    2.6 s
 100 blocks:  101 Vulkan partitions, partition() took    3.8 s
 200 blocks:  201 Vulkan partitions, partition() took   13.1 s
 400 blocks:  401 Vulkan partitions, partition() took   94.7 s
Repro 5 (needs ExecuTorch 1.6.0 built with Vulkan, and a GPU) 
import sys

import torch
from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner
from executorch.exir import to_edge_transform_and_lower
from executorch.runtime import Runtime


class Many(torch.nn.Module):
    def __init__(self, n):
        super().__init__()
        self.n = n

    def forward(self, x):
        for i in range(self.n):
            x = x * 1.0001 + 0.001  # 2 dispatches per step
        return x


for n in [int(a) for a in sys.argv[1:]] or [500, 1000, 2000]:
    x = torch.randn(1, 64)
    et = to_edge_transform_and_lower(torch.export.export(Many(n), (x,)), partitioner=[VulkanPartitioner()]).to_executorch()
    method = Runtime.get().load_program(et.buffer).load_method("forward")
    try:
        out = method.execute([x])[0]
        print(n, "steps: ok, max abs diff", (out - Many(n)(x)).abs().max().item())
    except RuntimeError as error:
        print(n, "steps: RuntimeError:", str(error).split("\n")[0][:200])

Output on an H100:

# built with the event tracer on (the default of the Python wheel's build preset)
500 steps: ok, max abs diff 0.0
1000 steps: ok, max abs diff 0.0
1500 steps: RuntimeError: Exception raised from write_timestamp at .../backends/vulkan/runtime/vk_api/QueryPool.cpp:79: (num_queries_ < config_.max_query_count) is false! ...
2000 steps: RuntimeError: (same)
# built with -DEXECUTORCH_ENABLE_EVENT_TRACER=OFF
1500 steps: ok, max abs diff 0.0
4000 steps: ok, max abs diff 0.0

LeRobot
  • Decide the rows marked "needs a decision".
Problem Policies Fix Status
lerobot-export has no Vulkan backend, so there is no Vulkan bundle for lerobot-rollout --bundle.path to run. all needs a decision: an executorch_vulkan backend in LeRobot, next to executorch_xnnpack and executorch_openvino Open
lerobot-export and the torch-free lerobot-rollout --bundle.path exist only on the lerobot-export-stack branch. all needs a decision: how to bring them upstream in reviewable pieces (same row as #24684) Open
With transformers 5.22, a pi0 or pi05 checkpoint saved with older SigLIP weight names does not load: model.safetensors does not hold model.paligemma_with_expert.paligemma.model.vision_tower.embeddings.patch_embedding.weight and 436 more. pi0, pi05 huggingface/lerobot#4775 Open

cc @SS-JIA @manuelcandales @digantdesai @cbilgin

Lenguaje dominante
Python
Estrellas
5.2k
Forks
1.2k
Merge medio
2 d 37 min
PR fusionados (30 d)
693

Preparar el entorno

Primeros pasos

  1. Lee el issue completo y luego la guía de contribución del proyecto.
  2. Comenta en el issue que vas a ocuparte — evita que dos personas hagan lo mismo.
  3. Haz un fork del repositorio y trabaja en una rama.
  4. Abre un pull request que haga referencia al número del issue.

Más de pytorch/executorch

Todos los issues de pytorch/executorch

Issues similares

Más issues de Python

Recibe los nuevos issues en tu correo

Un resumen breve de issues de GitHub para principiantes.