Add dense padded attention kernel and use unsafe rng key for generation

Author: coolkp

padded_flash_attn.py

@@ -57,18 +57,18 @@ jit_initializers: True
 from_pt: True
 split_head_dim: True
 attention: 'flash' # Supported attention: dot_product, flash, cudnn_flash_te, ring
-flash_min_seq_length: 4096
+flash_min_seq_length: 0
 dropout: 0.1
 
 flash_block_sizes: {
-  "block_q" : 1024,
-  "block_kv_compute" : 256,
-  "block_kv" : 1024,
-  "block_q_dkv" : 1024,
-  "block_kv_dkv" : 1024,
-  "block_kv_dkv_compute" : 256,
-  "block_q_dq" : 1024,
-  "block_kv_dq" : 1024
+  "block_q" : 3024,
+  "block_kv_compute" : 1024,
+  "block_kv" : 2048,
+  "block_q_dkv" : 3024,
+  "block_kv_dkv" : 2048,
+  "block_kv_dkv_compute" : 2048,
+  "block_q_dq" : 3024,
+  "block_kv_dq" : 2048
 }
 # Use on v6e
 # flash_block_sizes: {
@@ -82,16 +82,16 @@ flash_block_sizes: {
 #   "block_kv_dq" : 2048
 # }
 # Use on v5p
-flash_block_sizes: {
-  "block_q" : 1024,
-  "block_kv_compute" : 256,
-  "block_kv" : 3072,
-  "block_q_dkv" : 1024,
-  "block_kv_dkv" : 3072,
-  "block_kv_dkv_compute" : 256,
-  "block_q_dq" : 1024,
-  "block_kv_dq" : 3072
-}
+# flash_block_sizes: {
+#   "block_q" : 3024,
+#   "block_kv_compute" : 1024,
+#   "block_kv" : 2048,
+#   "block_q_dkv" : 1024,
+#   "block_kv_dkv" : 3072,
+#   "block_kv_dkv_compute" : 256,
+#   "block_q_dq" : 1024,
+#   "block_kv_dq" : 3072
+# }
 # GroupNorm groups
 norm_num_groups: 32
 
@@ -152,7 +152,7 @@ mesh_axes: ['data', 'fsdp', 'tensor']
 logical_axis_rules: [
                       ['batch', 'data'],
                       ['activation_batch', 'data'],
-                      ['activation_self_attn_heads', ['fsdp', 'tensor']],
+                      ['activation_self_attn_heads', ['fsdp', 'tensor']], 
                       ['activation_cross_attn_q_length', ['fsdp', 'tensor']],
                       ['activation_length', 'fsdp'],
                       ['activation_heads', 'tensor'],
@@ -284,7 +284,7 @@ flow_shift: 3.0
 # Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf
 guidance_rescale: 0.0
 num_inference_steps: 30
-fps: 24
+fps: 16
 save_final_checkpoint: False
 
 # SDXL Lightning parameters

splash_attn_benchmark.py

@@ -62,6 +62,14 @@ def delete_file(file_path: str):
 
 
 jax.config.update("jax_use_shardy_partitioner", True)
+jax.config.update("jax_default_prng_impl", "unsafe_rbg")
+  # TF allocates extraneous GPU memory when using TFDS data
+  # this leads to CUDA OOMs. WAR for now is to hide GPUs from TF
+  # tf.config.set_visible_devices([], "GPU")
+if "xla_tpu_spmd_rng_bit_generator_unsafe" not in os.environ.get("LIBTPU_INIT_ARGS", ""):
+  os.environ["LIBTPU_INIT_ARGS"] = (
+      os.environ.get("LIBTPU_INIT_ARGS", "") + " --xla_tpu_spmd_rng_bit_generator_unsafe=true"
+  )
 
 
 def inference_generate_video(config, pipeline, filename_prefix=""):
@@ -97,7 +105,6 @@ def inference_generate_video(config, pipeline, filename_prefix=""):
 def run(config, pipeline=None, filename_prefix=""):
   print("seed: ", config.seed)
   from maxdiffusion.checkpointing.wan_checkpointer import WanCheckpointer
-
   checkpoint_loader = WanCheckpointer(config, "WAN_CHECKPOINT")
   pipeline = checkpoint_loader.load_checkpoint()
   if pipeline is None:

src/maxdiffusion/configs/base_wan_14b.yml

@@ -498,11 +498,11 @@ def get_flash_block_sizes(config):
         block_q=int(config.flash_block_sizes["block_q"]),
         block_kv_compute=int(config.flash_block_sizes["block_kv_compute"]),
         block_kv=int(config.flash_block_sizes["block_kv"]),
-        block_q_dkv=int(config.flash_block_sizes["block_q_dkv"]),
-        block_kv_dkv=int(config.flash_block_sizes["block_kv_dkv"]),
-        block_kv_dkv_compute=int(config.flash_block_sizes["block_kv_dkv_compute"]),
-        block_q_dq=int(config.flash_block_sizes["block_q_dq"]),
-        block_kv_dq=int(config.flash_block_sizes["block_kv_dq"]),
+        block_q_dkv=config.flash_block_sizes.get("block_q_dkv"),
+        block_kv_dkv=config.flash_block_sizes.get("block_kv_dkv"),
+        block_kv_dkv_compute=config.flash_block_sizes.get("block_kv_dkv_compute"),
+        block_q_dq=config.flash_block_sizes.get("block_q_dq"),
+        block_kv_dq=config.flash_block_sizes.get("block_kv_dq"),
     )
   return flash_block_sizes
 

src/maxdiffusion/generate_wan.py

@@ -26,6 +26,7 @@
 from jax.experimental.pallas.ops.tpu.splash_attention import splash_attention_kernel
 from einops import rearrange
 from .. import common_types, max_logging
+from .padded_flash_attn import make_dense_padded_attention
 
 from . import quantizations
 
@@ -236,20 +237,23 @@ def wrap_flash_attention(query, key, value):
     kv_indices = jax.lax.broadcasted_iota(jnp.int32, (kv_padded_len,), 0)
     kv_segment_ids = (kv_indices < key_seq_len).astype(jnp.int32)
     segment_ids = splash_attention_kernel.SegmentIds(q=q_segment_ids, kv=kv_segment_ids)
-
-    # make_splash_mha is wrapped around shardmap and seq and head is already
-    # sharded based on in_specs, therefore setting head_shards=1 and q_seq_shards=1.
     splash_kernel = splash_attention_kernel.make_splash_mha(
         mask=multi_head_mask,
         head_shards=1,  # the sizes of the axis is sharding over heads
         q_seq_shards=1,  # the sizes of the axis is sharding over seq_len
         block_sizes=block_sizes,
         save_residuals=True if attention_kernel == "ring" else False,
     )
-    vmapped_splash = jax.vmap(splash_kernel, in_axes=(0, 0, 0, None))
+    vmapped_splash = jax.vmap(splash_kernel, in_axes=(0, 0, 0, None), out_axes=0)
 
     if attention_kernel == "flash":
+      # attention_output = vmapped_splash(query, key, value, segment_ids)
       attention_output = vmapped_splash(query, key, value, segment_ids)
+    elif attention_kernel == "dense_padded":
+      padded_kv_len = key.shape[1] - key_seq_len
+      dense_padded_attention_kernel = make_dense_padded_attention(block_sizes=block_sizes, kv_padding=padded_kv_len)
+      vmapped_splash = jax.vmap(dense_padded_attention_kernel, in_axes=(0, 0, 0), out_axes=0)
+      attention_output, _ = vmapped_splash(query, key, value)
     else:
       if num_fsdp_shards > 1:
         out, (lse,) = vmapped_splash(query, key, value, segment_ids)
@@ -458,6 +462,19 @@ def _apply_attention(
         dtype,
         attention_kernel,
     )
+  elif attention_kernel == "dense_padded":
+    return _tpu_flash_attention(
+        query,
+        key * scale,
+        value,
+        heads,
+        mesh,
+        axis_names_q,
+        axis_names_kv,
+        flash_block_sizes,
+        dtype,
+        attention_kernel,
+    )
   elif attention_kernel == "ring":
     return _tpu_flash_attention(
         query, key * scale, value, heads, mesh, axis_names_q, axis_names_kv, flash_block_sizes, dtype, attention_kernel
@@ -877,10 +894,10 @@ def __call__(
     dtype = hidden_states.dtype
     if encoder_hidden_states is None:
       encoder_hidden_states = hidden_states
-
-    query_proj = self.query(hidden_states)
-    key_proj = self.key(encoder_hidden_states)
-    value_proj = self.value(encoder_hidden_states)
+    with jax.named_scope("attention-projection"):
+      query_proj = self.query(hidden_states)
+      key_proj = self.key(encoder_hidden_states)
+      value_proj = self.value(encoder_hidden_states)
 
     if self.qk_norm:
       query_proj = self.norm_q(query_proj)
@@ -895,7 +912,8 @@ def __call__(
     query_proj = checkpoint_name(query_proj, "query_proj")
     key_proj = checkpoint_name(key_proj, "key_proj")
     value_proj = checkpoint_name(value_proj, "value_proj")
-    attn_output = self.attention_op.apply_attention(query_proj, key_proj, value_proj)
+    with jax.named_scope("attention-compute"):
+      attn_output = self.attention_op.apply_attention(query_proj, key_proj, value_proj)
 
     attn_output = attn_output.astype(dtype=dtype)
     attn_output = checkpoint_name(attn_output, "attn_output")