Fix sharding logic in tpu_custom_attention for 4D meshes

This fixes an issue with >=3-axis meshes where heads were not being appropriately sharded, causing asymmetric Q/KV sharding (replicating KV across context shards). The new logic uses name-based axis resolution and properly shards both Q and KV on the head-parallel axes. It also adds a safety guard to ensure num_heads is divisible by the head-parallelism size.

Author: Perseus14

src/maxdiffusion/kernels/custom_splash_attention.py

@@ -571,9 +571,11 @@ def _splash_attention(q, k, v):
 
 # ---------------------------------------------------------------------------
 # High-level attention function with shard_map
+# TODO: Move `tpu_custom_attention` and `make_custom_splash_sdpa` 
+# out of this file and into `attention_flax.py` where other orchestration 
+# functions live. This file should be strictly Pallas kernel definitions.
 # ---------------------------------------------------------------------------
 
-
 def tpu_custom_attention(
     query,
     key,
@@ -666,28 +668,55 @@ def _kernel_3d(q_3d, k_3d, v_3d):
     kv_partition_spec = P()
     out_constraint = P()
   else:
-    axis_names = mesh.axis_names
-    if len(axis_names) == 1:
-      tp_axis = axis_names[0]
-      q_partition_spec = P(None, tp_axis, None, None)
-      kv_partition_spec = P(None, tp_axis, None, None)
-      out_constraint = P(None, None, tp_axis, None)
-    elif len(axis_names) == 2:
-      dp_axis, tp_axis = axis_names[0], axis_names[1]
-      dp_size = mesh.shape[dp_axis]
-      if batch_size >= dp_size:
-        q_partition_spec = P(dp_axis, tp_axis, None, None)
-        kv_partition_spec = P(dp_axis, tp_axis, None, None)
-        out_constraint = P(dp_axis, None, tp_axis, None)
+    # Name-based axis resolution for the standard 4D mesh ("data", "fsdp", "context", "tensor").
+    # Axes are grouped by semantic role:
+    #   - Batch-parallel: "data", "fsdp"  (shard the batch dimension)
+    #   - Head-parallel:  "context", "tensor" (shard the head dimension)
+    # Only axes with size > 1 participate in sharding to avoid shard_map errors.
+    # Q and KV are sharded identically (symmetric head-parallel), matching the
+    # strategy that the legacy 2D (dp, tp) path used.
+    batch_axes = tuple(a for a in ("data", "fsdp") if a in mesh.shape and mesh.shape[a] > 1)
+    head_axes = tuple(a for a in ("context", "tensor") if a in mesh.shape and mesh.shape[a] > 1)
+
+    batch_parallel_size = 1
+    for a in batch_axes:
+      batch_parallel_size *= mesh.shape[a]
+    head_parallel_size = 1
+    for a in head_axes:
+      head_parallel_size *= mesh.shape[a]
+
+    # Flatten singleton tuples for clean PartitionSpecs.
+    batch_spec = batch_axes[0] if len(batch_axes) == 1 else (batch_axes if batch_axes else None)
+    head_spec = head_axes[0] if len(head_axes) == 1 else (head_axes if head_axes else None)
+
+    if (
+        batch_parallel_size > 1
+        and batch_size >= batch_parallel_size
+        and num_heads % head_parallel_size == 0
+    ):
+      # Head-parallel: batch on data/fsdp, heads on context/tensor.
+      q_partition_spec = P(batch_spec, head_spec, None, None)
+      kv_partition_spec = P(batch_spec, head_spec, None, None)
+      out_constraint = P(batch_spec, None, head_spec, None)
+    elif head_parallel_size > 1 or batch_parallel_size > 1:
+      # Batch too small for batch-parallel; shard all parallelism onto heads.
+      all_axes = batch_axes + head_axes
+      all_parallel_size = batch_parallel_size * head_parallel_size
+      all_spec = all_axes[0] if len(all_axes) == 1 else (all_axes if all_axes else None)
+      if num_heads % all_parallel_size == 0:
+        q_partition_spec = P(None, all_spec, None, None)
+        kv_partition_spec = P(None, all_spec, None, None)
+        out_constraint = P(None, None, all_spec, None)
       else:
-        all_axes = tuple(axis_names)
-        q_partition_spec = P(None, all_axes, None, None)
-        kv_partition_spec = P(None, all_axes, None, None)
-        out_constraint = P(None, None, all_axes, None)
+        # Heads not divisible by available parallelism; fall back to replicated.
+        q_partition_spec = P()
+        kv_partition_spec = P()
+        out_constraint = P()
     else:
-      q_partition_spec = P(axis_names[0], axis_names[1], axis_names[2], None)
-      kv_partition_spec = P(axis_names[0], axis_names[1], None, None)
-      out_constraint = P(axis_names[0], None, (axis_names[1], axis_names[2]), None)
+      # All axes are size-1; fully replicated (e.g. single-device).
+      q_partition_spec = P()
+      kv_partition_spec = P()
+      out_constraint = P()
 
   sharded_fn = shard_map(
       _attention_on_slices,