diff --git a/src/maxdiffusion/configuration_utils.py b/src/maxdiffusion/configuration_utils.py
index 091d6d05..396626e2 100644
--- a/src/maxdiffusion/configuration_utils.py
+++ b/src/maxdiffusion/configuration_utils.py
@@ -496,7 +496,9 @@ def extract_init_dict(cls, config_dict, **kwargs):
     diffusers_library = importlib.import_module(__name__.split(".")[0])
 
     if cls.has_compatibles:
-      compatible_classes = [c for c in cls._get_compatibles() if not isinstance(c, DummyObject)]  # pyrefly: ignore[missing-attribute]
+      compatible_classes = [
+          c for c in cls._get_compatibles() if not isinstance(c, DummyObject)
+      ]  # pyrefly: ignore[missing-attribute]
     else:
       compatible_classes = []
 
diff --git a/src/maxdiffusion/models/attention_flax.py b/src/maxdiffusion/models/attention_flax.py
index bdb738cf..d29eb26e 100644
--- a/src/maxdiffusion/models/attention_flax.py
+++ b/src/maxdiffusion/models/attention_flax.py
@@ -363,24 +363,25 @@ def _extract_custom_block_sizes(flash_block_sizes):
   if flash_block_sizes is not None:
     if isinstance(flash_block_sizes, dict):
       get = flash_block_sizes.get
-      bq = get("block_q", bq)
-      bkv = get("block_kv", bkv)
-      bkv_compute = get("block_kv_compute", bkv_compute)
-      bkv_compute_in = get("block_kv_compute_in", bkv_compute_in)
-      heads_per_tile = get("heads_per_tile", heads_per_tile)
-      vmem_limit_bytes = get("vmem_limit_bytes", vmem_limit_bytes)
+      bq = get("block_q", None) or bq
+      bkv = get("block_kv", None) or bkv
+      bkv_compute = get("block_kv_compute", None) or bkv_compute
+      bkv_compute_in = get("block_kv_compute_in", None) or bkv_compute_in
+      # A BlockSizes object carries heads_per_tile=None when the config dict omitted
+      # it; getattr then returns that None instead of the default, so coerce it back
+      # to 1 (the custom-kernel default) to keep the `heads_per_tile > 1` guards safe.
+      heads_per_tile = get("heads_per_tile", None) or heads_per_tile
+      vmem_limit_bytes = get("vmem_limit_bytes", None) or vmem_limit_bytes
     else:
-      bq = getattr(flash_block_sizes, "block_q", bq)
-      bkv = getattr(flash_block_sizes, "block_kv", bkv)
-      bkv_compute = getattr(flash_block_sizes, "block_kv_compute", bkv_compute)
-      bkv_compute_in = getattr(flash_block_sizes, "block_kv_compute_in", bkv_compute_in)
-      heads_per_tile = getattr(flash_block_sizes, "heads_per_tile", heads_per_tile)
-      vmem_limit_bytes = getattr(flash_block_sizes, "vmem_limit_bytes", vmem_limit_bytes)
-  # A BlockSizes object carries heads_per_tile=None when the config dict omitted
-  # it; getattr then returns that None instead of the default, so coerce it back
-  # to 1 (the custom-kernel default) to keep the `heads_per_tile > 1` guards safe.
-  if heads_per_tile is None:
-    heads_per_tile = 1
+      bq = getattr(flash_block_sizes, "block_q", None) or bq
+      bkv = getattr(flash_block_sizes, "block_kv", None) or bkv
+      bkv_compute = getattr(flash_block_sizes, "block_kv_compute", None) or bkv_compute
+      bkv_compute_in = getattr(flash_block_sizes, "block_kv_compute_in", None) or bkv_compute_in
+      # A BlockSizes object carries heads_per_tile=None when the config dict omitted
+      # it; getattr then returns that None instead of the default, so coerce it back
+      # to 1 (the custom-kernel default) to keep the `heads_per_tile > 1` guards safe.
+      heads_per_tile = getattr(flash_block_sizes, "heads_per_tile", None) or heads_per_tile
+      vmem_limit_bytes = getattr(flash_block_sizes, "vmem_limit_bytes", None) or vmem_limit_bytes
   return bq, bkv, bkv_compute, bkv_compute_in, heads_per_tile, vmem_limit_bytes