diff --git a/csrc/moe/moe_align_sum_kernels.cu b/csrc/moe/moe_align_sum_kernels.cu
index 5c9e47402408..9fcc88ef41bd 100644
--- a/csrc/moe/moe_align_sum_kernels.cu
+++ b/csrc/moe/moe_align_sum_kernels.cu
@@ -1,7 +1,8 @@
 #include <torch/all.h>
 #include <ATen/cuda/CUDAContext.h>
 #include <c10/cuda/CUDAGuard.h>
-#include <cub/cub.cuh>
+
+#include "../cub_helpers.h"
 
 #include <ATen/ATen.h>
 #include <ATen/cuda/Atomic.cuh>
diff --git a/csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h b/csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h
index 108091efbefa..69ac131bfa39 100644
--- a/csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h
+++ b/csrc/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h
@@ -5,9 +5,14 @@
 #include <c10/core/ScalarType.h>
 #include <torch/all.h>
 #include "dispatch.h"
-#include <cub/cub.cuh>
-#include <cub/device/device_radix_sort.cuh>
-#include <cub/util_type.cuh>
+
+#include "../cub_helpers.h"
+#ifndef USE_ROCM
+  #include <cub/device/device_radix_sort.cuh>
+  #include <cub/util_type.cuh>
+#else
+// hipcub includes these via hipcub.hpp
+#endif
 #include "cutlass/numeric_size.h"
 #include "cutlass/array.h"
 
diff --git a/csrc/sampler.cu b/csrc/sampler.cu
index d458f8e4c1d0..a2bd42813723 100644
--- a/csrc/sampler.cu
+++ b/csrc/sampler.cu
@@ -3,11 +3,7 @@
 #include <torch/cuda.h>
 #include <c10/cuda/CUDAGuard.h>
 
-#ifndef USE_ROCM
-  #include <cub/cub.cuh>
-#else
-  #include <hipcub/hipcub.hpp>
-#endif
+#include "cub_helpers.h"
 
 namespace vllm {