[DCU] fix 2.5 get_num_compute_streams

b4a2489f · yuguo · 2b05e121 · b4a2489f
Commit b4a2489f authored Jun 17, 2025 by yuguo
Hide whitespace changes
Inline Side-by-side

Showing with 4 additions and 0 deletions

transformer_engine/common/util/multi_stream.cpp transformer_engine/common/util/multi_stream.cpp +4 -0

No files found.
--- a/transformer_engine/common/util/multi_stream.cpp
+++ b/transformer_engine/common/util/multi_stream.cpp
@@ -50,7 +50,11 @@ cudaEvent_t get_compute_stream_event(int idx) {
 }

 int get_num_compute_streams() {
+#ifdef __HIP_PLATFORM_AMD__
+  static constexpr int num_compute_streams = 2;
+#else
  static constexpr int num_compute_streams = 4;
+#endif
  return num_compute_streams;
 }