GitLab

See also: https://stackoverflow.com/questions/55136414/constexpr-variable-captured-inside-lambda-loses-its-constexpr-ness

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

Follow xFormers's DISTPATCH_BOOL. Haven't tested it on Windows.

This is faster since we only need to do atomic adds on dq, instead of atomic
adds on both dk and dv.

Previously we were using block_size=256.

They don't have to have the same block size, number of threads, etc.

Theoretically this might have lower numerical error since the scaling is in
fp32 instead of fp16 (not sure, I haven't thought too carefully about it).
However, in practice, the numerical errors seem about the same.

This speeds up the fwd by 1.5x.

To reduce smem usage for SM75

From 4KB per buffer to 2KB per buffer. This saves us 8KB of smem (each Q and dO
have 2 buffers)