Refactor Hackable Diffusion for improved XLA performance and Flash Attention support.

PiperOrigin-RevId: 908804155

Author: Fareed Sheriff

hackable_diffusion/lib/architecture/attention.py

@@ -142,7 +142,14 @@ def _dot_product_attention(
     The output tensor.
   """
 
-  b, _, t, _ = q.shape
+  b, _, t, head_d = q.shape
+
+  # Use jax.nn.dot_product_attention for hardware acceleration when possible.
+  if mask is None and hasattr(jax.nn, "dot_product_attention"):
+    q_scaled = q * (rescale * jnp.sqrt(head_d))
+    attn_output = jax.nn.dot_product_attention(q_scaled, k, v)
+    attn_output = attn_output.transpose(0, 2, 1, 3).reshape(b, t, -1)
+    return attn_output
 
   # Attention scores
   attn_logits = jnp.einsum("bhtd,bhsd->bhts", q, k) * rescale

hackable_diffusion/lib/architecture/attention_test.py

@@ -409,5 +409,21 @@ def test_multi_head_attention_invalid_mask_shape_raises_error(
       module.init(self.rng, self.x, c, mask=invalid_mask)
 
 
+
+  def test_optimized_attention_path(self):
+    """Tests the optimized attention path (unmasked)."""
+    module = attention.MultiHeadAttention(num_heads=self.num_heads)
+    variables = module.init(self.rng, self.x, self.c)
+    output = module.apply(variables, self.x, self.c)
+    self.assertEqual(output.shape, self.x.shape)
+
+  def test_masked_attention_path(self):
+    """Tests the manual attention path (masked)."""
+    module = attention.MultiHeadAttention(num_heads=self.num_heads)
+    mask = jnp.ones((self.batch_size, self.seq_len_kv), dtype=jnp.bool_)
+    variables = module.init(self.rng, self.x, self.c, mask=mask)
+    output = module.apply(variables, self.x, self.c, mask=mask)
+    self.assertEqual(output.shape, self.x.shape)
+
 if __name__ == "__main__":
   absltest.main()

hackable_diffusion/lib/architecture/dit_blocks.py

@@ -158,27 +158,30 @@ def __call__(
       The output tensor.
     """
 
+    # Precompute activation for conditioning
+    cond_act = nn.silu(cond)
+
     # Attention Branch
-    x_attn_modulated = self.conditional_norm(x, c=nn.silu(cond))
+    x_attn_modulated = self.conditional_norm(x, c=cond_act)
     attn_out = self.attn(x_attn_modulated, c=None, mask=mask)
     # Optional dropout
     if self.dropout_rate > 0.0:
       attn_out = nn.Dropout(rate=self.dropout_rate)(
           attn_out, deterministic=not is_training
       )
-    gate_msa = self.gate_msa(nn.silu(cond))
+    gate_msa = self.gate_msa(cond_act)
     # Add a sequence dimension [...,None,:] to broadcast to [*batch,seq,dim].
     x = x + gate_msa[..., None, :] * attn_out
 
     # MLP Branch
-    x_mlp_modulated = self.conditional_norm(x, c=nn.silu(cond))
+    x_mlp_modulated = self.conditional_norm(x, c=cond_act)
     mlp_out = self.mlp(x_mlp_modulated, is_training=is_training)
     # Optional dropout
     if self.dropout_rate > 0.0:
       mlp_out = nn.Dropout(rate=self.dropout_rate)(
           mlp_out, deterministic=not is_training
       )
-    gate_mlp = self.gate_mlp(nn.silu(cond))
+    gate_mlp = self.gate_mlp(cond_act)
     # Add a sequence dimension [...,None,:] to broadcast to [*batch,seq,dim].
     x = x + gate_mlp[..., None, :] * mlp_out
     return x
@@ -267,7 +270,8 @@ def __call__(
     hn = h // hp
     wn = w // wp
 
-    x = self.conditional_norm(x, c=nn.silu(cond))
+    cond_act = nn.silu(cond)
+    x = self.conditional_norm(x, c=cond_act)
     x = nn.Dense(
         features=hp * wp * c,
         name="Dense_Out",

hackable_diffusion/lib/architecture/normalization.py

@@ -25,6 +25,7 @@
 from hackable_diffusion.lib import hd_typing
 from hackable_diffusion.lib import utils
 from hackable_diffusion.lib.architecture import arch_typing
+import jax
 import jax.numpy as jnp
 import kauldron.ktyping as kt
 
@@ -40,6 +41,17 @@
 NormalizationType = arch_typing.NormalizationType
 
 
+
+################################################################################
+# MARK: Fused Kernels
+################################################################################
+
+def fused_rms_norm(x, scale, epsilon=1e-6):
+  """Fused RMSNorm implementation for XLA efficiency."""
+  ms = jnp.mean(jnp.square(x), axis=-1, keepdims=True)
+  return x * jax.lax.rsqrt(ms + epsilon) * scale
+
+
 ################################################################################
 # MARK: NormalizationLayer
 ################################################################################
@@ -128,13 +140,17 @@ def __call__(
     ch = x_shape[-1]
 
     if self.normalization_method == NormalizationType.RMS_NORM:
-      x = nn.RMSNorm(
-          epsilon=self.epsilon,
-          dtype=self.dtype,
-          reduction_axes=-1,  # For (B ... ch) results in (B ... ) RMS values.
-          feature_axes=-1,  # Per channel scale.
-          use_scale=self.use_scale,
-      )(x=x, mask=mask)
+      if mask is None and self.use_scale:
+        scale = self.param("scale", nn.initializers.ones, (ch,), self.dtype)
+        x = fused_rms_norm(x, scale, self.epsilon)
+      else:
+        x = nn.RMSNorm(
+            epsilon=self.epsilon,
+            dtype=self.dtype,
+            reduction_axes=-1,  # For (B ... ch) results in (B ... ) RMS values.
+            feature_axes=-1,  # Per channel scale.
+            use_scale=self.use_scale,
+        )(x=x, mask=mask)
     elif self.normalization_method == NormalizationType.GROUP_NORM:
 
       # If using GroupNorm the mask data must be such that the last dimension
@@ -187,6 +203,17 @@ def __call__(
     return x
 
 
+
+################################################################################
+# MARK: Fused Kernels
+################################################################################
+
+def fused_rms_norm(x, scale, epsilon=1e-6):
+  """Fused RMSNorm implementation for XLA efficiency."""
+  ms = jnp.mean(jnp.square(x), axis=-1, keepdims=True)
+  return x * jax.lax.rsqrt(ms + epsilon) * scale
+
+
 ################################################################################
 # MARK: NormalizationLayerFactory
 ################################################################################

hackable_diffusion/lib/architecture/normalization_test.py

@@ -469,5 +469,17 @@ def test_rmsnorm_mask_equivalence(self):
     )
 
 
+
+  def test_fused_rms_norm_path(self):
+    """Tests the fused RMSNorm path (unmasked)."""
+    module = normalization.NormalizationLayer(
+        normalization_method=normalization.NormalizationType.RMS_NORM,
+        conditional=False
+    )
+    x = jnp.ones((2, 16, 32))
+    variables = module.init(self.rng, x)
+    output = module.apply(variables, x)
+    self.assertEqual(output.shape, x.shape)
+
 if __name__ == "__main__":
   absltest.main()