TP communication overlap: enable the overlap between GEMM chunk at Hopper BF16

Signed-off-by: Sangkug Lym <[email protected]>

Author: erhoo82

transformer_engine/pytorch/cpp_extensions/gemm.py

@@ -284,10 +284,13 @@ def gemm(
             assert (
                 extra_output_tensor is not None
             ), "SPLIT_PIPELINED_RS requires extra output tensor"
+            # Disable the overlap between GEMM chunks at ampere and below
+            major, _ = torch.cuda.get_device_capability()
+            overlap_gemm_chunks = True if major >= 9 else False
             args = tuple(
                 args
                 + (
-                    False,
+                    overlap_gemm_chunks,
                     extra_output_tensor,
                 )
             )