Support causal splash attention in Qwen3, enforce prompt_embeds sharding constraint, and fix unit test sharding handling

Author: amepas

src/maxdiffusion/generate_flux2klein.py

@@ -278,6 +278,7 @@ def main(argv):
       ulysses_shards=getattr(config, "ulysses_shards", -1),
       ulysses_attention_chunks=getattr(config, "ulysses_attention_chunks", 1),
       max_layer_to_run=getattr(config, "text_encoder_max_layer", 27),
+      is_causal=getattr(config, "text_encoder_is_causal", True),
   )
   qwen3_model = FlaxQwen3Model(qwen3_config)
 

src/maxdiffusion/models/attention_flax.py

@@ -127,8 +127,7 @@ def _reshape_batch_dim_to_heads(tensor, heads):
   tensor = tensor.reshape(batch_size // head_size, head_size, seq_len, dim)
   tensor = jnp.transpose(tensor, (0, 2, 1, 3))
   reshaped_tensor = tensor.reshape(batch_size // head_size, seq_len, dim * head_size)
-  axis_names = nn.logical_to_mesh_axes((BATCH, LENGTH, HEAD))
-  return jax.lax.with_sharding_constraint(reshaped_tensor, axis_names)
+  return nn.with_logical_constraint(reshaped_tensor, (BATCH, LENGTH, HEAD))
 
 
 def _reshape_heads_to_batch_dim(tensor, heads):
@@ -141,8 +140,7 @@ def _reshape_heads_to_batch_dim(tensor, heads):
   else:
     batch_size, head_size, seq_len, head_dim = tensor.shape
     reshaped_tensor = tensor.reshape(batch_size * head_size, seq_len, head_dim)
-  axis_names = nn.logical_to_mesh_axes((BATCH, LENGTH, HEAD))
-  return jax.lax.with_sharding_constraint(reshaped_tensor, axis_names)
+  return nn.with_logical_constraint(reshaped_tensor, (BATCH, LENGTH, HEAD))
 
 
 def _reshape_heads_to_head_dim(tensor):
@@ -151,8 +149,7 @@ def _reshape_heads_to_head_dim(tensor):
   b, h, s, d = tensor.shape
   tensor = jnp.transpose(tensor, axes=[0, 2, 1, 3])
   reshaped_tensor = jnp.reshape(tensor, (b, -1, h * d))
-  axis_names = nn.logical_to_mesh_axes((BATCH, LENGTH, HEAD))
-  return jax.lax.with_sharding_constraint(reshaped_tensor, axis_names)
+  return nn.with_logical_constraint(reshaped_tensor, (BATCH, LENGTH, HEAD))
 
 
 def _unflatten_heads(tensor, heads):
@@ -579,6 +576,7 @@ def _tpu_flash_attention(
     attention_mask: jax.Array = None,
     use_base2_exp: bool = False,
     use_experimental_scheduler: bool = False,
+    is_causal: bool = False,
 ) -> jax.Array:
   """TPU Flash Attention"""
 
@@ -656,8 +654,12 @@ def wrap_flash_attention(query, key, value, attention_mask):
     key, _, key_seq_len = _pad_data_for_flash(key, heads, block_kv)
     value, _, _ = _pad_data_for_flash(value, heads, block_kv)
 
-    mask = splash_attention_mask.FullMask(_shape=(query.shape[2], key.shape[2]))
-    multi_head_mask = splash_attention_mask.MultiHeadMask(masks=(mask,) * query.shape[1])
+    if is_causal:
+      mask = splash_attention_mask.CausalMask((query.shape[2], key.shape[2]))
+      multi_head_mask = splash_attention_mask.MultiHeadMask(masks=(mask,) * query.shape[1])
+    else:
+      mask = splash_attention_mask.FullMask(_shape=(query.shape[2], key.shape[2]))
+      multi_head_mask = splash_attention_mask.MultiHeadMask(masks=(mask,) * query.shape[1])
 
     segment_ids_cls = (
         tokamax_splash_base.SegmentIds if attention_kernel == "tokamax_ring" else splash_attention_kernel.SegmentIds
@@ -674,9 +676,14 @@ def wrap_flash_attention(query, key, value, attention_mask):
     # 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.
     if attention_kernel == "tokamax_flash":
-      mask = tokamax_splash_attention_mask.FullMask(
-          _shape=(query.shape[2], key.shape[2]),
-      )
+      if is_causal:
+        mask = tokamax_splash_attention_mask.CausalMask(
+            (query.shape[2], key.shape[2]),
+        )
+      else:
+        mask = tokamax_splash_attention_mask.FullMask(
+            _shape=(query.shape[2], key.shape[2]),
+        )
       splash_kernel = tokamax_splash_attention_kernel.make_splash_mha(
           mask=mask,
           q_seq_shards=1,  # the sizes of the axis is sharding over seq_len
@@ -1571,10 +1578,9 @@ def _cudnn_flash_attention(query: Array, key: Array, value: Array, heads: int, m
   key = _reshape_data_for_cudnn_flash(key, heads)
   value = _reshape_data_for_cudnn_flash(value, heads)
 
-  axis_names = nn.logical_to_mesh_axes((BATCH, LENGTH, HEAD, D_KV))
-  query = jax.lax.with_sharding_constraint(query, axis_names)
-  key = jax.lax.with_sharding_constraint(key, axis_names)
-  value = jax.lax.with_sharding_constraint(value, axis_names)
+  query = nn.with_logical_constraint(query, (BATCH, LENGTH, HEAD, D_KV))
+  key = nn.with_logical_constraint(key, (BATCH, LENGTH, HEAD, D_KV))
+  value = nn.with_logical_constraint(value, (BATCH, LENGTH, HEAD, D_KV))
 
   out = dpa_layer(query, key, value, mask=None)
   return _reshape_data_from_cudnn_flash(out)
@@ -1788,6 +1794,7 @@ def flash_kernel(q, k, v, context):
       attention_mask=context["attention_mask"],
       use_base2_exp=context["use_base2_exp"],
       use_experimental_scheduler=context["use_experimental_scheduler"],
+      is_causal=context.get("is_causal", False),
   )
 
 
@@ -1809,6 +1816,7 @@ def tokamax_flash_kernel(q, k, v, context):
       attention_mask=context["attention_mask"],
       use_base2_exp=context["use_base2_exp"],
       use_experimental_scheduler=context["use_experimental_scheduler"],
+      is_causal=context.get("is_causal", False),
   )
 
 
@@ -1830,6 +1838,7 @@ def tokamax_ring_kernel(q, k, v, context):
       attention_mask=context["attention_mask"],
       use_base2_exp=context["use_base2_exp"],
       use_experimental_scheduler=context["use_experimental_scheduler"],
+      is_causal=context.get("is_causal", False),
   )
 
 
@@ -1883,6 +1892,7 @@ def _apply_attention(
     use_experimental_scheduler: bool = False,
     ulysses_shards: int = -1,
     ulysses_attention_chunks: int = 1,
+    is_causal: bool = False,
 ):
   """Routes to different attention kernels using a module-level registry."""
 
@@ -1948,6 +1958,7 @@ def _apply_attention(
       "float32_qk_product": float32_qk_product,
       "use_memory_efficient_attention": use_memory_efficient_attention,
       "dpa_layer": dpa_layer,
+      "is_causal": is_causal,
   }
 
   # Module-level Registry lookup
@@ -2281,6 +2292,7 @@ class AttentionOp(nn.Module):
   use_experimental_scheduler: bool = False
   ulysses_shards: int = -1
   ulysses_attention_chunks: int = 1
+  is_causal: bool = False
 
   def setup(self):
     self.dpa_layer = None
@@ -2330,6 +2342,7 @@ def apply_attention(self, query: Array, key: Array, value: Array, attention_mask
         use_experimental_scheduler=self.use_experimental_scheduler,
         ulysses_shards=self.ulysses_shards,
         ulysses_attention_chunks=self.ulysses_attention_chunks,
+        is_causal=self.is_causal,
     )
 
 
@@ -2620,9 +2633,9 @@ def __call__(
       rngs: nnx.Rngs = None,
       cached_kv: Optional[Dict[str, Tuple[jax.Array, jax.Array]]] = None,
   ) -> jax.Array:
-    axis_names = nn.logical_to_mesh_axes((BATCH, LENGTH, HEAD))
-    hidden_states = jax.lax.with_sharding_constraint(hidden_states, axis_names)
-    encoder_hidden_states = jax.lax.with_sharding_constraint(encoder_hidden_states, axis_names)
+    hidden_states = nn.with_logical_constraint(hidden_states, (BATCH, LENGTH, HEAD))
+    if encoder_hidden_states is not None:
+      encoder_hidden_states = nn.with_logical_constraint(encoder_hidden_states, (BATCH, LENGTH, HEAD))
     dtype = hidden_states.dtype
     is_self_attention = encoder_hidden_states is None
     if encoder_hidden_states is None:

src/maxdiffusion/models/qwen3_flax.py

@@ -47,6 +47,7 @@ def __init__(
       ulysses_shards: int = -1,
       ulysses_attention_chunks: int = 1,
       max_layer_to_run: Optional[int] = 27,
+      is_causal: bool = True,
   ):
     self.vocab_size = vocab_size
     self.hidden_size = hidden_size
@@ -65,6 +66,7 @@ def __init__(
     self.ulysses_shards = ulysses_shards
     self.ulysses_attention_chunks = ulysses_attention_chunks
     self.max_layer_to_run = max_layer_to_run
+    self.is_causal = is_causal
 
 
 # -----------------------------------------------------------------------------
@@ -278,6 +280,7 @@ def __call__(
           dtype=self.config.dtype,
           ulysses_shards=self.config.ulysses_shards,
           ulysses_attention_chunks=self.config.ulysses_attention_chunks,
+          is_causal=getattr(self.config, "is_causal", True),
       )
       out = attn_op.apply_attention(q_3d, k_3d, v_3d, attention_mask=attention_mask)
     else:
@@ -293,9 +296,10 @@ def __call__(
 
       scores = jnp.matmul(q_f, jnp.transpose(k_f, (0, 1, 3, 2))) / math.sqrt(self.config.head_dim)
 
-      # 7. Apply causal attention mask
-      causal_mask = jnp.tril(jnp.ones((seq_len, seq_len), dtype=jnp.bool_))
-      scores = jnp.where(causal_mask, scores, -1e4)
+      # 7. Apply causal attention mask if configured
+      if getattr(self.config, "is_causal", True):
+        causal_mask = jnp.tril(jnp.ones((seq_len, seq_len), dtype=jnp.bool_))
+        scores = jnp.where(causal_mask, scores, -1e4)
 
       # 8. Apply padding attention mask if provided
       if attention_mask is not None:

src/maxdiffusion/pipelines/flux/flux2klein_pipeline.py

@@ -119,6 +119,8 @@ def qwen3_forward(q_params, ids, mask):
       h_27 = all_hidden_states[27]
       out = jnp.stack([h_9, h_18, h_27], axis=1)
       prompt_embeds = jnp.transpose(out, (0, 2, 1, 3)).reshape((ids.shape[0], ids.shape[1], -1))
+      context_spec = P(None, "context") if "context" in self.mesh.axis_names and self.mesh.shape["context"] > 1 else P()
+      prompt_embeds = jax.lax.with_sharding_constraint(prompt_embeds, jax.sharding.NamedSharding(self.mesh, context_spec))
       return prompt_embeds
 
     @jax.jit(static_argnums=(4, 5), donate_argnums=(1,))
@@ -489,6 +491,9 @@ def put_data_on_devices(x, sharding):
       max_logging.log(f"{host_prefix} Passed Phase A Sync Barrier (phase_a_complete) successfully! ✅")
 
       latents_jax = put_data_on_devices(latents_jax, data_sharding)
+      context_spec = P(None, "context") if "context" in self.mesh.axis_names and self.mesh.shape["context"] > 1 else P()
+      context_sharding = jax.sharding.NamedSharding(self.mesh, context_spec)
+      prompt_embeds_jax = put_data_on_devices(prompt_embeds_jax, context_sharding)
       txt_ids_val = put_data_on_devices(txt_ids_val, data_sharding)
       img_ids_val = put_data_on_devices(img_ids_val, data_sharding)