Fix: Plumb vmem limit and bkv_compute_in to the custom kernel

[eltsai]

Right now the vmem_limit_bytes and block_kv_compute_in from the argument was never passed into the kernel.

This PR plumbs vmem_limit_bytes from the flash_block_sizes config down to the Mosaic
CompilerParams so the custom splash kernel can run under a raised scoped-VMEM
limit.

• custom_splash_attention.py: vmem_limit_bytes param through both forwards ->
  pltpu.CompilerParams, make_splash_mha, tpu_custom_attention (+ flash_block_sizes
  read), make_custom_splash_sdpa (+ kwargs read).
• attention_flax.py: read vmem_limit_bytes in wrap_ulysses_attention, pass to
  make_splash_mha.
• max_utils.py: CustomFlashBlockSizes frozen carrier so vmem_limit_bytes (and
  block_kv_compute_in / heads_per_tile) survive config -> kernel instead of being
  dropped by the JAX BlockSizes dataclass.

By raising the vmem to the 64 MB limit, our grid search finds a better config:

export BQ=8448
export BKV=1024
export BKV_COMPUTE=1024
export BKV_COMPUTE_IN=1024
export HEADS_PER_TILE=1

E2E result:

==================================================
  TIMING SUMMARY
==================================================
  Load (checkpoint):     141.2s
  Compile:               183.3s
  Inference:             126.9s
  ────────────────────────────────────────
  Conditioning:            2.2s
    - VAE Encode:          0.0s
  Denoise Total:         122.6s
  VAE Decode:              2.1s
    - TPU Compute:         1.7s
    - Host Formatting:     0.4s
==================================================

Comparing to previous config (BQ=4864), the denoising steps is 3.7% faster (122.6 sec vs 127.3 sec).

[github-actions]

e2e testgrid: https://8bcf50593faf4ea38060e236169827e5-dot-us-central1.composer.googleusercontent.com/dags/maxdiffusion_tpu_e2e/grid

[Perseus14]

Pretty great find @eltsai, awesome work!

Few questions

• Could you clarify on the usage of vmem_limit? Do we add this param in the config file under flash_block_sizes, if so can we update the configs as well (if it doesn't affect other attention mechanisms)?
• Does setting the XLA flag for vmem_limit_bytes (eg: xla_tpu_scoped_vmem_limit_kib) affect this value?
• Can we update the PR title and description as block_kv_compute_in seems to be already passed into the kernel?

[github-actions]

🤖 Hi @Perseus14, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

[github-actions]

🤖 I'm sorry @Perseus14, but I was unable to process your request. Please see the logs for more details.

[github-actions]

🤖 Hi @Perseus14, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

[github-actions]

📋 Review Summary

This pull request correctly plumbs the vmem_limit_bytes configuration option and other block-level tuning parameters (such as block_kv_compute_in and heads_per_tile) all the way down to the custom TPU Pallas splash attention kernel. This allows executing under raised scoped VMEM limits, unlocking optimal layout sweeping and a ~3.7% performance improvement in denoising.

🔍 General Feedback

• Preservation of Configuration Parameters: Introducing the hashable, frozen CustomFlashBlockSizes carrier is an excellent design decision to prevent JAX from silently dropping these custom properties.
• Robust Config Extraction: The extraction of custom properties in attention_flax.py correctly handles both raw dicts and custom objects, which provides excellent resilience against different configuration entry points.
• Consistency Check: Please verify that tpu_custom_attention (often used by TorchAX or standalone tests) uses the same object-safe extraction logic to prevent AttributeError crashes, as detailed in the inline comment.

[github-actions]

🔴 With the introduction of the frozen dataclass CustomFlashBlockSizes in max_utils.py to preserve custom properties across JAX boundaries, passing flash_block_sizes to tpu_custom_attention (e.g., via the TorchAX SDPA wrapper path in make_custom_splash_sdpa) will result in a runtime crash. Since flash_block_sizes is no longer guaranteed to be a dictionary, calling .get() on a CustomFlashBlockSizes instance will raise an AttributeError.

We should safely extract values by checking if flash_block_sizes is a dictionary or an object, similar to how it is done in attention_flax.py.

The entire block starting from line 597 should be refactored as follows:

  if flash_block_sizes is not None:
    if isinstance(flash_block_sizes, dict):
      block_q = flash_block_sizes.get("block_q", block_q)
      block_kv = flash_block_sizes.get("block_kv", block_kv)
      block_kv_compute = flash_block_sizes.get("block_kv_compute", block_kv_compute)
      block_kv_compute_in = flash_block_sizes.get("block_kv_compute_in", block_kv_compute_in)
      heads_per_tile = flash_block_sizes.get("heads_per_tile", heads_per_tile)
      vmem_limit_bytes = flash_block_sizes.get("vmem_limit_bytes", vmem_limit_bytes)
    else:
      block_q = getattr(flash_block_sizes, "block_q", block_q)
      block_kv = getattr(flash_block_sizes, "block_kv", block_kv)
      block_kv_compute = getattr(flash_block_sizes, "block_kv_compute", block_kv_compute)
      block_kv_compute_in = getattr(flash_block_sizes, "block_kv_compute_in", block_kv_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)

[eltsai]

@Perseus14 (not sure why I can't directly reply to the comments)

• Could you clarify on the usage of vmem_limit? Do we add this param in the config file under flash_block_sizes, if so can we update the configs as well (if it doesn't affect other attention mechanisms)?

Yes, it's read from the flash_block_sizes dict in the config, for example.:

flash_block_sizes: {
  "block_q": 5888,
  "block_kv": 1024,
  "block_kv_compute": 1024,
  "block_kv_compute_in": 1024,
  "heads_per_tile": 1,
  "vmem_limit_bytes": 67108864   

• Does setting the XLA flag for vmem_limit_bytes (eg: xla_tpu_scoped_vmem_limit_kib) affect this value?

No, they're separate knobs at different layers and neither feeds the other. vmem_limit_bytes here is the Mosaic per-kernel VMEM budget for the splash kernel specifically (via CompilerParams).
xla_tpu_scoped_vmem_limit_kib is a global XLA-level scoped-VMEM limit for XLA-compiled HLO ops.

• Can we update the PR title and description as block_kv_compute_in seems to be already passed into the kernel?

before this commit, block_kv_compute_in (and heads_per_tile) were silently dropped by get_flash_block_sizes, because it converted the dict into the JAX BlockSizes dataclass, which has no field for them. The new CustomFlashBlockSizes carrier added here is what actually lets block_kv_compute_in, heads_per_tile, and vmem_limit_bytes survive to the kernel.