防止softmax的kernel中地址计算溢出 (#71935)

Qin-sx · web-flow · commit 8dd22d95f9ab · 2025-04-08T17:36:05.000+08:00
* ensure address do not overflow

	modified:   paddle/phi/kernels/gpudnn/softmax_gpudnn.h

* added a unit test

	modified:   test/legacy_test/test_log_softmax.py

* check run

	modified:   test_log_softmax.py
diff --git a/paddle/phi/kernels/gpudnn/softmax_gpudnn.h b/paddle/phi/kernels/gpudnn/softmax_gpudnn.h
@@ -445,8 +445,8 @@ __global__ void KeMatrixSoftmaxForward(T* softmax, const T* src, int dim_size) {
   using VecT = phi::AlignedVector<T, kVecSize>;
 
   int bid = blockIdx.x;
-  T* batch_input = const_cast<T*>(src) + bid * dim_size;
-  T* batch_output = softmax + bid * dim_size;
+  T* batch_input = const_cast<T*>(src) + (uint64_t)bid * dim_size;
+  T* batch_output = softmax + (uint64_t)bid * dim_size;
 
   const int input_align_shift =
       ((uint64_t)batch_input) % MATRIX_SOFTMAX_ALIGN_BYTES / sizeof(T);
diff --git a/test/legacy_test/test_log_softmax.py b/test/legacy_test/test_log_softmax.py
@@ -246,6 +246,23 @@ def test_errors(self):
             self.assertRaises(TypeError, F.log_softmax, x, dtype='int32')
 
 
+def _check_cuda_memory_20GB():
+    if not hasattr(paddle.device.cuda, 'get_device_properties'):
+        return False
+    gpu_info = paddle.device.cuda.get_device_properties(0)
+    return gpu_info.total_memory >= 20 * (1024**3)  # 20GB
+
+
+@unittest.skipIf(
+    not core.is_compiled_with_cuda() or not _check_cuda_memory_20GB(),
+    "Need CUDA support and at least 20GB GPU memory",
+)
+class TestLogSoftmaxLargeOp(unittest.TestCase):
+    def test_check_run(self):
+        x = paddle.randn([4, 4096, 131072 + 2048])  # 8GB+4*4096*2048
+        paddle.nn.functional.log_softmax(x, axis=-1)
+
+
 if __name__ == "__main__":
     paddle.enable_static()
     unittest.main()
