[ROCm] Enable fp_quantizer on ROCm (deepspeedai#7027)

This change is required to successfully build fp_quantizer extension on ROCm. --------- Co-authored-by: Logan Adams <[email protected]> Signed-off-by: gyou2021 <[email protected]>
gyou2021 · Feb 18, 2025 · 2ce885b · 2ce885b
1 parent f7e6f9b
commit 2ce885b
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/op_builder/fp_quantizer.py b/op_builder/fp_quantizer.py
@@ -94,7 +94,10 @@ def sources(self):
         ]
 
     def extra_ldflags(self):
-        return ['-lcurand']
+        if not self.is_rocm_pytorch():
+            return ['-lcurand']
+        else:
+            return []
 
     def include_paths(self):
         return ['csrc/fp_quantizer/includes', 'csrc/includes']