diff --git a/backends/mlu/tests/CMakeLists.txt b/backends/mlu/tests/CMakeLists.txt
index 5b44e3b603a..f77e6bcba7a 100644
--- a/backends/mlu/tests/CMakeLists.txt
+++ b/backends/mlu/tests/CMakeLists.txt
@@ -22,7 +22,7 @@ function(py_test_modules TARGET_NAME)
   add_test(
     NAME ${TARGET_NAME}
     COMMAND
-      ${CMAKE_COMMAND} -E env
+      ${CMAKE_COMMAND} -E env FLAGS_use_stream_safe_cuda_allocator=false
       CUSTOM_DEVICE_ROOT=${CMAKE_BINARY_DIR}/python/paddle_custom_device/
       PYTHONPATH=${PYTHON_SOURCE_DIR}:${PYTHON_SOURCE_DIR}/tests:$ENV{PYTHONPATH}
       ${py_test_modules_ENVS} python ${PYTHON_SOURCE_DIR}/tools/test_runner.py
@@ -37,7 +37,7 @@ endfunction()
 add_test(
   NAME test_LeNet_MNIST
   COMMAND
-    ${CMAKE_COMMAND} -E env
+    ${CMAKE_COMMAND} -E env FLAGS_use_stream_safe_cuda_allocator=false
     CUSTOM_DEVICE_ROOT=${CMAKE_BINARY_DIR}/python/paddle_custom_device/
     PYTHONPATH=${PYTHON_SOURCE_DIR}:${PYTHON_SOURCE_DIR}/tests:$ENV{PYTHONPATH}
     python test_LeNet_MNIST.py
diff --git a/backends/mlu/tests/unittests/test_multinomial_op_mlu.py b/backends/mlu/tests/unittests/test_multinomial_op_mlu.py
index 2170103b74d..2cb24914775 100644
--- a/backends/mlu/tests/unittests/test_multinomial_op_mlu.py
+++ b/backends/mlu/tests/unittests/test_multinomial_op_mlu.py
@@ -169,7 +169,11 @@ def test_static(self):
         train_program = base.Program()
         with base.program_guard(train_program, startup_program):
             x = paddle.static.data("x", shape=[4], dtype="float32")
-            out = paddle.multinomial(x, num_samples=100000, replacement=True)
+            outs = [
+                paddle.multinomial(x, num_samples=100000, replacement=True)
+                for _ in range(10)
+            ]
+            out = paddle.concat(outs, axis=0)
 
             place = base.CustomPlace("mlu", 0)
             exe = base.Executor(place)
diff --git a/backends/mlu/tools/dockerfile/Dockerfile.mlu.kylinv10.gcc82.py310 b/backends/mlu/tools/dockerfile/Dockerfile.mlu.kylinv10.gcc82.py310
index 20f27076259..78e2152560c 100644
--- a/backends/mlu/tools/dockerfile/Dockerfile.mlu.kylinv10.gcc82.py310
+++ b/backends/mlu/tools/dockerfile/Dockerfile.mlu.kylinv10.gcc82.py310
@@ -41,6 +41,7 @@ ENV FLAGS_use_stride_kernel=0
 ENV FLAGS_allocator_strategy=auto_growth
 ENV CNCL_MEM_POOL_MULTI_CLIQUE_ENABLE=1
 ENV PADDLE_XCCL_BACKEND=mlu
+ENV FLAGS_use_stream_safe_cuda_allocator=false
 
 # yum and pip clean
 RUN yum clean all && \
diff --git a/backends/mlu/tools/dockerfile/Dockerfile.mlu.ubuntu20.gcc84.py310 b/backends/mlu/tools/dockerfile/Dockerfile.mlu.ubuntu20.gcc84.py310
index 46b9c64223c..c3afa3c22ea 100644
--- a/backends/mlu/tools/dockerfile/Dockerfile.mlu.ubuntu20.gcc84.py310
+++ b/backends/mlu/tools/dockerfile/Dockerfile.mlu.ubuntu20.gcc84.py310
@@ -83,6 +83,7 @@ ENV FLAGS_use_stride_kernel=0
 ENV FLAGS_allocator_strategy=auto_growth
 ENV CNCL_MEM_POOL_MULTI_CLIQUE_ENABLE=1
 ENV PADDLE_XCCL_BACKEND=mlu
+ENV FLAGS_use_stream_safe_cuda_allocator=false
 
 # Clean
 RUN apt-get clean -y
diff --git a/backends/npu/tools/pr_ci_npu.sh b/backends/npu/tools/pr_ci_npu.sh
index fa7b36b5232..84ff04e083f 100644
--- a/backends/npu/tools/pr_ci_npu.sh
+++ b/backends/npu/tools/pr_ci_npu.sh
@@ -28,6 +28,7 @@ export CODE_ROOT
 
 # For paddle easy debugging
 export FLAGS_call_stack_level=2
+export FLAGS_enable_pir_api=1
 
 failed_test_lists=''
 tmp_dir=`mktemp -d`