Tag and save Pallas LSE statistics to eliminate attention recomputation entirely

Author: prishajain1

src/maxdiffusion/kernels/splash_attention/splash_attention_kernel.py

@@ -1264,6 +1264,7 @@ def _splash_attention_fwd(
       max_logit_value=max_logit_value,
   )
   logsumexp = stats["logsumexp"]  # save in the config base for the bwd pass
+  logsumexp = jax.checkpoint_name(logsumexp, "pallas_logsumexp")
   if config.use_base2_exp:  # for user, output values in natural base
     stats["logsumexp"] = stats["logsumexp"] / LOG2E
     stats["max_logits"] = stats["max_logits"] / LOG2E

src/maxdiffusion/models/flux/transformers/transformer_flux_flax.py

@@ -495,7 +495,7 @@ def setup(self):
     # 4. Force strict checkpointing on the Single Wrapper
     #RemattedSingleWrapper = nn.remat(ScannedSingleBlockWrapper, prevent_cse=True, policy=cp.checkpoint_dots_with_no_batch_dims)
     #RemattedSingleWrapper = nn.remat(ScannedSingleBlockWrapper, prevent_cse=True, policy=cp.offload_dot_with_no_batch_dims(offload_src="device", offload_dst="pinned_host"))
-    RemattedSingleWrapper = nn.remat(ScannedSingleBlockWrapper, prevent_cse=True, policy=cp.save_only_these_names("attn_output", "lin1_norm_hidden_states"))
+    RemattedSingleWrapper = nn.remat(ScannedSingleBlockWrapper, prevent_cse=True, policy=cp.save_only_these_names("attn_output", "lin1_norm_hidden_states", "pallas_logsumexp"))
 
     self.scanned_single_blocks = nn.scan(
         RemattedSingleWrapper,