[NPU] Fix dropout bug in downscale_in_infer mode (#1346)

Author: will-jl944

backends/npu/kernels/dropout_kernel.cc

@@ -400,14 +400,15 @@ void DropoutRawKernel(const Context& dev_ctx,
       EXEC_NPU_CMD(
           aclnnDropoutDoMask, dev_ctx, tmp_x, *mask, dropout_prob, tmp_out);
     }
+
+    if (!is_upscale) {
+      phi::Scalar revert_scale = static_cast<T>(1.0 - dropout_prob);
+      EXEC_NPU_CMD(aclnnInplaceMuls, dev_ctx, *out, revert_scale);
+    }
   } else {
     if (!is_upscale) {
-      const auto& muls_runner =
-          NpuOpRunner("Muls",
-                      {x},
-                      {*out},
-                      {{"value", static_cast<float>(1 - dropout_prob)}});
-      muls_runner.Run(stream);
+      phi::Scalar down_scale = static_cast<T>(1.0 - dropout_prob);
+      EXEC_NPU_CMD(aclnnMuls, dev_ctx, x, down_scale, *out);
       return;
     }
     TensorCopy(dev_ctx, x, false, out);
@@ -531,7 +532,6 @@ void DropoutGradRawKernel(const Context& dev_ctx,
     return;
   }
 
-  dropout_prob = is_upscale ? dropout_prob : 0.0;
   if (dx->dtype() == phi::DataType::FLOAT64) {
     // transform dx
     phi::DenseTensor tmp_dx;
@@ -563,6 +563,11 @@ void DropoutGradRawKernel(const Context& dev_ctx,
   } else {
     EXEC_NPU_CMD(aclnnDropoutDoMask, dev_ctx, dout, mask, dropout_prob, *dx);
   }
+
+  if (!is_upscale) {
+    phi::Scalar revert_scale = static_cast<T>(1.0 - dropout_prob);
+    EXEC_NPU_CMD(aclnnInplaceMuls, dev_ctx, *dx, revert_scale);
+  }
   return;
 }
 

backends/npu/tests/unittests/test_dropout_op_npu.py

@@ -228,6 +228,21 @@ def setUp(self):
         self.outputs = {"Out": self.inputs["X"]}
 
 
+@skip_check_grad_ci(reason="For inference, check_grad is not required.")
+class TestDropoutOpInferenceDownscale(TestDropoutOpInference):
+    def setUp(self):
+        self.op_type = "dropout"
+        self.set_npu()
+        self.init_dtype()
+        self.inputs = {"X": np.random.random((32, 64, 3)).astype(self.dtype)}
+        self.attrs = {
+            "dropout_prob": 0.75,
+            "is_test": True,
+            "dropout_implementation": "downgrade_in_infer",
+        }
+        self.outputs = {"Out": self.inputs["X"] * (1.0 - self.attrs["dropout_prob"])}
+
+
 class TestDropoutOpWithSeed(TestDropoutOp):
     # the seed is a Tensor
     def setUp(self):