Refactor bkv_compute_in into _BlockSizes

This change reduces parameter clutter by keeping all flash block size configuration grouped together in the _BlockSizes data class. make_splash_mha, _splash_attention_forward, and _splash_attention_forward_mhpt now extract it natively.

Author: Perseus14

src/maxdiffusion/kernels/custom_splash_attention.py

@@ -35,12 +35,13 @@
 
 
 class _BlockSizes:
-  __slots__ = ("block_q", "block_kv", "block_kv_compute")
+  __slots__ = ("block_q", "block_kv", "block_kv_compute", "block_kv_compute_in")
 
-  def __init__(self, block_q: int, block_kv: int, block_kv_compute: int | None = None):
+  def __init__(self, block_q: int, block_kv: int, block_kv_compute: int | None = None, block_kv_compute_in: int = DEFAULT_BKVCOMPUTEINSIZE):
     self.block_q = block_q
     self.block_kv = block_kv
     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
 
 
 def _flash_attention_kernel(
@@ -342,7 +343,6 @@ def _splash_attention_forward(
     k: jax.Array,
     v: jax.Array,
     block_sizes: _BlockSizes,
-    bkv_compute_in: int,
     q_seq_len: int | None = None,
     kv_seq_len: int | None = None,
     use_base2_exp: bool = True,
@@ -353,6 +353,7 @@ def _splash_attention_forward(
   head_dim_v = v.shape[-1]
   bq, bkv = block_sizes.block_q, block_sizes.block_kv
   bkv_compute = block_sizes.block_kv_compute
+  bkv_compute_in = block_sizes.block_kv_compute_in
   num_kv_heads = k.shape[0]
   padded_kv_seq_len = k.shape[1]
 
@@ -428,7 +429,6 @@ def _splash_attention_forward_mhpt(
     k: jax.Array,
     v: jax.Array,
     block_sizes: _BlockSizes,
-    bkv_compute_in: int,
     heads_per_tile: int,
     q_seq_len: int | None = None,
     kv_seq_len: int | None = None,
@@ -440,6 +440,7 @@ def _splash_attention_forward_mhpt(
   head_dim_v = v.shape[-1]
   bq, bkv = block_sizes.block_q, block_sizes.block_kv
   bkv_compute = block_sizes.block_kv_compute
+  bkv_compute_in = block_sizes.block_kv_compute_in
   num_kv_heads = k.shape[0]
   actual_q_seq_len = q_seq_len if q_seq_len is not None else padded_q_seq_len
   actual_kv_seq_len = kv_seq_len if kv_seq_len is not None else k.shape[1]
@@ -514,7 +515,6 @@ def out_index_map(h, i, j, *_):
 
 def make_splash_mha(
     block_sizes: _BlockSizes,
-    bkv_compute_in: int = DEFAULT_BKVCOMPUTEINSIZE,
     orig_q_seq_len: int | None = None,
     orig_kv_seq_len: int | None = None,
     heads_per_tile: int = 1,
@@ -529,7 +529,6 @@ def _splash_attention(q, k, v):
           k,
           v,
           block_sizes,
-          bkv_compute_in,
           heads_per_tile,
           q_seq_len=orig_q_seq_len,
           kv_seq_len=orig_kv_seq_len,
@@ -542,7 +541,6 @@ def _splash_attention(q, k, v):
         k,
         v,
         block_sizes,
-        bkv_compute_in,
         q_seq_len=orig_q_seq_len,
         kv_seq_len=orig_kv_seq_len,
         use_base2_exp=use_base2_exp,

src/maxdiffusion/models/attention_flax.py

@@ -656,11 +656,10 @@ def wrap_ulysses_attention(query, key, value):
       key, _, key_seq_len = _pad_data_for_flash(key, heads, bkv)
       value, _, _ = _pad_data_for_flash(value, heads, bkv)
 
-      bsizes = custom_splash._BlockSizes(block_q=bq, block_kv=bkv, block_kv_compute=bkv_compute)
+      bsizes = custom_splash._BlockSizes(block_q=bq, block_kv=bkv, block_kv_compute=bkv_compute, block_kv_compute_in=bkv_compute_in)
 
       splash_kernel = custom_splash.make_splash_mha(
           block_sizes=bsizes,
-          bkv_compute_in=bkv_compute_in,
           orig_q_seq_len=query_seq_len,
           orig_kv_seq_len=key_seq_len,
           heads_per_tile=heads_per_tile,