Integrate torchax custom attention kernel into ulysses

[eltsai]

Adding torchax path's custom kernel into ulysses (triggered when attention=ulysses_custom)

Inference time:

==================================================
  TIMING SUMMARY
==================================================
  Load (checkpoint):      80.6s
  Compile:               186.9s
  ────────────────────────────────────────
  Inference:             167.2s
==================================================

[github-actions]

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

[Perseus14]

Updating this to

bq = 4864
bkv = 1024
bkv_compute = 1024
bkv_compute_in = 1024
heads_per_tile = 1

and using this command gave me the following latency

  Load (checkpoint):     297.0s
  Compile:               219.8s
  ───────────────────────────────
  Inference:             147.4s

[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]

🤖 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]

🤖 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]

🤖 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]

🤖 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.

[eltsai]

Updated stats:

Accelerator	Sharding	E2E time	log	Video
v7x-8	dp2-context4-tp1	139.3s	log	Video
v7x-16	dp2-context8-tp1	70.2s	log	Video

[Perseus14]

@eltsai I have left a few comments on the code. I am hoping we can reuse the user defined configs like use_base2_exp, use_experimental_scheduler and flash_block_sizes for the custom kernel.

Let me know if we can incorporate these changes now or attempt them later

cc: @entrpn

[eltsai]

Updated the PR based on @Perseus14 's comments: passing use_base2_exp, use_experimental_scheduler and flash_block_sizes instead of hard coded.

Also moved the custom kernel file under src/maxdiffusion/kernels

Updated the profiling code to log perf with more granularity:

==================================================
  TIMING SUMMARY
==================================================
  Load (checkpoint):     130.9s
  Compile:               137.5s
  ────────────────────────────────────────
  Inference:             137.4s
  Conditioning:            2.4s
  Denoise Total:         121.8s
  VAE Decode:             13.3s
==================================================

[eltsai]

Updated perf logging code. It was not logging each component's time correctly. In WanPipeline2_2.__call__, we start the denoise_total timer, run the loop, and stop it. Because of jax's asynchronous execution, the timer stops before the TPU has actually finished executing all 40 steps.

New run:

==================================================
  TIMING SUMMARY
==================================================
  Load (checkpoint):     128.6s
  Compile:               137.6s
  ────────────────────────────────────────
  Inference:             136.9s
  Conditioning:            2.2s
  Denoise Total:         130.6s
  VAE Decode:              4.0s
==================================================

Also added the two new flags from @rishabhmanoj

disable_bounds_checks=True,
skip_device_barrier=True,

which give us about ~1 sec of gain