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>
This commit is contained in:
pepijn223
2026-07-09 13:36:48 +02:00
parent 1e94a4f62d
commit 33f0414733
+25
View File
@@ -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)