[sdpa] Fix bottom-right efficient attention when Q exceeds K (#193632)
## Human Note Part 1 of doing the acutal change I want which is aligning both sdpa and varlen to lower-right the bettter choice of mask semantic in the age of llms. This fixes some edge cases that would have borken mem-eff. Currently not user facing since we are still on upper left ## Agent note This commit was authored with an AI assistant. CUDA memory-efficient attention stored the bottom-right causal diagonal offset as an unsigned value. When the query sequence was longer than the key sequence, the negative offset wrapped and disabled masking in mixed query blocks. Entirely masked query blocks were also skipped without initializing forward output or query gradients. Store the offset as a signed value, clamp each block's key count at zero, and zero-initialize output and query gradients only when bottom-right masking may contain fully masked rows. This defines those rows as zero for both dense and packed variable-length attention while preserving the empty-allocation path for shapes where every query has at least one valid key. ## Test Plan ```bash source ~/.venvs/dev/bin/activate && source ~/dotfiles/scripts/_build_setup.sh -b200 && CCACHE_NOHASHDIR=true uv pip install -e . --no-build-isolation -v gpu-run auto -- env PYTHONPATH=$PWD /home/dev/.venvs/dev/bin/pytest -q -rs test/test_transformers.py -k "mem_efficient_bottom_right or fully_masked_out_rows or mem_efficient_attention_gqa_split_key" gpu-run auto -- env PYTHONPATH=$PWD /home/dev/.venvs/dev/bin/pytest -q test/test_varlen_attention.py gpu-run auto -- env PYTHONPATH=$PWD /home/dev/.venvs/dev/bin/pytest -q -rs test/test_flop_counter.py -k varlen git diff --check lintrunner -a aten/src/ATen/native/transformers/cuda/attention.cu aten/src/ATen/native/transformers/cuda/attention_backward.cu aten/src/ATen/native/transformers/cuda/mem_eff_attention/kernel_forward.h test/test_transformers.py ``` Pull Request resolved: https://github.com/pytorch/pytorch/pull/193632 Approved by: https://github.com/liangel-02 ghstack dependencies: #193630, #193631
D
drisspg committed
2509857d0e811629f1d34c3d187c172a7c8bc935
Parent: 7059aa3
Committed by PyTorch MergeBot <pytorchmergebot@users.noreply.github.com>
on 8/24/2026, 9:31:57 PM