__grid__constant is only suported on cuda version >= 11.8 (#121275)

drisspg · pytorchmergebot · commit f72eb5ae4c56 · 2024-03-06T03:44:59.000Z
Summary: Update the macros to exclude using __grid__constant on compiling for devices > sm80 but cuda version < 11.8. Test Plan: buck2 build --keep-going --config buck2.log_configured_graph_size=true --flagfile fbcode//mode/dev fbcode//sigrid/predictor/client/python:ig_sigrid_client_pybinding Differential Revision: D54556796 Co-authored-by: Driss Guessous <drisspg@meta.com> Pull Request resolved: #121275 Approved by: https://github.com/drisspg
diff --git a/aten/src/ATen/native/transformers/cuda/flash_attn/flash_bwd_launch_template.h b/aten/src/ATen/native/transformers/cuda/flash_attn/flash_bwd_launch_template.h
@@ -16,6 +16,10 @@ namespace pytorch_flash {
 // Determine if the architecture supports FLASH and define a macro to handle parameter modifiers
 #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
 #define ARCH_SUPPORTS_FLASH
+#endif
+
+#if defined(ARCH_SUPPORTS_FLASH) && defined(__CUDACC_VER_MAJOR__) && __CUDACC_VER_MAJOR__ >= 11 && \
+    defined(__CUDACC_VER_MINOR__) && __CUDACC_VER_MINOR__ >= 8
 #define KERNEL_PARAM_MODIFIER __grid_constant__
 #else
 #define KERNEL_PARAM_MODIFIER
diff --git a/aten/src/ATen/native/transformers/cuda/flash_attn/flash_fwd_launch_template.h b/aten/src/ATen/native/transformers/cuda/flash_attn/flash_fwd_launch_template.h
@@ -15,6 +15,10 @@ namespace pytorch_flash {
 // Determine if the architecture supports FLASH and define a macro to handle parameter modifiers
 #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
 #define ARCH_SUPPORTS_FLASH
+#endif
+
+#if defined(ARCH_SUPPORTS_FLASH) && defined(__CUDACC_VER_MAJOR__) && __CUDACC_VER_MAJOR__ >= 11 && \
+    defined(__CUDACC_VER_MINOR__) && __CUDACC_VER_MINOR__ >= 8
 #define KERNEL_PARAM_MODIFIER __grid_constant__
 #else
 #define KERNEL_PARAM_MODIFIER
