diff --git a/ex_engine/xllm_kernels/cuda/activation.cu b/ex_engine/xllm_kernels/cuda/activation.cu
index a7cafc63..409ca324 100644
--- a/ex_engine/xllm_kernels/cuda/activation.cu
+++ b/ex_engine/xllm_kernels/cuda/activation.cu
@@ -14,6 +14,7 @@ limitations under the License.
 ==============================================================================*/
 #include <c10/cuda/CUDAGuard.h>
 #include <torch/cuda.h>
+#include <torch/extension.h>
 
 #include <cstdint>
 
@@ -166,8 +167,8 @@ void act_and_mul(torch::Tensor out,
                  const std::string& act_mode) {
   if (act_mode != "silu" && act_mode != "gelu" && act_mode != "gelu_tanh" &&
       act_mode != "gelu_pytorch_tanh") {
-    LOG(FATAL) << "Unsupported act mode: " << act_mode
-               << ", only support silu, gelu, gelu_tanh, gelu_pytorch_tanh";
+    TORCH_CHECK(false, "Unsupported act mode: ", act_mode,
+                ", only support silu, gelu, gelu_tanh, gelu_pytorch_tanh");
   }
 
   // flashinfer act_and_mul ops
diff --git a/ex_engine/xllm_kernels/cuda/reshape_paged_cache.cu b/ex_engine/xllm_kernels/cuda/reshape_paged_cache.cu
index c835f95c..ab9591aa 100644
--- a/ex_engine/xllm_kernels/cuda/reshape_paged_cache.cu
+++ b/ex_engine/xllm_kernels/cuda/reshape_paged_cache.cu
@@ -14,6 +14,7 @@ limitations under the License.
 ==============================================================================*/
 
 #include <c10/cuda/CUDAStream.h>
+#include <torch/extension.h>
 
 
 #include "device_utils.cuh"