Add inference speed optimization

models/rfd3/configs/inference_engine/rfdiffusion3.yaml

@@ -66,3 +66,6 @@ dump_trajectories: False
 align_trajectory_structures: False
 prevalidate_inputs: False
 low_memory_mode: False # False for standard mode, True for memory efficient tokenization mode
+
+# torch.compile the hot diffusion submodules. Cuts steady-state rollout time by ~1.6x.
+compile_model: False

models/rfd3/src/rfd3/engine.py

@@ -72,6 +72,7 @@ class RFD3InferenceConfig:
     low_memory_mode: bool = (
         False  # False for standard mode, True for memory efficient tokenization mode
     )
+    compile_model: bool = False
 
     # Other:
     num_nodes: int = 1
@@ -160,6 +161,7 @@ def __init__(
         dump_trajectories: bool,
         align_trajectory_structures: bool,
         low_memory_mode: bool,
+        compile_model: bool = False,
         **kwargs,
     ):
         super().__init__(
@@ -202,6 +204,70 @@ def __init__(
             # HACK: Set attribute to the diffusion module
             os.environ["RFD3_LOW_MEMORY_MODE"] = "1"
 
+        self.compile_model = compile_model
+        self.compiled_ = False
+
+    # Submodules of the diffusion module that are pure tensor code and get re-entered
+    # once (encoder) or twice (the rest, via recycling) on every diffusion step.
+    _COMPILE_TARGETS = (
+        "encoder",
+        "diffusion_token_encoder",
+        "diffusion_transformer",
+        "decoder",
+    )
+
+    def initialize(self):
+        cfg = super().initialize()
+        if self.compile_model and not self.compiled_:
+            self._compile_diffusion_submodules()
+            self.compiled_ = True
+        return cfg
+
+    def _compile_diffusion_submodules(self) -> None:
+        """Wrap the hot diffusion submodules in `torch.compile`.
+
+        The rollout is dominated by many small kernels (~14k launches per diffusion
+        step), so inductor's fusion is worth roughly 1.6x on steady-state rollout time
+        at both small and large diffusion batch sizes. It costs a one-off warmup
+        (~85-90s warm cache, ~210s cold) charged to the first diffusion step, which is
+        why this is opt-in rather than the default: it is a loss for a single rollout
+        and a win from roughly the second onwards.
+        """
+        model = self.trainer.state["model"]
+
+        # Unwrap _FabricModule / DistributedDataParallel / EMA to reach the RFD3 net
+        net = getattr(model, "_forward_module", model)
+        for _ in range(5):
+            if hasattr(net, "diffusion_module"):
+                break
+            for attr in ("module", "shadow", "model"):
+                if hasattr(net, attr):
+                    net = getattr(net, attr)
+                    break
+        else:
+            ranked_logger.warning(
+                "Could not locate the diffusion module; skipping torch.compile."
+            )
+            return
+
+        diffusion_module = net.diffusion_module
+        for name in self._COMPILE_TARGETS:
+            submodule = getattr(diffusion_module, name, None)
+            if submodule is None:
+                continue
+            # dynamic=False: L and I are fixed for a given specification, so we want
+            # static-shape kernels rather than dynamic-shape guards.
+            setattr(
+                diffusion_module,
+                name,
+                torch.compile(submodule, dynamic=False),
+            )
+        ranked_logger.info(
+            "torch.compile enabled for diffusion submodules "
+            f"({', '.join(self._COMPILE_TARGETS)}). Expect a one-off warmup on the "
+            "first diffusion step."
+        )
+
     # The base `run` is positional (`inputs, *_`); this engine deliberately exposes a
     # richer keyword-only API, so the override is intentionally LSP-incompatible.
     def run(  # type: ignore[override]

models/rfd3/src/rfd3/model/layers/attention.py

@@ -1,4 +1,5 @@
 import math
+import os
 from math import sqrt
 
 import torch
@@ -336,17 +337,34 @@ def do_attention(Q_L, C_L, P_LL):
                 else:
                     # Original full P_LL path
                     b = self.to_b(P_LL)
-                    attn_out = sparse_pairbias_attention(
-                        Q=q,
-                        K=k,
-                        V=v,
-                        B=b,
-                        G=g,
-                        gather_bias=True,
-                        indices=indices,
-                        H=self.n_head,
-                        full=full,
-                    )  # [D, L, c]
+                    if use_dense_sdpa_pairbias(
+                        Q=q, indices=indices, full=full, H=self.n_head
+                    ):
+                        # Mathematically equivalent to the sparse path below, but avoids
+                        # materializing the (D, L, k, c) gathered K/V tensors. See the
+                        # function docstring for when this is (and isn't) the faster
+                        # choice.
+                        attn_out = dense_sdpa_pairbias_attention(
+                            Q=q,
+                            K=k,
+                            V=v,
+                            B=b,
+                            G=g,
+                            indices=indices,
+                            H=self.n_head,
+                        )  # [D, L, c]
+                    else:
+                        attn_out = sparse_pairbias_attention(
+                            Q=q,
+                            K=k,
+                            V=v,
+                            B=b,
+                            G=g,
+                            gather_bias=True,
+                            indices=indices,
+                            H=self.n_head,
+                            full=full,
+                        )  # [D, L, c]
 
             # Output projection (from adaLN-Zero)
             Q_L = self.to_o(attn_out)
@@ -373,6 +391,121 @@ def do_attention(Q_L, C_L, P_LL):
 ######################################################################################
 
 
+_DENSE_PATH_REPORTED = False
+
+
+def _report_attention_path(chosen: str, reason: str, shape: str) -> None:
+    """Log the attention path once per process so it is visible in run logs."""
+    global _DENSE_PATH_REPORTED
+    if _DENSE_PATH_REPORTED:
+        return
+    _DENSE_PATH_REPORTED = True
+    ranked_logger.info(
+        f"Atom attention path: {chosen} ({reason}) [{shape}]. "
+        "Set RFD3_DENSE_SDPA_ATTENTION=0 to force the original sparse path."
+    )
+
+
+def use_dense_sdpa_pairbias(Q, indices, full, H) -> bool:
+    """Decide whether to run the dense-SDPA path instead of the sparse gather path.
+
+    Only used at inference (the sparse path is kept for training, where its activation
+    memory scales as O(L * k) rather than O(L^2) and matters for the backward pass).
+
+    Set `RFD3_DENSE_SDPA_ATTENTION=0` to force the original sparse path.
+    """
+    D, L, _ = Q.shape
+    k = indices.shape[-1]
+    shape = f"D={D} L={L} k={k} H={H}"
+
+    if full:
+        _report_attention_path("SPARSE", "full=True", shape)
+        return False
+    if torch.is_grad_enabled():
+        _report_attention_path("SPARSE", "grad enabled (training)", shape)
+        return False
+    if os.environ.get("RFD3_DENSE_SDPA_ATTENTION", "1") != "1":
+        _report_attention_path("SPARSE", "disabled via env var", shape)
+        return False
+    if os.environ.get("RFD3_LOW_MEMORY_MODE", "0") == "1":
+        _report_attention_path("SPARSE", "low memory mode", shape)
+        return False
+    if not Q.is_cuda:
+        _report_attention_path("SPARSE", "not on CUDA", shape)
+        return False
+    # The dense path is only a win while the L x L attention/bias tensors stay small
+    # relative to the O(L * k) gathers it replaces, and while they fit comfortably in
+    # VRAM. Two (D, H, L, L) bf16 tensors are live at peak (bias + attention scores).
+    if L <= k:
+        _report_attention_path("SPARSE", f"L={L} <= k={k}", shape)
+        return False
+    free_bytes, _ = torch.cuda.mem_get_info(Q.device)
+    # Peak extra allocation is the (D, H, L, L) bf16 bias plus one transient of the same
+    # size from the masked_fill; SDPA itself streams the scores rather than storing them.
+    dense_bytes = 2 * D * H * L * L * 2
+    if dense_bytes >= 0.35 * free_bytes:
+        _report_attention_path(
+            "SPARSE",
+            f"needs {dense_bytes / 2**30:.1f} GiB of {free_bytes / 2**30:.1f} GiB free",
+            shape,
+        )
+        return False
+    _report_attention_path("DENSE-SDPA", f"{dense_bytes / 2**30:.1f} GiB bias", shape)
+    return True
+
+
+def dense_sdpa_pairbias_attention(Q, K, V, B, indices, H, G=None):
+    """Pair-bias attention over the same key set as `sparse_pairbias_attention`.
+
+    Rather than gathering K/V into (D, L, k, c), this builds the (D, H, L, L) additive
+    bias, sets non-selected keys to -inf, and defers to `scaled_dot_product_attention`.
+    The softmax gives zero weight to the masked keys, so the result matches the sparse
+    path up to floating-point summation order.
+
+    Q, K, V: (D, L, c) | B: (L, L, H) or (D, L, L, H) | G: (D, L, c) | indices: (D, L, k)
+    Returns (D, L, c).
+    """
+    D, L, c = Q.shape
+    d = c // H
+
+    # q/k are promoted to fp32 by the (auto-cast-excluded) kq RMSNorms; SDPA needs a
+    # single dtype for q/k/v, so settle on the value/bias dtype (bf16 under AMP).
+    dtype = V.dtype
+    q = Q.to(dtype).reshape(D, L, H, d).transpose(1, 2)  # (D, H, L, d)
+    k = K.to(dtype).reshape(D, L, H, d).transpose(1, 2)
+    v = V.reshape(D, L, H, d).transpose(1, 2)
+
+    # Additive pair bias -> (D, H, L, L)
+    if B.ndim == 3:  # (L, L, H), shared across the diffusion batch
+        bias = B.permute(2, 0, 1).unsqueeze(0).expand(D, -1, -1, -1)
+    elif B.ndim == 4:  # (D, L, L, H)
+        bias = B.permute(0, 3, 1, 2)
+    else:
+        raise ValueError(f"Unexpected pair-bias shape {tuple(B.shape)}")
+
+    # Mask out every key that the sparse path would not have gathered. The sparse path
+    # broadcasts indices over the batch via advanced indexing; scatter_ does not, so a
+    # shared (1, L, k) index set has to be expanded explicitly or batches 1.. would end
+    # up fully masked (-inf everywhere -> NaN out of the softmax).
+    indices = indices.to(torch.int64)
+    if indices.shape[0] != D:
+        if indices.shape[0] != 1:
+            raise ValueError(
+                f"indices batch dim {indices.shape[0]} is neither 1 nor {D}"
+            )
+        indices = indices.expand(D, -1, -1)
+    valid = torch.zeros((D, L, L), dtype=torch.bool, device=Q.device)
+    valid.scatter_(2, indices, True)
+    bias = bias.to(dtype).masked_fill(~valid.unsqueeze(1), float("-inf"))
+
+    attn_out = F.scaled_dot_product_attention(q, k, v, attn_mask=bias)  # (D, H, L, d)
+
+    if G is not None:
+        attn_out = attn_out * G.reshape(D, L, H, d).transpose(1, 2)
+
+    return attn_out.transpose(1, 2).reshape(D, L, c).contiguous()
+
+
 def sparse_pairbias_attention(
     Q, K, V, B, indices, H, gather_bias=True, G=None, full=False
 ):