From 33f0414733ef836e9cf1f632a394278ffcf93d80 Mon Sep 17 00:00:00 2001
From: pepijn223 <pepijn223@hf.co>
Date: Thu, 9 Jul 2026 13:36:48 +0200
Subject: [PATCH] feat(runtime): add --fp8 flag to enable PI052 FlashRT FP8
 MLPs

Wire the existing (but previously unreachable from the runtime) PI052
FlashRT FP8 MLP swap into the language runtime. --fp8 sets
config.use_flashrt_fp8_mlp before load; the policy calibrates and swaps
every Gemma + SigLIP MLP to fused FP8 on its first predict_action_chunk.
Ignored with a warning for policies without the flag (PI052 only).

Measured ~1.12x faster action-chunk inference (124 -> 111 ms) on an
RTX 5090; needs the `kernels` package (pin <0.13 for transformers) and
CUDA SM>=8.9, else it degrades to BF16.

Co-authored-by: Cursor <cursoragent@cursor.com>
---
 src/lerobot/runtime/cli.py | 25 +++++++++++++++++++++++++
 1 file changed, 25 insertions(+)

diff --git a/src/lerobot/runtime/cli.py b/src/lerobot/runtime/cli.py
index 746bf2bb4..eac66c619 100644
--- a/src/lerobot/runtime/cli.py
+++ b/src/lerobot/runtime/cli.py
@@ -383,6 +383,17 @@ def _parse_args(argv: list[str] | None = None, *, prog: str | None = None) -> ar
             "decode and a meaningless memory line."
         ),
     )
+    p.add_argument(
+        "--fp8",
+        action="store_true",
+        help=(
+            "PI052 only: enable FlashRT FP8 MLP kernels. Calibrates on the "
+            "first inference call and swaps every Gemma + SigLIP MLP to fused "
+            "FP8 (needs the `kernels` package and CUDA SM>=8.9; degrades to "
+            "BF16 otherwise). Speeds up the forward pass at a small accuracy "
+            "cost — see policies/pi052/flashrt_fp8.py."
+        ),
+    )
     p.add_argument(
         "--subtask_chunks_per_gen",
         type=int,
@@ -461,6 +472,7 @@ def _load_policy_and_preprocessor(
     dataset_repo_id: str | None,
     *,
     load_processors_from_checkpoint: bool = False,
+    fp8: bool = False,
 ) -> tuple[Any, Any, Any, Any]:
     """Load a policy checkpoint (local path or Hub repo id).
 
@@ -490,6 +502,18 @@ def _load_policy_and_preprocessor(
     if getattr(cfg, "gradient_checkpointing", False):
         cfg.gradient_checkpointing = False
 
+    # Opt-in FP8: only PI052 has the FlashRT MLP swap. The policy calibrates and
+    # swaps on its first predict_action_chunk when this flag is set.
+    if fp8:
+        if hasattr(cfg, "use_flashrt_fp8_mlp"):
+            cfg.use_flashrt_fp8_mlp = True
+            logging.info("[runtime] FP8 MLP kernels enabled (FlashRT) for %s", cfg.type)
+        else:
+            logging.warning(
+                "[runtime] --fp8 ignored: %s has no use_flashrt_fp8_mlp (PI052 only).",
+                cfg.type,
+            )
+
     ds_meta = None
     preprocessor = None
     postprocessor = None
@@ -1529,6 +1553,7 @@ def run(
         args.policy_path,
         args.dataset_repo_id,
         load_processors_from_checkpoint=sim_mode,
+        fp8=args.fp8,
     )
 
     policy_type = getattr(policy.config, "type", None)