[Release/2.6] Backport softmax fixes from 2.8dev (#2247)

This fixes OOB memory access for followng code
``` python
import torch
qk = torch.randn((9,1017), dtype=torch.float64, device='cuda')
smqk = torch.softmax(qk, dim=-1)
```

Upstream PR:
* pytorch#144009
* pytorch#154778

---------

Co-authored-by: eqy <[email protected]>

Author: xinyazhang

aten/src/ATen/native/cuda/SoftMax.cu

@@ -468,7 +468,7 @@ ilpReduce(index_t shift,
     if (offset >= shift && offset < size) {
       threadVal = r(threadVal, data[offset]);
     }
-    size -= blockDim.x;
+    size -= blockDim.x > size ? size : blockDim.x;
     data += blockDim.x;
   }
   index_t last = size % (ILP * blockDim.x);
@@ -518,7 +518,7 @@ WriteFpropResultsVectorized(
     if (offset >= shift && offset < size) {
       output[offset] = epilogue(input[offset]);
     }
-    size -= blockDim.x;
+    size -= blockDim.x > size ? size : blockDim.x;
     input += blockDim.x;
     output += blockDim.x;
   }
@@ -573,7 +573,7 @@ WriteBpropResultsVectorized(
     if (threadIdx.x >= shift) {
       gradInput[offset] = epilogue(gradOutput[offset], output[offset]);
     }
-    size -= blockDim.x;
+    size -= blockDim.x > size ? size : blockDim.x;
     gradInput += blockDim.x;
     output += blockDim.x;
     gradOutput += blockDim.x;

test/test_nn.py

@@ -10498,6 +10498,13 @@ def run_test(*shape):
         run_test(1100000000, 2)  # Illegal memory access https://github.com/pytorch/pytorch/issues/52715
         run_test(2200000000, 1)  # invalid configuration argument https://github.com/pytorch/pytorch/issues/52716
 
+    @onlyCUDA
+    @dtypes(torch.double)
+    def test_softmax_double(self, device, dtype):
+        logits = torch.randn(5, 513, dtype=dtype, device=device)
+        expected_ones = F.log_softmax(logits, dim=1).exp().sum(dim=1)
+        self.assertEqual(expected_ones, torch.ones_like(expected_ones))
+
     @onlyCUDA
     @dtypes(torch.half)
     @largeTensorTest("20GB")