Change:
return pad_bias<mem_eff_alignment>(attn_mask); } // Check and make the tensor contiguous if needed <del> if (attn_mask.sym_stride(0) % 16 != 0 || attn_mask.sym_stride(1) % 16 != 0 || attn_mask.sym_stride(2) % 16 != 0) { </del> <ins> auto needs_contig = [](const c10::SymInt& stride) { return (stride % 16 != 0) || (stride == 0); }; if (needs_contig(attn_mask.sym_stride(0)) || needs_contig(attn_mask.sym_stride(1)) || needs_contig(attn_mask.sym_stride(2)) || needs_contig(attn_mask.sym_stride(3))) { </ins> return attn_mask.contiguous(); }

Discussion:
Hi after more hours of debugging than I am comfortable to admit, I noticed the following breaking change between PyTorch 2.0.1 and PyTorch 2.1. The issue can be reproduced both with & . There is no issue on 2.0.1 For fp32 inputs to SDPA on CUDA device & passing a custom , we have the following: PyTorch 2.0.1: dispatches to math ( for other backends) PyTorch 2.1: dispatches to mem-efficient attention (likely following ). And it appears that the mem-efficient attention backend outputs wrong results when passing non-contigous key/value, while the math backend goes just fine. Reproduction: The result is: Note that these metrics are not super representative of how the drift is bad locally: ! The relevant tensors are attached in a zip here. Thank you! cc
So for this case the mask passed in is all Trues which will get converted to a "zeros" additive attention_bias mask.. in theory a no-op. When I run the same q,k,v input with no mask I get which leads me to believe that non-contiguous inputs alone are not the problem but non-contiguous + masking is no bueno
If I update the script to only call contiguous on either key or value but not both I still see the same results
Okayy I got it, so this happens when bias_mask is not contiguous
Weird - so not being contiguous was responsible of the issue?
I can confirm this issue is fixed on
Unfortunately this solutions makes the (1, 1, seqlenq, seqlenkv case) have to fully materialize the attention mask. I think that a third component of this bug is that it is occuring when seqlenq = 1, I tried with seqlen_q = 2 and it appears to not occur anymore
I was acting somewhat rash with the impeding 2.1.1 deadline and I am unhappy with my solution, and I didn't fully think through its ramifications. For a very hot path I.e. a mask of (1, 1, seqlenq, seqlenkv) this causes a substantial memory regression.

I am having a hard time figure out whats goin on with the repro mentioned here: The results of the repo appear to be dependent on order in which the ops are ran, which makes me thing there is some IMA happening, but alas I haven't found it.
Thanks a lot for looking into it sorry for having been a bit pushy about having this in the patch! I appear not to be able to reproduce the logits difference in using and A100-SXM4-80GB, is the output difference appearing using the memory patch
Tentatively marking hi-pri for the regression
Could you help double check if the memory issue is fixed in 2.1.2? Here are the results and myself gets using the reproduction example in the issue: This is the same as what we get from 2.1.0
So the memory issue was "fixed" in 2.1.1 via this PR: While this did change the result of the repro, it caused a memory usage "regression" for a very common hot path for SDPA. This PR: Was used to fix the memory regression. I did further analysis on the initial correctness report. I determined that the repro wasn't behaving as I expected.. e.g the order with which aten ops were called impacted. I am also not entirely sure that a contiguous vs non contiguous input should produce identical results (it likely should but I need to reason if there is some potential float addmulitple ops that get reordered cc ) I was leaving this open in case we need to revisit this but the the memory reduction should be visible via usual memory profiling techniques
I can also confirm that this issue is reintroduced in the 2.1.2 release, while it was fixed in 2.1.1. This is a bit bad for us as following 2.1.1 we SDPA by default in Transformers and a very common execution path has non-contiguous inputs: We'll likely have to do a patch and manually call , I guess. In 2.1.2, this codepath starts erroring out with . Here, silently SDPA starts outputting different results when inputs are not contiguous, which I think is not wishable. If the kernel requires contiguous inputs, why not check it on the C++ side and call if needed (maybe with a warning)?
do you have a repro for the illegal memory access?

I wasn't able to re-repro this with the existing one in this issue.
I don't have a minimal reproduction unfortunately, but you can try: and run with the following script: which errors at: This corresponds to this code: Note that on of Transformers you will not be able to reproduce as for now I just disabled the mem-efficient path for GPT BigCode due to this issue (fix at , removing the expand).
Is it possible to have minimal repro for this ? Running using your instructions I get following:
Oh sorry. You can try with instead that is non-gated. I'll try to come up with a minimal repro ASAP. I edited my message above. For the record using
Running above command produced for me: Notable Environment differences: CUDA-12.1 PyTorch: Main GPU: H100 `
While the repro isn't throwing a IMA compute-sanitizer is reporting errors:
Current gists with compute-sanitizer outputs:
I have a fix:
I rerun the script to confirm that this is fixed in 2.2.0
Okay I am going to close this issue for now