Description

Slice scaled_dot_product_attention into smaller chunks so that the SDPA of each chunk wouldn't request any allocation larger than the given limit.

This was initially designed to work around the 4GB single block allocation limitation of Intel compute-runtime (RuntimeError: Current platform can NOT allocate memory block with size larger than 4GB! Tried to allocate 8.00 GiB). Then I found out that setting a smaller limit would reduce the VRAM footprint during SDPA calculation. The current limit (VRAM // 8) was tuned for Intel Arc A770 16G and A750 8G without sacrificing performance.

With this change, A770 16G can generate 512x512 of batch size 32 and A750 8G can generate batch size 16

Test results:

Common settings: --use-ipex --opt-sdp-attention, txt2img, DPM++ 2M Karras, 20 steps, 512x512 resolution, batch count = 5

• Effective it/s == Batch size * Batch count * Steps / Total time taken
• RE in the table refers to RuntimeError: Current platform can NOT allocate memory block with size larger than 4GB!
• OOM in the table refers to RuntimeError: Allocation is out of device memory on current platform.

A770 16G (connected with two monitors [taking up ~1.1GB VRAM])

A750 8G (not connected with monitors)

Screenshots/videos:

Checklist:

• I have read contributing wiki page
• I have performed a self-review of my own code
• My code follows the style guidelines
• My code passes tests with --test-server --use-ipex --opt-sdp-attention