pass membermask to elect_one_sync

Signed-off-by: Varun Sundar Rabindranath <vsundarr@redhat.com>

Author: Varun Sundar Rabindranath

csrc/kernels/intranode.cu

@@ -596,10 +596,6 @@ combine(dtype_t* recv_x, float* recv_topk_weights,
 
     // TMA stuffs
 #ifndef DISABLE_SM90_FEATURES
-    auto set_n_bits = [](const int x) -> uint32_t {
-        static constexpr uint64_t one = 1; 
-        return static_cast<uint32_t>((one << x) - one);
-    };
     extern __shared__ __align__(1024) uint8_t smem_buffer[];
     auto tma_buffer = smem_buffer + (thread_id / 32) * kNumTMABytesPerWarp;
 #endif
@@ -841,11 +837,10 @@ combine(dtype_t* recv_x, float* recv_topk_weights,
                         out_dtypes[j] = static_cast<dtype_t>(values[j]);
 
 #ifndef DISABLE_SM90_FEATURES
-		    const int num_participating_threads = min(32, hidden_int4 - i);
-                    const unsigned int warp_mask = set_n_bits(num_participating_threads);
+		    auto const warp_mask = __activemask();
 
                     // Wait TMA arrival
-                    if (elect_one_sync())
+                    if (elect_one_sync(warp_mask))
                         tma_store_wait<kNumStages - 1>();
                     __syncwarp(warp_mask);
 
@@ -856,7 +851,7 @@ combine(dtype_t* recv_x, float* recv_topk_weights,
                     // Issue TMA
                     tma_store_fence();
                     __syncwarp(warp_mask);
-                    if (elect_one_sync()) {
+                    if (elect_one_sync(warp_mask)) {
                         auto tma_bytes = min(32, hidden_int4 - i) * static_cast<int>(sizeof(int4));
                         tma_store_1d(reinterpret_cast<int4*>(tma_buffer) + tma_stage_idx * 32,
                                      recv_int4 + token_idx * hidden_int4 + i, tma_bytes, false);

csrc/kernels/utils.cuh

@@ -310,7 +310,7 @@ __forceinline__ __device__ int get_lane_id() {
     return lane_id;
 }
 
-__device__ __forceinline__ uint32_t elect_one_sync() {
+__device__ __forceinline__ uint32_t elect_one_sync(const uint32_t membermask=0xffffffff) {
 #ifndef DISABLE_SM90_FEATURES
     uint32_t pred = 0;
     asm volatile(
@@ -321,7 +321,7 @@ __device__ __forceinline__ uint32_t elect_one_sync() {
       "@%%px mov.s32 %0, 1;\n"
       "}\n"
       : "+r"(pred)
-      : "r"(0xffffffff));
+      : "r"(membermask));
     return pred;
 #else
     return get_lane_id() == 0;