From 9ff418db7376900afb42e2d2ae22b95728cbaf38 Mon Sep 17 00:00:00 2001
From: chen <103103266+ckl117@users.noreply.github.com>
Date: Wed, 19 Nov 2025 16:02:21 +0800
Subject: [PATCH] check METAX_GPU (#5114)

---
 custom_ops/gpu_ops/helper.h | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/custom_ops/gpu_ops/helper.h b/custom_ops/gpu_ops/helper.h
index 9b54c1128..ef1c63022 100644
--- a/custom_ops/gpu_ops/helper.h
+++ b/custom_ops/gpu_ops/helper.h
@@ -725,6 +725,10 @@ inline void launchWithPdlWhenEnabled(KernelFn kernelFn,
                                      size_t dynamicShmSize,
                                      cudaStream_t stream,
                                      Args &&...args) {
+#ifdef PADDLE_WITH_CUSTOM_DEVICE_METAX_GPU
+  (*kernelFn)<<<grid, block, dynamicShmSize, stream>>>(
+      std::forward<Args>(args)...);
+#else
   cudaLaunchConfig_t kernelConfig;
   kernelConfig.gridDim = grid;
   kernelConfig.blockDim = block;
@@ -738,5 +742,6 @@ inline void launchWithPdlWhenEnabled(KernelFn kernelFn,
   kernelConfig.numAttrs = 1;
 
   cudaLaunchKernelEx(&kernelConfig, kernelFn, std::forward<Args>(args)...);
+#endif
 }
 #endif