From 820d04fae10b48c97d6bd73381cda814d09403d4 Mon Sep 17 00:00:00 2001
From: Josh Romero <joshr@nvidia.com>
Date: Mon, 29 Jun 2026 15:45:57 -0700
Subject: [PATCH] Avoid NVSHMEM device header parsing in non-RDC CUDA
 translation units.

Signed-off-by: Josh Romero <joshr@nvidia.com>
---
 include/internal/cudecomp_kernels.cuh | 13 ++++++++++++-
 src/cudecomp_kernels.cu               |  4 ++++
 2 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/include/internal/cudecomp_kernels.cuh b/include/internal/cudecomp_kernels.cuh
index 66138fd..e38abe3 100644
--- a/include/internal/cudecomp_kernels.cuh
+++ b/include/internal/cudecomp_kernels.cuh
@@ -22,12 +22,23 @@
 #include <cstddef>
 #include <cstdint>
 
+// Keep NVSHMEM device headers out of non-RDC CUDA translation units.
+#if defined(ENABLE_NVSHMEM) && defined(CUDECOMP_NVSHMEM_HOST_ONLY) && !defined(NVSHMEM_HOSTLIB_ONLY)
+#define CUDECOMP_UNDEF_NVSHMEM_HOSTLIB_ONLY
+#define NVSHMEM_HOSTLIB_ONLY
+#endif
+
 #ifdef ENABLE_NVSHMEM
 #include <nvshmem.h>
 #endif
 
 #include "internal/cudecomp_kernels.h"
 
+#ifdef CUDECOMP_UNDEF_NVSHMEM_HOSTLIB_ONLY
+#undef NVSHMEM_HOSTLIB_ONLY
+#undef CUDECOMP_UNDEF_NVSHMEM_HOSTLIB_ONLY
+#endif
+
 #define CUDECOMP_CUDA_NTHREADS (128)
 #define CUDECOMP_UNROLL_FACTOR (4)
 #define CUDECOMP_MIN_BLOCKS_PER_SM (16)
@@ -36,7 +47,7 @@
 
 namespace cudecomp {
 
-#ifdef ENABLE_NVSHMEM
+#if defined(ENABLE_NVSHMEM) && !defined(CUDECOMP_NVSHMEM_HOST_ONLY)
 template <typename T>
 __launch_bounds__(CUDECOMP_CUDA_NTHREADS) __global__
     void cudecomp_nvshmem_alltoallv_k(cudecompNvshmemA2AParams<T> params) {
diff --git a/src/cudecomp_kernels.cu b/src/cudecomp_kernels.cu
index 311ebb9..16a77b5 100644
--- a/src/cudecomp_kernels.cu
+++ b/src/cudecomp_kernels.cu
@@ -18,7 +18,11 @@
 #include <cuda/std/complex>
 
 #include "internal/checks.h"
+
+// This non-RDC translation unit only instantiates generic CUDA kernels.
+#define CUDECOMP_NVSHMEM_HOST_ONLY
 #include "internal/cudecomp_kernels.cuh"
+#undef CUDECOMP_NVSHMEM_HOST_ONLY
 
 namespace cudecomp {