perf(pi052): sync-free denoise loop + opt-in FlashRT FP8 MLP (#3870)

* perf(pi052): sync-free denoise loop (precompute timesteps, device masks, KV crop)

Remove the per-denoise-step CPU->GPU syncs and the per-step KV-cache deepcopy
from action sampling:

- precompute the timestep schedule once instead of rebuilding a tensor from a
  Python float every step (torch.tensor(time, device=cuda) is a host sync);
- build the constant [1, 0, ...] suffix attention mask on-device instead of
  torch.tensor(python_list, device=cuda);
- drop the per-step copy.deepcopy of the prefix KV cache: the expert forward
  appends the suffix K/V in place, so crop back to the prefix length afterwards
  (prefix K/V are read-only, so this is exact and the loop stays one graph).

Bit-exact: action max|delta|=0 vs the previous implementation; no API change.

* feat(pi052): optional FlashRT FP8 Gemma/SigLIP MLP swap (opt-in)

Opt-in (config use_flashrt_fp8_mlp) swap of the Gemma GeGLU + SigLIP GELU MLPs
to the FlashRT FP8 Hub kernels. When the flag is set, the first inference
calibrates static activation scales on that observation and swaps the MLP
modules in place (re-entry guarded); graceful BF16 fallback if the kernels are
unavailable.

Calibration follows the FlashRT contract: the FP8 modules are swapped in first,
then a single forward measures each GEMM's input/hidden amax on the
already-quantized (FP8-propagated) activations, with the preceding fixed
RMSNorm weight (1+w) folded into the GEMM and scale = amax/448 * 1.05.

On pi05_libero_pytorch (RTX 5090, torch.compile): ~1.91x end-to-end
(89.4 -> 46.7 ms) with the sync-free loop, action cos vs BF16 ~0.999
(maxdiff ~0.03) over real LIBERO frames.

src/lerobot/policies/pi052/configuration_pi052.py

@@ -206,6 +206,15 @@ class PI052Config(PI05Config):
     # at -4.5% step time on H100 (bench job 22161421); peak memory
     # unchanged. ``fused_linear_cross_entropy`` ships separately via
     # ``_shifted_lin_ce`` / ``_fast_lin_ce``.
+    use_flashrt_fp8_mlp: bool = False
+    """Opt-in: swap every Gemma GeGLU MLP (action expert + prefix LM) and the
+    SigLIP vision MLP to FlashRT fused FP8 kernels (Hugging Face Kernel Hub
+    ``flashrt/*``). The swap needs a one-time activation calibration on a real
+    observation, so it is applied explicitly via
+    ``PI052Policy.apply_flashrt_fp8_mlp(batch)`` after loading (not at build).
+    Degrades gracefully to BF16 if ``kernels`` / the FlashRT packages are
+    missing. Default off keeps behaviour identical to the BF16 path."""
+
     use_hf_kernels: bool = True
     """Deprecated. Liger HF kernels are patched unconditionally by
     ``_enable_hf_kernels`` — this field is retained as a no-op for

src/lerobot/policies/pi052/flashrt_fp8.py

@@ -0,0 +1,272 @@
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Optional FP8 MLP swap for PI052 using the FlashRT Hugging Face Kernel Hub.
+
+Replaces every Gemma GeGLU MLP (action expert + prefix language model) with the
+fused ``fp8_geglu_mlp_bf16`` kernel and the SigLIP vision-tower MLP with
+``fp8_gelu_mlp_bf16``. Static activation scales are calibrated once on a real
+observation; weights are quantized once. This is opt-in and degrades gracefully
+to the BF16 path if ``kernels`` or the FlashRT packages are unavailable.
+
+Use:
+    policy = PI052Policy.from_pretrained(...)
+    batch = preprocessor(observation)        # one representative observation
+    policy.apply_flashrt_fp8_mlp(batch)       # calibrate + swap in place
+"""
+
+from __future__ import annotations
+
+import logging
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+logger = logging.getLogger(__name__)
+
+_FP8_MAX = 448.0
+
+
+def _roundtrip_fp8(x: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
+    """Quantize->dequantize an activation through FP8 E4M3 at ``scale`` (f32)."""
+    q = torch.clamp(x.float() / scale.float(), -_FP8_MAX, _FP8_MAX).to(torch.float8_e4m3fn)
+    return q.float() * scale.float()
+_SWIGLU_REPO = "flashrt/flashrt-fp8-swiglu-ffn"
+_GELU_REPO = "flashrt/flashrt-fp8-ffn"
+_GEMM_REPO = "flashrt/flashrt-gemm-epilogues"
+
+
+def _get_kernel(repo: str):
+    """Load a FlashRT Hub package (cached). Returns None if unavailable."""
+    from kernels import get_kernel
+
+    return get_kernel(repo, version=1)
+
+
+def _quantize_fp8(weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    scale = max(weight.detach().float().abs().max().item(), 1e-12) / _FP8_MAX
+    fp8 = torch.clamp(weight.float() / scale, -_FP8_MAX, _FP8_MAX).to(torch.float8_e4m3fn)
+    return fp8.contiguous(), torch.tensor([scale], dtype=torch.float32)
+
+
+def _static_scale(amax: float, safety: float) -> torch.Tensor:
+    return torch.tensor([max(amax, 1e-12) / _FP8_MAX * safety], dtype=torch.float32)
+
+
+class _FlashRTGeGLU(nn.Module):
+    """FP8 drop-in for a Gemma GeGLU MLP (gate/up/down, gelu_pytorch_tanh, no bias)."""
+
+    def __init__(self, mlp, in_amax, hid_amax, ffn_ops, quant_ops, safety, fuse_weight=None):
+        super().__init__()
+        self.ffn_ops = ffn_ops
+        self.quant_ops = quant_ops
+        self.in_features = mlp.gate_proj.weight.shape[1]
+        device = mlp.gate_proj.weight.device
+        gate_up = torch.cat([mlp.gate_proj.weight, mlp.up_proj.weight], dim=0).float()
+        # Fold the preceding RMSNorm weight (1 + w) into the gate/up GEMM and feed
+        # the kernel channel_scale = 1/(1+w). This is exact (it just moves the
+        # per-channel (1+w) from the activation to the weight) and is what keeps
+        # FP8 accurate: the normed activation rms(x) is uniform, while
+        # rms(x)*(1+w) has per-channel outliers that per-tensor FP8 quantizes
+        # poorly. Mirrors the FlashRT runtime (norm runs with ones, weight folds
+        # 1+w). Only the fixed-weight (non-adaptive) norms fold; adaptive-RMSNorm
+        # layers pass fuse_weight=None (channel_scale = ones).
+        if fuse_weight is not None:
+            f = (1.0 + fuse_weight.detach().float())
+            gate_up = gate_up * f[None, :]
+            channel_scale = (1.0 / f).to(torch.bfloat16)
+        else:
+            channel_scale = torch.ones(self.in_features, dtype=torch.bfloat16)
+        gate_up_fp8, gate_up_scale = _quantize_fp8(gate_up)
+        down_fp8, down_scale = _quantize_fp8(mlp.down_proj.weight)
+        self.register_buffer("gate_up_fp8", gate_up_fp8.to(device))
+        self.register_buffer("down_fp8", down_fp8.to(device))
+        self.register_buffer("gate_up_scale", gate_up_scale.to(device))
+        self.register_buffer("down_scale", down_scale.to(device))
+        self.register_buffer("input_scale", _static_scale(in_amax, safety).to(device))
+        self.register_buffer("hidden_scale", _static_scale(hid_amax, safety).to(device))
+        self.register_buffer("channel_scale", channel_scale.to(device))
+        self.safety = safety
+        self.calibrating = False
+        self._ia = 0.0
+        self._ha = 0.0
+
+    def _calibrate_step(self, x):
+        # FP8-propagated calibration (FlashRT contract): measure input/hidden amax
+        # on the live (already-FP8-upstream) activation, running-max across steps.
+        flat = x.reshape(-1, self.in_features).to(torch.bfloat16)
+        xq = flat.float() * self.channel_scale.float()
+        self._ia = max(self._ia, xq.abs().max().item())
+        self.input_scale.copy_(_static_scale(self._ia, self.safety).to(self.input_scale.device))
+        xdq = _roundtrip_fp8(xq, self.input_scale)
+        wdq = self.gate_up_fp8.float() * self.gate_up_scale.float()
+        gate, up = (xdq @ wdq.t()).chunk(2, dim=-1)
+        hidden = F.gelu(gate, approximate="tanh") * up
+        self._ha = max(self._ha, hidden.abs().max().item())
+        self.hidden_scale.copy_(_static_scale(self._ha, self.safety).to(self.hidden_scale.device))
+
+    def forward(self, x):
+        if self.calibrating:
+            self._calibrate_step(x)
+        shape = x.shape
+        flat = x.reshape(-1, self.in_features).to(torch.bfloat16)
+        x_fp8 = self.quant_ops.channel_scale_quantize_fp8_static_bf16(flat, self.channel_scale, self.input_scale)
+        out = self.ffn_ops.fp8_geglu_mlp_bf16(
+            x_fp8, self.gate_up_fp8, self.down_fp8,
+            self.input_scale, self.gate_up_scale, self.hidden_scale, self.down_scale,
+        )
+        return out.reshape(shape)
+
+
+class _FlashRTGeluMLP(nn.Module):
+    """FP8 drop-in for a SigLIP MLP (fc1 -> gelu_tanh -> fc2, with bias)."""
+
+    def __init__(self, mlp, in_amax, hid_amax, ffn_ops, quant_ops, safety):
+        super().__init__()
+        self.ffn_ops = ffn_ops
+        self.quant_ops = quant_ops
+        self.in_features = mlp.fc1.weight.shape[1]
+        self.out_features = mlp.fc2.weight.shape[0]
+        device = mlp.fc1.weight.device
+        up_fp8, up_scale = _quantize_fp8(mlp.fc1.weight)
+        down_fp8, down_scale = _quantize_fp8(mlp.fc2.weight)
+        self.register_buffer("up_fp8", up_fp8.to(device))
+        self.register_buffer("down_fp8", down_fp8.to(device))
+        self.register_buffer("up_scale", up_scale.to(device))
+        self.register_buffer("down_scale", down_scale.to(device))
+        self.register_buffer("up_bias", mlp.fc1.bias.detach().to(torch.bfloat16))
+        self.register_buffer("down_bias", mlp.fc2.bias.detach().to(torch.bfloat16))
+        self.register_buffer("input_scale", _static_scale(in_amax, safety).to(device))
+        self.register_buffer("hidden_scale", _static_scale(hid_amax, safety).to(device))
+        self.register_buffer("channel_scale", torch.ones(self.in_features, device=device, dtype=torch.bfloat16))
+        self.safety = safety
+        self.calibrating = False
+        self._ia = 0.0
+        self._ha = 0.0
+
+    def _calibrate_step(self, x):
+        flat = x.reshape(-1, self.in_features).to(torch.bfloat16)
+        self._ia = max(self._ia, flat.float().abs().max().item())
+        self.input_scale.copy_(_static_scale(self._ia, self.safety).to(self.input_scale.device))
+        xdq = _roundtrip_fp8(flat.float(), self.input_scale)
+        hid = (xdq @ (self.up_fp8.float() * self.up_scale.float()).t()) + self.up_bias.float()
+        hid = F.gelu(hid, approximate="tanh")
+        self._ha = max(self._ha, hid.abs().max().item())
+        self.hidden_scale.copy_(_static_scale(self._ha, self.safety).to(self.hidden_scale.device))
+
+    def forward(self, x):
+        if self.calibrating:
+            self._calibrate_step(x)
+        shape = x.shape
+        dtype = x.dtype
+        flat = x.reshape(-1, self.in_features).to(torch.bfloat16)
+        x_fp8 = self.quant_ops.channel_scale_quantize_fp8_static_bf16(flat, self.channel_scale, self.input_scale)
+        out = self.ffn_ops.fp8_gelu_mlp_bf16(
+            x_fp8, self.up_fp8, self.up_bias, self.down_fp8, self.down_bias,
+            self.input_scale, self.up_scale, self.hidden_scale, self.down_scale,
+        )
+        return out.reshape(*shape[:-1], self.out_features).to(dtype)
+
+
+def _siglip_mlps(model) -> list:
+    tower = model.paligemma_with_expert.paligemma.model.vision_tower
+    return [m for _, m in tower.named_modules() if type(m).__name__ == "SiglipMLP"]
+
+
+def _run_forward(policy, batches) -> None:
+    """Run predict_action_chunk in eager mode (a compiled ``sample_actions``
+    bypasses the Python module forwards, so drop it for the calibration pass)."""
+    model = policy.model
+    saved = {name: vars(model).pop(name) for name in ("sample_actions", "forward") if name in vars(model)}
+    with torch.inference_mode():
+        for batch in batches:
+            policy.predict_action_chunk({k: (v.clone() if torch.is_tensor(v) else v) for k, v in batch.items()})
+    torch.cuda.synchronize()
+    vars(model).update(saved)
+
+
+def _fixed_norm_weight(norm):
+    """RMSNorm (1+w) fold weight if ``norm`` is fixed-weight; None if adaptive."""
+    return norm.weight if getattr(norm, "dense", None) is None else None
+
+
+def _fp8_supported(device) -> bool:
+    """FP8 E4M3 tensor cores require CUDA SM >= 8.9 (Ada / Hopper / Blackwell).
+    Older GPUs (e.g. A100 SM 8.0) and CPU have no FP8 path, so the kernels would
+    fail at runtime — gate here and keep BF16."""
+    if device.type != "cuda" or not torch.cuda.is_available():
+        return False
+    major, minor = torch.cuda.get_device_capability(device)
+    return (major, minor) >= (8, 9)
+
+
+def apply_fp8_mlp(policy, batch, *, safety: float = 1.05) -> bool:
+    """Swap every Gemma GeGLU MLP and SigLIP GELU MLP to the FlashRT FP8 kernels.
+
+    Calibration is FP8-propagated (the FlashRT contract): the FP8 modules are
+    swapped in first, then a single forward on one representative frame measures
+    each GEMM's input/hidden amax on the *already-quantized* activations it sees
+    at runtime (not on a clean BF16 forward), running-max across denoise steps.
+    The preceding fixed RMSNorm weight (1+w) is folded into the GEMM so the
+    quantized activation is the uniform rms(x); adaptive-RMSNorm inputs (action
+    expert) do not fold. ``scale = amax/448 * safety``. A single frame is enough.
+    Returns True on success, False (no-op, BF16 kept) if the device lacks FP8
+    support or the kernels are unavailable.
+    """
+    device = next(policy.parameters()).device
+    if not _fp8_supported(device):
+        logger.warning(
+            "PI052: device %s has no FP8 (E4M3) support (needs CUDA SM>=8.9); keeping BF16.",
+            device,
+        )
+        return False
+    batches = batch if isinstance(batch, (list, tuple)) else [batch]
+    try:
+        ffn_ops = _get_kernel(_SWIGLU_REPO)
+        gelu_ops = _get_kernel(_GELU_REPO)
+        quant_ops = _get_kernel(_GEMM_REPO)
+    except Exception as exc:  # noqa: BLE001
+        logger.warning("PI052: FlashRT FP8 kernels unavailable (%s); keeping BF16.", exc)
+        return False
+
+    model = policy.model
+    calibrating = []
+
+    gemma_layers = (
+        list(model.paligemma_with_expert.gemma_expert.model.layers)
+        + list(model.paligemma_with_expert.paligemma.model.language_model.layers)
+    )
+    for layer in gemma_layers:
+        fw = _fixed_norm_weight(layer.post_attention_layernorm)
+        layer.mlp = _FlashRTGeGLU(layer.mlp, 1.0, 1.0, ffn_ops, quant_ops, safety, fuse_weight=fw).to(device)
+        calibrating.append(layer.mlp)
+
+    siglip = _siglip_mlps(model)
+    for mlp_parent in model.paligemma_with_expert.paligemma.model.vision_tower.vision_model.encoder.layers:
+        mlp_parent.mlp = _FlashRTGeluMLP(mlp_parent.mlp, 1.0, 1.0, gelu_ops, quant_ops, safety).to(device)
+        calibrating.append(mlp_parent.mlp)
+
+    # FP8-propagated calibration: one forward with every module in calibrate mode.
+    for m in calibrating:
+        m.calibrating = True
+    _run_forward(policy, batches)
+    for m in calibrating:
+        m.calibrating = False
+
+    logger.info(
+        "PI052: FlashRT FP8 enabled (%d Gemma + %d SigLIP MLPs).",
+        len(gemma_layers), len(siglip),
+    )
+    return True

src/lerobot/policies/pi052/modeling_pi052.py

@@ -38,7 +38,6 @@ for the LM head.
 from __future__ import annotations
 
 import builtins
-import copy
 import logging
 import math
 import types
@@ -416,8 +415,13 @@ class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
 
         embs = torch.cat(embs, dim=1)
         pad_masks = torch.cat(pad_masks, dim=1)
-        att_masks = torch.tensor(att_masks, dtype=embs.dtype, device=embs.device)
-        att_masks = att_masks[None, :].expand(bsize, len(att_masks))
+        # The suffix mask is the constant [1, 0, ..., 0]; build it on-device
+        # rather than via torch.tensor(python_list, device=cuda), which is a
+        # host->device sync on every denoise step.
+        n = len(att_masks)
+        att_masks = torch.zeros(n, dtype=embs.dtype, device=embs.device)
+        att_masks[0] = 1
+        att_masks = att_masks[None, :].expand(bsize, n)
 
         return embs, pad_masks, att_masks, adarms_cond
 
@@ -517,10 +521,17 @@ class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
 
         dt = -1.0 / num_steps
 
+        # Precompute the whole timestep schedule on-device once, instead of
+        # rebuilding a tensor from a Python float every step
+        # (``torch.tensor(time, device=cuda)`` is a host->device sync ×num_steps).
+        times = torch.tensor(
+            [1.0 + s * dt for s in range(num_steps)], dtype=torch.float32, device=device
+        )
+
         x_t = noise
         for step in range(num_steps):
-            time = 1.0 + step * dt
-            time_tensor = torch.tensor(time, dtype=torch.float32, device=device).expand(bsize)
+            time = 1.0 + step * dt  # Python float kept for the RTC branch below
+            time_tensor = times[step].expand(bsize)
 
             def denoise_step_partial_call(input_x_t, current_timestep=time_tensor):
                 return self.denoise_step(
@@ -579,7 +590,12 @@ class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
         )
         self.paligemma_with_expert.gemma_expert.model.config._attn_implementation = "eager"  # noqa: SLF001
 
-        past_key_values = copy.deepcopy(past_key_values)
+        # The expert forward appends the suffix K/V to the prefix cache in-place
+        # (GemmaAttention.update runs even with use_cache=False), so each step
+        # must start from a prefix-only cache. Instead of deep-copying the whole
+        # cache every step, let it append and crop back to the prefix length
+        # afterwards (the prefix K/V are read-only, so this is exact and keeps
+        # the loop a single graph).
         outputs_embeds, _ = self.paligemma_with_expert.forward(
             attention_mask=full_att_2d_masks_4d,
             position_ids=position_ids,
@@ -588,6 +604,7 @@ class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
             use_cache=False,
             adarms_cond=[None, adarms_cond],
         )
+        past_key_values.crop(prefix_len)
 
         suffix_out = outputs_embeds[1]
         suffix_out = suffix_out[:, -self.config.chunk_size :]
@@ -1063,6 +1080,18 @@ class PI052Policy(PreTrainedPolicy):
         # subtask can be held across several chunks (see subtask_replan_steps).
         self._subtask_chunk_counter = 0
 
+    def apply_flashrt_fp8_mlp(self, batch: dict[str, Tensor], *, safety: float = 1.05) -> bool:
+        """Opt-in: swap every Gemma + SigLIP MLP to FlashRT fused FP8 kernels.
+
+        Calibrates static activation scales once on ``batch`` (one representative
+        observation, already through the preprocessor) and swaps the MLP modules
+        in place. Returns False (no-op, BF16 kept) if the kernels are missing.
+        Gated by ``config.use_flashrt_fp8_mlp`` — see flashrt_fp8.py.
+        """
+        from .flashrt_fp8 import apply_fp8_mlp  # noqa: PLC0415
+
+        return apply_fp8_mlp(self, batch, safety=safety)
+
     # ------------------------------------------------------------------
     # Head unfreeze helper
     # ------------------------------------------------------------------
@@ -2360,6 +2389,13 @@ class PI052Policy(PreTrainedPolicy):
         """Predict a chunk of actions given environment observations."""
         self.eval()
 
+        # Opt-in FlashRT FP8: calibrate static scales on the first real observation
+        # and swap the MLPs in place. Guard set before the call so the calibration
+        # forward (which re-enters predict_action_chunk) does not recurse.
+        if self.config.use_flashrt_fp8_mlp and not getattr(self, "_fp8_applied", False):
+            self._fp8_applied = True
+            self.apply_flashrt_fp8_mlp(batch)
+
         # Prepare inputs
         images, img_masks = self._preprocess_images(batch)
         tokens, masks = batch[f"{OBS_LANGUAGE_TOKENS}"], batch[f"{OBS_LANGUAGE_ATTENTION_MASK}"]