Fix 2CTA TMEM Free (#189)

* Fix 2CTA TMEM

* Fix BF16 as well

Author: zheanxu

deep_gemm/include/deep_gemm/impls/sm100_bf16_gemm.cuh

@@ -411,12 +411,13 @@ sm100_bf16_gemm_impl(int* grouped_layout,
                  tensor_map_cd);
             }
         }
-
-        // Deallocate tensor memory by the last UMMA store warp
-        // NOTES: warp 0 is waiting TMA store
-        if (epilogue_warp_idx == kNumUMMAStoreThreads / 32 - 1)
-            Allocator().free(0, kNumTmemCols);
     }
+
+    // Deallocate tensor memory
+    kNumMulticast > 1 ? cute::cluster_sync() : __syncthreads();
+    if (warp_idx == 0)
+        Allocator().free(0, kNumTmemCols);
+
 #else
     if (blockIdx.x == 0 and threadIdx.x == 0)
         DG_DEVICE_ASSERT(false and "This kernel only support sm_100f");

deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_gemm_1d1d.cuh

@@ -489,12 +489,13 @@ sm100_fp8_fp4_gemm_1d1d_impl(int* grouped_layout,
                  tensor_map_cd);
             }
         }
-
-        // Deallocate tensor memory by the last UMMA store warp
-        // NOTES: warp 0 is waiting TMA store
-        if (epilogue_warp_idx == kNumUMMAStoreThreads / 32 - 1)
-            Allocator().free(0, kNumTmemCols);
     }
+
+    // Deallocate tensor memory
+    kNumMulticast > 1 ? cute::cluster_sync() : __syncthreads();
+    if (warp_idx == 0)
+        Allocator().free(0, kNumTmemCols);
+
 #else
     if (blockIdx.x == 0 and threadIdx.x == 0)
         DG_DEVICE_ASSERT(false and "This kernel only support sm_100f");