From 3fa1e4f7541ff59eb07c063b9fc032f49b270f22 Mon Sep 17 00:00:00 2001 From: YanjingLi Date: Thu, 6 Aug 2026 23:30:20 -0700 Subject: [PATCH] dense attention and compile --- .../inference_engine/rfdiffusion3.yaml | 3 + models/rfd3/src/rfd3/engine.py | 66 ++++++++ .../rfd3/src/rfd3/model/layers/attention.py | 155 ++++++++++++++++-- 3 files changed, 213 insertions(+), 11 deletions(-) diff --git a/models/rfd3/configs/inference_engine/rfdiffusion3.yaml b/models/rfd3/configs/inference_engine/rfdiffusion3.yaml index cd468d5f..6cadde8b 100644 --- a/models/rfd3/configs/inference_engine/rfdiffusion3.yaml +++ b/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 diff --git a/models/rfd3/src/rfd3/engine.py b/models/rfd3/src/rfd3/engine.py index 1bd16cad..dea76ef6 100644 --- a/models/rfd3/src/rfd3/engine.py +++ b/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] diff --git a/models/rfd3/src/rfd3/model/layers/attention.py b/models/rfd3/src/rfd3/model/layers/attention.py index 31021c47..7690639e 100644 --- a/models/rfd3/src/rfd3/model/layers/attention.py +++ b/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 ):