fix: enable ulysses sharding for custom kernels and improve scaling precision

[Perseus14]

Description

This PR introduces two small but important fixes to the Ulysses attention implementation:

1. Enable Ulysses sharding for custom kernels: In src/maxdiffusion/pyconfig.py, the check for Ulysses attention is changed from attention == "ulysses" to "ulysses" in attention. This ensures that custom attention implementations that include "ulysses" in their identifier (e.g., custom_ulysses) will correctly trigger Ulysses sequence sharding instead of falling back to default sharding strategies.
2. Ensure correct padding: In src/maxdiffusion/models/attention_flax.py, padding was done for query variable but query_scaled was used in the attention calculation. This fixes it and ensures the padded variable is used in the attention calculation.
3. Improve scaling precision: In src/maxdiffusion/models/attention_flax.py, the hardcoded constant 1.44269504 used to scale queries for base-2 exponentiation is replaced with math.log2(math.e). This provides better precision and makes the intent of the code clearer.

Generation Time

• Main: ~420s
• Branch: ~140s

[github-actions]

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

[eltsai]

Thanks for fixing this @Perseus14 ! I tested on v6e (don't have v7x for now due to capacity) and the generation time is 204 sec:

==================================================
  TIMING SUMMARY
==================================================
  Load (checkpoint):      97.3s
  Compile:               229.8s
  ────────────────────────────────────────
  Inference:             204.0s
  Conditioning:           11.8s
  Denoise Total:         183.6s
  VAE Decode:              8.6s
==================================================

I think the generation time is expected, because on v7x we are seeing 28% speed boost (140 vs 194.4). From go/wan-dashboard the e2e gen time is 322 sec, the speed gain is about 36% percent.