pytorch

2800a04a - Add device range helper and remove sm86 specific check for memory efficient attention (#102985)

Add device range helper and remove sm86 specific check for memory efficient attention (#102985) # Summary Since we have upstreamed the latest changes of memory efficient attetnion we can remove the sm86/sm89 specific check. All head_sizes (assuming correctly alignment) should work for sm86 and sm89 size and don't have a max capability. If head_size > 96 there will be a big drop in performance but should not error and still maintain memory savings by not materializing attention weights. Pull Request resolved: https://github.com/pytorch/pytorch/pull/102985 Approved by: https://github.com/cpuhrsch