fix https://github.com/NVIDIA/TensorRT-Edge-LLM/issues/204: pass the CUDA include path to the XQA NVRTC JIT

cpp/CMakeLists.txt

@@ -182,6 +182,31 @@ target_include_directories(edgellmPluginJit PRIVATE ${COMMON_INCLUDE_DIRS}
                                                     ${NVRTC_INCLUDE_DIR})
 target_link_libraries(
   edgellmPluginJit PRIVATE ${NVRTC_LIB} ${CUDA_DRIVER_LINK_LIB} ${CUDART_LIB})
+# The XQA JIT compiles its kernel through NVRTC at load time and needs the CUDA
+# runtime headers on the include path (cuda_fp16.h -> vector_types.h). Those are
+# not necessarily where nvrtc.h is: in a full toolkit they sit together, but a
+# standalone NVRTC package ships nvrtc.h alone. Search for the header that is
+# actually included, across the roots FindNVRTC already honours and in the same
+# precedence, so a non-default NVRTC_ROOT_DIR or CUDA_TARGET_DIR cannot leave
+# the JIT reading a different toolkit's headers than the one it links against.
+find_path(
+  EDGELLM_JIT_CUDA_INCLUDE_DIR vector_types.h
+  HINTS ${NVRTC_INCLUDE_DIR} ${NVRTC_ROOT_DIR}/include
+        ${CUDA_TARGET_DIR}/include ${CUDA_DIR}/include
+  NO_DEFAULT_PATH)
+if(NOT EDGELLM_JIT_CUDA_INCLUDE_DIR)
+  message(
+    WARNING
+      "No vector_types.h under the NVRTC/CUDA roots; the XQA JIT will fall back "
+      "to ${CUDA_DIR}/include. Set NVRTC_ROOT_DIR or CUDA_DIR to a CUDA Toolkit "
+      "with the runtime headers, or set EDGELLM_NVRTC_INCLUDE at run time.")
+  set(EDGELLM_JIT_CUDA_INCLUDE_DIR "${CUDA_DIR}/include")
+endif()
+# Cross builds bake a path that is correct on the BUILD host; where the toolkit
+# sits elsewhere on the target, EDGELLM_NVRTC_INCLUDE overrides it at run time.
+target_compile_definitions(
+  edgellmPluginJit
+  PRIVATE EDGELLM_CUDA_INCLUDE_DIR="${EDGELLM_JIT_CUDA_INCLUDE_DIR}")
 add_dependencies(edgellmPluginJit generatePluginJitEmbeddedSources)
 edgellm_apply_qnx_warning_suppressions(edgellmPluginJit)
 edgellm_set_hidden_visibility(edgellmPluginJit)

cpp/kernels/decodeAttentionKernels/decoderXQAJitCompiler.cpp

@@ -24,11 +24,18 @@
 
 #include <algorithm>
 #include <array>
+#include <cstdlib>
 #include <stdexcept>
 #include <string>
 #include <type_traits>
 #include <vector>
 
+// Set by the build system to the CUDA toolkit this build was configured against. The
+// conventional install path is the fallback for builds that do not define it.
+#ifndef EDGELLM_CUDA_INCLUDE_DIR
+#define EDGELLM_CUDA_INCLUDE_DIR "/usr/local/cuda/include"
+#endif
+
 namespace trt_edgellm
 {
 namespace
@@ -105,7 +112,18 @@ std::vector<std::string> buildNvrtcOptions(XQAJitKey const& key)
     options.emplace_back("--use_fast_math");
     options.emplace_back("--device-as-default-execution-space");
     options.emplace_back(getGpuArchitectureOption(key.sm));
-    // No -I flags needed: all headers are passed as virtual includes to nvrtcCreateProgram.
+    // The kernel's own headers are virtual includes on nvrtcCreateProgram, but cuda_fp16.h
+    // includes vector_types.h from the CUDA toolkit, which NVRTC does not carry as a builtin
+    // on CUDA 12.x for x86. Without this include path the compile fails with
+    //   cuda_fp16.h(129): catastrophic error: cannot open source file "vector_types.h"
+    // and it surfaces during ONNX parsing as a plugin creation failure, which points away
+    // from the cause. EDGELLM_NVRTC_INCLUDE overrides the configured path for installs that
+    // are relocated away from the toolkit they were built against.
+    {
+        char const* const cudaIncludeOverride = std::getenv("EDGELLM_NVRTC_INCLUDE");
+        options.emplace_back(
+            std::string("-I") + (cudaIncludeOverride != nullptr ? cudaIncludeOverride : EDGELLM_CUDA_INCLUDE_DIR));
+    }
     options.emplace_back("-DGENERATE_CUBIN=1");
     options.emplace_back("-DNDEBUG");
     options.emplace_back("-DINFINITY=__int_as_float(0x7f800000)");