From 77ecfb9b29649dcd263d6df6fb53477fc289ff9c Mon Sep 17 00:00:00 2001 From: Benson Ma Date: Fri, 24 May 2024 12:21:39 -0700 Subject: [PATCH] [fbgemm_gpu] Enable NCCL code - Enable NCCL code and tests for multiple GPU car --- .../experimental/example/test/triton_example_test.py | 2 +- fbgemm_gpu/experimental/gen_ai/CMakeLists.txt | 10 +++++----- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/fbgemm_gpu/experimental/example/test/triton_example_test.py b/fbgemm_gpu/experimental/example/test/triton_example_test.py index 8e043656ad..3e134835cc 100644 --- a/fbgemm_gpu/experimental/example/test/triton_example_test.py +++ b/fbgemm_gpu/experimental/example/test/triton_example_test.py @@ -14,7 +14,7 @@ @triton.jit # fmt: off def triton_add_kernel(x_ptr, y_ptr, z_ptr, n_elements, BLOCK_SIZE: tl.constexpr) -> None: -# fmt: on +# fmt: on # noqa E115 # We use a 1D launch grid so axis is 0. pid = tl.program_id(axis=0) diff --git a/fbgemm_gpu/experimental/gen_ai/CMakeLists.txt b/fbgemm_gpu/experimental/gen_ai/CMakeLists.txt index a82236cbba..a95759befb 100644 --- a/fbgemm_gpu/experimental/gen_ai/CMakeLists.txt +++ b/fbgemm_gpu/experimental/gen_ai/CMakeLists.txt @@ -35,14 +35,14 @@ set(quantize_ops_sources src/quantize/quantize.cu src/quantize/quantize.cpp) -# set(comm_ops_sources -# src/comm/car.cu -# src/comm/car.cpp) +set(comm_ops_sources + src/comm/car.cu + src/comm/car.cpp) set(experimental_gen_ai_cpp_source_files ${attention_ops_sources} - ${quantize_ops_sources}) -# ${comm_ops_sources}) + ${quantize_ops_sources} + ${comm_ops_sources}) set_source_files_properties(${experimental_gen_ai_cpp_source_files} PROPERTIES INCLUDE_DIRECTORIES