Refactor attention subsystem into 4-layer Onion Architecture (Phases 1-6)

- Extract legacy UNet transformer blocks from attention_flax.py into src/maxdiffusion/models/unet_transformer_blocks_flax.py (-68% LOC reduction in attention_flax.py).
- Create src/maxdiffusion/models/attention_strategies/ package with communication-only strategies (SingleShardStrategy, RingAttentionStrategy, UlyssesStrategy, CustomRingAttentionStrategy) and Protocol definition.
- Introduce extensible KERNEL_REGISTRY in attention_utils.py and attention_dispatch.py with @register_kernel decorators for all 13 attention backends.
- Implement robust memory-efficient attention chunking and ValueError tensor validation in attention_dispatch.py.
- Implement scan-based ring communication loops (length=ring_size-1) without static loop unrolling in custom_ring.py to keep compile times and HLO graph size minimal.
- Ensure exact cross-attention sequence length and value head dimension handling across Ulysses and Ring strategies.
- Maintain 100% mathematical and operational parity with origin/main across all attention kernels (dot_product, Pallas splash, TokaMax ring, Ulysses, and custom dense splash kernels).
- Add unit test suite src/maxdiffusion/tests/attention_strategies_test.py covering strategy instantiation, hooks, block size adapters, and kernel registry validation.

Author: Perseus14

src/maxdiffusion/kernels/custom_splash_attention.py

@@ -46,6 +46,22 @@ def __init__(
     self.block_kv_compute = block_kv_compute if block_kv_compute is not None else block_kv
     self.block_kv_compute_in = block_kv_compute_in
 
+  @property
+  def bq(self) -> int:
+    return self.block_q
+
+  @property
+  def bkv(self) -> int:
+    return self.block_kv
+
+  @property
+  def bkv_compute(self) -> int:
+    return self.block_kv_compute
+
+  @property
+  def bkv_compute_in(self) -> int:
+    return self.block_kv_compute_in
+
 
 # Fixed-m softmax-bound constants. Instead of tracking the online-softmax
 # running max per KV block, eligible heads subtract a precomputed per-query
@@ -63,7 +79,8 @@ def __init__(
 # logits differently, breaking the cross-shard merge). Without k-smoothing the
 # per-row max logit has no >=0 guarantee, so the safe bound halves (calibrated
 # for ring_size=2, matching DiffusionServing's ring gate).
-_FIXED_M_RING_SAFE_BOUND = _FIXED_M_SAFE_BOUND / 2.0
+FIXED_M_RING_SAFE_BOUND = _FIXED_M_SAFE_BOUND / 2.0
+_FIXED_M_RING_SAFE_BOUND = FIXED_M_RING_SAFE_BOUND
 
 
 def _flash_attention_kernel(
@@ -815,3 +832,9 @@ def _splash_attention(q, k, v, mk=None):
     )
 
   return _splash_attention
+
+
+# Public aliases for first-class attention strategies
+BlockSizes = _BlockSizes
+FIXED_M_SAFE_BOUND = _FIXED_M_SAFE_BOUND
+splash_attention_forward_ring = _splash_attention_forward_ring