Softmax docstrings and type fixes (#37)

* Softmax docs and type fixes
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

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

* change API, better naming, const fixes
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

docs/api/c/index.rst

@@ -18,5 +18,6 @@ directly from C/C++, without Python.
    cast.h <cast>
    gemm.h <gemm>
    layer_norm.h <layer_norm>
+   softmax.h <softmax>
    transformer_engine.h <transformer_engine>
    transpose.h <transpose>

docs/api/c/softmax.rst

+..
+    Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+
+    See LICENSE for license information.
+
+softmax.h
+======
+
+.. doxygenfile:: softmax.h

transformer_engine/common/fused_softmax/scaled_masked_softmax.cu

@@ -349,7 +349,7 @@ __global__ void scaled_masked_softmax_warp_forward(
 template <typename input_t, typename output_t, typename acc_t, int log2_elements>
 __global__ void scaled_masked_softmax_warp_backward(
     output_t *gradInput,
-    input_t *grad,
+    const input_t *grad,
     const input_t *output,
     acc_t scale,
     int micro_batch_size,
@@ -773,7 +773,7 @@ void dispatch_scaled_masked_softmax_forward(
 template<typename input_t, typename output_t, typename acc_t>
 void dispatch_scaled_masked_softmax_backward(
     output_t *grad_input,
-    input_t *grad,
+    const input_t *grad,
     const input_t *output,
     const acc_t scale,
     int query_seq_len,
@@ -968,7 +968,8 @@ void scaled_softmax_forward(
 }
 
 void scaled_softmax_backward(
-    const Tensor output_grads,
+    Tensor output_grads,
+    const Tensor incoming_grads,
     const Tensor softmax_results,
     float scale_factor,
     cudaStream_t stream) {
@@ -983,7 +984,7 @@ void scaled_softmax_backward(
     TRANSFORMER_ENGINE_TYPE_SWITCH_16BIT(output_grads.dtype, softmax_type,
         dispatch_scaled_masked_softmax_backward<softmax_type, softmax_type, float>(
             reinterpret_cast<softmax_type*>(output_grads.dptr),
-            reinterpret_cast<softmax_type*>(output_grads.dptr),
+            reinterpret_cast<softmax_type const*>(incoming_grads.dptr),
             reinterpret_cast<softmax_type const*>(softmax_results.dptr),
             scale_factor,
             query_seq_len,
@@ -1023,7 +1024,8 @@ void scaled_masked_softmax_forward(
 
 
 void scaled_masked_softmax_backward(
-    const Tensor output_grads,
+    Tensor output_grads,
+    const Tensor incoming_grads,
     const Tensor softmax_results,
     float scale_factor,
     cudaStream_t stream
@@ -1038,7 +1040,7 @@ void scaled_masked_softmax_backward(
     TRANSFORMER_ENGINE_TYPE_SWITCH_16BIT(output_grads.dtype, softmax_type,
         dispatch_scaled_masked_softmax_backward<softmax_type, softmax_type, float>(
             reinterpret_cast<softmax_type*>(output_grads.dptr),
-            reinterpret_cast<softmax_type*>(output_grads.dptr),
+            reinterpret_cast<softmax_type const*>(incoming_grads.dptr),
             reinterpret_cast<softmax_type const*>(softmax_results.dptr),
             scale_factor,
             query_seq_len,
@@ -1068,14 +1070,16 @@ void nvte_scaled_softmax_forward(
 
 
 void nvte_scaled_softmax_backward(
-    const NVTETensor output_grads,
+    const NVTETensor incoming_grads,
     const NVTETensor softmax_results,
+    NVTETensor output_grads,
     float scale_factor,
     cudaStream_t stream
 ) {
     using namespace transformer_engine;
     scaled_softmax_backward(
-        *reinterpret_cast<const Tensor*>(output_grads),
+        *reinterpret_cast<Tensor*>(output_grads),
+        *reinterpret_cast<const Tensor*>(incoming_grads),
         *reinterpret_cast<const Tensor*>(softmax_results),
         scale_factor,
         stream);
@@ -1100,15 +1104,17 @@ void nvte_scaled_masked_softmax_forward(
 
 
 void nvte_scaled_masked_softmax_backward(
-    const NVTETensor input,
-    NVTETensor softmax_results,
+    const NVTETensor incoming_grads,
+    const NVTETensor softmax_results,
+    NVTETensor output_grads,
     float scale_factor,
     cudaStream_t stream
 ) {
     using namespace transformer_engine;
     scaled_masked_softmax_backward(
-        *reinterpret_cast<const Tensor*>(input),
-        *reinterpret_cast<Tensor*>(softmax_results),
+        *reinterpret_cast<Tensor*>(output_grads),
+        *reinterpret_cast<const Tensor*>(incoming_grads),
+        *reinterpret_cast<const Tensor*>(softmax_results),
         scale_factor,
         stream);
 }

transformer_engine/common/fused_softmax/scaled_upper_triang_masked_softmax.cu

@@ -248,7 +248,7 @@ __global__ void scaled_upper_triang_masked_softmax_warp_forward(
 template <typename input_t, typename output_t, typename acc_t, int log2_elements>
 __global__ void scaled_upper_triang_masked_softmax_warp_backward(
     output_t *gradInput,
-    input_t *grad,
+    const input_t *grad,
     const input_t *output,
     acc_t scale,
     int micro_batch_size,
@@ -509,7 +509,7 @@ void dispatch_scaled_upper_triang_masked_softmax_forward(
 template<typename input_t, typename output_t, typename acc_t>
 void dispatch_scaled_upper_triang_masked_softmax_backward(
     output_t *grad_input,
-    input_t *grad,
+    const input_t *grad,
     const input_t *output,
     const acc_t scale,
     int softmax_elements,
@@ -683,7 +683,8 @@ void scaled_upper_triang_masked_softmax_forward(
 
 
 void scaled_upper_triang_masked_softmax_backward(
-    const Tensor output_grads,
+    Tensor output_grads,
+    const Tensor incoming_grads,
     const Tensor softmax_results,
     float scale_factor,
     cudaStream_t stream)  {
@@ -695,7 +696,7 @@ void scaled_upper_triang_masked_softmax_backward(
     TRANSFORMER_ENGINE_TYPE_SWITCH_16BIT(output_grads.dtype, softmax_type,
         dispatch_scaled_upper_triang_masked_softmax_backward<softmax_type, softmax_type, float>(
             reinterpret_cast<softmax_type*>(output_grads.dptr),
-            reinterpret_cast<softmax_type*>(output_grads.dptr),
+            reinterpret_cast<softmax_type const*>(incoming_grads.dptr),
             reinterpret_cast<softmax_type const*>(softmax_results.dptr),
             scale_factor,
             seq_len,
@@ -723,14 +724,16 @@ void nvte_scaled_upper_triang_masked_softmax_forward(
 
 
 void nvte_scaled_upper_triang_masked_softmax_backward(
-    const NVTETensor output_grads,
+    const NVTETensor incoming_grads,
     const NVTETensor softmax_results,
+    NVTETensor output_grads,
     float scale_factor,
     cudaStream_t stream
 ) {
     using namespace transformer_engine;
     scaled_upper_triang_masked_softmax_backward(
-        *reinterpret_cast<const Tensor*>(output_grads),
+        *reinterpret_cast<Tensor*>(output_grads),
+        *reinterpret_cast<const Tensor*>(incoming_grads),
         *reinterpret_cast<const Tensor*>(softmax_results),
         scale_factor,
         stream);

transformer_engine/common/include/transformer_engine/softmax.h

@@ -15,7 +15,13 @@
 extern "C" {
 #endif
 
-
+/*! \brief Compute scaled softmax activation on the input.
+ *
+ *  \param[in]     input           Input tensor for softmax.
+ *  \param[out]    softmax_results Output tensor.
+ *  \param[in]     scale_factor    Scalar for the input tensor.
+ *  \param[in]     stream          CUDA stream used for the operation.
+ */
 void nvte_scaled_softmax_forward(
     const NVTETensor input,
     NVTETensor softmax_results,
@@ -24,14 +30,35 @@ void nvte_scaled_softmax_forward(
 );
 
 
+/*! \brief Compute the backward of the scaled softmax activation.
+ *
+ *  - `incoming_grads` is the input tensor containing the gradients received from the following layer.
+ *  - `softmax_results` is the output tensor of the corresponding forward softmax operation.
+ *  - `output_grads` is the output tensor containing the computed gradients.
+ *
+ *  \param[in]     incoming_grads  Input gradient tensor for backward.
+ *  \param[in]     softmax_results Output tensor of softmax forward.
+ *  \param[out]    output_grads    Output tensor.
+ *  \param[in]     scale_factor    Scalar for the output tensor.
+ *  \param[in]     stream          CUDA stream used for the operation.
+ */
 void nvte_scaled_softmax_backward(
-    const NVTETensor output_grads,
+    const NVTETensor incoming_grads,
     const NVTETensor softmax_results,
+    NVTETensor output_grads,
     float scale_factor,
     cudaStream_t stream
 );
 
 
+/*! \brief Compute scaled masked softmax activation on the input.
+ *
+ *  \param[in]     input           Input tensor for softmax.
+ *  \param[in]     mask            Mask for the input tensor.
+ *  \param[out]    softmax_results Output tensor.
+ *  \param[in]     scale_factor    Scalar for the input tensor.
+ *  \param[in]     stream          CUDA stream used for the operation.
+ */
 void nvte_scaled_masked_softmax_forward(
     const NVTETensor input,
     const NVTETensor mask,
@@ -41,14 +68,34 @@ void nvte_scaled_masked_softmax_forward(
 );
 
 
+/*! \brief Compute the backward of the scaled masked softmax activation.
+ *
+ *  - `incoming_grads` is the input tensor containing the gradients received from the following layer.
+ *  - `softmax_results` is the output tensor of the corresponding forward softmax operation.
+ *  - `output_grads` is the output tensor containing the computed gradients.
+ *
+ *  \param[in]     incoming_grads  Input gradient tensor for backward.
+ *  \param[in]     softmax_results Output tensor of softmax forward.
+ *  \param[out]    output_grads    Output tensor.
+ *  \param[in]     scale_factor    Scalar for the output tensor.
+ *  \param[in]     stream          CUDA stream used for the operation.
+ */
 void nvte_scaled_masked_softmax_backward(
-    const NVTETensor input,
-    NVTETensor softmax_results,
+    const NVTETensor incoming_grads,
+    const NVTETensor softmax_results,
+    NVTETensor output_grads,
     float scale_factor,
     cudaStream_t stream
 );
 
 
+/*! \brief Compute scaled softmax activation using a 2D upper triangular mask on the input.
+ *
+ *  \param[in]     input           Input tensor for softmax.
+ *  \param[out]    softmax_results Output tensor.
+ *  \param[in]     scale_factor    Scalar for the input tensor.
+ *  \param[in]     stream          CUDA stream used for the operation.
+ */
 void nvte_scaled_upper_triang_masked_softmax_forward(
     const NVTETensor input,
     NVTETensor softmax_results,
@@ -57,9 +104,22 @@ void nvte_scaled_upper_triang_masked_softmax_forward(
 );
 
 
+/*! \brief Compute the backward of the scaled softmax activation using a 2D upper triangular mask.
+ *
+ *  - `incoming_grads` is the input tensor containing the gradients received from the following layer.
+ *  - `softmax_results` is the output tensor of the corresponding forward softmax operation.
+ *  - `output_grads` is the output tensor containing the computed gradients.
+ *
+ *  \param[in]     incoming_grads  Input gradient tensor for backward.
+ *  \param[in]     softmax_results Output tensor of softmax forward.
+ *  \param[out]    output_grads    Output tensor.
+ *  \param[in]     scale_factor    Scalar for the output tensor.
+ *  \param[in]     stream          CUDA stream used for the operation.
+ */
 void nvte_scaled_upper_triang_masked_softmax_backward(
-    const NVTETensor output_grads,
+    const NVTETensor incoming_grads,
     const NVTETensor softmax_results,
+    NVTETensor output_grads,
     float scale_factor,
     cudaStream_t stream
 );

transformer_engine/pytorch/csrc/extensions.cu

@@ -537,8 +537,9 @@ at::Tensor scaled_softmax_backward(at::Tensor output_grad_,
     auto output_grads_cu = makeTransformerEngineTensor(output_grads);
     auto softmax_results_cu = makeTransformerEngineTensor(softmax_results);
 
+    // Produce gradients in place.
     nvte_scaled_softmax_backward(
-          output_grads_cu.data(), softmax_results_cu.data(),
+          output_grads_cu.data(), softmax_results_cu.data(), output_grads_cu.data(),
           scale_factor, at::cuda::getCurrentCUDAStream());
 
     return output_grads;
@@ -608,8 +609,9 @@ at::Tensor scaled_masked_softmax_backward(at::Tensor output_grad_,
     auto output_grads_cu = makeTransformerEngineTensor(output_grads);
     auto softmax_results_cu = makeTransformerEngineTensor(softmax_results);
 
+    // Produce gradients in place.
     nvte_scaled_softmax_backward(
-          output_grads_cu.data(), softmax_results_cu.data(),
+          output_grads_cu.data(), softmax_results_cu.data(), output_grads_cu.data(),
           scale_factor, at::cuda::getCurrentCUDAStream());
 
     return output_grads;
@@ -671,8 +673,10 @@ at::Tensor scaled_upper_triang_masked_softmax_backward(at::Tensor output_grads_,
     auto output_grads_cu = makeTransformerEngineTensor(output_grads);
     auto softmax_results_cu = makeTransformerEngineTensor(softmax_results);
 
+    // Produce gradients in place.
     nvte_scaled_upper_triang_masked_softmax_backward(output_grads_cu.data(),
                                                      softmax_results_cu.data(),
+                                                     output_grads_cu.data(),
                                                      scale_factor,
                                                      at::cuda::getCurrentCUDAStream());