From c27483933d50a693fef9c48418d2664cf6a6a6f8 Mon Sep 17 00:00:00 2001
From: Omar Elayan <142979319+oelayan7@users.noreply.github.com>
Date: Tue, 10 Sep 2024 19:08:50 +0300
Subject: [PATCH] wrap include cuda_bf16.h with ifdef BF16_AVAILABLE (#6520)

---
 csrc/fp_quantizer/fp_quantize.cu         | 2 ++
 csrc/fp_quantizer/includes/fp_quantize.h | 2 ++
 2 files changed, 4 insertions(+)

diff --git a/csrc/fp_quantizer/fp_quantize.cu b/csrc/fp_quantizer/fp_quantize.cu
index cca063956167..66ea7392e011 100644
--- a/csrc/fp_quantizer/fp_quantize.cu
+++ b/csrc/fp_quantizer/fp_quantize.cu
@@ -15,7 +15,9 @@
 #include <cuda_fp16.h>
 #include <curand_kernel.h>
 
+#ifdef BF16_AVAILABLE
 #include <cuda_bf16.h>
+#endif
 #include <cuda_runtime_api.h>
 
 using ROp = reduce::ROpType;
diff --git a/csrc/fp_quantizer/includes/fp_quantize.h b/csrc/fp_quantizer/includes/fp_quantize.h
index 507252d6e722..60c75541f603 100644
--- a/csrc/fp_quantizer/includes/fp_quantize.h
+++ b/csrc/fp_quantizer/includes/fp_quantize.h
@@ -10,7 +10,9 @@
 
 #include <cuda_fp16.h>
 
+#ifdef BF16_AVAILABLE
 #include <cuda_bf16.h>
+#endif
 #include <cuda_runtime_api.h>
 #include <stdio.h>